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 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 pub fn n_workers(&self) -> usize {
251 self.threads.len()
252 }
253
254 pub fn run_rows(&self, rows: usize, f: &(dyn Fn(usize, usize) + Sync)) {
265 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 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 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; 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 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 f(nw, n);
339
340 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#[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 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 let mut seen = 0usize;
427 loop {
428 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 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 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
467pub 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
501pub 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 #[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 assert_eq!(
587 P::cores_from_capacities(&[1024, 1024, 1024, 1024, 768, 768, 768, 768]),
588 Some(8)
589 );
590 assert_eq!(
592 P::cores_from_capacities(&[1024, 1024, 1024, 1024, 350, 350, 350, 350]),
593 Some(4)
594 );
595 assert_eq!(
597 P::cores_from_capacities(&[1024, 800, 800, 800, 800, 300, 300, 300]),
598 Some(5)
599 );
600 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 assert_eq!(counter.load(Ordering::Relaxed), 400);
659 }
660
661 #[test]
662 fn pool_wakes_after_park() {
663 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 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}