src/cpu/thread_pool.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #include <brotensor/detail/cpu/thread_pool.h> | ||
| 2 | |||
| 3 | namespace brotensor::detail::cpu { | ||
| 4 | |||
| 5 | namespace { | ||
| 6 | // Set while the calling thread (original caller or a worker) is inside a | ||
| 7 | // job's fn_ calls. Guards against the reentrancy hazard where fn_ itself | ||
| 8 | // invokes parallel_for again: without this, every thread currently | ||
| 9 | // executing fn_ would race to claim the singleton's shared job state for | ||
| 10 | // the nested call. A nested call instead runs as a plain sequential loop | ||
| 11 | // on whichever thread hits it. | ||
| 12 | thread_local bool t_active = false; | ||
| 13 | } // namespace | ||
| 14 | |||
| 15 | 56 | ThreadPool::ThreadPool() { | |
| 16 | 28 | unsigned hw = std::thread::hardware_concurrency(); | |
| 17 |
1/2✓ Branch 0 taken 28 times.
✗ Branch 1 not taken.
|
28 | if (hw == 0) hw = 1; |
| 18 | // Reserve one logical core for the calling thread, which also drains | ||
| 19 | // work instead of idling while it waits for the pool. | ||
| 20 |
1/2✓ Branch 0 taken 28 times.
✗ Branch 1 not taken.
|
28 | const unsigned num_workers = hw > 1 ? hw - 1 : 0; |
| 21 |
1/2✓ Branch 0 taken 28 times.
✗ Branch 1 not taken.
|
28 | workers_.reserve(num_workers); |
| 22 |
2/2✓ Branch 0 taken 56 times.
✓ Branch 1 taken 28 times.
|
84 | for (unsigned i = 0; i < num_workers; ++i) { |
| 23 |
1/2✓ Branch 0 taken 56 times.
✗ Branch 1 not taken.
|
112 | workers_.emplace_back([this] { worker_loop(); }); |
| 24 | 56 | } | |
| 25 | 28 | } | |
| 26 | |||
| 27 | 56 | ThreadPool::~ThreadPool() { | |
| 28 |
1/2✓ Branch 0 taken 28 times.
✗ Branch 1 not taken.
|
28 | shutdown(); |
| 29 | 56 | } | |
| 30 | |||
| 31 | 28 | void ThreadPool::shutdown() { | |
| 32 | 28 | shutdown_.store(true, std::memory_order_relaxed); | |
| 33 | 28 | generation_.fetch_add(1, std::memory_order_release); | |
| 34 | 28 | generation_.notify_all(); | |
| 35 |
2/2✓ Branch 0 taken 56 times.
✓ Branch 1 taken 28 times.
|
84 | for (auto& t : workers_) { |
| 36 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 56 times.
|
56 | if (t.joinable()) t.join(); |
| 37 | } | ||
| 38 | // Idempotent: an empty vector makes a second call (e.g. the eventual | ||
| 39 | // destructor, after an explicit early shutdown()) a no-op. | ||
| 40 | 28 | workers_.clear(); | |
| 41 | 28 | } | |
| 42 | |||
| 43 | 2998 | ThreadPool& ThreadPool::instance() { | |
| 44 |
4/6✓ Branch 0 taken 28 times.
✓ Branch 1 taken 2970 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 28 times.
✓ Branch 4 taken 28 times.
✗ Branch 5 not taken.
|
2998 | static ThreadPool pool; |
| 45 | 2998 | return pool; | |
| 46 | ✗ | } | |
| 47 | |||
| 48 | 56 | void ThreadPool::worker_loop() { | |
| 49 | 56 | int seen = 0; | |
| 50 | 5318 | for (;;) { | |
| 51 | 5318 | generation_.wait(seen); | |
| 52 | 5318 | seen = generation_.load(std::memory_order_acquire); | |
| 53 |
2/2✓ Branch 0 taken 5262 times.
✓ Branch 1 taken 56 times.
|
5318 | if (shutdown_.load(std::memory_order_relaxed)) return; |
| 54 | |||
| 55 | 5262 | t_active = true; | |
| 56 | 567500 | for (;;) { | |
| 57 | 567500 | const std::size_t i = cursor_.fetch_add(1, std::memory_order_relaxed); | |
| 58 |
2/2✓ Branch 0 taken 5262 times.
✓ Branch 1 taken 562238 times.
|
567500 | if (i >= n_) break; |
| 59 | 562238 | (*fn_)(i); | |
| 60 | } | ||
| 61 | 5262 | t_active = false; | |
| 62 |
2/2✓ Branch 0 taken 2625 times.
✓ Branch 1 taken 2637 times.
|
5262 | if (outstanding_.fetch_sub(1, std::memory_order_acq_rel) == 1) { |
| 63 | 2637 | outstanding_.notify_all(); | |
| 64 | 2637 | } | |
| 65 | } | ||
| 66 | } | ||
| 67 | |||
| 68 | 2998 | void ThreadPool::run(std::size_t n, const std::function<void(std::size_t)>& fn) { | |
| 69 |
2/2✓ Branch 0 taken 2997 times.
✓ Branch 1 taken 1 time.
|
2998 | if (n == 0) return; |
| 70 |
5/6✓ Branch 0 taken 2989 times.
✓ Branch 1 taken 8 times.
✓ Branch 2 taken 2989 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 352 times.
✓ Branch 5 taken 2637 times.
|
2997 | if (t_active || workers_.empty() || n == 1) { |
| 71 | // Either nested (this thread is already draining a job's fn_, so | ||
| 72 | // the shared job state is in use — see class comment) or not worth | ||
| 73 | // dispatching (no spare cores, or a single item). | ||
| 74 |
2/2✓ Branch 0 taken 360 times.
✓ Branch 1 taken 452 times.
|
812 | for (std::size_t i = 0; i < n; ++i) fn(i); |
| 75 | 360 | return; | |
| 76 | } | ||
| 77 | |||
| 78 | 2637 | fn_ = &fn; | |
| 79 | 2637 | n_ = n; | |
| 80 | 2637 | cursor_.store(0, std::memory_order_relaxed); | |
| 81 | 2637 | outstanding_.store(static_cast<int>(workers_.size()), std::memory_order_relaxed); | |
| 82 | |||
| 83 | 2637 | generation_.fetch_add(1, std::memory_order_release); | |
| 84 | 2637 | generation_.notify_all(); | |
| 85 | |||
| 86 | // The calling thread drains work too instead of idling. | ||
| 87 | 2637 | t_active = true; | |
| 88 | 405913 | for (;;) { | |
| 89 | 405913 | const std::size_t i = cursor_.fetch_add(1, std::memory_order_relaxed); | |
| 90 |
2/2✓ Branch 0 taken 403276 times.
✓ Branch 1 taken 2637 times.
|
405913 | if (i >= n_) break; |
| 91 | 403276 | fn(i); | |
| 92 | } | ||
| 93 | 2637 | t_active = false; | |
| 94 | |||
| 95 | 5127 | for (;;) { | |
| 96 | 5127 | const int rem = outstanding_.load(std::memory_order_acquire); | |
| 97 |
2/2✓ Branch 0 taken 2490 times.
✓ Branch 1 taken 2637 times.
|
5127 | if (rem == 0) break; |
| 98 | 2490 | outstanding_.wait(rem); | |
| 99 | } | ||
| 100 | 2998 | } | |
| 101 | |||
| 102 | } // namespace brotensor::detail::cpu | ||
| 103 |