1use std::cell::UnsafeCell;
31use std::sync::Arc;
32use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
33
34pub static FORCED_THREADS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
38
39pub static WORKER_TIDS: std::sync::Mutex<Vec<i32>> = std::sync::Mutex::new(Vec::new());
43
44
45#[derive(Clone, Copy)]
49struct TaskPtr(*const (dyn Fn(usize, usize) + Sync));
50unsafe impl Send for TaskPtr {}
51
52struct Inner {
53 epoch: AtomicUsize,
55 remaining: AtomicUsize,
57 slot: UnsafeCell<Option<(TaskPtr, usize)>>,
61 shutdown: AtomicBool,
62 spin_budget: usize,
64 parked: Box<[AtomicBool]>,
67}
68
69unsafe impl Sync for Inner {}
73
74static DISPATCHES: AtomicUsize = AtomicUsize::new(0);
77
78pub fn dispatch_count() -> usize {
80 DISPATCHES.load(Ordering::Relaxed)
81}
82
83pub struct Pool {
86 inner: Arc<Inner>,
87 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 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 #[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 #[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 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 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 pub fn n_workers(&self) -> usize {
256 self.threads.len()
257 }
258
259 pub fn run_rows(&self, rows: usize, f: &(dyn Fn(usize, usize) + Sync)) {
270 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 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 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; 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 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 f(nw, n);
344
345 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#[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 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 let mut seen = 0usize;
432 loop {
433 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 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 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
472pub 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
506pub 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 #[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 assert_eq!(
592 P::cores_from_capacities(&[1024, 1024, 1024, 1024, 768, 768, 768, 768]),
593 Some(8)
594 );
595 assert_eq!(
597 P::cores_from_capacities(&[1024, 1024, 1024, 1024, 350, 350, 350, 350]),
598 Some(4)
599 );
600 assert_eq!(
602 P::cores_from_capacities(&[1024, 800, 800, 800, 800, 300, 300, 300]),
603 Some(5)
604 );
605 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 assert_eq!(counter.load(Ordering::Relaxed), 400);
664 }
665
666 #[test]
667 fn pool_wakes_after_park() {
668 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 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}