When running parallel tasks with fewer than the maximum number of threads, ensure that the thread_id given to the caller's thread function is

src/pthreads.c

@@ -240,7 +240,7 @@ static int32_t wait_on_num_active_threads(pthreadpool_t threadpool,
         // First increase the `num_waiting_threads` counter and re-check
         // `num_active_threads` thereafter to avoid slipping past calls to
         // `signal_num_active_threads`.
-        pthreadpool_fetch_add_acquire_release_uint32_t(
+        pthreadpool_fetch_add_sequentially_consistent_uint32_t(
             &threadpool->num_waiting_threads, 1);
         if ((curr_active_threads = pthreadpool_load_consume_int32_t(
                  &threadpool->num_active_threads)) <= 0) {
@@ -324,7 +324,8 @@ static void signal_num_active_threads(pthreadpool_t threadpool,
                                       uint32_t max_num_waiting) {
 #if PTHREADPOOL_USE_FUTEX
   const uint32_t num_waiting_threads =
-      pthreadpool_load_consume_uint32_t(&threadpool->num_waiting_threads);
+      pthreadpool_load_sequentially_consistent_uint32_t(
+          &threadpool->num_waiting_threads);
   if (num_waiting_threads > max_num_waiting) {
     futex_wake_n((pthreadpool_atomic_uint32_t*)&threadpool->num_active_threads,
                  num_waiting_threads - max_num_waiting);
@@ -547,10 +548,18 @@ static pthreadpool_thread_return_t thread_main(void* arg) {
         const uint32_t max_active_threads =
             pthreadpool_load_acquire_size_t(&threadpool->threads_count);
         if (curr_active_threads < max_active_threads) {
+          // If we have more threads than requested, spoof the `thread_id` to be
+          // `curr_active_threads`, which will be unique for each thread in
+          // this task.
+          const uint32_t curr_thread_id =
+              (max_active_threads < threadpool->max_num_threads)
+                  ? curr_active_threads
+                  : thread_id;
+          pthreadpool_log_debug("thread %u working on job %u as thread %u.",
+                                thread_id, threadpool->job_id, curr_thread_id);
+
           // Do the needful.
-          pthreadpool_log_debug("thread %u working on job %u.", thread_id,
-                                threadpool->job_id);
-          run_thread_function(threadpool, thread_id);
+          run_thread_function(threadpool, curr_thread_id);
         }
 
         // Ring the bell on the way out.

src/threadpool-atomics.h

@@ -308,6 +308,11 @@ static inline int32_t pthreadpool_fetch_add_sequentially_consistent_int32_t(
   return atomic_fetch_add_explicit(address, value, memory_order_seq_cst);
 }
 
+static inline int32_t pthreadpool_fetch_add_sequentially_consistent_uint32_t(
+    pthreadpool_atomic_uint32_t* address, uint32_t value) {
+  return atomic_fetch_add_explicit(address, value, memory_order_seq_cst);
+}
+
 static inline uint32_t pthreadpool_exchange_acquire_uint32_t(
     pthreadpool_atomic_uint32_t* address, uint32_t value) {
   return atomic_exchange_explicit(address, value, memory_order_acquire);

test/pthreadpool_v2.cc

@@ -10830,15 +10830,31 @@ TEST(Parallelize6DTile2D, MultiThreadPoolWorkStealing) {
 
 struct CheckThreadIDData {
   CheckThreadIDData(size_t num_threads) : num_threads(num_threads) {}
-  size_t num_threads;
+  const size_t num_threads;
   std::set<size_t> thread_ids;
 };
 
 static void CheckThreadID(CheckThreadIDData* data, size_t thread_id, size_t) {
-  static std::mutex mutex;                  // NOLINT(build/c++11)
-  std::lock_guard<std::mutex> lock(mutex);  // NOLINT(build/c++11)
+  static std::mutex mutex;  // NOLINT(build/c++11)
+  static std::condition_variable cond_var;
+  std::unique_lock<std::mutex> lock(mutex);  // NOLINT(build/c++11)
+
+  // Make sure we have a valit `thread_id`.
+  ASSERT_LT(thread_id, data->num_threads);
+
+  // Make sure this thread is here for the first time.
+  ASSERT_FALSE(data->thread_ids.contains(thread_id));
   data->thread_ids.insert(thread_id);
-  ASSERT_LE(data->thread_ids.size(), data->num_threads);
+
+  // Wait until everyone has arrived before returning to avoid work-stealing.
+  if (data->thread_ids.size() == data->num_threads) {
+    cond_var.notify_all();
+  } else {
+    /* Wait until all items are computed */
+    while (data->thread_ids.size() < data->num_threads) {
+      cond_var.wait(lock);
+    }
+  }
 }
 
 TEST(SetNumThreads, ValidRange) {
@@ -10871,7 +10887,7 @@ TEST(SetNumThreads, ValidRange) {
     pthreadpool_parallelize_1d_with_thread(
         threadpool.get(),
         reinterpret_cast<pthreadpool_task_1d_with_thread_t>(CheckThreadID),
-        (void*)&data, kParallelize1DRange, /*flags=*/0);
+        (void*)&data, num_threads, /*flags=*/0);
   }
 }
 
@@ -10907,7 +10923,7 @@ TEST(SetNumThreads, Maximum) {
     pthreadpool_parallelize_1d_with_thread(
         threadpool.get(),
         reinterpret_cast<pthreadpool_task_1d_with_thread_t>(CheckThreadID),
-        (void*)&data1, kParallelize1DRange, /*flags=*/0);
+        (void*)&data1, num_threads, /*flags=*/0);
 
     // Set the maximum of threads ((kNumThreadpoolThreads + 1)).
     ASSERT_EQ(pthreadpool_set_threads_count(threadpool.get(), 0),
@@ -10919,7 +10935,7 @@ TEST(SetNumThreads, Maximum) {
     pthreadpool_parallelize_1d_with_thread(
         threadpool.get(),
         reinterpret_cast<pthreadpool_task_1d_with_thread_t>(CheckThreadID),
-        (void*)&data2, kParallelize1DRange, /*flags=*/0);
+        (void*)&data2, max_num_threads, /*flags=*/0);
   }
 }
 
@@ -10950,7 +10966,7 @@ TEST(SetNumThreads, TooHigh) {
     pthreadpool_parallelize_1d_with_thread(
         threadpool.get(),
         reinterpret_cast<pthreadpool_task_1d_with_thread_t>(CheckThreadID),
-        (void*)&data, kParallelize1DRange, /*flags=*/0);
+        (void*)&data, max_num_threads, /*flags=*/0);
   }
 }