include/brotensor/detail/cpu/thread_pool.h
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #pragma once | ||
| 2 | |||
| 3 | // Minimal lock-free thread pool for parallelizing embarrassingly-parallel | ||
| 4 | // CPU ops (matmul's row axis, a norm's batch*channel axis, attention's | ||
| 5 | // batch*head axis, ...). Per this project's convention (atomics / single- | ||
| 6 | // owner / queues, never a lock), work handoff uses std::atomic<T>::wait / | ||
| 7 | // notify (C++20) instead of a condition_variable + mutex, and work items | ||
| 8 | // are claimed via an atomic fetch_add cursor rather than a guarded queue. | ||
| 9 | // | ||
| 10 | // Usage: brotensor::detail::cpu::parallel_for(n, [&](std::size_t i) { ... }); | ||
| 11 | // blocks the calling thread until fn(i) has run for every i in [0, n) — | ||
| 12 | // across the pool's workers plus the calling thread itself, which also | ||
| 13 | // drains work instead of idling. Falls back to a plain sequential loop | ||
| 14 | // when n <= 1 or the host has no spare cores, so tiny problem sizes don't | ||
| 15 | // pay thread hand-off cost. | ||
| 16 | // | ||
| 17 | // Single job in flight at a time: ThreadPool::instance() is a process-wide | ||
| 18 | // singleton, and run() assumes it is not re-entered from a second | ||
| 19 | // concurrent application thread while a call is outstanding (matches this | ||
| 20 | // library's usage — CPU ops are driven from one calling thread; run() | ||
| 21 | // parallelizes *within* one op call, not across concurrent callers). | ||
| 22 | // | ||
| 23 | // Nested calls (an outer parallel_for's callback invoking parallel_for | ||
| 24 | // again — e.g. one op calling another that also parallelizes internally) | ||
| 25 | // are safe but not further parallelized: a per-thread flag detects that the | ||
| 26 | // calling thread is already draining a pool job and runs the nested call | ||
| 27 | // as a plain sequential loop instead of touching the shared job state a | ||
| 28 | // second time, on whichever thread hits it (the original caller or a | ||
| 29 | // worker). | ||
| 30 | |||
| 31 | #include <atomic> | ||
| 32 | #include <cstddef> | ||
| 33 | #include <functional> | ||
| 34 | #include <thread> | ||
| 35 | #include <vector> | ||
| 36 | |||
| 37 | namespace brotensor::detail::cpu { | ||
| 38 | |||
| 39 | class ThreadPool { | ||
| 40 | public: | ||
| 41 | static ThreadPool& instance(); | ||
| 42 | // The pool if instance() has made it, else null: shutdown asks this, so | ||
| 43 | // a process that never ran a CPU kernel does not start a thread per core | ||
| 44 | // only to join them. | ||
| 45 | static ThreadPool* existing(); | ||
| 46 | |||
| 47 | void run(std::size_t n, const std::function<void(std::size_t)>& fn); | ||
| 48 | |||
| 49 | // Workers plus the calling thread, i.e. the effective parallelism a | ||
| 50 | // run() call can spread work across. | ||
| 51 | 197 | int num_threads() const { return static_cast<int>(workers_.size()) + 1; } | |
| 52 | |||
| 53 | // Explicitly joins every worker thread. Callers should invoke this | ||
| 54 | // deterministically during their own shutdown sequence — the pool is | ||
| 55 | // a Meyers singleton, so its destructor only runs during the process's | ||
| 56 | // static-destruction phase, by which point every *other* thread has | ||
| 57 | // already been suspended by RtlExitUserProcess. A worker suspended | ||
| 58 | // mid-operation while holding some global lock (e.g. the Debug CRT's | ||
| 59 | // iterator-checking mutex, taken by any std::vector destructor) can | ||
| 60 | // then deadlock the main thread's own exit-time TLS destructors | ||
| 61 | // waiting on that same lock forever. Calling shutdown() early — before | ||
| 62 | // any of that starts — avoids the whole class of hazard. Idempotent: | ||
| 63 | // safe to call more than once (the eventual destructor call included). | ||
| 64 | void shutdown(); | ||
| 65 | |||
| 66 | ThreadPool(const ThreadPool&) = delete; | ||
| 67 | ThreadPool& operator=(const ThreadPool&) = delete; | ||
| 68 | |||
| 69 | private: | ||
| 70 | ThreadPool(); | ||
| 71 | ~ThreadPool(); | ||
| 72 | |||
| 73 | void worker_loop(); | ||
| 74 | |||
| 75 | std::vector<std::thread> workers_; | ||
| 76 | |||
| 77 | // Job state, valid only while a run() call is outstanding (single job | ||
| 78 | // in flight at a time — see class comment). | ||
| 79 | const std::function<void(std::size_t)>* fn_ = nullptr; | ||
| 80 | std::size_t n_ = 0; | ||
| 81 | std::atomic<std::size_t> cursor_{0}; | ||
| 82 | std::atomic<int> outstanding_{0}; // workers still draining this job | ||
| 83 | |||
| 84 | // Bumped once per run() call; workers block in generation_.wait(seen) | ||
| 85 | // until it changes, then pick up the new job (or exit, on shutdown). | ||
| 86 | std::atomic<int> generation_{0}; | ||
| 87 | std::atomic<bool> shutdown_{false}; | ||
| 88 | }; | ||
| 89 | |||
| 90 | 3522 | inline void parallel_for(std::size_t n, const std::function<void(std::size_t)>& fn) { | |
| 91 | 3522 | ThreadPool::instance().run(n, fn); | |
| 92 | 3522 | } | |
| 93 | |||
| 94 | } // namespace brotensor::detail::cpu | ||
| 95 |