diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 491316f74912..db717894a81b 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -489,6 +489,10 @@ struct ggml_threadpool { atomic_int GGML_CACHE_ALIGN n_barrier; atomic_int GGML_CACHE_ALIGN n_barrier_passed; atomic_int GGML_CACHE_ALIGN current_chunk; // currently processing chunk during Mat_Mul, shared between all the threads. + // Chunk counter dedicated to ggml_compute_forward_mul_mat(). It is kept at n_threads between two + // matmuls (see the chunk loop there), which removes the need to re-initialize it - and therefore + // the need for a barrier - at the beginning of every matmul. + atomic_int GGML_CACHE_ALIGN mul_mat_chunk; // these are atomic as an annotation for thread-sanitizer atomic_bool stop; // Used for stopping the threadpool altogether @@ -1354,15 +1358,14 @@ UseGgmlGemm1:; } } #endif - } - if (ith == 0) { - // Every thread starts at ith, so the first unprocessed chunk is nth. This save a bit of coordination right at the start. - atomic_store_explicit(¶ms->threadpool->current_chunk, nth, memory_order_relaxed); + // the converted rows of src1 are shared by all the threads, so they have to be + // published before the matmul itself can start. when no conversion is needed the + // threads only read src0/src1 and write to disjoint parts of dst, so no + // synchronization is required here at all. + ggml_barrier(params->threadpool); } - ggml_barrier(params->threadpool); - #if GGML_USE_LLAMAFILE if (src1->type != vec_dot_type) { const void* wdata = (src1->type == vec_dot_type) ? src1->data : params->wdata; @@ -1420,10 +1423,12 @@ UseGgmlGemm2:; const int64_t dr0 = (nr0 + nchunk0 - 1) / nchunk0; const int64_t dr1 = (nr1 + nchunk1 - 1) / nchunk1; + const int64_t nchunk = nchunk0 * nchunk1; + // The first chunk comes from our thread_id, the rest will get auto-assigned. int current_chunk = ith; - while (current_chunk < nchunk0 * nchunk1) { + while (current_chunk < nchunk) { const int64_t ith0 = current_chunk % nchunk0; const int64_t ith1 = current_chunk / nchunk0; @@ -1443,11 +1448,18 @@ UseGgmlGemm2:; } ggml_compute_forward_mul_mat_one_chunk(params, dst, src0->type, num_rows_per_vec_dot, ir0_start, ir0_end, ir1_start, ir1_end); - if (nth >= nchunk0 * nchunk1) { + if (nth >= nchunk) { break; } - current_chunk = atomic_fetch_add_explicit(¶ms->threadpool->current_chunk, 1, memory_order_relaxed); + // mul_mat_chunk starts at nth, so exactly nchunk tickets are handed out here: the + // nchunk - nth remaining chunks plus one out-of-range ticket per thread, which is how + // each thread leaves the loop. The thread that draws the highest ticket is therefore the + // last one to touch the counter and restores it to nth for the next matmul. + current_chunk = atomic_fetch_add_explicit(¶ms->threadpool->mul_mat_chunk, 1, memory_order_relaxed); + if (current_chunk == nth + nchunk - 1) { + atomic_store_explicit(¶ms->threadpool->mul_mat_chunk, nth, memory_order_relaxed); + } } } @@ -3284,6 +3296,7 @@ static struct ggml_threadpool * ggml_threadpool_new_impl( threadpool->n_barrier = 0; threadpool->n_barrier_passed = 0; threadpool->current_chunk = 0; + threadpool->mul_mat_chunk = tpp->n_threads; threadpool->stop = false; threadpool->pause = tpp->paused; threadpool->abort = -1; @@ -3384,6 +3397,7 @@ enum ggml_status ggml_graph_compute(struct ggml_cgraph * cgraph, struct ggml_cpl // update the number of threads from the actual number of threads that we got from OpenMP n_threads = omp_get_num_threads(); atomic_store_explicit(&threadpool->n_graph, n_threads, memory_order_relaxed); + atomic_store_explicit(&threadpool->mul_mat_chunk, n_threads, memory_order_relaxed); } // Apply thread CPU mask and priority @@ -3397,6 +3411,7 @@ enum ggml_status ggml_graph_compute(struct ggml_cgraph * cgraph, struct ggml_cpl } } else { atomic_store_explicit(&threadpool->n_graph, 1, memory_order_relaxed); + atomic_store_explicit(&threadpool->mul_mat_chunk, 1, memory_order_relaxed); ggml_graph_compute_thread(&threadpool->workers[0]); } #else @@ -3405,6 +3420,8 @@ enum ggml_status ggml_graph_compute(struct ggml_cgraph * cgraph, struct ggml_cpl n_threads = threadpool->n_threads; } + atomic_store_explicit(&threadpool->mul_mat_chunk, n_threads, memory_order_relaxed); + // Kick all threads to start the new graph ggml_graph_compute_kickoff(threadpool, n_threads);