Skip to main content

ferrox_core/
threads.rs

1//! CPU worker-pool policy, shared by `ferrox` (CLI) and `ferrox-server`.
2//!
3//! Three things this module exists to control, all of which were measured
4//! to matter far more than any kernel change on Apple Silicon:
5//!
6//! 1. **Thread count.** `available_parallelism()` on an M2 Pro reports 10
7//!    (6 performance + 4 efficiency cores). Splitting a decode GEMV
8//!    across all 10 makes every fork-join wait on the slowest E-core
9//!    slice. llama.cpp defaults to `hw.perflevel0.physicalcpu` for
10//!    exactly this reason (`common_cpu_get_num_math`), and collapses when
11//!    forced above it -- 346 -> 176 tok/s on SmolLM2-135M Q8_0 going from
12//!    `-t 4` to `-t 10`. So default to the performance-core count, not
13//!    the logical-core count.
14//!
15//! 2. **Thread QoS.** macOS schedules threads onto E-cores based on their
16//!    Quality-of-Service class, and QoS is *inherited* from whichever
17//!    thread spawned them. Rayon builds its global pool lazily, on first
18//!    use -- which inside `ferrox-server` is a Tokio `spawn_blocking`
19//!    task, not the main thread. If that blocking thread carries a
20//!    demoted QoS, every rayon worker inherits it and the whole matvec
21//!    runs on efficiency cores. [`init_cpu_pool`] pins the workers to
22//!    `USER_INTERACTIVE` explicitly so the pool's placement does not
23//!    depend on who happened to touch rayon first.
24//!
25//! 3. **SMT siblings.** The same argument as (1), for the other kind of
26//!    fake core. On a 16C/32T host `available_parallelism()` reports 32,
27//!    so an auto-sized pool puts two workers on every physical core.
28//!    MoE decode is memory-bandwidth-bound: a sibling adds no bandwidth,
29//!    contends for the same core's load ports, and turns a spin barrier
30//!    into a livelock-grade tax once the pool is oversubscribed. So the
31//!    non-macOS width comes from [`physical_core_count`], which
32//!    deduplicates `thread_siblings_list` across this process's affinity
33//!    mask; sysfs being unreadable degrades to the
34//!    `available_parallelism` answer rather than failing.
35
36use rayon::prelude::*;
37use std::collections::HashSet;
38use std::path::Path;
39
40/// macOS QoS classes (`sys/qos.h`). Only the ones we name are listed.
41#[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
66/// The calling thread's macOS QoS class, as a human-readable name.
67/// `None` off macOS, where the concept does not exist.
68pub 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
79/// Number of *performance* cores, which is the useful width for a decode
80/// GEMV. On macOS this is `hw.perflevel0.physicalcpu` (llama.cpp reads
81/// the same key). Elsewhere it is [`physical_core_count`] — one logical
82/// CPU per physical core inside this process's affinity mask — and only
83/// a host whose sysfs topology is unreadable falls all the way back to
84/// `available_parallelism`.
85pub 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    // SAFETY: `key` is NUL-terminated, and `out`/`len` describe a live
113    // i32 of exactly `len` bytes, which is what these keys return.
114    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
130// ---------------------------------------------------------------------
131// SMT topology: one worker per physical core
132//
133// Ported from FreeToken's `freetoken/moe/cpu_executor.py`
134// (`physical_core_cpus`, `resolve_threads_and_affinity`, and the pool
135// sizing in `CpuMoeExecutor.__init__`). Apache-2.0; see
136// `docs/THIRD_PARTY_NOTICES.md`.
137//
138// Every rule below is a pure function over an injected [`CpuTopology`],
139// with a thin wrapper that reads the real `/sys`. That split is not
140// cosmetic: the rules only *matter* on an SMT host with a restricted
141// affinity mask, and no CI machine can be assumed to be one, so a
142// sysfs-reading implementation would be a set of policies that are never
143// actually exercised until a production box gets them wrong.
144// ---------------------------------------------------------------------
145
146/// Where Linux publishes per-CPU topology. [`CpuTopology::detect`] reads
147/// `{root}/cpu{n}/topology/thread_siblings_list` under this directory.
148pub const SYSFS_CPU_ROOT: &str = "/sys/devices/system/cpu";
149
150/// Parse one `thread_siblings_list` line into the logical CPU ids it
151/// names.
152///
153/// The kernel prints a cpulist, and both spellings occur on real
154/// hardware: `0-1` on a box that numbers a core's siblings adjacently,
155/// `0,64` on one that numbers every first-sibling before any second.
156/// Handling only one of the two silently collapses or explodes the
157/// deduplicated core count on the other, which is exactly the sizing
158/// mistake this module exists to prevent — so both are parsed, and mixed
159/// `0-1,64-65` forms with them. Tokens that are neither are dropped
160/// rather than poisoning the list with a bogus CPU id.
161pub 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                // A cpulist range may carry a `:used/group` stride
171                // suffix; sibling lists never do, so take the plain
172                // prefix and ignore anything past it.
173                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
190/// The logical CPUs this process may actually run on, ascending.
191///
192/// This is the affinity mask, not the machine: a pool sized to the
193/// machine inside a `taskset`- or cpuset-confined container
194/// oversubscribes the slice it was given by exactly the factor it was
195/// confined by, and every worker then time-slices against every other.
196/// Linux reads `sched_getaffinity`; everywhere else (macOS included,
197/// which has no equivalent) this is `0..available_parallelism()`, i.e.
198/// today's answer.
199pub 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    // SAFETY: a zeroed `cpu_set_t` is a valid (empty) mask; the kernel
215    // writes at most `cpusetsize` bytes into it, and we pass exactly the
216    // size of the live local. `CPU_ISSET` only reads that same local.
217    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/// The SMT layout of the logical CPUs this process may run on: for each
235/// allowed CPU, in ascending id order, the set of CPUs sharing its
236/// physical core.
237///
238/// An **empty** sibling list means "sysfs did not answer for this CPU".
239/// Such a CPU is treated as a physical core of its own and is never
240/// merged with another unknown one — guessing that two unreadable CPUs
241/// are siblings would silently halve the pool on any host that simply
242/// has no `topology/` directory (a container with a masked `/sys`, some
243/// hypervisors, anything non-Linux), and that regression arrives with no
244/// error attached to it.
245#[derive(Debug, Clone, PartialEq, Eq, Default)]
246pub struct CpuTopology {
247    /// `(cpu id, its thread siblings)`, ascending by cpu id, unique.
248    entries: Vec<(usize, Vec<usize>)>,
249}
250
251impl CpuTopology {
252    /// Build a topology from explicit sibling lists — the injection
253    /// point that makes every rule in this module testable on a host
254    /// with no SMT, no root, and no particular CPU count.
255    ///
256    /// Entries are sorted by CPU id and deduplicated, because the rules
257    /// pick the *first* allowed CPU of each physical core as that core's
258    /// representative and that choice must not depend on the order the
259    /// caller happened to enumerate in.
260    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    /// Read `{root}/cpu{n}/topology/thread_siblings_list` for each CPU in
271    /// `allowed`. A CPU whose file is missing or unreadable gets an empty
272    /// sibling list (see the type docs: it becomes its own core), so a
273    /// host without sysfs topology degrades to one worker per allowed CPU
274    /// instead of erroring out mid-boot.
275    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    /// This host's topology: [`SYSFS_CPU_ROOT`] restricted to
289    /// [`process_affinity_cpus`].
290    pub fn detect() -> Self {
291        Self::read_from(Path::new(SYSFS_CPU_ROOT), &process_affinity_cpus())
292    }
293
294    /// The allowed logical CPUs, ascending.
295    pub fn allowed_cpus(&self) -> Vec<usize> {
296        self.entries.iter().map(|(cpu, _)| *cpu).collect()
297    }
298
299    /// Number of allowed logical CPUs, SMT siblings included.
300    pub fn len(&self) -> usize {
301        self.entries.len()
302    }
303
304    /// Whether no CPU at all is described.
305    pub fn is_empty(&self) -> bool {
306        self.entries.is_empty()
307    }
308}
309
310/// One logical CPU per physical core, restricted to `topology`.
311///
312/// MoE decode is memory-bandwidth-bound, so SMT siblings only contend
313/// for the same core's load ports without adding bandwidth; one logical
314/// CPU per physical core is the fastest and, more importantly, the most
315/// *stable* width. Siblings are deduplicated by the set their
316/// `thread_siblings_list` names, and the lowest-numbered allowed CPU of
317/// each core wins.
318///
319/// Never returns an empty list: with no topology it degrades to the
320/// allowed CPUs, and with nothing allowed at all to `[0]`, because every
321/// caller downstream divides work by this length.
322pub 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            // Unknown topology for this CPU: a core of its own, never
328            // merged with another unknown CPU.
329            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
350/// [`physical_core_cpus_in`] against this host's real topology.
351pub fn physical_core_cpus() -> Vec<usize> {
352    physical_core_cpus_in(&CpuTopology::detect())
353}
354
355/// How many physical cores this process may run on.
356///
357/// Clamped by `available_parallelism`, which is the only figure that
358/// accounts for a cgroup CPU *quota* — a container pinned to 8 CPUs but
359/// throttled to two cores' worth of runtime reports 8 in its affinity
360/// mask and 2 here, and a pool built for 8 spends the difference being
361/// throttled mid-GEMV. A host with no readable topology lands on
362/// `available_parallelism` outright, which is the pre-topology answer.
363pub 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
370/// `(num_threads, core_ids)` for a pinned worker pool.
371///
372/// `requested == 0` means one worker per physical core, pinned to it.
373/// An explicit count is honoured exactly, spreading across physical-core
374/// **representatives first** and only then across the remaining logical
375/// CPUs, so distinct hardware threads are used before any core is
376/// doubled up — filling CPUs in numeric order instead would put workers
377/// 0 and 1 on one core's two siblings on every host that numbers
378/// siblings adjacently, i.e. half the pool contending before a second
379/// core has been touched at all. A count larger than the allowed CPU set
380/// wraps, deliberately: the caller asked for that width.
381pub 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
405/// [`resolve_threads_and_affinity_in`] against this host's real topology.
406pub fn resolve_threads_and_affinity(requested: usize) -> (usize, Vec<usize>) {
407    resolve_threads_and_affinity_in(requested, &CpuTopology::detect())
408}
409
410/// A sized worker pool: how many workers, which logical CPU each pins
411/// to, and the CPU set aside for a coordinator, if any.
412#[derive(Debug, Clone, PartialEq, Eq, Default)]
413pub struct CpuPoolPlan {
414    /// Worker count. Always equal to `core_ids.len()`.
415    pub num_threads: usize,
416    /// The logical CPU each worker pins to, in worker order.
417    pub core_ids: Vec<usize>,
418    /// The logical CPU donated to the coordinator, or `None` when no
419    /// coordinator was asked for or the pool was too small to donate.
420    pub coordinator_cpu: Option<usize>,
421}
422
423/// Size a pinned pool, optionally reserving a core for a coordinator
424/// thread (the one that polls the device doorbell and drives the pool).
425///
426/// The donation rule is the point: under *auto* sizing only, and only
427/// while more than two workers survive it, the coordinator takes the
428/// last physical core and the pool drops from N to N-1 workers. A
429/// coordinator that instead time-slices against a full-width pool
430/// measurably destabilizes throughput on a fully-subscribed box — it is
431/// a spinner, so the worker sharing its core is the one the fork-join
432/// barrier waits for, every single step. An explicit thread count is
433/// never silently reduced: the operator asked for that width.
434pub 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
452/// [`plan_cpu_pool_in`] against this host's real topology.
453pub fn plan_cpu_pool(requested: usize, reserve_coordinator: bool) -> CpuPoolPlan {
454    plan_cpu_pool_in(requested, reserve_coordinator, &CpuTopology::detect())
455}
456
457/// The intra-op width a *second* thread pool may still use once `plan`
458/// has claimed its cores: `physical_cores - workers - coordinator - 1`,
459/// clamped into `1..=configured`. The trailing `-1` is the calling
460/// thread itself, which is running the surrounding forward.
461///
462/// Without this clamp the framework's own pool defaults to the full core
463/// count and each of its threads lands on a core a pinned, spinning
464/// worker already owns — the pinned pool cannot yield, so the
465/// oversubscription is paid as scheduler latency on the decode critical
466/// path instead of buying parallelism anywhere. Never returns 0: a
467/// zero-width intra-op pool is not a degraded configuration, it is a
468/// broken one.
469pub 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
482/// How many rayon workers to run: `FERROX_CPU_THREADS`, else
483/// `RAYON_NUM_THREADS`, else [`perf_core_count`].
484pub 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
497/// Prefer serial when fork-join overhead exceeds the matvec work.
498/// ~256k element-ops matches the previous `prefer_serial_matvec` gate.
499pub fn should_parallelize(n_rows: usize, n_cols: usize) -> bool {
500    n_rows > 1 && n_rows.saturating_mul(n_cols) >= 256_000
501}
502
503/// One output row per slot; parallel when [`should_parallelize`] says the
504/// matvec is big enough to pay for the fork-join. Rows are never split.
505pub 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    // Global rayon, the same pool act-quant uses. A second dedicated
517    // matvec pool of P-core width was measured to regress CPU pp512 on
518    // Host B (~40 -> ~14 tok/s), so there is only one pool.
519    let rows = &mut output[..n];
520    rows.par_iter_mut()
521        .enumerate()
522        .for_each(|(row, out)| row_fn(row, out));
523}
524
525/// Chunk-parallel sibling of [`for_each_row`]; chunks are never split.
526pub 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
560/// Builds the global rayon pool with an explicit width and an explicit
561/// QoS, so neither depends on which thread first touched rayon. Safe to
562/// call more than once and from either binary; a pool that already
563/// exists is left alone.
564///
565/// Returns the thread count the pool was built with, or `None` if the
566/// global pool already existed.
567pub 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                // SAFETY: sets only the calling thread's QoS class.
577                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        // `resolve_cpu_threads` reads process-global env, and tests share
622        // a process, so assert the fallback only when nothing is set.
623        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    /// An 8-logical/4-physical SMT host, siblings numbered adjacently.
640    /// Every CPU is in the affinity mask.
641    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    /// THE central test, and it FAILS against the pre-topology
655    /// implementation: on this host `available_parallelism()` reports 8
656    /// while there are only 4 physical cores, so the old
657    /// `perf_core_count` fallback would have sized the pool to 8 and put
658    /// two bandwidth-bound workers on every core. The auto width must be
659    /// the physical-core count, and the chosen CPUs must be one per core.
660    #[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        // AMD/POWER style: all first siblings, then all second siblings.
681        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        // Garbage tokens are dropped, not turned into CPU 0.
698        assert_eq!(parse_thread_siblings_list(""), Vec::<usize>::new());
699        assert_eq!(parse_thread_siblings_list("x,-,7"), vec![7]);
700        // A descending range names nothing rather than panicking.
701        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        // Empty sibling lists model an unreadable `topology/` directory:
707        // the answer must be today's `available_parallelism`-shaped one.
708        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        // Same 8T/4C host, but only CPUs 1, 3, 4, 5 are allowed. Cores
729        // {0,1} and {2,3} contribute their high sibling (the only allowed
730        // one); core {4,5} contributes 4 once, not twice.
731        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        // Four workers land on four distinct cores...
752        assert_eq!(
753            resolve_threads_and_affinity_in(4, &topology),
754            (4, vec![0, 2, 4, 6])
755        );
756        // ...and only the fifth onward touches a sibling.
757        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        // An operator-supplied width is never silently reduced.
793        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        // 2 physical cores: donating would leave a single worker, so the
801        // coordinator shares instead.
802        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        // 16 physical cores, 3 workers + coordinator -> 16-3-1-1 = 11.
816        let plan = plan_cpu_pool_in(0, true, &smt_8t_4c());
817        assert_eq!(clamp_intra_op_threads(16, &plan, 16), 11);
818        // A configured width below the spare is left alone.
819        assert_eq!(clamp_intra_op_threads(4, &plan, 16), 4);
820        // A fully claimed machine still gets a usable intra-op width.
821        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        // cpu0/cpu1 are siblings; cpu2 has no topology directory at all.
834        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}