1use 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
36pub static FORCED_THREADS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
40
41pub static WORKER_TIDS: std::sync::Mutex<Vec<i32>> = std::sync::Mutex::new(Vec::new());
46
47#[derive(Clone, Copy)]
51struct TaskPtr(*const (dyn Fn(usize, usize) + Sync));
52unsafe impl Send for TaskPtr {}
53
54struct Inner {
55 epoch: AtomicUsize,
57 remaining: AtomicUsize,
59 slot: UnsafeCell<Option<(TaskPtr, usize, usize, usize)>>,
73 shutdown: AtomicBool,
74 spin_budget: AtomicUsize,
76 parked: Box<[AtomicBool]>,
79 #[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
88unsafe impl Sync for Inner {}
92
93static DISPATCHES: AtomicUsize = AtomicUsize::new(0);
96
97pub fn dispatch_count() -> usize {
99 DISPATCHES.load(Ordering::Relaxed)
100}
101
102pub struct Pool {
105 inner: Arc<Inner>,
106 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
118pub(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 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 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 #[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 #[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 #[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 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 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 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 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 pub fn n_workers(&self) -> usize {
332 self.threads.len()
333 }
334
335 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 pub(crate) fn set_spin_budget(&self, spins: usize) {
389 self.inner.spin_budget.store(spins, Ordering::Relaxed);
390 }
391
392 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 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 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 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 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; 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 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 f(nw, n);
520
521 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#[cfg(target_os = "linux")]
548mod numa {
549 use std::sync::Mutex;
550 use std::sync::atomic::{AtomicUsize, Ordering};
551
552 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 #[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 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 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 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 let cached = unsafe { libc::mincore(p, PAGE, &mut vec) } == 0 && vec & 1 == 1;
679 if cached {
680 incore += 1;
681 unsafe { std::ptr::read_volatile(p as *const u8) };
683 }
684 }
685 let mut status = vec![-1i32; pages.len()];
686 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 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 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 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 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 pub(super) fn set_affinity(tid: i32, cpus: &[usize]) -> bool {
830 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#[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 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 #[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 let mut seen = 0usize;
915 loop {
916 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 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 let (task, n, dev, limit) =
949 unsafe { (*inner.slot.get()).expect("job published with epoch") };
950 if idx >= limit {
951 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
963pub 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
997pub 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
1051pub(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 #[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 #[inline]
1084 pub(crate) fn new(p: *mut f32) -> Self {
1085 Self(p)
1086 }
1087
1088 #[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 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 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 assert_eq!(
1188 P::cores_from_capacities(&[1024, 1024, 1024, 1024, 768, 768, 768, 768]),
1189 Some(8)
1190 );
1191 assert_eq!(
1193 P::cores_from_capacities(&[1024, 1024, 1024, 1024, 350, 350, 350, 350]),
1194 Some(4)
1195 );
1196 assert_eq!(
1198 P::cores_from_capacities(&[1024, 800, 800, 800, 800, 300, 300, 300]),
1199 Some(5)
1200 );
1201 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 assert_eq!(counter.load(Ordering::Relaxed), 400);
1260 }
1261
1262 #[test]
1263 fn pool_wakes_after_park() {
1264 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 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 assert_eq!(grain_for(24, 49), 1);
1303 assert_eq!(grain_for(4096, 49), 32);
1305 assert_eq!(grain_for(32768, 49), 83);
1306 assert_eq!(grain_for(0, 49), 1);
1308 assert_eq!(grain_for(7, 1), 7);
1309 assert!(grain_for(1, 49) >= 1);
1310 }
1311}