From b56a3c2dc98fa591eabf2a5b63aad8551984585c Mon Sep 17 00:00:00 2001
From: Kawrakow <iwankawrakow@gmail.com>
Date: Thu, 19 Mar 2026 07:33:00 +0100
Subject: [PATCH] Better barrier (#1456)

---
 ggml/include/ggml.h         |  1 +
 ggml/src/ggml.c             | 31 +++++++++++++++++++++----------
 src/llama-build-context.cpp |  4 +++-
 3 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h
index 3c103e6a..ec45e96d 100644
--- a/ggml/include/ggml.h
+++ b/ggml/include/ggml.h
@@ -846,6 +846,7 @@ extern "C" {
         int size;
         int n_nodes;
         int n_leafs;
+        int n_batch;
 
         struct ggml_tensor ** nodes;
         struct ggml_tensor ** grads;
diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c
index 03ff48a5..549a6127 100644
--- a/ggml/src/ggml.c
+++ b/ggml/src/ggml.c
@@ -2783,6 +2783,7 @@ struct ggml_compute_state_shared {
     const struct ggml_cplan * cplan;
 
     int n_threads;
+    int n_batch;
 
     // synchronization primitives
     atomic_int n_barrier;
@@ -4497,16 +4498,7 @@ inline static void ggml_critical_section_start(void) {
     }
 }
 
-#ifdef GGML_USE_OPENMP
-static void ggml_barrier(struct ggml_compute_state_shared * shared) {
-    if (shared->n_threads == 1) {
-        return;
-    }
-
-    #pragma omp barrier
-}
-#else
-static void ggml_barrier(struct ggml_compute_state_shared * shared) {
+static inline void ggml_barrier_impl(struct ggml_compute_state_shared * shared) {
     if (shared->n_threads == 1) {
         return;
     }
@@ -4539,6 +4531,22 @@ static void ggml_barrier(struct ggml_compute_state_shared * shared) {
         }
     }
 }
+
+#ifdef GGML_USE_OPENMP
+static void ggml_barrier(struct ggml_compute_state_shared * shared) {
+    if (shared->n_threads == 1) {
+        return;
+    }
+    if (shared && shared->n_batch > 32) {
+        ggml_barrier_impl(shared);
+        return;
+    }
+    #pragma omp barrier
+}
+#else
+static void ggml_barrier(struct ggml_compute_state_shared * shared) {
+    ggml_barrier_impl(shared);
+}
 #endif
 
 // TODO: make this somehow automatically executed
@@ -25878,6 +25886,7 @@ struct ggml_cgraph * ggml_new_graph_custom(struct ggml_context * ctx, size_t siz
         /*.size         =*/ size,
         /*.n_nodes      =*/ 0,
         /*.n_leafs      =*/ 0,
+        /*.n_batch      =*/ 0,
         /*.nodes        =*/ nodes_ptr,
         /*.grads        =*/ grads_ptr,
         /*.leafs        =*/ leafs_ptr,
@@ -25899,6 +25908,7 @@ struct ggml_cgraph ggml_graph_view(struct ggml_cgraph * cgraph0, int i0, int i1)
         /*.size         =*/ 0,
         /*.n_nodes      =*/ i1 - i0,
         /*.n_leafs      =*/ 0,
+        /*.n_batch      =*/ cgraph0->n_batch,
         /*.nodes        =*/ cgraph0->nodes + i0,
         /*.grads        =*/ cgraph0->grads ? cgraph0->grads + i0 : NULL,
         /*.leafs        =*/ NULL,
@@ -26638,6 +26648,7 @@ enum ggml_status ggml_graph_compute(struct ggml_cgraph * cgraph, struct ggml_cpl
         /*.cgraph                  =*/ cgraph,
         /*.cgraph_plan             =*/ cplan,
         /*.n_threads               =*/ n_threads,
+        /*.n_batch                 =*/ cgraph->n_batch,
         /*.n_barrier               =*/ 0,
         /*.n_barrier_passed        =*/ 0,
         /*.abort_callback          =*/ NULL,
diff --git a/src/llama-build-context.cpp b/src/llama-build-context.cpp
index 010f5d4e..7e5d7207 100644
--- a/src/llama-build-context.cpp
+++ b/src/llama-build-context.cpp
@@ -9988,8 +9988,10 @@ ggml_cgraph * llm_build_context::llama_build_graph(
             GGML_ABORT("fatal error");
     }
 
+    result->n_batch = llm.n_tokens;
+
     // add on pooling layer
-    if (lctx.cparams.mtp_op_type == MTP_OP_NONE && (lctx.cparams.embeddings || 
+    if (lctx.cparams.mtp_op_type == MTP_OP_NONE && (lctx.cparams.embeddings ||
         (lctx.model.hparams.nextn_predict_layers > 0 || lctx.model.mtp))) {
         result = llm.append_pooling(result);
     }