Skip to main content

cortiq_engine/
pool.rs

1//! Persistent worker pool for row-parallel matvecs.
2//!
3//! Threads are spawned once and spin-then-park between calls — vmfcore
4//! measured spawn-per-matvec at ~+27% decode cost versus a persistent
5//! pool. Parallelism is by disjoint row ranges, so results are
6//! bit-identical to the serial path (each row's dot product is computed
7//! the same way).
8//!
9//! Dispatch is a single shared job slot + atomic epoch (roadmap §3 P0):
10//! the caller publishes one pointer, bumps the epoch and JOINS THE WORK
11//! as the extra worker instead of blocking on a latch. The previous
12//! design allocated an `Arc<Latch>` and pushed a message into every
13//! worker's mpsc channel for every matvec (~200 dispatches/token) —
14//! with decode-grade matvecs that synchronization was its own budget.
15//! Workers spin for `CMF_POOL_SPIN` iterations before parking.
16//! Default 4000: at ~39 dispatches/token, park-immediately pays the
17//! unpark syscall on every worker for every dispatch — measured on an
18//! M4 (interleaved A/B, current epoch dispatch + parked-flag design):
19//! Qwen-0.5B q8 decode 101→115 tok/s, q4t 117→149, the 50M bench model
20//! 549→954 at spin=4000 vs spin=0. An early measurement that showed
21//! spinning LOSING (−25% on q8) predates the parked-flag skip and the
22//! multi-matrix dispatch cuts; it no longer reproduces. Over-spinning
23//! still hurts (200k: −15% vs 4k — spinners steal the caller's serial
24//! cycles), so the budget stays bounded. `CMF_POOL_SPIN=0` restores
25//! park-immediately for share-the-box serving.
26//!
27//! `CMF_THREADS` env: 0/1 = serial, N = worker count
28//! (default: available_parallelism − 1, capped at 8).
29
30use std::cell::UnsafeCell;
31use std::sync::Arc;
32#[cfg(any(target_os = "android", target_os = "linux"))]
33use std::sync::Mutex;
34use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
35
36/// Embedder override for the pool size (C ABI `cortiq_set_threads`):
37/// 0 = unset, consult CMF_THREADS / topology as before. Read once at
38/// pool construction, so set it before the load.
39pub static FORCED_THREADS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
40
41/// Kernel thread ids of the last fully constructed pool's workers
42/// (Android/Linux) — what ADPF's PerformanceHintManager needs to attribute
43/// work to the governor. Published as one complete snapshot after that
44/// pool's per-instance registration barrier; empty elsewhere.
45pub static WORKER_TIDS: std::sync::Mutex<Vec<i32>> = std::sync::Mutex::new(Vec::new());
46
47/// A `*const dyn Fn` that may cross a thread boundary. Safety is
48/// provided by `Pool::run`: the caller blocks until every worker has
49/// finished, so the borrow outlives all uses.
50#[derive(Clone, Copy)]
51struct TaskPtr(*const (dyn Fn(usize, usize) + Sync));
52unsafe impl Send for TaskPtr {}
53
54struct Inner {
55    /// Bumped once per published job; workers watch it.
56    epoch: AtomicUsize,
57    /// Workers still running the current job (excludes the caller).
58    remaining: AtomicUsize,
59    /// The published job: closure pointer + total participant count.
60    /// Written by the caller BEFORE the epoch bump, read by workers
61    /// AFTER they observe the new epoch (acquire/release pairing).
62    /// (task, worker count, publisher's GPU device, worker limit). The
63    /// device rides along because a dispatch begun on card 1 must not
64    /// finish on card 0: worker threads have their own thread-locals,
65    /// and the engine resolves its wgpu context through one. The limit
66    /// is how many workers PARTICIPATE: a job with eight grains has no
67    /// use for three hundred workers — the unpark syscalls and the
68    /// remaining-drain would BE the job (measured: 361 pool dispatches
69    /// per DeepSeek-V4 token, and CMF_THREADS=64 vs 380 was 1.3 vs 2.4
70    /// tok/s with no other change). Workers at or past the limit skip
71    /// the job entirely and never touch `remaining`.
72    slot: UnsafeCell<Option<(TaskPtr, usize, usize, usize)>>,
73    shutdown: AtomicBool,
74    /// Spin iterations before a worker parks (0 = park immediately).
75    spin_budget: AtomicUsize,
76    /// Per-worker "I am parked" flags — lets the caller skip the unpark
77    /// syscall for workers that are still spinning.
78    parked: Box<[AtomicBool]>,
79    /// Per-pool registration state. `WORKER_TIDS` is a process-wide
80    /// snapshot for ADPF and cannot be a construction barrier: another
81    /// pool may clear and republish that snapshot concurrently.
82    #[cfg(any(target_os = "android", target_os = "linux"))]
83    registered: AtomicUsize,
84    #[cfg(any(target_os = "android", target_os = "linux"))]
85    worker_tids: Mutex<Vec<i32>>,
86}
87
88// SAFETY: `slot` is only written while no job is in flight (run()
89// returns after `remaining` hits 0) and only read after the epoch
90// publication that follows the write.
91unsafe impl Sync for Inner {}
92
93/// Process-wide dispatch counter (roadmap §3 P0 «измерения»): one tick
94/// per published job. `bench --json` reports dispatches/token from it.
95static DISPATCHES: AtomicUsize = AtomicUsize::new(0);
96
97/// Total pool jobs published since process start (all pools).
98pub fn dispatch_count() -> usize {
99    DISPATCHES.load(Ordering::Relaxed)
100}
101
102/// Persistent thread pool: shared job slot, epoch dispatch, caller
103/// participation.
104pub struct Pool {
105    inner: Arc<Inner>,
106    /// Thread handles for `unpark` (same order as `parked`).
107    threads: Vec<std::thread::Thread>,
108    joins: Vec<std::thread::JoinHandle<()>>,
109}
110
111fn spin_budget_from_env() -> usize {
112    std::env::var("CMF_POOL_SPIN")
113        .ok()
114        .and_then(|v| v.parse::<usize>().ok())
115        .unwrap_or(4000)
116}
117
118/// Rows per chunk: enough chunks to balance, large enough to keep the SDOT
119/// inner loop and the prefetcher in their stride — and never so coarse that
120/// ONE worker takes the whole job.
121///
122/// That last clause was missing. The floor was a flat 32, so any job with
123/// fewer than 32 rows went entirely to whichever worker grabbed the cursor
124/// first while the other 48 were woken, found nothing, and left. The
125/// hyper-connection projection has 24 rows and is called 86 times a token:
126/// it paid the full price of a fan-out and ran single-threaded.
127pub(crate) fn grain_for(rows: usize, workers: usize) -> usize {
128    if rows == 0 || workers <= 1 {
129        return rows.max(1);
130    }
131    let balanced = (rows / (workers * 8)).max(32);
132    // One chunk per worker at the very least.
133    balanced.min(rows.div_ceil(workers)).max(1)
134}
135
136impl Pool {
137    pub fn new(n_workers: usize) -> Self {
138        Self::with_spin(n_workers, spin_budget_from_env())
139    }
140
141    /// Explicit spin budget (tests pin it without touching the env).
142    pub fn with_spin(n_workers: usize, spin_budget: usize) -> Self {
143        let inner = Arc::new(Inner {
144            epoch: AtomicUsize::new(0),
145            remaining: AtomicUsize::new(0),
146            slot: UnsafeCell::new(None),
147            shutdown: AtomicBool::new(false),
148            spin_budget: AtomicUsize::new(spin_budget),
149            parked: (0..n_workers).map(|_| AtomicBool::new(false)).collect(),
150            #[cfg(any(target_os = "android", target_os = "linux"))]
151            registered: AtomicUsize::new(0),
152            #[cfg(any(target_os = "android", target_os = "linux"))]
153            worker_tids: Mutex::new(Vec::with_capacity(n_workers)),
154        });
155        let mut joins = Vec::with_capacity(n_workers);
156        for w in 0..n_workers {
157            let inner = inner.clone();
158            let h = std::thread::Builder::new()
159                .name(format!("cmf-pool-{w}"))
160                .spawn(move || {
161                    #[cfg(any(target_os = "android", target_os = "linux"))]
162                    {
163                        let tid = unsafe { libc::gettid() } as i32;
164                        if let Ok(mut tids) = inner.worker_tids.lock() {
165                            tids.push(tid);
166                        }
167                        inner.registered.fetch_add(1, Ordering::Release);
168                    }
169                    worker_loop(&inner, w)
170                })
171                .expect("spawn pool worker");
172            joins.push(h);
173        }
174        // Per-pool registration barrier: `spawn` returns before the closure runs,
175        // and the embedder reads `cortiq_worker_tids` right after load —
176        // on a phone only the first worker had registered by then (the
177        // '· 1 threads' About line that misled the cmfmobile device
178        // investigation twice). Thread start is milliseconds; wait for
179        // every worker has registered before construction returns.
180        #[cfg(any(target_os = "android", target_os = "linux"))]
181        while inner.registered.load(Ordering::Acquire) < n_workers {
182            std::thread::yield_now();
183        }
184        #[cfg(any(target_os = "android", target_os = "linux"))]
185        if let (Ok(mut global), Ok(local)) = (WORKER_TIDS.lock(), inner.worker_tids.lock()) {
186            *global = local.clone();
187        }
188        let threads = joins.iter().map(|h| h.thread().clone()).collect();
189        Self {
190            inner,
191            threads,
192            joins,
193        }
194    }
195
196    /// Big-core count on heterogeneous ARM (big.LITTLE): the kernel
197    /// exposes per-core capacity on Android and most ARM Linux; efficiency
198    /// cores in the pool DRAG the big ones on our row-parallel jobs (the
199    /// same cliff llama.cpp hits at -t 10 on an M4: 163 → 112 tok/s).
200    /// None = capacities absent or homogeneous.
201    #[cfg(all(
202        target_arch = "aarch64",
203        any(target_os = "linux", target_os = "android")
204    ))]
205    fn big_cores() -> Option<usize> {
206        Self::cores_from_capacities(&core_capacities())
207    }
208
209    /// How many cores the pool should use, from the kernel's per-core
210    /// capacity values. Capacity folds µarch × clock into one number,
211    /// and the two need different treatment: cores of ANOTHER µarch
212    /// (A5xx efficiency cluster next to A7xx/X: capacity ratio ≥ ~2)
213    /// drag row-parallel work down and are excluded; cores of the SAME
214    /// µarch merely clock-binned (JLQ JR510: 8×A55 as 4×2.0 + 4×1.5 GHz,
215    /// ratio 1.33) pull their weight and must ALL be used. The 1.6
216    /// threshold splits the two regimes: on a Snapdragon 8-class part
217    /// it keeps X + A7xx mid cores and drops A5xx.
218    #[cfg_attr(
219        not(all(
220            target_arch = "aarch64",
221            any(target_os = "linux", target_os = "android")
222        )),
223        allow(dead_code)
224    )]
225    fn cores_from_capacities(caps: &[u64]) -> Option<usize> {
226        let max = *caps.iter().max()?;
227        let min = *caps.iter().min()?;
228        if caps.len() < 2 || max == min {
229            return None;
230        }
231        Some(caps.iter().filter(|&&c| c * 8 >= max * 5).count())
232    }
233
234    #[cfg(target_os = "macos")]
235    fn big_cores() -> Option<usize> {
236        // Apple silicon: the P-only default measured WORSE than mixing the
237        // efficiency cores in — the grain-pulling dispatch absorbs the
238        // speed skew exactly as designed, and decode is memory-bound
239        // enough that E-cores add real serviceable work (M4, dense 3B:
240        // 4 threads 8.4 tok/s, 6-9 threads 9.6-10.7). Fall through to
241        // available_parallelism - 1; CMF_THREADS still pins by hand.
242        // The sysctl probe stays for introspection tooling.
243        if true {
244            return None;
245        }
246        #[allow(unreachable_code)]
247        unsafe extern "C" {
248            fn sysctlbyname(
249                name: *const std::ffi::c_char,
250                oldp: *mut std::ffi::c_void,
251                oldlenp: *mut usize,
252                newp: *mut std::ffi::c_void,
253                newlen: usize,
254            ) -> std::ffi::c_int;
255        }
256        unsafe {
257            let name = std::ffi::CString::new("hw.perflevel0.physicalcpu").ok()?;
258            let mut count: i32 = 0;
259            let mut size = std::mem::size_of::<i32>();
260            let ret = sysctlbyname(
261                name.as_ptr(),
262                &mut count as *mut i32 as *mut std::ffi::c_void,
263                &mut size,
264                std::ptr::null_mut(),
265                0,
266            );
267            if ret == 0 && count > 0 {
268                Some(count as usize)
269            } else {
270                None
271            }
272        }
273    }
274
275    #[cfg(not(any(
276        all(
277            target_arch = "aarch64",
278            any(target_os = "linux", target_os = "android")
279        ),
280        target_os = "macos"
281    )))]
282    fn big_cores() -> Option<usize> {
283        None
284    }
285
286    /// The thread count `from_env` would use RIGHT NOW: forced (C ABI)
287    /// > CMF_THREADS > big-core topology > available_parallelism−1.
288    /// > ≤1 means the model runs serial (no pool). Introspection
289    /// > (`execution_mode`, status endpoints) must report THIS, not
290    /// > available_parallelism.
291    pub fn effective_threads() -> usize {
292        let forced = FORCED_THREADS.load(std::sync::atomic::Ordering::Relaxed);
293        if forced > 0 {
294            return forced;
295        }
296        match std::env::var("CMF_THREADS") {
297            Ok(v) => v.parse::<usize>().unwrap_or(0),
298            Err(_) => match Self::big_cores() {
299                Some(big) => big,
300                None => {
301                    // The cap was 8, which left big machines idle: on a
302                    // 256-core EPYC, Nanbeige 4.2 decoded at 7.4 tok/s on
303                    // the default 8 threads and 14.8 at 32, with prefill
304                    // 12 -> ~16 over the same move. Past ~32 it falls off
305                    // hard (5.5 at 64, 1.6 at 256) — decode is
306                    // memory-bound and the extra threads only add
307                    // dispatch barriers — so 32 is a ceiling, not a
308                    // target. Machines with 9 cores or fewer are
309                    // unaffected: avail-1 already bounds them.
310                    let avail = std::thread::available_parallelism()
311                        .map(|n| n.get())
312                        .unwrap_or(1);
313                    avail.saturating_sub(1).min(32)
314                }
315            },
316        }
317    }
318
319    /// Pool sized from `CMF_THREADS` (see module docs). `None` = serial.
320    /// Without the env, heterogeneous ARM defaults to its BIG cores.
321    pub fn from_env() -> Option<Arc<Self>> {
322        let n = Self::effective_threads();
323        if n <= 1 {
324            None
325        } else {
326            Some(Arc::new(Self::new(n)))
327        }
328    }
329
330    /// Spawned worker threads (the caller joins each job on top).
331    pub fn n_workers(&self) -> usize {
332        self.threads.len()
333    }
334
335    /// Keep the pool on the NUMA node that holds `regions` (the model's
336    /// weight bytes). Linux with two or more nodes only; `CMF_NUMA=0`
337    /// turns it off, `CMF_NUMA=node:<n>` forces a node.
338    ///
339    /// WHY: decode streams every weight once per token, and on a
340    /// two-socket host the page cache holds a file on whichever node
341    /// read it. Unpinned, the scheduler spreads the workers over both
342    /// sockets and half the matvec rows cross the socket link. Measured
343    /// on a 2×EPYC 7763 pod with the model's pages all on node 0 (31 CPUs
344    /// of cgroup quota): a STREAM-style read over a node-0 buffer gives
345    /// 42 GB/s from 31 unpinned threads and 74 GB/s from 31 threads kept
346    /// on node 0. The mask is the node's physical cores (first SMT
347    /// sibling) when there are enough of them for the pool, else the
348    /// whole node; never narrower than the pool, so nothing oversubscribes.
349    /// Threads are bound to a SET of cores, not to one core each: the
350    /// scheduler still balances inside the node. The calling thread
351    /// adopts the same mask on its next dispatch.
352    pub fn bind_numa(&self, regions: &[&[u8]]) {
353        #[cfg(target_os = "linux")]
354        {
355            let Some((node, cpus)) = numa::choose(regions, self.threads.len() + 1) else {
356                return;
357            };
358            let mut applied = 0usize;
359            if let Ok(tids) = self.inner.worker_tids.lock() {
360                for &tid in tids.iter() {
361                    if numa::set_affinity(tid, &cpus) {
362                        applied += 1;
363                    }
364                }
365            }
366            numa::publish(cpus.clone());
367            numa::adopt_caller();
368            tracing::info!(
369                "numa: pool bound to node {node} ({} cpus, {applied}/{} workers)",
370                cpus.len(),
371                self.threads.len()
372            );
373            if std::env::var("CMF_NUMA_TRACE").is_ok_and(|v| v != "0") {
374                eprintln!(
375                    "numa: pool bound to node {node}: {} cpus, {applied}/{} workers",
376                    cpus.len(),
377                    self.threads.len()
378                );
379            }
380        }
381        #[cfg(not(target_os = "linux"))]
382        let _ = regions;
383    }
384
385    /// Retune an already-created pool for an architecture with a measured
386    /// dispatch cadence. The environment remains the operator override; this
387    /// hook only changes the automatic default after model geometry is known.
388    pub(crate) fn set_spin_budget(&self, spins: usize) {
389        self.inner.spin_budget.store(spins, Ordering::Relaxed);
390    }
391
392    /// Run `f(row_start, row_end)` over `0..rows`, self-balancing.
393    ///
394    /// One dispatch, but workers pull row-ranges from a shared cursor
395    /// instead of each taking a fixed 1/n slice. On a heterogeneous CPU
396    /// (Apple Silicon: 4 P-cores + 6 E-cores here) a static split makes
397    /// every matvec end at the SLOWEST core's pace while the fast ones
398    /// idle at the barrier; pulling by grain lets a P-core take several
399    /// chunks for each one an E-core takes, so skew collapses to a
400    /// single grain. Row ranges stay disjoint and each row's dot is
401    /// computed exactly as in the serial path → bit-identical output.
402    pub fn run_rows(&self, rows: usize, f: &(dyn Fn(usize, usize) + Sync)) {
403        let grain = grain_for(rows, self.threads.len() + 1);
404        let chunks = rows.div_ceil(grain.max(1));
405        let next = AtomicUsize::new(0);
406        self.run_limited(chunks, &|_w, _n| loop {
407            let start = next.fetch_add(grain, Ordering::Relaxed);
408            if start >= rows {
409                break;
410            }
411            f(start, (start + grain).min(rows));
412        });
413    }
414
415    /// `run`, waking at most `max_workers` workers. Same grain, same
416    /// row split, bit-identical results — only the number of threads
417    /// woken changes, so an 8-grain job stops paying 380 unparks. Only
418    /// cursor-style closures (which ignore their (idx, n) arguments)
419    /// come through here: the caller identifies itself as `limit`,
420    /// which under a cap is NOT `n_workers()`.
421    fn run_limited(&self, max_workers: usize, f: &(dyn Fn(usize, usize) + Sync)) {
422        #[cfg(target_os = "linux")]
423        numa::adopt_caller();
424        let nw = self.threads.len().min(max_workers);
425        if nw == self.threads.len() {
426            return self.run(f);
427        }
428        DISPATCHES.fetch_add(1, Ordering::Relaxed);
429        let ptr: *const (dyn Fn(usize, usize) + Sync) = f;
430        let ptr: *const (dyn Fn(usize, usize) + Sync + 'static) =
431            unsafe { std::mem::transmute(ptr) };
432        let dev = crate::gpu::current_device();
433        // SAFETY: same contract as `run` — no job in flight, and the
434        // wait below outlives every borrow of `f`.
435        unsafe { *self.inner.slot.get() = Some((TaskPtr(ptr), nw + 1, dev, nw)) };
436        self.inner.remaining.store(nw, Ordering::Relaxed);
437        self.inner.epoch.fetch_add(1, Ordering::SeqCst);
438        for (i, t) in self.threads.iter().enumerate().take(nw) {
439            if self.inner.parked[i].load(Ordering::SeqCst) {
440                t.unpark();
441            }
442        }
443        f(nw, nw + 1);
444        let mut spins = 0usize;
445        while self.inner.remaining.load(Ordering::Acquire) != 0 {
446            spins += 1;
447            if spins < 10_000 {
448                std::hint::spin_loop();
449            } else {
450                std::thread::yield_now();
451            }
452        }
453    }
454
455    /// Multi-matrix job: one dispatch serves SEVERAL row spaces
456    /// (roadmap §3 P0 — «одна внешняя публикация job на слой»). Parts
457    /// are laid out back-to-back in a virtual row space and pulled by
458    /// grain from one shared cursor, so QKV or gate+up cost a single
459    /// barrier instead of one each. Each part's `f(start, end)` sees its
460    /// OWN row indices — per-row math and outputs are bit-identical to
461    /// separate `run_rows` calls.
462    pub fn run_many(&self, parts: &[(usize, &(dyn Fn(usize, usize) + Sync))]) {
463        let total: usize = parts.iter().map(|p| p.0).sum();
464        if total == 0 {
465            return;
466        }
467        let grain = grain_for(total, self.threads.len() + 1);
468        let chunks = total.div_ceil(grain.max(1));
469        let next = AtomicUsize::new(0);
470        self.run_limited(chunks, &|_w, _n| loop {
471            let s = next.fetch_add(grain, Ordering::Relaxed);
472            if s >= total {
473                break;
474            }
475            let e = (s + grain).min(total);
476            let mut base = 0usize;
477            for &(rows, f) in parts {
478                let a = s.max(base);
479                let b = e.min(base + rows);
480                if a < b {
481                    f(a - base, b - base);
482                }
483                base += rows;
484                if base >= e {
485                    break;
486                }
487            }
488        });
489    }
490
491    /// Run `f(worker_idx, n_participants)` on every worker AND the
492    /// calling thread (`worker_idx = n_workers()` for the caller);
493    /// returns when all participants have finished.
494    pub fn run(&self, f: &(dyn Fn(usize, usize) + Sync)) {
495        #[cfg(target_os = "linux")]
496        numa::adopt_caller();
497        DISPATCHES.fetch_add(1, Ordering::Relaxed);
498        let nw = self.threads.len();
499        let n = nw + 1; // caller participates
500        // SAFETY: the wait loop below blocks until every worker is done,
501        // so extending the borrow to 'static never outlives the call.
502        let ptr: *const (dyn Fn(usize, usize) + Sync) = f;
503        let ptr: *const (dyn Fn(usize, usize) + Sync + 'static) =
504            unsafe { std::mem::transmute(ptr) };
505        // SAFETY: no job in flight (previous run() drained `remaining`),
506        // so the slot is not being read.
507        let dev = crate::gpu::current_device();
508        unsafe { *self.inner.slot.get() = Some((TaskPtr(ptr), n, dev, nw)) };
509        self.inner.remaining.store(nw, Ordering::Relaxed);
510        self.inner.epoch.fetch_add(1, Ordering::SeqCst);
511        for (i, t) in self.threads.iter().enumerate() {
512            if self.inner.parked[i].load(Ordering::SeqCst) {
513                t.unpark();
514            }
515        }
516
517        // The caller's share — the barrier costs nothing while there is
518        // real work to do.
519        f(nw, n);
520
521        // Wait for the stragglers (bounded by one worker's chunk).
522        let mut spins = 0usize;
523        while self.inner.remaining.load(Ordering::Acquire) != 0 {
524            spins += 1;
525            if spins < 10_000 {
526                std::hint::spin_loop();
527            } else {
528                std::thread::yield_now();
529            }
530        }
531    }
532}
533
534impl Drop for Pool {
535    fn drop(&mut self) {
536        self.inner.shutdown.store(true, Ordering::SeqCst);
537        for t in &self.threads {
538            t.unpark();
539        }
540        for h in self.joins.drain(..) {
541            let _ = h.join();
542        }
543    }
544}
545
546/// NUMA placement for the pool (see `Pool::bind_numa`).
547#[cfg(target_os = "linux")]
548mod numa {
549    use std::sync::Mutex;
550    use std::sync::atomic::{AtomicUsize, Ordering};
551
552    /// The published mask; `EPOCH` bumps on every publish so a calling
553    /// thread re-adopts at most once per bind.
554    static MASK: Mutex<Vec<usize>> = Mutex::new(Vec::new());
555    static EPOCH: AtomicUsize = AtomicUsize::new(0);
556    thread_local! {
557        static SEEN: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
558    }
559
560    pub(super) fn publish(cpus: Vec<usize>) {
561        if let Ok(mut m) = MASK.lock() {
562            *m = cpus;
563        }
564        EPOCH.fetch_add(1, Ordering::Release);
565    }
566
567    /// One relaxed load + one TLS read per dispatch when nothing changed.
568    #[inline]
569    pub(super) fn adopt_caller() {
570        let e = EPOCH.load(Ordering::Acquire);
571        if e == 0 || SEEN.with(|c| c.get()) == e {
572            return;
573        }
574        SEEN.with(|c| c.set(e));
575        if let Ok(m) = MASK.lock() {
576            if !m.is_empty() {
577                set_affinity(0, &m);
578            }
579        }
580    }
581
582    pub(super) fn parse_list(s: &str) -> Vec<usize> {
583        let mut out = Vec::new();
584        for part in s.trim().split(',') {
585            let part = part.trim();
586            if part.is_empty() {
587                continue;
588            }
589            match part.split_once('-') {
590                Some((a, b)) => {
591                    if let (Ok(a), Ok(b)) = (a.parse::<usize>(), b.parse::<usize>()) {
592                        out.extend(a..=b);
593                    }
594                }
595                None => {
596                    if let Ok(a) = part.parse() {
597                        out.push(a);
598                    }
599                }
600            }
601        }
602        out
603    }
604
605    fn nodes() -> Vec<(usize, Vec<usize>)> {
606        let mut v = Vec::new();
607        let Ok(rd) = std::fs::read_dir("/sys/devices/system/node") else {
608            return v;
609        };
610        for e in rd.flatten() {
611            let name = e.file_name().to_string_lossy().to_string();
612            let Some(id) = name.strip_prefix("node").and_then(|x| x.parse::<usize>().ok()) else {
613                continue;
614            };
615            if let Ok(l) = std::fs::read_to_string(e.path().join("cpulist")) {
616                let cpus = parse_list(&l);
617                if !cpus.is_empty() {
618                    v.push((id, cpus));
619                }
620            }
621        }
622        v.sort();
623        v
624    }
625
626    fn allowed() -> Vec<usize> {
627        // SAFETY: plain syscall into a zeroed, correctly sized set.
628        unsafe {
629            let mut set: libc::cpu_set_t = std::mem::zeroed();
630            if libc::sched_getaffinity(0, std::mem::size_of::<libc::cpu_set_t>(), &mut set) != 0 {
631                return Vec::new();
632            }
633            (0..libc::CPU_SETSIZE as usize)
634                .filter(|&c| libc::CPU_ISSET(c, &set))
635                .collect()
636        }
637    }
638
639    /// First SMT sibling of its core (or no topology info: count it).
640    fn primary(cpu: usize) -> bool {
641        let p = format!("/sys/devices/system/cpu/cpu{cpu}/topology/thread_siblings_list");
642        match std::fs::read_to_string(p) {
643            Ok(l) => parse_list(&l).first().is_none_or(|&f| f == cpu),
644            Err(_) => true,
645        }
646    }
647
648    /// Where the weights live: (pages sampled, sampled pages in the page
649    /// cache, mapped pages per node). The sample is ≤ 4096 pages spread
650    /// over `regions`; every sampled page that is already cached is mapped
651    /// here with one read (a minor fault — `mincore` says it is cached, so
652    /// no disk I/O), because both node queries below only see pages mapped
653    /// into THIS process. Per-node counts come from `move_pages` in query
654    /// mode, or — where a container's seccomp profile refuses that syscall
655    /// (EPERM on the RunPod image) — from `/proc/self/numa_maps` for the
656    /// mappings that hold the regions.
657    fn page_nodes(regions: &[&[u8]]) -> (usize, usize, Vec<usize>) {
658        const PAGE: usize = 4096;
659        let total: usize = regions.iter().map(|r| r.len() / PAGE).sum();
660        if total == 0 {
661            return (0, 0, Vec::new());
662        }
663        let stride = total.div_ceil(4096).max(1);
664        let mut pages: Vec<*mut libc::c_void> = Vec::new();
665        for r in regions {
666            let base = (r.as_ptr() as usize).div_ceil(PAGE) * PAGE;
667            let end = r.as_ptr() as usize + r.len();
668            let mut a = base;
669            while a + PAGE <= end {
670                pages.push(a as *mut libc::c_void);
671                a += PAGE * stride;
672            }
673        }
674        let mut incore = 0usize;
675        for &p in &pages {
676            let mut vec = 0u8;
677            // SAFETY: `p` is a page-aligned address inside a live mapping.
678            let cached = unsafe { libc::mincore(p, PAGE, &mut vec) } == 0 && vec & 1 == 1;
679            if cached {
680                incore += 1;
681                // SAFETY: readable mapped byte; volatile so it is not elided.
682                unsafe { std::ptr::read_volatile(p as *const u8) };
683            }
684        }
685        let mut status = vec![-1i32; pages.len()];
686        // SAFETY: query-only move_pages on our own mappings; `nodes` is
687        // NULL so nothing moves, `status` has one slot per page.
688        let rc = unsafe {
689            libc::syscall(
690                libc::SYS_move_pages,
691                0,
692                pages.len() as libc::c_ulong,
693                pages.as_mut_ptr(),
694                std::ptr::null::<libc::c_int>(),
695                status.as_mut_ptr(),
696                0,
697            )
698        };
699        let mut by = Vec::new();
700        if rc == 0 {
701            for &st in &status {
702                if st >= 0 {
703                    let n = st as usize;
704                    if by.len() <= n {
705                        by.resize(n + 1, 0);
706                    }
707                    by[n] += 1;
708                }
709            }
710        } else {
711            by = numa_maps_nodes(regions);
712        }
713        (pages.len(), incore, by)
714    }
715
716    /// Mapped pages per node of every mapping that overlaps `regions`,
717    /// from `/proc/self/maps` (ranges) + `/proc/self/numa_maps` (`N<k>=`).
718    fn numa_maps_nodes(regions: &[&[u8]]) -> Vec<usize> {
719        let (Ok(maps), Ok(nm)) = (
720            std::fs::read_to_string("/proc/self/maps"),
721            std::fs::read_to_string("/proc/self/numa_maps"),
722        ) else {
723            return Vec::new();
724        };
725        let spans: Vec<(usize, usize)> = regions
726            .iter()
727            .map(|r| (r.as_ptr() as usize, r.as_ptr() as usize + r.len()))
728            .collect();
729        let mut starts = std::collections::HashSet::new();
730        for line in maps.lines() {
731            let Some((range, _)) = line.split_once(' ') else {
732                continue;
733            };
734            let Some((a, b)) = range.split_once('-') else {
735                continue;
736            };
737            let (Ok(a), Ok(b)) = (usize::from_str_radix(a, 16), usize::from_str_radix(b, 16)) else {
738                continue;
739            };
740            if spans.iter().any(|&(s, e)| s < b && a < e) {
741                starts.insert(a);
742            }
743        }
744        let mut by = Vec::new();
745        for line in nm.lines() {
746            let mut it = line.split_whitespace();
747            let Some(a) = it.next().and_then(|a| usize::from_str_radix(a, 16).ok()) else {
748                continue;
749            };
750            if !starts.contains(&a) {
751                continue;
752            }
753            for f in it {
754                let Some((k, v)) = f.split_once('=') else {
755                    continue;
756                };
757                let (Some(n), Ok(v)) = (
758                    k.strip_prefix('N').and_then(|n| n.parse::<usize>().ok()),
759                    v.parse::<usize>(),
760                ) else {
761                    continue;
762                };
763                if by.len() <= n {
764                    by.resize(n + 1, 0);
765                }
766                by[n] += v;
767            }
768        }
769        by
770    }
771
772    /// (node, cpu mask) for a pool of `threads` participants, or None.
773    pub(super) fn choose(regions: &[&[u8]], threads: usize) -> Option<(usize, Vec<usize>)> {
774        let env = std::env::var("CMF_NUMA").ok();
775        if matches!(env.as_deref(), Some("0") | Some("off")) {
776            return None;
777        }
778        let trace = std::env::var("CMF_NUMA_TRACE").is_ok_and(|v| v != "0");
779        let nodes = nodes();
780        if nodes.len() < 2 {
781            if trace {
782                eprintln!("numa: {} node(s) visible — nothing to bind", nodes.len());
783            }
784            return None;
785        }
786        // `CMF_NUMA=node:<n>` forces a node (plain "0" means OFF).
787        let forced = env
788            .as_deref()
789            .and_then(|v| v.strip_prefix("node:"))
790            .and_then(|v| v.parse::<usize>().ok());
791        let node = match forced {
792            Some(n) => n,
793            None => {
794                // Auto: only when the weights already sit on ONE node
795                // (≥ 90% of the resident sample, and most of the sample
796                // resident). A file spread over both nodes is better
797                // served by both sockets; a cold file has no home yet.
798                let (sampled, incore, by) = page_nodes(regions);
799                let resident: usize = by.iter().sum();
800                if trace {
801                    eprintln!(
802                        "numa: sampled {sampled} weight pages, {incore} cached, mapped by node {by:?}"
803                    );
804                }
805                if sampled == 0 || incore * 2 < sampled || resident == 0 {
806                    return None;
807                }
808                let (n, &cnt) = by.iter().enumerate().max_by_key(|(_, c)| **c)?;
809                if cnt * 10 < resident * 9 {
810                    return None;
811                }
812                n
813            }
814        };
815        let cpus = &nodes.iter().find(|(id, _)| *id == node)?.1;
816        let allowed = allowed();
817        let usable: Vec<usize> = cpus.iter().copied().filter(|c| allowed.contains(c)).collect();
818        let prim: Vec<usize> = usable.iter().copied().filter(|&c| primary(c)).collect();
819        if prim.len() >= threads {
820            Some((node, prim))
821        } else if usable.len() >= threads {
822            Some((node, usable))
823        } else {
824            None
825        }
826    }
827
828    /// Bind thread `tid` (0 = the calling thread) to `cpus`.
829    pub(super) fn set_affinity(tid: i32, cpus: &[usize]) -> bool {
830        // SAFETY: plain syscall with a zeroed, correctly sized set.
831        unsafe {
832            let mut set: libc::cpu_set_t = std::mem::zeroed();
833            for &c in cpus {
834                if c < libc::CPU_SETSIZE as usize {
835                    libc::CPU_SET(c, &mut set);
836                }
837            }
838            libc::sched_setaffinity(tid, std::mem::size_of::<libc::cpu_set_t>(), &set) == 0
839        }
840    }
841}
842
843/// Per-core capacity: the kernel's `cpu_capacity` (µarch × clock) when
844/// EAS exposes it, else `cpufreq/cpuinfo_max_freq` — same cluster
845/// ordering, so the 62.5% big-core rule keeps working on EAS-less
846/// kernels (TUNING.md open item: pinning silently did nothing there).
847#[cfg(any(
848    target_os = "android",
849    all(target_arch = "aarch64", target_os = "linux")
850))]
851fn core_capacities() -> Vec<u64> {
852    let read_all = |leaf: &str| -> Vec<u64> {
853        let mut vals = Vec::new();
854        for cpu in 0.. {
855            let path = format!("/sys/devices/system/cpu/cpu{cpu}/{leaf}");
856            match std::fs::read_to_string(&path) {
857                Ok(v) => match v.trim().parse() {
858                    Ok(x) => vals.push(x),
859                    Err(_) => break,
860                },
861                Err(_) => break,
862            }
863        }
864        vals
865    };
866    let caps = read_all("cpu_capacity");
867    if caps.len() >= 2 {
868        return caps;
869    }
870    read_all("cpufreq/cpuinfo_max_freq")
871}
872
873#[cfg(target_os = "android")]
874fn pin_thread_to_big_cores() {
875    use std::mem;
876    let caps = core_capacities();
877    let max = caps.iter().copied().max().unwrap_or(0);
878    let min = caps.iter().copied().min().unwrap_or(0);
879
880    // Only pin if heterogeneous
881    if caps.len() < 2 || max == min {
882        return;
883    }
884
885    unsafe {
886        let mut set: libc::cpu_set_t = mem::zeroed();
887        for (i, &c) in caps.iter().enumerate() {
888            if c * 8 >= max * 5 {
889                libc::CPU_SET(i, &mut set);
890            }
891        }
892        libc::sched_setaffinity(0, mem::size_of::<libc::cpu_set_t>(), &set);
893    }
894}
895
896fn worker_loop(inner: &Inner, idx: usize) {
897    #[cfg(target_os = "android")]
898    pin_thread_to_big_cores();
899    // Apple silicon: ask for the performance cores. Threads spawned
900    // without a QoS class land on the efficiency cores when the
901    // scheduler feels like it — a user's video-VAE encode on an M4 sat
902    // on the E-cores at 100% with the P-cores asleep for 140 s (HF
903    // discussion #4). USER_INITIATED is the class an interactive tool's
904    // work belongs to; the ~4 P-cores then take the pool's grains.
905    #[cfg(target_os = "macos")]
906    unsafe {
907        libc::pthread_set_qos_class_self_np(libc::qos_class_t::QOS_CLASS_USER_INITIATED, 0);
908    }
909
910    // The pool is created at epoch 0; baseline MUST be 0, not a fresh
911    // epoch read — if the caller publishes a job before the OS actually
912    // starts this thread, reading the live epoch would adopt that job's
913    // epoch as "already seen", skip it, and deadlock the caller's wait.
914    let mut seen = 0usize;
915    loop {
916        // Wait for a new epoch: spin first (decode publishes the next
917        // matvec within microseconds), park only when idle for real.
918        let mut spins = 0usize;
919        loop {
920            let e = inner.epoch.load(Ordering::Acquire);
921            if e != seen {
922                seen = e;
923                break;
924            }
925            if inner.shutdown.load(Ordering::Relaxed) {
926                return;
927            }
928            if spins < inner.spin_budget.load(Ordering::Relaxed) {
929                spins += 1;
930                std::hint::spin_loop();
931            } else {
932                inner.parked[idx].store(true, Ordering::SeqCst);
933                // Re-check under SeqCst: the caller bumps the epoch
934                // BEFORE reading `parked`, so either it sees our flag
935                // (and unparks) or we see its epoch here — a missed
936                // wakeup is impossible. Spurious unparks just loop.
937                if inner.epoch.load(Ordering::SeqCst) == seen
938                    && !inner.shutdown.load(Ordering::Relaxed)
939                {
940                    std::thread::park();
941                }
942                inner.parked[idx].store(false, Ordering::SeqCst);
943            }
944        }
945        // SAFETY: the slot was written before the epoch bump we just
946        // observed (release/acquire), and stays valid until `remaining`
947        // drops to zero — which happens only after `f` returns below.
948        let (task, n, dev, limit) =
949            unsafe { (*inner.slot.get()).expect("job published with epoch") };
950        if idx >= limit {
951            // Not invited: a bounded dispatch (run_rows with few grains)
952            // counted only `limit` workers into `remaining`. Executing —
953            // or decrementing — here would corrupt the barrier.
954            continue;
955        }
956        let f = unsafe { &*task.0 };
957        crate::gpu::set_current_device(dev);
958        f(idx, n);
959        inner.remaining.fetch_sub(1, Ordering::AcqRel);
960    }
961}
962
963/// Row-parallel dense matvec: `out[o] = Σ_j w[o·in + j]·x[j]`.
964/// Bit-identical to the serial loop (row order does not change math).
965pub fn matvec_rows(pool: Option<&Pool>, w: &[f32], x: &[f32], out: &mut [f32]) {
966    let in_dim = x.len();
967    let out_dim = out.len();
968    debug_assert!(w.len() >= out_dim * in_dim);
969
970    let row_dot = |o: usize| -> f32 {
971        let row = &w[o * in_dim..(o + 1) * in_dim];
972        let mut sum = 0.0f32;
973        for j in 0..in_dim {
974            sum += row[j] * x[j];
975        }
976        sum
977    };
978
979    match pool {
980        Some(pool) if out_dim >= 256 => {
981            let out_addr = SendMut(out.as_mut_ptr());
982            let run_range = move |start: usize, end: usize| {
983                for o in start..end {
984                    unsafe { *out_addr.at(o) = row_dot(o) };
985                }
986            };
987            pool.run_rows(out_dim, &run_range);
988        }
989        _ => {
990            for (o, dst) in out.iter_mut().enumerate() {
991                *dst = row_dot(o);
992            }
993        }
994    }
995}
996
997/// Two-input row matvec: one pass over the weight rows serves BOTH
998/// inputs — CPU decode is memory-bound, so the second position costs a
999/// fraction of the first (this is where MTP speculative verify wins).
1000/// Per-output accumulation order matches the single-input path exactly
1001/// → bit-identical results.
1002pub fn matvec_rows2(
1003    pool: Option<&Pool>,
1004    w: &[f32],
1005    x1: &[f32],
1006    x2: &[f32],
1007    out1: &mut [f32],
1008    out2: &mut [f32],
1009) {
1010    let in_dim = x1.len();
1011    debug_assert_eq!(x2.len(), in_dim);
1012    let out_dim = out1.len();
1013    debug_assert_eq!(out2.len(), out_dim);
1014    debug_assert!(w.len() >= out_dim * in_dim);
1015
1016    let row_dots = |o: usize| -> (f32, f32) {
1017        let row = &w[o * in_dim..(o + 1) * in_dim];
1018        let (mut s1, mut s2) = (0.0f32, 0.0f32);
1019        for j in 0..in_dim {
1020            s1 += row[j] * x1[j];
1021            s2 += row[j] * x2[j];
1022        }
1023        (s1, s2)
1024    };
1025
1026    match pool {
1027        Some(pool) if out_dim >= 256 => {
1028            let o1 = SendMut(out1.as_mut_ptr());
1029            let o2 = SendMut(out2.as_mut_ptr());
1030            let run_range = move |start: usize, end: usize| {
1031                for o in start..end {
1032                    let (s1, s2) = row_dots(o);
1033                    unsafe {
1034                        *o1.at(o) = s1;
1035                        *o2.at(o) = s2;
1036                    }
1037                }
1038            };
1039            pool.run_rows(out_dim, &run_range);
1040        }
1041        _ => {
1042            for o in 0..out_dim {
1043                let (s1, s2) = row_dots(o);
1044                out1[o] = s1;
1045                out2[o] = s2;
1046            }
1047        }
1048    }
1049}
1050
1051/// `SendMut` for any element type — the sampler's sparse chain writes
1052/// per-grain candidate lists.
1053pub(crate) struct SendMutT<T>(*mut T);
1054unsafe impl<T> Send for SendMutT<T> {}
1055unsafe impl<T> Sync for SendMutT<T> {}
1056impl<T> Clone for SendMutT<T> {
1057    fn clone(&self) -> Self {
1058        *self
1059    }
1060}
1061impl<T> Copy for SendMutT<T> {}
1062impl<T> SendMutT<T> {
1063    #[inline]
1064    pub(crate) fn new(p: *mut T) -> Self {
1065        Self(p)
1066    }
1067    /// Same contract as `SendMut::at`: disjoint indices, pointee outlives
1068    /// the joined dispatch.
1069    #[inline]
1070    pub(crate) fn at(self, i: usize) -> *mut T {
1071        unsafe { self.0.add(i) }
1072    }
1073}
1074
1075#[derive(Clone, Copy)]
1076pub(crate) struct SendMut(*mut f32);
1077unsafe impl Send for SendMut {}
1078unsafe impl Sync for SendMut {}
1079
1080impl SendMut {
1081    /// The caller promises the threads it hands this to write disjoint
1082    /// indices, and that the pointee outlives them.
1083    #[inline]
1084    pub(crate) fn new(p: *mut f32) -> Self {
1085        Self(p)
1086    }
1087
1088    /// Method receiver forces the closure to capture the whole (Sync)
1089    /// wrapper, not the bare `*mut f32` field (edition-2021 precise capture).
1090    #[inline]
1091    pub(crate) fn at(self, i: usize) -> *mut f32 {
1092        unsafe { self.0.add(i) }
1093    }
1094}
1095
1096#[cfg(test)]
1097mod tests {
1098    #[test]
1099    #[cfg(target_os = "linux")]
1100    fn numa_cpulist_parses_ranges_and_singles() {
1101        assert_eq!(super::numa::parse_list("0-3,8,10-11\n"), vec![0, 1, 2, 3, 8, 10, 11]);
1102        assert_eq!(super::numa::parse_list(""), Vec::<usize>::new());
1103    }
1104
1105    #[test]
1106    #[cfg(any(target_os = "android", target_os = "linux"))]
1107    fn worker_tids_registered_before_new_returns() {
1108        // WORKER_TIDS is only the last completed pool's process-wide
1109        // snapshot; another test can publish a different valid snapshot
1110        // immediately after `new` returns. Check this pool's private
1111        // registration state instead.
1112        use std::collections::HashSet;
1113        let p = super::Pool::new(3);
1114        let local: Vec<_> = p.inner.worker_tids.lock().unwrap().clone();
1115        let registered = p.inner.registered.load(Ordering::Acquire);
1116        let unique: HashSet<_> = local.iter().copied().collect();
1117        assert!(
1118            registered == 3
1119                && local.len() == 3
1120                && unique.len() == 3
1121                && local.iter().all(|&tid| tid > 0),
1122            "all worker tids must be privately registered before new returns \
1123             (registered {registered}, local {}, unique {})",
1124            local.len(),
1125            unique.len()
1126        );
1127    }
1128
1129    #[test]
1130    fn forced_threads_overrides_env_and_topology() {
1131        use std::sync::atomic::Ordering;
1132        super::FORCED_THREADS.store(3, Ordering::Relaxed);
1133        let pool = super::Pool::from_env().expect("forced 3 → pool");
1134        assert_eq!(pool.n_workers(), 3);
1135        super::FORCED_THREADS.store(1, Ordering::Relaxed);
1136        assert!(super::Pool::from_env().is_none(), "forced 1 → serial");
1137        super::FORCED_THREADS.store(0, Ordering::Relaxed);
1138    }
1139
1140    #[test]
1141    #[cfg(any(target_os = "android", target_os = "linux"))]
1142    fn concurrent_pool_constructors_complete_without_registration_race() {
1143        // WORKER_TIDS is a process-wide publication target. Before the
1144        // per-pool counter, a larger constructor could have all its workers
1145        // append, then a concurrent one could clear that vector; the larger
1146        // constructor would wait forever for a length that could never return.
1147        // Start unlike-sized constructors together so that regression is
1148        // exercised without relying on the test harness' scheduling.
1149        use std::sync::{Barrier, mpsc};
1150        use std::time::Duration;
1151
1152        for round in 0..16 {
1153            let start = Arc::new(Barrier::new(3));
1154            let (done_tx, done_rx) = mpsc::channel();
1155            let mut joins = Vec::new();
1156            for workers in [8usize, 1usize] {
1157                let start = start.clone();
1158                let done_tx = done_tx.clone();
1159                joins.push(std::thread::spawn(move || {
1160                    start.wait();
1161                    let pool = Pool::with_spin(workers, 0);
1162                    done_tx.send(pool.n_workers()).unwrap();
1163                }));
1164            }
1165            drop(done_tx);
1166            start.wait();
1167            let mut sizes = Vec::with_capacity(2);
1168            for _ in 0..2 {
1169                sizes.push(
1170                    done_rx
1171                        .recv_timeout(Duration::from_secs(10))
1172                        .unwrap_or_else(|_| panic!("pool constructor stalled in round {round}")),
1173                );
1174            }
1175            sizes.sort_unstable();
1176            assert_eq!(sizes, [1, 8]);
1177            for join in joins {
1178                join.join().unwrap();
1179            }
1180        }
1181    }
1182
1183    #[test]
1184    fn capacity_split_clock_bins_vs_microarch() {
1185        type P = super::Pool;
1186        // JR510: all-A55, two clock bins — use every core.
1187        assert_eq!(
1188            P::cores_from_capacities(&[1024, 1024, 1024, 1024, 768, 768, 768, 768]),
1189            Some(8)
1190        );
1191        // Classic big.LITTLE (A78 + A55) — big only.
1192        assert_eq!(
1193            P::cores_from_capacities(&[1024, 1024, 1024, 1024, 350, 350, 350, 350]),
1194            Some(4)
1195        );
1196        // Three-tier flagship: X + A7xx mids stay, A5xx littles go.
1197        assert_eq!(
1198            P::cores_from_capacities(&[1024, 800, 800, 800, 800, 300, 300, 300]),
1199            Some(5)
1200        );
1201        // Uniform: no signal, caller falls back.
1202        assert_eq!(P::cores_from_capacities(&[1024; 8]), None);
1203        assert_eq!(P::cores_from_capacities(&[]), None);
1204    }
1205
1206    use super::*;
1207
1208    #[test]
1209    fn parallel_matvec_equals_serial_bitexact() {
1210        let (out_dim, in_dim) = (512, 64);
1211        let w: Vec<f32> = (0..out_dim * in_dim)
1212            .map(|i| (i as f32 * 0.013).sin())
1213            .collect();
1214        let x: Vec<f32> = (0..in_dim).map(|i| (i as f32 * 0.07).cos()).collect();
1215
1216        let mut serial = vec![0.0f32; out_dim];
1217        matvec_rows(None, &w, &x, &mut serial);
1218
1219        let pool = Pool::new(4);
1220        let mut parallel = vec![0.0f32; out_dim];
1221        matvec_rows(Some(&pool), &w, &x, &mut parallel);
1222
1223        assert_eq!(serial, parallel, "row-parallel must be bit-identical");
1224    }
1225
1226    #[test]
1227    fn fused_pair_equals_two_singles_bitexact() {
1228        let (out_dim, in_dim) = (300, 48);
1229        let w: Vec<f32> = (0..out_dim * in_dim)
1230            .map(|i| (i as f32 * 0.011).sin())
1231            .collect();
1232        let x1: Vec<f32> = (0..in_dim).map(|i| (i as f32 * 0.03).cos()).collect();
1233        let x2: Vec<f32> = (0..in_dim).map(|i| (i as f32 * 0.09).sin()).collect();
1234
1235        let mut a1 = vec![0.0f32; out_dim];
1236        let mut a2 = vec![0.0f32; out_dim];
1237        matvec_rows(None, &w, &x1, &mut a1);
1238        matvec_rows(None, &w, &x2, &mut a2);
1239
1240        for pool in [None, Some(Pool::new(3))] {
1241            let mut b1 = vec![0.0f32; out_dim];
1242            let mut b2 = vec![0.0f32; out_dim];
1243            matvec_rows2(pool.as_ref(), &w, &x1, &x2, &mut b1, &mut b2);
1244            assert_eq!(a1, b1, "fused lane 1 must be bit-identical");
1245            assert_eq!(a2, b2, "fused lane 2 must be bit-identical");
1246        }
1247    }
1248
1249    #[test]
1250    fn pool_survives_many_runs() {
1251        let pool = Pool::new(3);
1252        let counter = AtomicUsize::new(0);
1253        for _ in 0..100 {
1254            pool.run(&|_, _| {
1255                counter.fetch_add(1, Ordering::Relaxed);
1256            });
1257        }
1258        // 3 workers + the participating caller = 4 executions per run.
1259        assert_eq!(counter.load(Ordering::Relaxed), 400);
1260    }
1261
1262    #[test]
1263    fn pool_wakes_after_park() {
1264        // Force immediate parking (no spin) — the epoch/parked handshake
1265        // must still never miss a wakeup.
1266        let pool = Pool::with_spin(2, 0);
1267        let counter = AtomicUsize::new(0);
1268        for _ in 0..50 {
1269            pool.run(&|_, _| {
1270                counter.fetch_add(1, Ordering::Relaxed);
1271            });
1272            // Give workers time to actually park between jobs.
1273            std::thread::sleep(std::time::Duration::from_micros(200));
1274        }
1275        assert_eq!(counter.load(Ordering::Relaxed), 150);
1276    }
1277
1278    #[test]
1279    fn worker_indices_are_distinct_and_cover_range() {
1280        let pool = Pool::new(3);
1281        let hits: Vec<AtomicUsize> = (0..4).map(|_| AtomicUsize::new(0)).collect();
1282        for _ in 0..20 {
1283            pool.run(&|widx, n| {
1284                assert_eq!(n, 4);
1285                hits[widx].fetch_add(1, Ordering::Relaxed);
1286            });
1287        }
1288        for (i, h) in hits.iter().enumerate() {
1289            assert_eq!(h.load(Ordering::Relaxed), 20, "participant {i} missed runs");
1290        }
1291    }
1292}
1293
1294#[cfg(test)]
1295mod grain_tests {
1296    use super::grain_for;
1297
1298    #[test]
1299    fn a_short_job_still_reaches_every_worker() {
1300        // 24 rows, 49 workers: the old flat floor of 32 handed all 24 to the
1301        // first worker and woke the rest for nothing.
1302        assert_eq!(grain_for(24, 49), 1);
1303        // Wide jobs keep the stride the SDOT loop wants.
1304        assert_eq!(grain_for(4096, 49), 32);
1305        assert_eq!(grain_for(32768, 49), 83);
1306        // Degenerate shapes must not divide by zero or return zero.
1307        assert_eq!(grain_for(0, 49), 1);
1308        assert_eq!(grain_for(7, 1), 7);
1309        assert!(grain_for(1, 49) >= 1);
1310    }
1311}