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;
32use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
33
34/// Embedder override for the pool size (C ABI `cortiq_set_threads`):
35/// 0 = unset, consult CMF_THREADS / topology as before. Read once at
36/// pool construction, so set it before the load.
37pub static FORCED_THREADS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
38
39/// Kernel thread ids of the CURRENT pool's workers (Android/Linux) —
40/// what ADPF's PerformanceHintManager needs to attribute work to the
41/// governor. Refilled on every pool construction; empty elsewhere.
42pub static WORKER_TIDS: std::sync::Mutex<Vec<i32>> = std::sync::Mutex::new(Vec::new());
43
44
45/// A `*const dyn Fn` that may cross a thread boundary. Safety is
46/// provided by `Pool::run`: the caller blocks until every worker has
47/// finished, so the borrow outlives all uses.
48#[derive(Clone, Copy)]
49struct TaskPtr(*const (dyn Fn(usize, usize) + Sync));
50unsafe impl Send for TaskPtr {}
51
52struct Inner {
53    /// Bumped once per published job; workers watch it.
54    epoch: AtomicUsize,
55    /// Workers still running the current job (excludes the caller).
56    remaining: AtomicUsize,
57    /// The published job: closure pointer + total participant count.
58    /// Written by the caller BEFORE the epoch bump, read by workers
59    /// AFTER they observe the new epoch (acquire/release pairing).
60    slot: UnsafeCell<Option<(TaskPtr, usize)>>,
61    shutdown: AtomicBool,
62    /// Spin iterations before a worker parks (0 = park immediately).
63    spin_budget: usize,
64    /// Per-worker "I am parked" flags — lets the caller skip the unpark
65    /// syscall for workers that are still spinning.
66    parked: Box<[AtomicBool]>,
67}
68
69// SAFETY: `slot` is only written while no job is in flight (run()
70// returns after `remaining` hits 0) and only read after the epoch
71// publication that follows the write.
72unsafe impl Sync for Inner {}
73
74/// Process-wide dispatch counter (roadmap §3 P0 «измерения»): one tick
75/// per published job. `bench --json` reports dispatches/token from it.
76static DISPATCHES: AtomicUsize = AtomicUsize::new(0);
77
78/// Total pool jobs published since process start (all pools).
79pub fn dispatch_count() -> usize {
80    DISPATCHES.load(Ordering::Relaxed)
81}
82
83/// Persistent thread pool: shared job slot, epoch dispatch, caller
84/// participation.
85pub struct Pool {
86    inner: Arc<Inner>,
87    /// Thread handles for `unpark` (same order as `parked`).
88    threads: Vec<std::thread::Thread>,
89    joins: Vec<std::thread::JoinHandle<()>>,
90}
91
92fn spin_budget_from_env() -> usize {
93    std::env::var("CMF_POOL_SPIN")
94        .ok()
95        .and_then(|v| v.parse::<usize>().ok())
96        .unwrap_or(4000)
97}
98
99impl Pool {
100    pub fn new(n_workers: usize) -> Self {
101        Self::with_spin(n_workers, spin_budget_from_env())
102    }
103
104    /// Explicit spin budget (tests pin it without touching the env).
105    pub fn with_spin(n_workers: usize, spin_budget: usize) -> Self {
106        let inner = Arc::new(Inner {
107            epoch: AtomicUsize::new(0),
108            remaining: AtomicUsize::new(0),
109            slot: UnsafeCell::new(None),
110            shutdown: AtomicBool::new(false),
111            spin_budget,
112            parked: (0..n_workers).map(|_| AtomicBool::new(false)).collect(),
113        });
114        let mut joins = Vec::with_capacity(n_workers);
115        if let Ok(mut tids) = WORKER_TIDS.lock() {
116            tids.clear();
117        }
118        for w in 0..n_workers {
119            let inner = inner.clone();
120            let h = std::thread::Builder::new()
121                .name(format!("cmf-pool-{w}"))
122                .spawn(move || {
123                    #[cfg(any(target_os = "android", target_os = "linux"))]
124                    if let Ok(mut tids) = WORKER_TIDS.lock() {
125                        tids.push(unsafe { libc::gettid() } as i32);
126                    }
127                    worker_loop(&inner, w)
128                })
129                .expect("spawn pool worker");
130            joins.push(h);
131        }
132        let threads = joins.iter().map(|h| h.thread().clone()).collect();
133        Self {
134            inner,
135            threads,
136            joins,
137        }
138    }
139
140    /// Big-core count on heterogeneous ARM (big.LITTLE): the kernel
141    /// exposes per-core capacity on Android and most ARM Linux; efficiency
142    /// cores in the pool DRAG the big ones on our row-parallel jobs (the
143    /// same cliff llama.cpp hits at -t 10 on an M4: 163 → 112 tok/s).
144    /// None = capacities absent or homogeneous.
145    #[cfg(all(
146        target_arch = "aarch64",
147        any(target_os = "linux", target_os = "android")
148    ))]
149    fn big_cores() -> Option<usize> {
150        Self::cores_from_capacities(&core_capacities())
151    }
152
153    /// How many cores the pool should use, from the kernel's per-core
154    /// capacity values. Capacity folds µarch × clock into one number,
155    /// and the two need different treatment: cores of ANOTHER µarch
156    /// (A5xx efficiency cluster next to A7xx/X: capacity ratio ≥ ~2)
157    /// drag row-parallel work down and are excluded; cores of the SAME
158    /// µarch merely clock-binned (JLQ JR510: 8×A55 as 4×2.0 + 4×1.5 GHz,
159    /// ratio 1.33) pull their weight and must ALL be used. The 1.6
160    /// threshold splits the two regimes: on a Snapdragon 8-class part
161    /// it keeps X + A7xx mid cores and drops A5xx.
162    #[cfg_attr(
163        not(all(
164            target_arch = "aarch64",
165            any(target_os = "linux", target_os = "android")
166        )),
167        allow(dead_code)
168    )]
169    fn cores_from_capacities(caps: &[u64]) -> Option<usize> {
170        let max = *caps.iter().max()?;
171        let min = *caps.iter().min()?;
172        if caps.len() < 2 || max == min {
173            return None;
174        }
175        Some(caps.iter().filter(|&&c| c * 8 >= max * 5).count())
176    }
177
178    #[cfg(target_os = "macos")]
179    fn big_cores() -> Option<usize> {
180        unsafe extern "C" {
181            fn sysctlbyname(
182                name: *const std::ffi::c_char,
183                oldp: *mut std::ffi::c_void,
184                oldlenp: *mut usize,
185                newp: *mut std::ffi::c_void,
186                newlen: usize,
187            ) -> std::ffi::c_int;
188        }
189        unsafe {
190            let name = std::ffi::CString::new("hw.perflevel0.physicalcpu").ok()?;
191            let mut count: i32 = 0;
192            let mut size = std::mem::size_of::<i32>();
193            let ret = sysctlbyname(
194                name.as_ptr(),
195                &mut count as *mut i32 as *mut std::ffi::c_void,
196                &mut size,
197                std::ptr::null_mut(),
198                0,
199            );
200            if ret == 0 && count > 0 {
201                Some(count as usize)
202            } else {
203                None
204            }
205        }
206    }
207
208    #[cfg(not(any(
209        all(
210            target_arch = "aarch64",
211            any(target_os = "linux", target_os = "android")
212        ),
213        target_os = "macos"
214    )))]
215    fn big_cores() -> Option<usize> {
216        None
217    }
218
219    /// The thread count `from_env` would use RIGHT NOW: forced (C ABI)
220    /// > CMF_THREADS > big-core topology > available_parallelism−1.
221    /// ≤1 means the model runs serial (no pool). Introspection
222    /// (`execution_mode`, status endpoints) must report THIS, not
223    /// available_parallelism.
224    pub fn effective_threads() -> usize {
225        let forced = FORCED_THREADS.load(std::sync::atomic::Ordering::Relaxed);
226        if forced > 0 {
227            return forced;
228        }
229        match std::env::var("CMF_THREADS") {
230            Ok(v) => v.parse::<usize>().unwrap_or(0),
231            Err(_) => match Self::big_cores() {
232                Some(big) => big,
233                None => {
234                    let avail = std::thread::available_parallelism()
235                        .map(|n| n.get())
236                        .unwrap_or(1);
237                    avail.saturating_sub(1).min(8)
238                }
239            },
240        }
241    }
242
243    /// Pool sized from `CMF_THREADS` (see module docs). `None` = serial.
244    /// Without the env, heterogeneous ARM defaults to its BIG cores.
245    pub fn from_env() -> Option<Arc<Self>> {
246        let n = Self::effective_threads();
247        if n <= 1 {
248            None
249        } else {
250            Some(Arc::new(Self::new(n)))
251        }
252    }
253
254    /// Spawned worker threads (the caller joins each job on top).
255    pub fn n_workers(&self) -> usize {
256        self.threads.len()
257    }
258
259    /// Run `f(row_start, row_end)` over `0..rows`, self-balancing.
260    ///
261    /// One dispatch, but workers pull row-ranges from a shared cursor
262    /// instead of each taking a fixed 1/n slice. On a heterogeneous CPU
263    /// (Apple Silicon: 4 P-cores + 6 E-cores here) a static split makes
264    /// every matvec end at the SLOWEST core's pace while the fast ones
265    /// idle at the barrier; pulling by grain lets a P-core take several
266    /// chunks for each one an E-core takes, so skew collapses to a
267    /// single grain. Row ranges stay disjoint and each row's dot is
268    /// computed exactly as in the serial path → bit-identical output.
269    pub fn run_rows(&self, rows: usize, f: &(dyn Fn(usize, usize) + Sync)) {
270        // Enough chunks to balance, large enough to keep the SDOT inner
271        // loop and the hardware prefetcher in their stride.
272        let grain = (rows / ((self.threads.len() + 1) * 8)).max(32);
273        let next = AtomicUsize::new(0);
274        self.run(&|_w, _n| loop {
275            let start = next.fetch_add(grain, Ordering::Relaxed);
276            if start >= rows {
277                break;
278            }
279            f(start, (start + grain).min(rows));
280        });
281    }
282
283    /// Multi-matrix job: one dispatch serves SEVERAL row spaces
284    /// (roadmap §3 P0 — «одна внешняя публикация job на слой»). Parts
285    /// are laid out back-to-back in a virtual row space and pulled by
286    /// grain from one shared cursor, so QKV or gate+up cost a single
287    /// barrier instead of one each. Each part's `f(start, end)` sees its
288    /// OWN row indices — per-row math and outputs are bit-identical to
289    /// separate `run_rows` calls.
290    pub fn run_many(&self, parts: &[(usize, &(dyn Fn(usize, usize) + Sync))]) {
291        let total: usize = parts.iter().map(|p| p.0).sum();
292        if total == 0 {
293            return;
294        }
295        let grain = (total / ((self.threads.len() + 1) * 8)).max(32);
296        let next = AtomicUsize::new(0);
297        self.run(&|_w, _n| loop {
298            let s = next.fetch_add(grain, Ordering::Relaxed);
299            if s >= total {
300                break;
301            }
302            let e = (s + grain).min(total);
303            let mut base = 0usize;
304            for &(rows, f) in parts {
305                let a = s.max(base);
306                let b = e.min(base + rows);
307                if a < b {
308                    f(a - base, b - base);
309                }
310                base += rows;
311                if base >= e {
312                    break;
313                }
314            }
315        });
316    }
317
318    /// Run `f(worker_idx, n_participants)` on every worker AND the
319    /// calling thread (`worker_idx = n_workers()` for the caller);
320    /// returns when all participants have finished.
321    pub fn run(&self, f: &(dyn Fn(usize, usize) + Sync)) {
322        DISPATCHES.fetch_add(1, Ordering::Relaxed);
323        let nw = self.threads.len();
324        let n = nw + 1; // caller participates
325        // SAFETY: the wait loop below blocks until every worker is done,
326        // so extending the borrow to 'static never outlives the call.
327        let ptr: *const (dyn Fn(usize, usize) + Sync) = f;
328        let ptr: *const (dyn Fn(usize, usize) + Sync + 'static) =
329            unsafe { std::mem::transmute(ptr) };
330        // SAFETY: no job in flight (previous run() drained `remaining`),
331        // so the slot is not being read.
332        unsafe { *self.inner.slot.get() = Some((TaskPtr(ptr), n)) };
333        self.inner.remaining.store(nw, Ordering::Relaxed);
334        self.inner.epoch.fetch_add(1, Ordering::SeqCst);
335        for (i, t) in self.threads.iter().enumerate() {
336            if self.inner.parked[i].load(Ordering::SeqCst) {
337                t.unpark();
338            }
339        }
340
341        // The caller's share — the barrier costs nothing while there is
342        // real work to do.
343        f(nw, n);
344
345        // Wait for the stragglers (bounded by one worker's chunk).
346        let mut spins = 0usize;
347        while self.inner.remaining.load(Ordering::Acquire) != 0 {
348            spins += 1;
349            if spins < 10_000 {
350                std::hint::spin_loop();
351            } else {
352                std::thread::yield_now();
353            }
354        }
355    }
356}
357
358impl Drop for Pool {
359    fn drop(&mut self) {
360        self.inner.shutdown.store(true, Ordering::SeqCst);
361        for t in &self.threads {
362            t.unpark();
363        }
364        for h in self.joins.drain(..) {
365            let _ = h.join();
366        }
367    }
368}
369
370/// Per-core capacity: the kernel's `cpu_capacity` (µarch × clock) when
371/// EAS exposes it, else `cpufreq/cpuinfo_max_freq` — same cluster
372/// ordering, so the 62.5% big-core rule keeps working on EAS-less
373/// kernels (TUNING.md open item: pinning silently did nothing there).
374#[cfg(any(
375    target_os = "android",
376    all(target_arch = "aarch64", target_os = "linux")
377))]
378fn core_capacities() -> Vec<u64> {
379    let read_all = |leaf: &str| -> Vec<u64> {
380        let mut vals = Vec::new();
381        for cpu in 0.. {
382            let path = format!("/sys/devices/system/cpu/cpu{cpu}/{leaf}");
383            match std::fs::read_to_string(&path) {
384                Ok(v) => match v.trim().parse() {
385                    Ok(x) => vals.push(x),
386                    Err(_) => break,
387                },
388                Err(_) => break,
389            }
390        }
391        vals
392    };
393    let caps = read_all("cpu_capacity");
394    if caps.len() >= 2 {
395        return caps;
396    }
397    read_all("cpufreq/cpuinfo_max_freq")
398}
399
400#[cfg(target_os = "android")]
401fn pin_thread_to_big_cores() {
402    use std::mem;
403    let caps = core_capacities();
404    let max = caps.iter().copied().max().unwrap_or(0);
405    let min = caps.iter().copied().min().unwrap_or(0);
406
407    // Only pin if heterogeneous
408    if caps.len() < 2 || max == min {
409        return;
410    }
411
412    unsafe {
413        let mut set: libc::cpu_set_t = mem::zeroed();
414        for (i, &c) in caps.iter().enumerate() {
415            if c * 8 >= max * 5 {
416                libc::CPU_SET(i, &mut set);
417            }
418        }
419        libc::sched_setaffinity(0, mem::size_of::<libc::cpu_set_t>(), &set);
420    }
421}
422
423fn worker_loop(inner: &Inner, idx: usize) {
424    #[cfg(target_os = "android")]
425    pin_thread_to_big_cores();
426
427    // The pool is created at epoch 0; baseline MUST be 0, not a fresh
428    // epoch read — if the caller publishes a job before the OS actually
429    // starts this thread, reading the live epoch would adopt that job's
430    // epoch as "already seen", skip it, and deadlock the caller's wait.
431    let mut seen = 0usize;
432    loop {
433        // Wait for a new epoch: spin first (decode publishes the next
434        // matvec within microseconds), park only when idle for real.
435        let mut spins = 0usize;
436        loop {
437            let e = inner.epoch.load(Ordering::Acquire);
438            if e != seen {
439                seen = e;
440                break;
441            }
442            if inner.shutdown.load(Ordering::Relaxed) {
443                return;
444            }
445            if spins < inner.spin_budget {
446                spins += 1;
447                std::hint::spin_loop();
448            } else {
449                inner.parked[idx].store(true, Ordering::SeqCst);
450                // Re-check under SeqCst: the caller bumps the epoch
451                // BEFORE reading `parked`, so either it sees our flag
452                // (and unparks) or we see its epoch here — a missed
453                // wakeup is impossible. Spurious unparks just loop.
454                if inner.epoch.load(Ordering::SeqCst) == seen
455                    && !inner.shutdown.load(Ordering::Relaxed)
456                {
457                    std::thread::park();
458                }
459                inner.parked[idx].store(false, Ordering::SeqCst);
460            }
461        }
462        // SAFETY: the slot was written before the epoch bump we just
463        // observed (release/acquire), and stays valid until `remaining`
464        // drops to zero — which happens only after `f` returns below.
465        let (task, n) = unsafe { (*inner.slot.get()).expect("job published with epoch") };
466        let f = unsafe { &*task.0 };
467        f(idx, n);
468        inner.remaining.fetch_sub(1, Ordering::AcqRel);
469    }
470}
471
472/// Row-parallel dense matvec: `out[o] = Σ_j w[o·in + j]·x[j]`.
473/// Bit-identical to the serial loop (row order does not change math).
474pub fn matvec_rows(pool: Option<&Pool>, w: &[f32], x: &[f32], out: &mut [f32]) {
475    let in_dim = x.len();
476    let out_dim = out.len();
477    debug_assert!(w.len() >= out_dim * in_dim);
478
479    let row_dot = |o: usize| -> f32 {
480        let row = &w[o * in_dim..(o + 1) * in_dim];
481        let mut sum = 0.0f32;
482        for j in 0..in_dim {
483            sum += row[j] * x[j];
484        }
485        sum
486    };
487
488    match pool {
489        Some(pool) if out_dim >= 256 => {
490            let out_addr = SendMut(out.as_mut_ptr());
491            let run_range = move |start: usize, end: usize| {
492                for o in start..end {
493                    unsafe { *out_addr.at(o) = row_dot(o) };
494                }
495            };
496            pool.run_rows(out_dim, &run_range);
497        }
498        _ => {
499            for (o, dst) in out.iter_mut().enumerate() {
500                *dst = row_dot(o);
501            }
502        }
503    }
504}
505
506/// Two-input row matvec: one pass over the weight rows serves BOTH
507/// inputs — CPU decode is memory-bound, so the second position costs a
508/// fraction of the first (this is where MTP speculative verify wins).
509/// Per-output accumulation order matches the single-input path exactly
510/// → bit-identical results.
511pub fn matvec_rows2(
512    pool: Option<&Pool>,
513    w: &[f32],
514    x1: &[f32],
515    x2: &[f32],
516    out1: &mut [f32],
517    out2: &mut [f32],
518) {
519    let in_dim = x1.len();
520    debug_assert_eq!(x2.len(), in_dim);
521    let out_dim = out1.len();
522    debug_assert_eq!(out2.len(), out_dim);
523    debug_assert!(w.len() >= out_dim * in_dim);
524
525    let row_dots = |o: usize| -> (f32, f32) {
526        let row = &w[o * in_dim..(o + 1) * in_dim];
527        let (mut s1, mut s2) = (0.0f32, 0.0f32);
528        for j in 0..in_dim {
529            s1 += row[j] * x1[j];
530            s2 += row[j] * x2[j];
531        }
532        (s1, s2)
533    };
534
535    match pool {
536        Some(pool) if out_dim >= 256 => {
537            let o1 = SendMut(out1.as_mut_ptr());
538            let o2 = SendMut(out2.as_mut_ptr());
539            let run_range = move |start: usize, end: usize| {
540                for o in start..end {
541                    let (s1, s2) = row_dots(o);
542                    unsafe {
543                        *o1.at(o) = s1;
544                        *o2.at(o) = s2;
545                    }
546                }
547            };
548            pool.run_rows(out_dim, &run_range);
549        }
550        _ => {
551            for o in 0..out_dim {
552                let (s1, s2) = row_dots(o);
553                out1[o] = s1;
554                out2[o] = s2;
555            }
556        }
557    }
558}
559
560#[derive(Clone, Copy)]
561struct SendMut(*mut f32);
562unsafe impl Send for SendMut {}
563unsafe impl Sync for SendMut {}
564
565impl SendMut {
566    /// Method receiver forces the closure to capture the whole (Sync)
567    /// wrapper, not the bare `*mut f32` field (edition-2021 precise capture).
568    #[inline]
569    fn at(self, i: usize) -> *mut f32 {
570        unsafe { self.0.add(i) }
571    }
572}
573
574#[cfg(test)]
575mod tests {
576    #[test]
577    fn forced_threads_overrides_env_and_topology() {
578        use std::sync::atomic::Ordering;
579        super::FORCED_THREADS.store(3, Ordering::Relaxed);
580        let pool = super::Pool::from_env().expect("forced 3 → pool");
581        assert_eq!(pool.n_workers(), 3);
582        super::FORCED_THREADS.store(1, Ordering::Relaxed);
583        assert!(super::Pool::from_env().is_none(), "forced 1 → serial");
584        super::FORCED_THREADS.store(0, Ordering::Relaxed);
585    }
586
587    #[test]
588    fn capacity_split_clock_bins_vs_microarch() {
589        type P = super::Pool;
590        // JR510: all-A55, two clock bins — use every core.
591        assert_eq!(
592            P::cores_from_capacities(&[1024, 1024, 1024, 1024, 768, 768, 768, 768]),
593            Some(8)
594        );
595        // Classic big.LITTLE (A78 + A55) — big only.
596        assert_eq!(
597            P::cores_from_capacities(&[1024, 1024, 1024, 1024, 350, 350, 350, 350]),
598            Some(4)
599        );
600        // Three-tier flagship: X + A7xx mids stay, A5xx littles go.
601        assert_eq!(
602            P::cores_from_capacities(&[1024, 800, 800, 800, 800, 300, 300, 300]),
603            Some(5)
604        );
605        // Uniform: no signal, caller falls back.
606        assert_eq!(P::cores_from_capacities(&[1024; 8]), None);
607        assert_eq!(P::cores_from_capacities(&[]), None);
608    }
609
610    use super::*;
611
612    #[test]
613    fn parallel_matvec_equals_serial_bitexact() {
614        let (out_dim, in_dim) = (512, 64);
615        let w: Vec<f32> = (0..out_dim * in_dim)
616            .map(|i| (i as f32 * 0.013).sin())
617            .collect();
618        let x: Vec<f32> = (0..in_dim).map(|i| (i as f32 * 0.07).cos()).collect();
619
620        let mut serial = vec![0.0f32; out_dim];
621        matvec_rows(None, &w, &x, &mut serial);
622
623        let pool = Pool::new(4);
624        let mut parallel = vec![0.0f32; out_dim];
625        matvec_rows(Some(&pool), &w, &x, &mut parallel);
626
627        assert_eq!(serial, parallel, "row-parallel must be bit-identical");
628    }
629
630    #[test]
631    fn fused_pair_equals_two_singles_bitexact() {
632        let (out_dim, in_dim) = (300, 48);
633        let w: Vec<f32> = (0..out_dim * in_dim)
634            .map(|i| (i as f32 * 0.011).sin())
635            .collect();
636        let x1: Vec<f32> = (0..in_dim).map(|i| (i as f32 * 0.03).cos()).collect();
637        let x2: Vec<f32> = (0..in_dim).map(|i| (i as f32 * 0.09).sin()).collect();
638
639        let mut a1 = vec![0.0f32; out_dim];
640        let mut a2 = vec![0.0f32; out_dim];
641        matvec_rows(None, &w, &x1, &mut a1);
642        matvec_rows(None, &w, &x2, &mut a2);
643
644        for pool in [None, Some(Pool::new(3))] {
645            let mut b1 = vec![0.0f32; out_dim];
646            let mut b2 = vec![0.0f32; out_dim];
647            matvec_rows2(pool.as_ref(), &w, &x1, &x2, &mut b1, &mut b2);
648            assert_eq!(a1, b1, "fused lane 1 must be bit-identical");
649            assert_eq!(a2, b2, "fused lane 2 must be bit-identical");
650        }
651    }
652
653    #[test]
654    fn pool_survives_many_runs() {
655        let pool = Pool::new(3);
656        let counter = AtomicUsize::new(0);
657        for _ in 0..100 {
658            pool.run(&|_, _| {
659                counter.fetch_add(1, Ordering::Relaxed);
660            });
661        }
662        // 3 workers + the participating caller = 4 executions per run.
663        assert_eq!(counter.load(Ordering::Relaxed), 400);
664    }
665
666    #[test]
667    fn pool_wakes_after_park() {
668        // Force immediate parking (no spin) — the epoch/parked handshake
669        // must still never miss a wakeup.
670        let pool = Pool::with_spin(2, 0);
671        let counter = AtomicUsize::new(0);
672        for _ in 0..50 {
673            pool.run(&|_, _| {
674                counter.fetch_add(1, Ordering::Relaxed);
675            });
676            // Give workers time to actually park between jobs.
677            std::thread::sleep(std::time::Duration::from_micros(200));
678        }
679        assert_eq!(counter.load(Ordering::Relaxed), 150);
680    }
681
682    #[test]
683    fn worker_indices_are_distinct_and_cover_range() {
684        let pool = Pool::new(3);
685        let hits: Vec<AtomicUsize> = (0..4).map(|_| AtomicUsize::new(0)).collect();
686        for _ in 0..20 {
687            pool.run(&|widx, n| {
688                assert_eq!(n, 4);
689                hits[widx].fetch_add(1, Ordering::Relaxed);
690            });
691        }
692        for (i, h) in hits.iter().enumerate() {
693            assert_eq!(h.load(Ordering::Relaxed), 20, "participant {i} missed runs");
694        }
695    }
696}