From 3008b31b17ba768485d9d67156964f859b5597a7 Mon Sep 17 00:00:00 2001
From: Faisal Zaghloul <quic_fzaghlou@quicinc.com>
Date: Wed, 31 Jul 2024 12:42:31 -0400
Subject: [PATCH] fix deadlock for cases where cgraph.n_nodes == 1

and fix --poll case
---
 ggml/src/ggml.c | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c
index d6aebd37b..5375bdd31 100644
--- a/ggml/src/ggml.c
+++ b/ggml/src/ggml.c
@@ -19114,7 +19114,14 @@ static thread_ret_t ggml_graph_compute_thread(void * data) {
         }
     }
 
+    if (cgraph->n_nodes == 1) {
+        // We need a barrier before disabling new_work in case we have a trivial graph
+        ggml_barrier(state->threadpool);
+    }
+
     if (!state->threadpool->disposable && state->ith == 0) {
+        // Don't need a lock, because there is a barrier after this, and only after that
+        // do the secondary threads go into standby
         state->threadpool->new_work = false;
     }
 
@@ -19356,6 +19363,8 @@ enum ggml_status ggml_graph_compute(struct ggml_cgraph * cgraph, struct ggml_cpl
             threadpool->new_work = true;
             ggml_cond_broadcast(&threadpool->cond);
             ggml_mutex_unlock(&threadpool->mutex);
+        } else {
+            threadpool->new_work = true;
         }
     }
     // this is a work thread too