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    /// Pool sized from `CMF_THREADS` (see module docs). `None` = serial.
220    /// Without the env, heterogeneous ARM defaults to its BIG cores.
221    pub fn from_env() -> Option<Arc<Self>> {
222        let forced = FORCED_THREADS.load(std::sync::atomic::Ordering::Relaxed);
223        if forced > 0 {
224            return if forced <= 1 {
225                None
226            } else {
227                Some(Arc::new(Self::new(forced)))
228            };
229        }
230        let n = match std::env::var("CMF_THREADS") {
231            Ok(v) => v.parse::<usize>().unwrap_or(0),
232            Err(_) => match Self::big_cores() {
233                Some(big) => big,
234                None => {
235                    let avail = std::thread::available_parallelism()
236                        .map(|n| n.get())
237                        .unwrap_or(1);
238                    avail.saturating_sub(1).min(8)
239                }
240            },
241        };
242        if n <= 1 {
243            None
244        } else {
245            Some(Arc::new(Self::new(n)))
246        }
247    }
248
249    /// Spawned worker threads (the caller joins each job on top).
250    pub fn n_workers(&self) -> usize {
251        self.threads.len()
252    }
253
254    /// Run `f(row_start, row_end)` over `0..rows`, self-balancing.
255    ///
256    /// One dispatch, but workers pull row-ranges from a shared cursor
257    /// instead of each taking a fixed 1/n slice. On a heterogeneous CPU
258    /// (Apple Silicon: 4 P-cores + 6 E-cores here) a static split makes
259    /// every matvec end at the SLOWEST core's pace while the fast ones
260    /// idle at the barrier; pulling by grain lets a P-core take several
261    /// chunks for each one an E-core takes, so skew collapses to a
262    /// single grain. Row ranges stay disjoint and each row's dot is
263    /// computed exactly as in the serial path → bit-identical output.
264    pub fn run_rows(&self, rows: usize, f: &(dyn Fn(usize, usize) + Sync)) {
265        // Enough chunks to balance, large enough to keep the SDOT inner
266        // loop and the hardware prefetcher in their stride.
267        let grain = (rows / ((self.threads.len() + 1) * 8)).max(32);
268        let next = AtomicUsize::new(0);
269        self.run(&|_w, _n| loop {
270            let start = next.fetch_add(grain, Ordering::Relaxed);
271            if start >= rows {
272                break;
273            }
274            f(start, (start + grain).min(rows));
275        });
276    }
277
278    /// Multi-matrix job: one dispatch serves SEVERAL row spaces
279    /// (roadmap §3 P0 — «одна внешняя публикация job на слой»). Parts
280    /// are laid out back-to-back in a virtual row space and pulled by
281    /// grain from one shared cursor, so QKV or gate+up cost a single
282    /// barrier instead of one each. Each part's `f(start, end)` sees its
283    /// OWN row indices — per-row math and outputs are bit-identical to
284    /// separate `run_rows` calls.
285    pub fn run_many(&self, parts: &[(usize, &(dyn Fn(usize, usize) + Sync))]) {
286        let total: usize = parts.iter().map(|p| p.0).sum();
287        if total == 0 {
288            return;
289        }
290        let grain = (total / ((self.threads.len() + 1) * 8)).max(32);
291        let next = AtomicUsize::new(0);
292        self.run(&|_w, _n| loop {
293            let s = next.fetch_add(grain, Ordering::Relaxed);
294            if s >= total {
295                break;
296            }
297            let e = (s + grain).min(total);
298            let mut base = 0usize;
299            for &(rows, f) in parts {
300                let a = s.max(base);
301                let b = e.min(base + rows);
302                if a < b {
303                    f(a - base, b - base);
304                }
305                base += rows;
306                if base >= e {
307                    break;
308                }
309            }
310        });
311    }
312
313    /// Run `f(worker_idx, n_participants)` on every worker AND the
314    /// calling thread (`worker_idx = n_workers()` for the caller);
315    /// returns when all participants have finished.
316    pub fn run(&self, f: &(dyn Fn(usize, usize) + Sync)) {
317        DISPATCHES.fetch_add(1, Ordering::Relaxed);
318        let nw = self.threads.len();
319        let n = nw + 1; // caller participates
320        // SAFETY: the wait loop below blocks until every worker is done,
321        // so extending the borrow to 'static never outlives the call.
322        let ptr: *const (dyn Fn(usize, usize) + Sync) = f;
323        let ptr: *const (dyn Fn(usize, usize) + Sync + 'static) =
324            unsafe { std::mem::transmute(ptr) };
325        // SAFETY: no job in flight (previous run() drained `remaining`),
326        // so the slot is not being read.
327        unsafe { *self.inner.slot.get() = Some((TaskPtr(ptr), n)) };
328        self.inner.remaining.store(nw, Ordering::Relaxed);
329        self.inner.epoch.fetch_add(1, Ordering::SeqCst);
330        for (i, t) in self.threads.iter().enumerate() {
331            if self.inner.parked[i].load(Ordering::SeqCst) {
332                t.unpark();
333            }
334        }
335
336        // The caller's share — the barrier costs nothing while there is
337        // real work to do.
338        f(nw, n);
339
340        // Wait for the stragglers (bounded by one worker's chunk).
341        let mut spins = 0usize;
342        while self.inner.remaining.load(Ordering::Acquire) != 0 {
343            spins += 1;
344            if spins < 10_000 {
345                std::hint::spin_loop();
346            } else {
347                std::thread::yield_now();
348            }
349        }
350    }
351}
352
353impl Drop for Pool {
354    fn drop(&mut self) {
355        self.inner.shutdown.store(true, Ordering::SeqCst);
356        for t in &self.threads {
357            t.unpark();
358        }
359        for h in self.joins.drain(..) {
360            let _ = h.join();
361        }
362    }
363}
364
365/// Per-core capacity: the kernel's `cpu_capacity` (µarch × clock) when
366/// EAS exposes it, else `cpufreq/cpuinfo_max_freq` — same cluster
367/// ordering, so the 62.5% big-core rule keeps working on EAS-less
368/// kernels (TUNING.md open item: pinning silently did nothing there).
369#[cfg(any(
370    target_os = "android",
371    all(target_arch = "aarch64", target_os = "linux")
372))]
373fn core_capacities() -> Vec<u64> {
374    let read_all = |leaf: &str| -> Vec<u64> {
375        let mut vals = Vec::new();
376        for cpu in 0.. {
377            let path = format!("/sys/devices/system/cpu/cpu{cpu}/{leaf}");
378            match std::fs::read_to_string(&path) {
379                Ok(v) => match v.trim().parse() {
380                    Ok(x) => vals.push(x),
381                    Err(_) => break,
382                },
383                Err(_) => break,
384            }
385        }
386        vals
387    };
388    let caps = read_all("cpu_capacity");
389    if caps.len() >= 2 {
390        return caps;
391    }
392    read_all("cpufreq/cpuinfo_max_freq")
393}
394
395#[cfg(target_os = "android")]
396fn pin_thread_to_big_cores() {
397    use std::mem;
398    let caps = core_capacities();
399    let max = caps.iter().copied().max().unwrap_or(0);
400    let min = caps.iter().copied().min().unwrap_or(0);
401
402    // Only pin if heterogeneous
403    if caps.len() < 2 || max == min {
404        return;
405    }
406
407    unsafe {
408        let mut set: libc::cpu_set_t = mem::zeroed();
409        for (i, &c) in caps.iter().enumerate() {
410            if c * 8 >= max * 5 {
411                libc::CPU_SET(i, &mut set);
412            }
413        }
414        libc::sched_setaffinity(0, mem::size_of::<libc::cpu_set_t>(), &set);
415    }
416}
417
418fn worker_loop(inner: &Inner, idx: usize) {
419    #[cfg(target_os = "android")]
420    pin_thread_to_big_cores();
421
422    // The pool is created at epoch 0; baseline MUST be 0, not a fresh
423    // epoch read — if the caller publishes a job before the OS actually
424    // starts this thread, reading the live epoch would adopt that job's
425    // epoch as "already seen", skip it, and deadlock the caller's wait.
426    let mut seen = 0usize;
427    loop {
428        // Wait for a new epoch: spin first (decode publishes the next
429        // matvec within microseconds), park only when idle for real.
430        let mut spins = 0usize;
431        loop {
432            let e = inner.epoch.load(Ordering::Acquire);
433            if e != seen {
434                seen = e;
435                break;
436            }
437            if inner.shutdown.load(Ordering::Relaxed) {
438                return;
439            }
440            if spins < inner.spin_budget {
441                spins += 1;
442                std::hint::spin_loop();
443            } else {
444                inner.parked[idx].store(true, Ordering::SeqCst);
445                // Re-check under SeqCst: the caller bumps the epoch
446                // BEFORE reading `parked`, so either it sees our flag
447                // (and unparks) or we see its epoch here — a missed
448                // wakeup is impossible. Spurious unparks just loop.
449                if inner.epoch.load(Ordering::SeqCst) == seen
450                    && !inner.shutdown.load(Ordering::Relaxed)
451                {
452                    std::thread::park();
453                }
454                inner.parked[idx].store(false, Ordering::SeqCst);
455            }
456        }
457        // SAFETY: the slot was written before the epoch bump we just
458        // observed (release/acquire), and stays valid until `remaining`
459        // drops to zero — which happens only after `f` returns below.
460        let (task, n) = unsafe { (*inner.slot.get()).expect("job published with epoch") };
461        let f = unsafe { &*task.0 };
462        f(idx, n);
463        inner.remaining.fetch_sub(1, Ordering::AcqRel);
464    }
465}
466
467/// Row-parallel dense matvec: `out[o] = Σ_j w[o·in + j]·x[j]`.
468/// Bit-identical to the serial loop (row order does not change math).
469pub fn matvec_rows(pool: Option<&Pool>, w: &[f32], x: &[f32], out: &mut [f32]) {
470    let in_dim = x.len();
471    let out_dim = out.len();
472    debug_assert!(w.len() >= out_dim * in_dim);
473
474    let row_dot = |o: usize| -> f32 {
475        let row = &w[o * in_dim..(o + 1) * in_dim];
476        let mut sum = 0.0f32;
477        for j in 0..in_dim {
478            sum += row[j] * x[j];
479        }
480        sum
481    };
482
483    match pool {
484        Some(pool) if out_dim >= 256 => {
485            let out_addr = SendMut(out.as_mut_ptr());
486            let run_range = move |start: usize, end: usize| {
487                for o in start..end {
488                    unsafe { *out_addr.at(o) = row_dot(o) };
489                }
490            };
491            pool.run_rows(out_dim, &run_range);
492        }
493        _ => {
494            for (o, dst) in out.iter_mut().enumerate() {
495                *dst = row_dot(o);
496            }
497        }
498    }
499}
500
501/// Two-input row matvec: one pass over the weight rows serves BOTH
502/// inputs — CPU decode is memory-bound, so the second position costs a
503/// fraction of the first (this is where MTP speculative verify wins).
504/// Per-output accumulation order matches the single-input path exactly
505/// → bit-identical results.
506pub fn matvec_rows2(
507    pool: Option<&Pool>,
508    w: &[f32],
509    x1: &[f32],
510    x2: &[f32],
511    out1: &mut [f32],
512    out2: &mut [f32],
513) {
514    let in_dim = x1.len();
515    debug_assert_eq!(x2.len(), in_dim);
516    let out_dim = out1.len();
517    debug_assert_eq!(out2.len(), out_dim);
518    debug_assert!(w.len() >= out_dim * in_dim);
519
520    let row_dots = |o: usize| -> (f32, f32) {
521        let row = &w[o * in_dim..(o + 1) * in_dim];
522        let (mut s1, mut s2) = (0.0f32, 0.0f32);
523        for j in 0..in_dim {
524            s1 += row[j] * x1[j];
525            s2 += row[j] * x2[j];
526        }
527        (s1, s2)
528    };
529
530    match pool {
531        Some(pool) if out_dim >= 256 => {
532            let o1 = SendMut(out1.as_mut_ptr());
533            let o2 = SendMut(out2.as_mut_ptr());
534            let run_range = move |start: usize, end: usize| {
535                for o in start..end {
536                    let (s1, s2) = row_dots(o);
537                    unsafe {
538                        *o1.at(o) = s1;
539                        *o2.at(o) = s2;
540                    }
541                }
542            };
543            pool.run_rows(out_dim, &run_range);
544        }
545        _ => {
546            for o in 0..out_dim {
547                let (s1, s2) = row_dots(o);
548                out1[o] = s1;
549                out2[o] = s2;
550            }
551        }
552    }
553}
554
555#[derive(Clone, Copy)]
556struct SendMut(*mut f32);
557unsafe impl Send for SendMut {}
558unsafe impl Sync for SendMut {}
559
560impl SendMut {
561    /// Method receiver forces the closure to capture the whole (Sync)
562    /// wrapper, not the bare `*mut f32` field (edition-2021 precise capture).
563    #[inline]
564    fn at(self, i: usize) -> *mut f32 {
565        unsafe { self.0.add(i) }
566    }
567}
568
569#[cfg(test)]
570mod tests {
571    #[test]
572    fn forced_threads_overrides_env_and_topology() {
573        use std::sync::atomic::Ordering;
574        super::FORCED_THREADS.store(3, Ordering::Relaxed);
575        let pool = super::Pool::from_env().expect("forced 3 → pool");
576        assert_eq!(pool.n_workers(), 3);
577        super::FORCED_THREADS.store(1, Ordering::Relaxed);
578        assert!(super::Pool::from_env().is_none(), "forced 1 → serial");
579        super::FORCED_THREADS.store(0, Ordering::Relaxed);
580    }
581
582    #[test]
583    fn capacity_split_clock_bins_vs_microarch() {
584        type P = super::Pool;
585        // JR510: all-A55, two clock bins — use every core.
586        assert_eq!(
587            P::cores_from_capacities(&[1024, 1024, 1024, 1024, 768, 768, 768, 768]),
588            Some(8)
589        );
590        // Classic big.LITTLE (A78 + A55) — big only.
591        assert_eq!(
592            P::cores_from_capacities(&[1024, 1024, 1024, 1024, 350, 350, 350, 350]),
593            Some(4)
594        );
595        // Three-tier flagship: X + A7xx mids stay, A5xx littles go.
596        assert_eq!(
597            P::cores_from_capacities(&[1024, 800, 800, 800, 800, 300, 300, 300]),
598            Some(5)
599        );
600        // Uniform: no signal, caller falls back.
601        assert_eq!(P::cores_from_capacities(&[1024; 8]), None);
602        assert_eq!(P::cores_from_capacities(&[]), None);
603    }
604
605    use super::*;
606
607    #[test]
608    fn parallel_matvec_equals_serial_bitexact() {
609        let (out_dim, in_dim) = (512, 64);
610        let w: Vec<f32> = (0..out_dim * in_dim)
611            .map(|i| (i as f32 * 0.013).sin())
612            .collect();
613        let x: Vec<f32> = (0..in_dim).map(|i| (i as f32 * 0.07).cos()).collect();
614
615        let mut serial = vec![0.0f32; out_dim];
616        matvec_rows(None, &w, &x, &mut serial);
617
618        let pool = Pool::new(4);
619        let mut parallel = vec![0.0f32; out_dim];
620        matvec_rows(Some(&pool), &w, &x, &mut parallel);
621
622        assert_eq!(serial, parallel, "row-parallel must be bit-identical");
623    }
624
625    #[test]
626    fn fused_pair_equals_two_singles_bitexact() {
627        let (out_dim, in_dim) = (300, 48);
628        let w: Vec<f32> = (0..out_dim * in_dim)
629            .map(|i| (i as f32 * 0.011).sin())
630            .collect();
631        let x1: Vec<f32> = (0..in_dim).map(|i| (i as f32 * 0.03).cos()).collect();
632        let x2: Vec<f32> = (0..in_dim).map(|i| (i as f32 * 0.09).sin()).collect();
633
634        let mut a1 = vec![0.0f32; out_dim];
635        let mut a2 = vec![0.0f32; out_dim];
636        matvec_rows(None, &w, &x1, &mut a1);
637        matvec_rows(None, &w, &x2, &mut a2);
638
639        for pool in [None, Some(Pool::new(3))] {
640            let mut b1 = vec![0.0f32; out_dim];
641            let mut b2 = vec![0.0f32; out_dim];
642            matvec_rows2(pool.as_ref(), &w, &x1, &x2, &mut b1, &mut b2);
643            assert_eq!(a1, b1, "fused lane 1 must be bit-identical");
644            assert_eq!(a2, b2, "fused lane 2 must be bit-identical");
645        }
646    }
647
648    #[test]
649    fn pool_survives_many_runs() {
650        let pool = Pool::new(3);
651        let counter = AtomicUsize::new(0);
652        for _ in 0..100 {
653            pool.run(&|_, _| {
654                counter.fetch_add(1, Ordering::Relaxed);
655            });
656        }
657        // 3 workers + the participating caller = 4 executions per run.
658        assert_eq!(counter.load(Ordering::Relaxed), 400);
659    }
660
661    #[test]
662    fn pool_wakes_after_park() {
663        // Force immediate parking (no spin) — the epoch/parked handshake
664        // must still never miss a wakeup.
665        let pool = Pool::with_spin(2, 0);
666        let counter = AtomicUsize::new(0);
667        for _ in 0..50 {
668            pool.run(&|_, _| {
669                counter.fetch_add(1, Ordering::Relaxed);
670            });
671            // Give workers time to actually park between jobs.
672            std::thread::sleep(std::time::Duration::from_micros(200));
673        }
674        assert_eq!(counter.load(Ordering::Relaxed), 150);
675    }
676
677    #[test]
678    fn worker_indices_are_distinct_and_cover_range() {
679        let pool = Pool::new(3);
680        let hits: Vec<AtomicUsize> = (0..4).map(|_| AtomicUsize::new(0)).collect();
681        for _ in 0..20 {
682            pool.run(&|widx, n| {
683                assert_eq!(n, 4);
684                hits[widx].fetch_add(1, Ordering::Relaxed);
685            });
686        }
687        for (i, h) in hits.iter().enumerate() {
688            assert_eq!(h.load(Ordering::Relaxed), 20, "participant {i} missed runs");
689        }
690    }
691}