use rayon::prelude::*;
use std::collections::HashSet;
use std::path::Path;
#[cfg(target_os = "macos")]
mod qos {
pub const QOS_CLASS_USER_INTERACTIVE: u32 = 0x21;
pub const QOS_CLASS_USER_INITIATED: u32 = 0x19;
pub const QOS_CLASS_DEFAULT: u32 = 0x15;
pub const QOS_CLASS_UTILITY: u32 = 0x11;
pub const QOS_CLASS_BACKGROUND: u32 = 0x09;
extern "C" {
pub fn pthread_set_qos_class_self_np(qos: u32, relative_priority: i32) -> i32;
pub fn qos_class_self() -> u32;
}
pub fn name(class: u32) -> &'static str {
match class {
QOS_CLASS_USER_INTERACTIVE => "user-interactive",
QOS_CLASS_USER_INITIATED => "user-initiated",
QOS_CLASS_DEFAULT => "default",
QOS_CLASS_UTILITY => "utility",
QOS_CLASS_BACKGROUND => "background",
_ => "unspecified",
}
}
}
pub fn current_qos_name() -> Option<&'static str> {
#[cfg(target_os = "macos")]
{
Some(qos::name(unsafe { qos::qos_class_self() }))
}
#[cfg(not(target_os = "macos"))]
{
None
}
}
pub fn perf_core_count() -> usize {
#[cfg(target_os = "macos")]
{
if let Some(n) = sysctl_usize("hw.perflevel0.physicalcpu") {
if n > 0 {
return n;
}
}
}
physical_core_count()
}
#[cfg(target_os = "macos")]
fn sysctl_usize(name: &str) -> Option<usize> {
use std::ffi::CString;
extern "C" {
fn sysctlbyname(
name: *const std::os::raw::c_char,
oldp: *mut std::ffi::c_void,
oldlenp: *mut usize,
newp: *mut std::ffi::c_void,
newlen: usize,
) -> std::os::raw::c_int;
}
let key = CString::new(name).ok()?;
let mut out: i32 = 0;
let mut len = std::mem::size_of::<i32>();
let rc = unsafe {
sysctlbyname(
key.as_ptr(),
&mut out as *mut i32 as *mut std::ffi::c_void,
&mut len,
std::ptr::null_mut(),
0,
)
};
if rc == 0 && out > 0 {
Some(out as usize)
} else {
None
}
}
pub const SYSFS_CPU_ROOT: &str = "/sys/devices/system/cpu";
pub fn parse_thread_siblings_list(text: &str) -> Vec<usize> {
let mut out = Vec::new();
for token in text.trim().split(',') {
let token = token.trim();
if token.is_empty() {
continue;
}
match token.split_once('-') {
Some((lo, hi)) => {
let hi = hi.split(':').next().unwrap_or(hi);
if let (Ok(lo), Ok(hi)) = (lo.trim().parse::<usize>(), hi.trim().parse::<usize>()) {
if lo <= hi {
out.extend(lo..=hi);
}
}
}
None => {
if let Ok(cpu) = token.parse::<usize>() {
out.push(cpu);
}
}
}
}
out
}
pub fn process_affinity_cpus() -> Vec<usize> {
#[cfg(target_os = "linux")]
{
if let Some(cpus) = sched_affinity_cpus() {
return cpus;
}
}
let n = std::thread::available_parallelism()
.map(|n| n.get())
.unwrap_or(1);
(0..n).collect()
}
#[cfg(target_os = "linux")]
fn sched_affinity_cpus() -> Option<Vec<usize>> {
unsafe {
let mut set: libc::cpu_set_t = std::mem::zeroed();
let rc = libc::sched_getaffinity(0, std::mem::size_of::<libc::cpu_set_t>(), &mut set);
if rc != 0 {
return None;
}
let cpus: Vec<usize> = (0..libc::CPU_SETSIZE as usize)
.filter(|&cpu| libc::CPU_ISSET(cpu, &set))
.collect();
if cpus.is_empty() {
None
} else {
Some(cpus)
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct CpuTopology {
entries: Vec<(usize, Vec<usize>)>,
}
impl CpuTopology {
pub fn from_sibling_lists<I>(entries: I) -> Self
where
I: IntoIterator<Item = (usize, Vec<usize>)>,
{
let mut entries: Vec<(usize, Vec<usize>)> = entries.into_iter().collect();
entries.sort_by_key(|(cpu, _)| *cpu);
entries.dedup_by_key(|(cpu, _)| *cpu);
Self { entries }
}
pub fn read_from(root: &Path, allowed: &[usize]) -> Self {
Self::from_sibling_lists(allowed.iter().map(|&cpu| {
let path = root
.join(format!("cpu{cpu}"))
.join("topology")
.join("thread_siblings_list");
let siblings = std::fs::read_to_string(&path)
.map(|text| parse_thread_siblings_list(&text))
.unwrap_or_default();
(cpu, siblings)
}))
}
pub fn detect() -> Self {
Self::read_from(Path::new(SYSFS_CPU_ROOT), &process_affinity_cpus())
}
pub fn allowed_cpus(&self) -> Vec<usize> {
self.entries.iter().map(|(cpu, _)| *cpu).collect()
}
pub fn len(&self) -> usize {
self.entries.len()
}
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
}
pub fn physical_core_cpus_in(topology: &CpuTopology) -> Vec<usize> {
let mut reps: Vec<usize> = Vec::new();
let mut seen: HashSet<Vec<usize>> = HashSet::new();
for (cpu, siblings) in &topology.entries {
if siblings.is_empty() {
reps.push(*cpu);
continue;
}
let mut key = siblings.clone();
key.sort_unstable();
key.dedup();
if seen.insert(key) {
reps.push(*cpu);
}
}
if !reps.is_empty() {
return reps;
}
let allowed = topology.allowed_cpus();
if allowed.is_empty() {
vec![0]
} else {
allowed
}
}
pub fn physical_core_cpus() -> Vec<usize> {
physical_core_cpus_in(&CpuTopology::detect())
}
pub fn physical_core_count() -> usize {
let logical = std::thread::available_parallelism()
.map(|n| n.get())
.unwrap_or(1);
physical_core_cpus().len().clamp(1, logical.max(1))
}
pub fn resolve_threads_and_affinity_in(
requested: usize,
topology: &CpuTopology,
) -> (usize, Vec<usize>) {
let reps = physical_core_cpus_in(topology);
if requested == 0 {
let n = reps.len();
return (n, reps);
}
let rep_set: HashSet<usize> = reps.iter().copied().collect();
let mut order = reps;
order.extend(
topology
.allowed_cpus()
.into_iter()
.filter(|cpu| !rep_set.contains(cpu)),
);
if order.is_empty() {
order.push(0);
}
let core_ids = (0..requested).map(|i| order[i % order.len()]).collect();
(requested, core_ids)
}
pub fn resolve_threads_and_affinity(requested: usize) -> (usize, Vec<usize>) {
resolve_threads_and_affinity_in(requested, &CpuTopology::detect())
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct CpuPoolPlan {
pub num_threads: usize,
pub core_ids: Vec<usize>,
pub coordinator_cpu: Option<usize>,
}
pub fn plan_cpu_pool_in(
requested: usize,
reserve_coordinator: bool,
topology: &CpuTopology,
) -> CpuPoolPlan {
let (mut num_threads, mut core_ids) = resolve_threads_and_affinity_in(requested, topology);
let mut coordinator_cpu = None;
if reserve_coordinator && requested == 0 && num_threads > 2 {
coordinator_cpu = core_ids.pop();
num_threads -= 1;
}
CpuPoolPlan {
num_threads,
core_ids,
coordinator_cpu,
}
}
pub fn plan_cpu_pool(requested: usize, reserve_coordinator: bool) -> CpuPoolPlan {
plan_cpu_pool_in(requested, reserve_coordinator, &CpuTopology::detect())
}
pub fn clamp_intra_op_threads(
configured: usize,
plan: &CpuPoolPlan,
physical_cores: usize,
) -> usize {
let coordinator = usize::from(plan.coordinator_cpu.is_some());
let spare = physical_cores
.saturating_sub(plan.num_threads)
.saturating_sub(coordinator)
.saturating_sub(1);
configured.min(spare).max(1)
}
pub fn resolve_cpu_threads() -> usize {
for key in ["FERROX_CPU_THREADS", "RAYON_NUM_THREADS"] {
if let Ok(v) = std::env::var(key) {
if let Ok(n) = v.trim().parse::<usize>() {
if n > 0 {
return n;
}
}
}
}
perf_core_count()
}
pub fn should_parallelize(n_rows: usize, n_cols: usize) -> bool {
n_rows > 1 && n_rows.saturating_mul(n_cols) >= 256_000
}
pub fn for_each_row<F>(output: &mut [f32], n_rows: usize, n_cols: usize, row_fn: F)
where
F: Fn(usize, &mut f32) + Send + Sync,
{
let n = n_rows.min(output.len());
if !should_parallelize(n, n_cols) {
for (row, out) in output.iter_mut().enumerate().take(n) {
row_fn(row, out);
}
return;
}
let rows = &mut output[..n];
rows.par_iter_mut()
.enumerate()
.for_each(|(row, out)| row_fn(row, out));
}
pub fn for_each_chunk_init<S, I, F>(
output: &mut [f32],
chunk_len: usize,
work_per_chunk: usize,
init: I,
f: F,
) where
I: Fn() -> S + Send + Sync,
S: Send,
F: Fn(&mut S, usize, &mut [f32]) + Send + Sync,
{
if chunk_len == 0 {
return;
}
let n_chunks = output.len() / chunk_len;
if !should_parallelize(n_chunks, work_per_chunk) {
let mut state = init();
for (i, chunk) in output[..n_chunks * chunk_len]
.chunks_mut(chunk_len)
.enumerate()
{
f(&mut state, i, chunk);
}
return;
}
let chunks = &mut output[..n_chunks * chunk_len];
let init = &init;
let f = &f;
chunks
.par_chunks_mut(chunk_len)
.enumerate()
.for_each_init(init, |state, (i, c)| f(state, i, c));
}
pub fn init_cpu_pool() -> Option<usize> {
let threads = resolve_cpu_threads();
let log = std::env::var_os("FERROX_QOS_LOG").is_some();
let built = rayon::ThreadPoolBuilder::new()
.num_threads(threads)
.start_handler(move |idx| {
#[cfg(target_os = "macos")]
{
let before = unsafe { qos::qos_class_self() };
let rc = unsafe {
qos::pthread_set_qos_class_self_np(qos::QOS_CLASS_USER_INTERACTIVE, 0)
};
if log {
eprintln!(
"ferrox: rayon worker {idx} qos {} -> {} (rc={rc})",
qos::name(before),
qos::name(unsafe { qos::qos_class_self() }),
);
}
}
#[cfg(not(target_os = "macos"))]
{
let _ = (idx, log);
}
})
.build_global()
.is_ok();
if built {
Some(threads)
} else {
None
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn perf_core_count_is_at_least_one_and_no_more_than_logical_cores() {
let logical = std::thread::available_parallelism()
.map(|n| n.get())
.unwrap_or(1);
let perf = perf_core_count();
assert!(perf >= 1, "perf core count must be positive, got {perf}");
assert!(
perf <= logical,
"perf cores ({perf}) cannot exceed logical cores ({logical})"
);
}
#[test]
fn resolved_thread_count_falls_back_to_perf_cores_without_env_overrides() {
if std::env::var_os("FERROX_CPU_THREADS").is_none()
&& std::env::var_os("RAYON_NUM_THREADS").is_none()
{
assert_eq!(resolve_cpu_threads(), perf_core_count());
}
}
#[test]
fn current_qos_name_is_reported_on_macos_and_absent_elsewhere() {
let qos = current_qos_name();
#[cfg(target_os = "macos")]
assert!(qos.is_some(), "macOS must report a QoS class");
#[cfg(not(target_os = "macos"))]
assert!(qos.is_none(), "QoS is a macOS-only concept");
}
fn smt_8t_4c() -> CpuTopology {
CpuTopology::from_sibling_lists([
(0, vec![0, 1]),
(1, vec![0, 1]),
(2, vec![2, 3]),
(3, vec![2, 3]),
(4, vec![4, 5]),
(5, vec![4, 5]),
(6, vec![6, 7]),
(7, vec![6, 7]),
])
}
#[test]
fn an_smt_host_is_sized_to_its_physical_cores_not_its_logical_cpus() {
let topology = smt_8t_4c();
assert_eq!(
topology.len(),
8,
"the fixture must have twice as many logical CPUs as cores"
);
assert_eq!(
physical_core_cpus_in(&topology),
vec![0, 2, 4, 6],
"one representative per physical core, lowest sibling first"
);
let (threads, core_ids) = resolve_threads_and_affinity_in(0, &topology);
assert_eq!(threads, 4, "auto sizing must not count SMT siblings");
assert_eq!(core_ids, vec![0, 2, 4, 6]);
}
#[test]
fn siblings_numbered_apart_are_deduplicated_the_same_as_adjacent_ones() {
let topology = CpuTopology::from_sibling_lists([
(0, vec![0, 64]),
(1, vec![1, 65]),
(64, vec![0, 64]),
(65, vec![1, 65]),
]);
assert_eq!(physical_core_cpus_in(&topology), vec![0, 1]);
}
#[test]
fn thread_siblings_lists_parse_as_ranges_comma_lists_and_mixtures() {
assert_eq!(parse_thread_siblings_list("0-1\n"), vec![0, 1]);
assert_eq!(parse_thread_siblings_list("0,64\n"), vec![0, 64]);
assert_eq!(parse_thread_siblings_list(" 3 "), vec![3]);
assert_eq!(parse_thread_siblings_list("0-1,64-65"), vec![0, 1, 64, 65]);
assert_eq!(parse_thread_siblings_list("2-4"), vec![2, 3, 4]);
assert_eq!(parse_thread_siblings_list(""), Vec::<usize>::new());
assert_eq!(parse_thread_siblings_list("x,-,7"), vec![7]);
assert_eq!(parse_thread_siblings_list("5-1"), Vec::<usize>::new());
}
#[test]
fn a_host_without_sysfs_topology_degrades_to_one_worker_per_allowed_cpu() {
let topology =
CpuTopology::from_sibling_lists((0..6).map(|cpu| (cpu, Vec::<usize>::new())));
assert_eq!(physical_core_cpus_in(&topology), vec![0, 1, 2, 3, 4, 5]);
assert_eq!(resolve_threads_and_affinity_in(0, &topology).0, 6);
}
#[test]
fn an_empty_topology_still_yields_one_usable_cpu() {
let topology = CpuTopology::default();
assert!(topology.is_empty());
assert_eq!(physical_core_cpus_in(&topology), vec![0]);
assert_eq!(resolve_threads_and_affinity_in(0, &topology), (1, vec![0]));
assert_eq!(
resolve_threads_and_affinity_in(2, &topology),
(2, vec![0, 0])
);
}
#[test]
fn cores_outside_the_affinity_mask_are_never_used_as_representatives() {
let full = smt_8t_4c();
let allowed = [1usize, 3, 4, 5];
let topology = CpuTopology::from_sibling_lists(
full.allowed_cpus()
.into_iter()
.filter(|cpu| allowed.contains(cpu))
.map(|cpu| {
(
cpu,
parse_thread_siblings_list(&format!("{}-{}", cpu & !1, cpu | 1)),
)
}),
);
assert_eq!(physical_core_cpus_in(&topology), vec![1, 3, 4]);
assert_eq!(resolve_threads_and_affinity_in(0, &topology).0, 3);
}
#[test]
fn an_explicit_count_fills_physical_cores_before_doubling_up_siblings() {
let topology = smt_8t_4c();
assert_eq!(
resolve_threads_and_affinity_in(4, &topology),
(4, vec![0, 2, 4, 6])
);
assert_eq!(
resolve_threads_and_affinity_in(6, &topology),
(6, vec![0, 2, 4, 6, 1, 3])
);
assert_eq!(
resolve_threads_and_affinity_in(8, &topology),
(8, vec![0, 2, 4, 6, 1, 3, 5, 7])
);
}
#[test]
fn an_explicit_count_larger_than_the_machine_wraps_instead_of_truncating() {
let topology = smt_8t_4c();
let (threads, core_ids) = resolve_threads_and_affinity_in(10, &topology);
assert_eq!(threads, 10, "an explicit width is honoured exactly");
assert_eq!(core_ids.len(), 10);
assert_eq!(&core_ids[8..], &[0, 2], "wraps back to the representatives");
}
#[test]
fn auto_sizing_donates_the_last_physical_core_to_the_coordinator() {
let plan = plan_cpu_pool_in(0, true, &smt_8t_4c());
assert_eq!(plan.num_threads, 3, "workers drop from N to N-1");
assert_eq!(plan.core_ids, vec![0, 2, 4]);
assert_eq!(plan.coordinator_cpu, Some(6));
assert_eq!(plan.num_threads, plan.core_ids.len());
}
#[test]
fn no_core_is_donated_without_a_coordinator_or_for_an_explicit_count() {
let topology = smt_8t_4c();
let no_coordinator = plan_cpu_pool_in(0, false, &topology);
assert_eq!(no_coordinator.num_threads, 4);
assert_eq!(no_coordinator.coordinator_cpu, None);
let explicit = plan_cpu_pool_in(4, true, &topology);
assert_eq!(explicit.num_threads, 4);
assert_eq!(explicit.coordinator_cpu, None);
}
#[test]
fn a_pool_of_two_or_fewer_keeps_its_workers_rather_than_donating() {
let dual = CpuTopology::from_sibling_lists([
(0, vec![0, 1]),
(1, vec![0, 1]),
(2, vec![2, 3]),
(3, vec![2, 3]),
]);
let plan = plan_cpu_pool_in(0, true, &dual);
assert_eq!(plan.num_threads, 2);
assert_eq!(plan.coordinator_cpu, None);
}
#[test]
fn the_intra_op_clamp_leaves_a_core_for_the_calling_thread() {
let plan = plan_cpu_pool_in(0, true, &smt_8t_4c());
assert_eq!(clamp_intra_op_threads(16, &plan, 16), 11);
assert_eq!(clamp_intra_op_threads(4, &plan, 16), 4);
assert_eq!(clamp_intra_op_threads(16, &plan, 4), 1);
assert_eq!(clamp_intra_op_threads(16, &plan, 0), 1);
}
#[test]
fn sibling_lists_are_read_from_a_sysfs_layout_on_disk() {
let root = std::env::temp_dir().join(format!(
"ferrox-threads-sysfs-{}-{:?}",
std::process::id(),
std::thread::current().id()
));
let _ = std::fs::remove_dir_all(&root);
for cpu in [0usize, 1] {
let dir = root.join(format!("cpu{cpu}")).join("topology");
std::fs::create_dir_all(&dir).expect("temp sysfs tree must be creatable");
std::fs::write(dir.join("thread_siblings_list"), "0-1\n")
.expect("temp sibling list must be writable");
}
let topology = CpuTopology::read_from(&root, &[0, 1, 2]);
assert_eq!(topology.allowed_cpus(), vec![0, 1, 2]);
assert_eq!(
physical_core_cpus_in(&topology),
vec![0, 2],
"cpu2 is unreadable, so it counts as a core of its own"
);
let _ = std::fs::remove_dir_all(&root);
}
#[test]
fn this_hosts_physical_core_count_is_positive_and_within_its_logical_cpus() {
let logical = std::thread::available_parallelism()
.map(|n| n.get())
.unwrap_or(1);
let physical = physical_core_count();
assert!(physical >= 1, "physical core count must be positive");
assert!(
physical <= logical,
"physical cores ({physical}) cannot exceed logical cores ({logical})"
);
assert_eq!(
physical_core_cpus().len().clamp(1, logical.max(1)),
physical
);
}
#[test]
fn this_hosts_affinity_mask_is_non_empty_and_ascending() {
let cpus = process_affinity_cpus();
assert!(!cpus.is_empty(), "a running process may run somewhere");
assert!(
cpus.windows(2).all(|w| w[0] < w[1]),
"affinity CPUs must be ascending and unique: {cpus:?}"
);
}
#[test]
fn for_each_row_parallel_matches_serial() {
let n = 4097usize;
let f = |row: usize| ((row % 97) as f32) * 0.25 - 3.0;
let mut par = vec![0.0f32; n];
for_each_row(&mut par, n, 4096, |row, slot| *slot = f(row));
let mut serial = vec![0.0f32; n];
for (row, slot) in serial.iter_mut().enumerate() {
*slot = f(row);
}
assert_eq!(par, serial);
}
}