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 | 52 | ThreadPool::ThreadPool() { | |
| 16 | 26 | unsigned hw = std::thread::hardware_concurrency(); | |
| 17 |
1/2✓ Branch 0 taken 26 times.
✗ Branch 1 not taken.
|
26 | 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 26 times.
✗ Branch 1 not taken.
|
26 | const unsigned num_workers = hw > 1 ? hw - 1 : 0; |
| 21 |
1/2✓ Branch 0 taken 26 times.
✗ Branch 1 not taken.
|
26 | workers_.reserve(num_workers); |
| 22 |
2/2✓ Branch 0 taken 52 times.
✓ Branch 1 taken 26 times.
|
78 | for (unsigned i = 0; i < num_workers; ++i) { |
| 23 |
1/2✓ Branch 0 taken 52 times.
✗ Branch 1 not taken.
|
104 | workers_.emplace_back([this] { worker_loop(); }); |
| 24 | 52 | } | |
| 25 | 26 | } | |
| 26 | |||
| 27 | 52 | ThreadPool::~ThreadPool() { | |
| 28 |
1/2✓ Branch 0 taken 26 times.
✗ Branch 1 not taken.
|
26 | shutdown(); |
| 29 | 52 | } | |
| 30 | |||
| 31 | 26 | void ThreadPool::shutdown() { | |
| 32 | 26 | shutdown_.store(true, std::memory_order_relaxed); | |
| 33 | 26 | generation_.fetch_add(1, std::memory_order_release); | |
| 34 | 26 | generation_.notify_all(); | |
| 35 |
2/2✓ Branch 0 taken 52 times.
✓ Branch 1 taken 26 times.
|
78 | for (auto& t : workers_) { |
| 36 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 52 times.
|
52 | 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 | 26 | workers_.clear(); | |
| 41 | 26 | } | |
| 42 | |||
| 43 | 1692 | ThreadPool& ThreadPool::instance() { | |
| 44 |
4/6✓ Branch 0 taken 26 times.
✓ Branch 1 taken 1666 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 26 times.
✓ Branch 4 taken 26 times.
✗ Branch 5 not taken.
|
1692 | static ThreadPool pool; |
| 45 | 1692 | return pool; | |
| 46 | ✗ | } | |
| 47 | |||
| 48 | 52 | void ThreadPool::worker_loop() { | |
| 49 | 52 | int seen = 0; | |
| 50 | 2685 | for (;;) { | |
| 51 | 2685 | generation_.wait(seen); | |
| 52 | 2685 | seen = generation_.load(std::memory_order_acquire); | |
| 53 |
2/2✓ Branch 0 taken 2633 times.
✓ Branch 1 taken 52 times.
|
2685 | if (shutdown_.load(std::memory_order_relaxed)) return; |
| 54 | |||
| 55 | 2633 | t_active = true; | |
| 56 | 566054 | for (;;) { | |
| 57 | 566054 | const std::size_t i = cursor_.fetch_add(1, std::memory_order_relaxed); | |
| 58 |
2/2✓ Branch 0 taken 2633 times.
✓ Branch 1 taken 563421 times.
|
566054 | if (i >= n_) break; |
| 59 | 563421 | (*fn_)(i); | |
| 60 | } | ||
| 61 | 2633 | t_active = false; | |
| 62 |
2/2✓ Branch 0 taken 1303 times.
✓ Branch 1 taken 1330 times.
|
2633 | if (outstanding_.fetch_sub(1, std::memory_order_acq_rel) == 1) { |
| 63 | 1330 | outstanding_.notify_all(); | |
| 64 | 1330 | } | |
| 65 | } | ||
| 66 | } | ||
| 67 | |||
| 68 | 1691 | void ThreadPool::run(std::size_t n, const std::function<void(std::size_t)>& fn) { | |
| 69 |
2/2✓ Branch 0 taken 1690 times.
✓ Branch 1 taken 1 time.
|
1691 | if (n == 0) return; |
| 70 |
5/6✓ Branch 0 taken 1682 times.
✓ Branch 1 taken 8 times.
✓ Branch 2 taken 1682 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 352 times.
✓ Branch 5 taken 1330 times.
|
1690 | 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 457 times.
|
817 | for (std::size_t i = 0; i < n; ++i) fn(i); |
| 75 | 360 | return; | |
| 76 | } | ||
| 77 | |||
| 78 | 1330 | fn_ = &fn; | |
| 79 | 1330 | n_ = n; | |
| 80 | 1330 | cursor_.store(0, std::memory_order_relaxed); | |
| 81 | 1330 | outstanding_.store(static_cast<int>(workers_.size()), std::memory_order_relaxed); | |
| 82 | |||
| 83 | 1330 | generation_.fetch_add(1, std::memory_order_release); | |
| 84 | 1330 | generation_.notify_all(); | |
| 85 | |||
| 86 | // The calling thread drains work too instead of idling. | ||
| 87 | 1330 | t_active = true; | |
| 88 | 375107 | for (;;) { | |
| 89 | 375107 | const std::size_t i = cursor_.fetch_add(1, std::memory_order_relaxed); | |
| 90 |
2/2✓ Branch 0 taken 373777 times.
✓ Branch 1 taken 1330 times.
|
375107 | if (i >= n_) break; |
| 91 | 373777 | fn(i); | |
| 92 | } | ||
| 93 | 1330 | t_active = false; | |
| 94 | |||
| 95 | 2962 | for (;;) { | |
| 96 | 2962 | const int rem = outstanding_.load(std::memory_order_acquire); | |
| 97 |
2/2✓ Branch 0 taken 1632 times.
✓ Branch 1 taken 1330 times.
|
2962 | if (rem == 0) break; |
| 98 | 1632 | outstanding_.wait(rem); | |
| 99 | } | ||
| 100 | 1691 | } | |
| 101 | |||
| 102 | } // namespace brotensor::detail::cpu | ||
| 103 |