diff --git a/src/pthreads.c b/src/pthreads.c index 0ea30be..f0b25ae 100644 --- a/src/pthreads.c +++ b/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. diff --git a/src/threadpool-atomics.h b/src/threadpool-atomics.h index cda87c5..67d4870 100644 --- a/src/threadpool-atomics.h +++ b/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); diff --git a/test/pthreadpool_v2.cc b/test/pthreadpool_v2.cc index 72ca4bc..a1a8b89 100644 --- a/test/pthreadpool_v2.cc +++ b/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 thread_ids; }; static void CheckThreadID(CheckThreadIDData* data, size_t thread_id, size_t) { - static std::mutex mutex; // NOLINT(build/c++11) - std::lock_guard lock(mutex); // NOLINT(build/c++11) + static std::mutex mutex; // NOLINT(build/c++11) + static std::condition_variable cond_var; + std::unique_lock 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(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(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(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(CheckThreadID), - (void*)&data, kParallelize1DRange, /*flags=*/0); + (void*)&data, max_num_threads, /*flags=*/0); } }