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 | 90 | ThreadPool::ThreadPool() { | |
| 16 | 45 | unsigned hw = std::thread::hardware_concurrency(); | |
| 17 |
1/2✓ Branch 0 taken 45 times.
✗ Branch 1 not taken.
|
45 | 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 45 times.
✗ Branch 1 not taken.
|
45 | const unsigned num_workers = hw > 1 ? hw - 1 : 0; |
| 21 |
1/2✓ Branch 0 taken 45 times.
✗ Branch 1 not taken.
|
45 | workers_.reserve(num_workers); |
| 22 |
2/2✓ Branch 0 taken 90 times.
✓ Branch 1 taken 45 times.
|
135 | for (unsigned i = 0; i < num_workers; ++i) { |
| 23 |
1/2✓ Branch 0 taken 90 times.
✗ Branch 1 not taken.
|
180 | workers_.emplace_back([this] { worker_loop(); }); |
| 24 | 90 | } | |
| 25 | 45 | } | |
| 26 | |||
| 27 | 90 | ThreadPool::~ThreadPool() { | |
| 28 |
1/2✓ Branch 0 taken 45 times.
✗ Branch 1 not taken.
|
45 | shutdown(); |
| 29 | 90 | } | |
| 30 | |||
| 31 | 45 | void ThreadPool::shutdown() { | |
| 32 | 45 | shutdown_.store(true, std::memory_order_relaxed); | |
| 33 | 45 | generation_.fetch_add(1, std::memory_order_release); | |
| 34 | 45 | generation_.notify_all(); | |
| 35 |
2/2✓ Branch 0 taken 90 times.
✓ Branch 1 taken 45 times.
|
135 | for (auto& t : workers_) { |
| 36 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 90 times.
|
90 | 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 | 45 | workers_.clear(); | |
| 41 | 45 | } | |
| 42 | |||
| 43 | namespace { | ||
| 44 | std::atomic<ThreadPool*> g_made{nullptr}; | ||
| 45 | } // namespace | ||
| 46 | |||
| 47 | 3719 | ThreadPool& ThreadPool::instance() { | |
| 48 |
4/6✓ Branch 0 taken 45 times.
✓ Branch 1 taken 3674 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 45 times.
✓ Branch 4 taken 45 times.
✗ Branch 5 not taken.
|
3719 | static ThreadPool pool; |
| 49 | 3719 | g_made.store(&pool, std::memory_order_release); | |
| 50 | 3719 | return pool; | |
| 51 | ✗ | } | |
| 52 | |||
| 53 | ✗ | ThreadPool* ThreadPool::existing() { | |
| 54 | ✗ | return g_made.load(std::memory_order_acquire); | |
| 55 | } | ||
| 56 | |||
| 57 | 90 | void ThreadPool::worker_loop() { | |
| 58 | 90 | int seen = 0; | |
| 59 | 6327 | for (;;) { | |
| 60 | 6327 | generation_.wait(seen); | |
| 61 | 6327 | seen = generation_.load(std::memory_order_acquire); | |
| 62 |
2/2✓ Branch 0 taken 6237 times.
✓ Branch 1 taken 90 times.
|
6327 | if (shutdown_.load(std::memory_order_relaxed)) return; |
| 63 | |||
| 64 | 6237 | t_active = true; | |
| 65 | 563969 | for (;;) { | |
| 66 | 563969 | const std::size_t i = cursor_.fetch_add(1, std::memory_order_relaxed); | |
| 67 |
2/2✓ Branch 0 taken 6237 times.
✓ Branch 1 taken 557732 times.
|
563969 | if (i >= n_) break; |
| 68 | 557732 | (*fn_)(i); | |
| 69 | } | ||
| 70 | 6237 | t_active = false; | |
| 71 |
2/2✓ Branch 0 taken 3100 times.
✓ Branch 1 taken 3137 times.
|
6237 | if (outstanding_.fetch_sub(1, std::memory_order_acq_rel) == 1) { |
| 72 | 3137 | outstanding_.notify_all(); | |
| 73 | 3137 | } | |
| 74 | } | ||
| 75 | } | ||
| 76 | |||
| 77 | 3522 | void ThreadPool::run(std::size_t n, const std::function<void(std::size_t)>& fn) { | |
| 78 |
2/2✓ Branch 0 taken 3521 times.
✓ Branch 1 taken 1 time.
|
3522 | if (n == 0) return; |
| 79 |
5/6✓ Branch 0 taken 3513 times.
✓ Branch 1 taken 8 times.
✓ Branch 2 taken 3513 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 376 times.
✓ Branch 5 taken 3137 times.
|
3521 | if (t_active || workers_.empty() || n == 1) { |
| 80 | // Either nested (this thread is already draining a job's fn_, so | ||
| 81 | // the shared job state is in use — see class comment) or not worth | ||
| 82 | // dispatching (no spare cores, or a single item). | ||
| 83 |
2/2✓ Branch 0 taken 384 times.
✓ Branch 1 taken 499 times.
|
883 | for (std::size_t i = 0; i < n; ++i) fn(i); |
| 84 | 384 | return; | |
| 85 | } | ||
| 86 | |||
| 87 | 3137 | fn_ = &fn; | |
| 88 | 3137 | n_ = n; | |
| 89 | 3137 | cursor_.store(0, std::memory_order_relaxed); | |
| 90 | 3137 | outstanding_.store(static_cast<int>(workers_.size()), std::memory_order_relaxed); | |
| 91 | |||
| 92 | 3137 | generation_.fetch_add(1, std::memory_order_release); | |
| 93 | 3137 | generation_.notify_all(); | |
| 94 | |||
| 95 | // The calling thread drains work too instead of idling. | ||
| 96 | 3137 | t_active = true; | |
| 97 | 451371 | for (;;) { | |
| 98 | 451371 | const std::size_t i = cursor_.fetch_add(1, std::memory_order_relaxed); | |
| 99 |
2/2✓ Branch 0 taken 448234 times.
✓ Branch 1 taken 3137 times.
|
451371 | if (i >= n_) break; |
| 100 | 448234 | fn(i); | |
| 101 | } | ||
| 102 | 3137 | t_active = false; | |
| 103 | |||
| 104 | 6194 | for (;;) { | |
| 105 | 6194 | const int rem = outstanding_.load(std::memory_order_acquire); | |
| 106 |
2/2✓ Branch 0 taken 3057 times.
✓ Branch 1 taken 3137 times.
|
6194 | if (rem == 0) break; |
| 107 | 3057 | outstanding_.wait(rem); | |
| 108 | } | ||
| 109 | 3522 | } | |
| 110 | |||
| 111 | } // namespace brotensor::detail::cpu | ||
| 112 |