From e8a391afb29f15092237933340edddd186ec226c Mon Sep 17 00:00:00 2001
From: Sam Schweigel <sam.schweigel@juliacomputing.com>
Date: Thu, 23 Jan 2025 14:33:28 -0800
Subject: [PATCH] Add JL_THREADPOOL_ID_*, fix missing a thread when setting
 affinities

---
 src/init.c      |  4 ++--
 src/julia.h     |  3 +++
 src/threading.c | 31 ++++++++++++++++++-------------
 3 files changed, 23 insertions(+), 15 deletions(-)

diff --git a/src/init.c b/src/init.c
index 7b41e63e98455..ff07c3a9e26ab 100644
--- a/src/init.c
+++ b/src/init.c
@@ -902,8 +902,8 @@ static NOINLINE void _finish_julia_init(JL_IMAGE_SEARCH rel, jl_ptls_t ptls, jl_
         jl_n_markthreads = 0;
         jl_n_sweepthreads = 0;
         jl_n_gcthreads = 0;
-        jl_n_threads_per_pool[0] = 0; // Interactive threadpool
-        jl_n_threads_per_pool[1] = 1; // Default threadpool
+        jl_n_threads_per_pool[JL_THREADPOOL_ID_INTERACTIVE] = 0;
+        jl_n_threads_per_pool[JL_THREADPOOL_ID_DEFAULT] = 1;
     } else {
         post_image_load_hooks();
     }
diff --git a/src/julia.h b/src/julia.h
index b5416568b7ae9..b2d72a7763d2b 100644
--- a/src/julia.h
+++ b/src/julia.h
@@ -2058,6 +2058,9 @@ extern JL_DLLIMPORT _Atomic(int) jl_n_threads;
 extern JL_DLLIMPORT int jl_n_gcthreads;
 extern int jl_n_markthreads;
 extern int jl_n_sweepthreads;
+
+#define JL_THREADPOOL_ID_INTERACTIVE 0
+#define JL_THREADPOOL_ID_DEFAULT 1
 extern JL_DLLIMPORT int *jl_n_threads_per_pool;
 
 // environment entries
diff --git a/src/threading.c b/src/threading.c
index caf5b86e0b208..e0d0e04625b8f 100644
--- a/src/threading.c
+++ b/src/threading.c
@@ -778,9 +778,9 @@ void jl_init_threading(void)
     }
 
     jl_all_tls_states_size = nthreads + nthreadsi + ngcthreads;
-    jl_n_threads_per_pool = (int*)malloc_s(2 * sizeof(int));
-    jl_n_threads_per_pool[0] = nthreadsi;
-    jl_n_threads_per_pool[1] = nthreads;
+    jl_n_threads_per_pool = (int*)malloc_s(jl_n_threadpools * sizeof(int));
+    jl_n_threads_per_pool[JL_THREADPOOL_ID_INTERACTIVE] = nthreadsi;
+    jl_n_threads_per_pool[JL_THREADPOOL_ID_DEFAULT] = nthreads;
     assert(jl_all_tls_states_size > 0);
     jl_atomic_store_release(&jl_all_tls_states, (jl_ptls_t*)calloc(jl_all_tls_states_size, sizeof(jl_ptls_t)));
     jl_atomic_store_release(&jl_n_threads, jl_all_tls_states_size);
@@ -793,9 +793,9 @@ uv_barrier_t thread_init_done;
 void jl_start_threads(void)
 {
     int nthreads = jl_atomic_load_relaxed(&jl_n_threads);
-    int ngcthreads = jl_n_gcthreads;
-    int nthreadsi = jl_n_threads_per_pool[0];
-    int nmutator_threads = nthreads - ngcthreads;
+    int ninteractive_threads = jl_n_threads_per_pool[JL_THREADPOOL_ID_INTERACTIVE];
+    int ndefault_threads = jl_n_threads_per_pool[JL_THREADPOOL_ID_DEFAULT];
+    int nmutator_threads = nthreads - jl_n_gcthreads;
 
     int cpumasksize = uv_cpumask_size();
     char *cp;
@@ -811,17 +811,19 @@ void jl_start_threads(void)
     if (cp && strcmp(cp, "0") != 0)
         exclusive = 1;
 
-    // exclusive use: affinitize threads, master thread on proc 0, rest
-    // according to a 'compact' policy
+    // exclusive use: affinitize threads, master thread on proc 0, threads in
+    // default pool according to a 'compact' policy
     // non-exclusive: no affinity settings; let the kernel move threads about
     if (exclusive) {
-        if (nmutator_threads - nthreadsi > jl_cpu_threads()) {
+        if (ndefault_threads > jl_cpu_threads()) {
             jl_printf(JL_STDERR, "ERROR: Too many threads requested for %s option.\n", MACHINE_EXCLUSIVE_NAME);
             exit(1);
         }
         memset(mask, 0, cpumasksize);
 
-        if (nthreadsi == 0) {
+        // If there are no interactive threads, the master thread is in the
+        // default pool and we must affinitize it
+        if (ninteractive_threads == 0) {
             mask[0] = 1;
             uvtid = uv_thread_self();
             uv_thread_setaffinity(&uvtid, mask, NULL, cpumasksize);
@@ -838,10 +840,13 @@ void jl_start_threads(void)
         t->tid = i;
         t->barrier = &thread_init_done;
         uv_thread_create(&uvtid, jl_threadfun, t);
-        if (exclusive && i > nthreadsi) {
-            mask[i - nthreadsi] = 1;
+
+        // Interactive pool threads get the low IDs, so check if this is a
+        // default pool thread.  The master thread is already on CPU 0.
+        if (exclusive && i >= ninteractive_threads) {
+            mask[i - ninteractive_threads] = 1;
             uv_thread_setaffinity(&uvtid, mask, NULL, cpumasksize);
-            mask[i - nthreadsi] = 0;
+            mask[i - ninteractive_threads] = 0;
         }
         uv_thread_detach(&uvtid);
     }