1use rayon::prelude::*;
37use std::collections::HashSet;
38use std::path::Path;
39
40#[cfg(target_os = "macos")]
42mod qos {
43 pub const QOS_CLASS_USER_INTERACTIVE: u32 = 0x21;
44 pub const QOS_CLASS_USER_INITIATED: u32 = 0x19;
45 pub const QOS_CLASS_DEFAULT: u32 = 0x15;
46 pub const QOS_CLASS_UTILITY: u32 = 0x11;
47 pub const QOS_CLASS_BACKGROUND: u32 = 0x09;
48
49 extern "C" {
50 pub fn pthread_set_qos_class_self_np(qos: u32, relative_priority: i32) -> i32;
51 pub fn qos_class_self() -> u32;
52 }
53
54 pub fn name(class: u32) -> &'static str {
55 match class {
56 QOS_CLASS_USER_INTERACTIVE => "user-interactive",
57 QOS_CLASS_USER_INITIATED => "user-initiated",
58 QOS_CLASS_DEFAULT => "default",
59 QOS_CLASS_UTILITY => "utility",
60 QOS_CLASS_BACKGROUND => "background",
61 _ => "unspecified",
62 }
63 }
64}
65
66pub fn current_qos_name() -> Option<&'static str> {
69 #[cfg(target_os = "macos")]
70 {
71 Some(qos::name(unsafe { qos::qos_class_self() }))
72 }
73 #[cfg(not(target_os = "macos"))]
74 {
75 None
76 }
77}
78
79pub fn perf_core_count() -> usize {
86 #[cfg(target_os = "macos")]
87 {
88 if let Some(n) = sysctl_usize("hw.perflevel0.physicalcpu") {
89 if n > 0 {
90 return n;
91 }
92 }
93 }
94 physical_core_count()
95}
96
97#[cfg(target_os = "macos")]
98fn sysctl_usize(name: &str) -> Option<usize> {
99 use std::ffi::CString;
100 extern "C" {
101 fn sysctlbyname(
102 name: *const std::os::raw::c_char,
103 oldp: *mut std::ffi::c_void,
104 oldlenp: *mut usize,
105 newp: *mut std::ffi::c_void,
106 newlen: usize,
107 ) -> std::os::raw::c_int;
108 }
109 let key = CString::new(name).ok()?;
110 let mut out: i32 = 0;
111 let mut len = std::mem::size_of::<i32>();
112 let rc = unsafe {
115 sysctlbyname(
116 key.as_ptr(),
117 &mut out as *mut i32 as *mut std::ffi::c_void,
118 &mut len,
119 std::ptr::null_mut(),
120 0,
121 )
122 };
123 if rc == 0 && out > 0 {
124 Some(out as usize)
125 } else {
126 None
127 }
128}
129
130pub const SYSFS_CPU_ROOT: &str = "/sys/devices/system/cpu";
149
150pub fn parse_thread_siblings_list(text: &str) -> Vec<usize> {
162 let mut out = Vec::new();
163 for token in text.trim().split(',') {
164 let token = token.trim();
165 if token.is_empty() {
166 continue;
167 }
168 match token.split_once('-') {
169 Some((lo, hi)) => {
170 let hi = hi.split(':').next().unwrap_or(hi);
174 if let (Ok(lo), Ok(hi)) = (lo.trim().parse::<usize>(), hi.trim().parse::<usize>()) {
175 if lo <= hi {
176 out.extend(lo..=hi);
177 }
178 }
179 }
180 None => {
181 if let Ok(cpu) = token.parse::<usize>() {
182 out.push(cpu);
183 }
184 }
185 }
186 }
187 out
188}
189
190pub fn process_affinity_cpus() -> Vec<usize> {
200 #[cfg(target_os = "linux")]
201 {
202 if let Some(cpus) = sched_affinity_cpus() {
203 return cpus;
204 }
205 }
206 let n = std::thread::available_parallelism()
207 .map(|n| n.get())
208 .unwrap_or(1);
209 (0..n).collect()
210}
211
212#[cfg(target_os = "linux")]
213fn sched_affinity_cpus() -> Option<Vec<usize>> {
214 unsafe {
218 let mut set: libc::cpu_set_t = std::mem::zeroed();
219 let rc = libc::sched_getaffinity(0, std::mem::size_of::<libc::cpu_set_t>(), &mut set);
220 if rc != 0 {
221 return None;
222 }
223 let cpus: Vec<usize> = (0..libc::CPU_SETSIZE as usize)
224 .filter(|&cpu| libc::CPU_ISSET(cpu, &set))
225 .collect();
226 if cpus.is_empty() {
227 None
228 } else {
229 Some(cpus)
230 }
231 }
232}
233
234#[derive(Debug, Clone, PartialEq, Eq, Default)]
246pub struct CpuTopology {
247 entries: Vec<(usize, Vec<usize>)>,
249}
250
251impl CpuTopology {
252 pub fn from_sibling_lists<I>(entries: I) -> Self
261 where
262 I: IntoIterator<Item = (usize, Vec<usize>)>,
263 {
264 let mut entries: Vec<(usize, Vec<usize>)> = entries.into_iter().collect();
265 entries.sort_by_key(|(cpu, _)| *cpu);
266 entries.dedup_by_key(|(cpu, _)| *cpu);
267 Self { entries }
268 }
269
270 pub fn read_from(root: &Path, allowed: &[usize]) -> Self {
276 Self::from_sibling_lists(allowed.iter().map(|&cpu| {
277 let path = root
278 .join(format!("cpu{cpu}"))
279 .join("topology")
280 .join("thread_siblings_list");
281 let siblings = std::fs::read_to_string(&path)
282 .map(|text| parse_thread_siblings_list(&text))
283 .unwrap_or_default();
284 (cpu, siblings)
285 }))
286 }
287
288 pub fn detect() -> Self {
291 Self::read_from(Path::new(SYSFS_CPU_ROOT), &process_affinity_cpus())
292 }
293
294 pub fn allowed_cpus(&self) -> Vec<usize> {
296 self.entries.iter().map(|(cpu, _)| *cpu).collect()
297 }
298
299 pub fn len(&self) -> usize {
301 self.entries.len()
302 }
303
304 pub fn is_empty(&self) -> bool {
306 self.entries.is_empty()
307 }
308}
309
310pub fn physical_core_cpus_in(topology: &CpuTopology) -> Vec<usize> {
323 let mut reps: Vec<usize> = Vec::new();
324 let mut seen: HashSet<Vec<usize>> = HashSet::new();
325 for (cpu, siblings) in &topology.entries {
326 if siblings.is_empty() {
327 reps.push(*cpu);
330 continue;
331 }
332 let mut key = siblings.clone();
333 key.sort_unstable();
334 key.dedup();
335 if seen.insert(key) {
336 reps.push(*cpu);
337 }
338 }
339 if !reps.is_empty() {
340 return reps;
341 }
342 let allowed = topology.allowed_cpus();
343 if allowed.is_empty() {
344 vec![0]
345 } else {
346 allowed
347 }
348}
349
350pub fn physical_core_cpus() -> Vec<usize> {
352 physical_core_cpus_in(&CpuTopology::detect())
353}
354
355pub fn physical_core_count() -> usize {
364 let logical = std::thread::available_parallelism()
365 .map(|n| n.get())
366 .unwrap_or(1);
367 physical_core_cpus().len().clamp(1, logical.max(1))
368}
369
370pub fn resolve_threads_and_affinity_in(
382 requested: usize,
383 topology: &CpuTopology,
384) -> (usize, Vec<usize>) {
385 let reps = physical_core_cpus_in(topology);
386 if requested == 0 {
387 let n = reps.len();
388 return (n, reps);
389 }
390 let rep_set: HashSet<usize> = reps.iter().copied().collect();
391 let mut order = reps;
392 order.extend(
393 topology
394 .allowed_cpus()
395 .into_iter()
396 .filter(|cpu| !rep_set.contains(cpu)),
397 );
398 if order.is_empty() {
399 order.push(0);
400 }
401 let core_ids = (0..requested).map(|i| order[i % order.len()]).collect();
402 (requested, core_ids)
403}
404
405pub fn resolve_threads_and_affinity(requested: usize) -> (usize, Vec<usize>) {
407 resolve_threads_and_affinity_in(requested, &CpuTopology::detect())
408}
409
410#[derive(Debug, Clone, PartialEq, Eq, Default)]
413pub struct CpuPoolPlan {
414 pub num_threads: usize,
416 pub core_ids: Vec<usize>,
418 pub coordinator_cpu: Option<usize>,
421}
422
423pub fn plan_cpu_pool_in(
435 requested: usize,
436 reserve_coordinator: bool,
437 topology: &CpuTopology,
438) -> CpuPoolPlan {
439 let (mut num_threads, mut core_ids) = resolve_threads_and_affinity_in(requested, topology);
440 let mut coordinator_cpu = None;
441 if reserve_coordinator && requested == 0 && num_threads > 2 {
442 coordinator_cpu = core_ids.pop();
443 num_threads -= 1;
444 }
445 CpuPoolPlan {
446 num_threads,
447 core_ids,
448 coordinator_cpu,
449 }
450}
451
452pub fn plan_cpu_pool(requested: usize, reserve_coordinator: bool) -> CpuPoolPlan {
454 plan_cpu_pool_in(requested, reserve_coordinator, &CpuTopology::detect())
455}
456
457pub fn clamp_intra_op_threads(
470 configured: usize,
471 plan: &CpuPoolPlan,
472 physical_cores: usize,
473) -> usize {
474 let coordinator = usize::from(plan.coordinator_cpu.is_some());
475 let spare = physical_cores
476 .saturating_sub(plan.num_threads)
477 .saturating_sub(coordinator)
478 .saturating_sub(1);
479 configured.min(spare).max(1)
480}
481
482pub fn resolve_cpu_threads() -> usize {
485 for key in ["FERROX_CPU_THREADS", "RAYON_NUM_THREADS"] {
486 if let Ok(v) = std::env::var(key) {
487 if let Ok(n) = v.trim().parse::<usize>() {
488 if n > 0 {
489 return n;
490 }
491 }
492 }
493 }
494 perf_core_count()
495}
496
497pub fn should_parallelize(n_rows: usize, n_cols: usize) -> bool {
500 n_rows > 1 && n_rows.saturating_mul(n_cols) >= 256_000
501}
502
503pub fn for_each_row<F>(output: &mut [f32], n_rows: usize, n_cols: usize, row_fn: F)
506where
507 F: Fn(usize, &mut f32) + Send + Sync,
508{
509 let n = n_rows.min(output.len());
510 if !should_parallelize(n, n_cols) {
511 for (row, out) in output.iter_mut().enumerate().take(n) {
512 row_fn(row, out);
513 }
514 return;
515 }
516 let rows = &mut output[..n];
520 rows.par_iter_mut()
521 .enumerate()
522 .for_each(|(row, out)| row_fn(row, out));
523}
524
525pub fn for_each_chunk_init<S, I, F>(
527 output: &mut [f32],
528 chunk_len: usize,
529 work_per_chunk: usize,
530 init: I,
531 f: F,
532) where
533 I: Fn() -> S + Send + Sync,
534 S: Send,
535 F: Fn(&mut S, usize, &mut [f32]) + Send + Sync,
536{
537 if chunk_len == 0 {
538 return;
539 }
540 let n_chunks = output.len() / chunk_len;
541 if !should_parallelize(n_chunks, work_per_chunk) {
542 let mut state = init();
543 for (i, chunk) in output[..n_chunks * chunk_len]
544 .chunks_mut(chunk_len)
545 .enumerate()
546 {
547 f(&mut state, i, chunk);
548 }
549 return;
550 }
551 let chunks = &mut output[..n_chunks * chunk_len];
552 let init = &init;
553 let f = &f;
554 chunks
555 .par_chunks_mut(chunk_len)
556 .enumerate()
557 .for_each_init(init, |state, (i, c)| f(state, i, c));
558}
559
560pub fn init_cpu_pool() -> Option<usize> {
568 let threads = resolve_cpu_threads();
569 let log = std::env::var_os("FERROX_QOS_LOG").is_some();
570 let built = rayon::ThreadPoolBuilder::new()
571 .num_threads(threads)
572 .start_handler(move |idx| {
573 #[cfg(target_os = "macos")]
574 {
575 let before = unsafe { qos::qos_class_self() };
576 let rc = unsafe {
578 qos::pthread_set_qos_class_self_np(qos::QOS_CLASS_USER_INTERACTIVE, 0)
579 };
580 if log {
581 eprintln!(
582 "ferrox: rayon worker {idx} qos {} -> {} (rc={rc})",
583 qos::name(before),
584 qos::name(unsafe { qos::qos_class_self() }),
585 );
586 }
587 }
588 #[cfg(not(target_os = "macos"))]
589 {
590 let _ = (idx, log);
591 }
592 })
593 .build_global()
594 .is_ok();
595 if built {
596 Some(threads)
597 } else {
598 None
599 }
600}
601
602#[cfg(test)]
603mod tests {
604 use super::*;
605
606 #[test]
607 fn perf_core_count_is_at_least_one_and_no_more_than_logical_cores() {
608 let logical = std::thread::available_parallelism()
609 .map(|n| n.get())
610 .unwrap_or(1);
611 let perf = perf_core_count();
612 assert!(perf >= 1, "perf core count must be positive, got {perf}");
613 assert!(
614 perf <= logical,
615 "perf cores ({perf}) cannot exceed logical cores ({logical})"
616 );
617 }
618
619 #[test]
620 fn resolved_thread_count_falls_back_to_perf_cores_without_env_overrides() {
621 if std::env::var_os("FERROX_CPU_THREADS").is_none()
624 && std::env::var_os("RAYON_NUM_THREADS").is_none()
625 {
626 assert_eq!(resolve_cpu_threads(), perf_core_count());
627 }
628 }
629
630 #[test]
631 fn current_qos_name_is_reported_on_macos_and_absent_elsewhere() {
632 let qos = current_qos_name();
633 #[cfg(target_os = "macos")]
634 assert!(qos.is_some(), "macOS must report a QoS class");
635 #[cfg(not(target_os = "macos"))]
636 assert!(qos.is_none(), "QoS is a macOS-only concept");
637 }
638
639 fn smt_8t_4c() -> CpuTopology {
642 CpuTopology::from_sibling_lists([
643 (0, vec![0, 1]),
644 (1, vec![0, 1]),
645 (2, vec![2, 3]),
646 (3, vec![2, 3]),
647 (4, vec![4, 5]),
648 (5, vec![4, 5]),
649 (6, vec![6, 7]),
650 (7, vec![6, 7]),
651 ])
652 }
653
654 #[test]
661 fn an_smt_host_is_sized_to_its_physical_cores_not_its_logical_cpus() {
662 let topology = smt_8t_4c();
663 assert_eq!(
664 topology.len(),
665 8,
666 "the fixture must have twice as many logical CPUs as cores"
667 );
668 assert_eq!(
669 physical_core_cpus_in(&topology),
670 vec![0, 2, 4, 6],
671 "one representative per physical core, lowest sibling first"
672 );
673 let (threads, core_ids) = resolve_threads_and_affinity_in(0, &topology);
674 assert_eq!(threads, 4, "auto sizing must not count SMT siblings");
675 assert_eq!(core_ids, vec![0, 2, 4, 6]);
676 }
677
678 #[test]
679 fn siblings_numbered_apart_are_deduplicated_the_same_as_adjacent_ones() {
680 let topology = CpuTopology::from_sibling_lists([
682 (0, vec![0, 64]),
683 (1, vec![1, 65]),
684 (64, vec![0, 64]),
685 (65, vec![1, 65]),
686 ]);
687 assert_eq!(physical_core_cpus_in(&topology), vec![0, 1]);
688 }
689
690 #[test]
691 fn thread_siblings_lists_parse_as_ranges_comma_lists_and_mixtures() {
692 assert_eq!(parse_thread_siblings_list("0-1\n"), vec![0, 1]);
693 assert_eq!(parse_thread_siblings_list("0,64\n"), vec![0, 64]);
694 assert_eq!(parse_thread_siblings_list(" 3 "), vec![3]);
695 assert_eq!(parse_thread_siblings_list("0-1,64-65"), vec![0, 1, 64, 65]);
696 assert_eq!(parse_thread_siblings_list("2-4"), vec![2, 3, 4]);
697 assert_eq!(parse_thread_siblings_list(""), Vec::<usize>::new());
699 assert_eq!(parse_thread_siblings_list("x,-,7"), vec![7]);
700 assert_eq!(parse_thread_siblings_list("5-1"), Vec::<usize>::new());
702 }
703
704 #[test]
705 fn a_host_without_sysfs_topology_degrades_to_one_worker_per_allowed_cpu() {
706 let topology =
709 CpuTopology::from_sibling_lists((0..6).map(|cpu| (cpu, Vec::<usize>::new())));
710 assert_eq!(physical_core_cpus_in(&topology), vec![0, 1, 2, 3, 4, 5]);
711 assert_eq!(resolve_threads_and_affinity_in(0, &topology).0, 6);
712 }
713
714 #[test]
715 fn an_empty_topology_still_yields_one_usable_cpu() {
716 let topology = CpuTopology::default();
717 assert!(topology.is_empty());
718 assert_eq!(physical_core_cpus_in(&topology), vec![0]);
719 assert_eq!(resolve_threads_and_affinity_in(0, &topology), (1, vec![0]));
720 assert_eq!(
721 resolve_threads_and_affinity_in(2, &topology),
722 (2, vec![0, 0])
723 );
724 }
725
726 #[test]
727 fn cores_outside_the_affinity_mask_are_never_used_as_representatives() {
728 let full = smt_8t_4c();
732 let allowed = [1usize, 3, 4, 5];
733 let topology = CpuTopology::from_sibling_lists(
734 full.allowed_cpus()
735 .into_iter()
736 .filter(|cpu| allowed.contains(cpu))
737 .map(|cpu| {
738 (
739 cpu,
740 parse_thread_siblings_list(&format!("{}-{}", cpu & !1, cpu | 1)),
741 )
742 }),
743 );
744 assert_eq!(physical_core_cpus_in(&topology), vec![1, 3, 4]);
745 assert_eq!(resolve_threads_and_affinity_in(0, &topology).0, 3);
746 }
747
748 #[test]
749 fn an_explicit_count_fills_physical_cores_before_doubling_up_siblings() {
750 let topology = smt_8t_4c();
751 assert_eq!(
753 resolve_threads_and_affinity_in(4, &topology),
754 (4, vec![0, 2, 4, 6])
755 );
756 assert_eq!(
758 resolve_threads_and_affinity_in(6, &topology),
759 (6, vec![0, 2, 4, 6, 1, 3])
760 );
761 assert_eq!(
762 resolve_threads_and_affinity_in(8, &topology),
763 (8, vec![0, 2, 4, 6, 1, 3, 5, 7])
764 );
765 }
766
767 #[test]
768 fn an_explicit_count_larger_than_the_machine_wraps_instead_of_truncating() {
769 let topology = smt_8t_4c();
770 let (threads, core_ids) = resolve_threads_and_affinity_in(10, &topology);
771 assert_eq!(threads, 10, "an explicit width is honoured exactly");
772 assert_eq!(core_ids.len(), 10);
773 assert_eq!(&core_ids[8..], &[0, 2], "wraps back to the representatives");
774 }
775
776 #[test]
777 fn auto_sizing_donates_the_last_physical_core_to_the_coordinator() {
778 let plan = plan_cpu_pool_in(0, true, &smt_8t_4c());
779 assert_eq!(plan.num_threads, 3, "workers drop from N to N-1");
780 assert_eq!(plan.core_ids, vec![0, 2, 4]);
781 assert_eq!(plan.coordinator_cpu, Some(6));
782 assert_eq!(plan.num_threads, plan.core_ids.len());
783 }
784
785 #[test]
786 fn no_core_is_donated_without_a_coordinator_or_for_an_explicit_count() {
787 let topology = smt_8t_4c();
788 let no_coordinator = plan_cpu_pool_in(0, false, &topology);
789 assert_eq!(no_coordinator.num_threads, 4);
790 assert_eq!(no_coordinator.coordinator_cpu, None);
791
792 let explicit = plan_cpu_pool_in(4, true, &topology);
794 assert_eq!(explicit.num_threads, 4);
795 assert_eq!(explicit.coordinator_cpu, None);
796 }
797
798 #[test]
799 fn a_pool_of_two_or_fewer_keeps_its_workers_rather_than_donating() {
800 let dual = CpuTopology::from_sibling_lists([
803 (0, vec![0, 1]),
804 (1, vec![0, 1]),
805 (2, vec![2, 3]),
806 (3, vec![2, 3]),
807 ]);
808 let plan = plan_cpu_pool_in(0, true, &dual);
809 assert_eq!(plan.num_threads, 2);
810 assert_eq!(plan.coordinator_cpu, None);
811 }
812
813 #[test]
814 fn the_intra_op_clamp_leaves_a_core_for_the_calling_thread() {
815 let plan = plan_cpu_pool_in(0, true, &smt_8t_4c());
817 assert_eq!(clamp_intra_op_threads(16, &plan, 16), 11);
818 assert_eq!(clamp_intra_op_threads(4, &plan, 16), 4);
820 assert_eq!(clamp_intra_op_threads(16, &plan, 4), 1);
822 assert_eq!(clamp_intra_op_threads(16, &plan, 0), 1);
823 }
824
825 #[test]
826 fn sibling_lists_are_read_from_a_sysfs_layout_on_disk() {
827 let root = std::env::temp_dir().join(format!(
828 "ferrox-threads-sysfs-{}-{:?}",
829 std::process::id(),
830 std::thread::current().id()
831 ));
832 let _ = std::fs::remove_dir_all(&root);
833 for cpu in [0usize, 1] {
835 let dir = root.join(format!("cpu{cpu}")).join("topology");
836 std::fs::create_dir_all(&dir).expect("temp sysfs tree must be creatable");
837 std::fs::write(dir.join("thread_siblings_list"), "0-1\n")
838 .expect("temp sibling list must be writable");
839 }
840 let topology = CpuTopology::read_from(&root, &[0, 1, 2]);
841 assert_eq!(topology.allowed_cpus(), vec![0, 1, 2]);
842 assert_eq!(
843 physical_core_cpus_in(&topology),
844 vec![0, 2],
845 "cpu2 is unreadable, so it counts as a core of its own"
846 );
847 let _ = std::fs::remove_dir_all(&root);
848 }
849
850 #[test]
851 fn this_hosts_physical_core_count_is_positive_and_within_its_logical_cpus() {
852 let logical = std::thread::available_parallelism()
853 .map(|n| n.get())
854 .unwrap_or(1);
855 let physical = physical_core_count();
856 assert!(physical >= 1, "physical core count must be positive");
857 assert!(
858 physical <= logical,
859 "physical cores ({physical}) cannot exceed logical cores ({logical})"
860 );
861 assert_eq!(
862 physical_core_cpus().len().clamp(1, logical.max(1)),
863 physical
864 );
865 }
866
867 #[test]
868 fn this_hosts_affinity_mask_is_non_empty_and_ascending() {
869 let cpus = process_affinity_cpus();
870 assert!(!cpus.is_empty(), "a running process may run somewhere");
871 assert!(
872 cpus.windows(2).all(|w| w[0] < w[1]),
873 "affinity CPUs must be ascending and unique: {cpus:?}"
874 );
875 }
876
877 #[test]
878 fn for_each_row_parallel_matches_serial() {
879 let n = 4097usize;
880 let f = |row: usize| ((row % 97) as f32) * 0.25 - 3.0;
881
882 let mut par = vec![0.0f32; n];
883 for_each_row(&mut par, n, 4096, |row, slot| *slot = f(row));
884
885 let mut serial = vec![0.0f32; n];
886 for (row, slot) in serial.iter_mut().enumerate() {
887 *slot = f(row);
888 }
889 assert_eq!(par, serial);
890 }
891}