diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index d6807b6dd47a..946a282eec09 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2848,6 +2848,8 @@ extern "C" { uint32_t poll; // polling level (0 - no polling, 100 - aggressive polling) bool strict_cpu; // strict cpu placement bool paused; // start in paused state + void (*thread_create_callback)(void); // callback invoked when thread is created + void (*thread_destroy_callback)(void); // callback invoked when thread is destroyed }; struct ggml_threadpool; // forward declaration, see ggml.c diff --git a/ggml/src/ggml-c.c b/ggml/src/ggml-c.c index b43016c87d21..51a1b717a58e 100644 --- a/ggml/src/ggml-c.c +++ b/ggml/src/ggml-c.c @@ -7800,6 +7800,8 @@ void ggml_threadpool_params_init(struct ggml_threadpool_params * p, int n_thread p->poll = 50; // hybrid-polling enabled p->strict_cpu = false; // no strict placement (all threads share same cpumask) p->paused = false; // threads are ready to go + p->thread_create_callback = 0; + p->thread_destroy_callback = 0; memset(p->cpumask, 0, GGML_MAX_N_THREADS); // all-zero means use the default affinity (usually inherited) } diff --git a/ggml/src/ggml-cpu/ggml-cpu-c.c b/ggml/src/ggml-cpu/ggml-cpu-c.c index eb8341c9aecc..b6969d72eb82 100644 --- a/ggml/src/ggml-cpu/ggml-cpu-c.c +++ b/ggml/src/ggml-cpu/ggml-cpu-c.c @@ -491,6 +491,9 @@ struct ggml_threadpool { int32_t prio; // Scheduling priority uint32_t poll; // Polling level (0 - no polling) + void (*thread_create_callback)(void); + void (*thread_destroy_callback)(void); + enum ggml_status ec; }; @@ -3160,6 +3163,10 @@ static thread_ret_t ggml_graph_compute_secondary_thread(void* data) { struct ggml_compute_state * state = (struct ggml_compute_state *) data; struct ggml_threadpool * threadpool = state->threadpool; + if (threadpool->thread_create_callback) { + threadpool->thread_create_callback(); + } + ggml_thread_apply_priority(threadpool->prio); if (ggml_thread_cpumask_is_valid(state->cpumask)) { ggml_thread_apply_affinity(state->cpumask); @@ -3190,6 +3197,10 @@ static thread_ret_t ggml_graph_compute_secondary_thread(void* data) { } } + if (threadpool->thread_destroy_callback) { + threadpool->thread_destroy_callback(); + } + return (thread_ret_t) 0; } @@ -3249,6 +3260,8 @@ static struct ggml_threadpool * ggml_threadpool_new_impl( threadpool->n_threads = tpp->n_threads; threadpool->poll = tpp->poll; threadpool->prio = tpp->prio; + threadpool->thread_create_callback = tpp->thread_create_callback; + threadpool->thread_destroy_callback = tpp->thread_destroy_callback; threadpool->ec = GGML_STATUS_SUCCESS; }