use std::sync::OnceLock;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum CpuTier {
Scalar,
Avx2,
Avx512,
Neon,
NeonDotprod,
NeonI8mm,
}
impl CpuTier {
pub fn label(self) -> &'static str {
match self {
CpuTier::Scalar => "scalar",
CpuTier::Avx2 => "avx2",
CpuTier::Avx512 => "avx512",
CpuTier::Neon => "neon",
CpuTier::NeonDotprod => "neon+dotprod",
CpuTier::NeonI8mm => "neon+i8mm",
}
}
fn parse(s: &str) -> Option<CpuTier> {
match s.trim().to_ascii_lowercase().as_str() {
"scalar" | "none" | "off" => Some(CpuTier::Scalar),
#[cfg(target_arch = "x86_64")]
"avx2" => Some(CpuTier::Avx2),
#[cfg(target_arch = "x86_64")]
"avx512" => Some(CpuTier::Avx512),
#[cfg(target_arch = "aarch64")]
"neon" => Some(CpuTier::Neon),
#[cfg(target_arch = "aarch64")]
"dotprod" | "neon+dotprod" | "neon,dotprod" => Some(CpuTier::NeonDotprod),
#[cfg(target_arch = "aarch64")]
"i8mm" | "neon+i8mm" | "neon,i8mm" => Some(CpuTier::NeonI8mm),
_ => None,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct CpuFeatures {
pub tier: CpuTier,
pub avx2: bool,
pub fma: bool,
pub avx512f: bool,
pub avx512bw: bool,
pub avx512vnni: bool,
pub neon: bool,
pub dotprod: bool,
pub i8mm: bool,
}
impl CpuFeatures {
const NONE: CpuFeatures = CpuFeatures {
tier: CpuTier::Scalar,
avx2: false,
fma: false,
avx512f: false,
avx512bw: false,
avx512vnni: false,
neon: false,
dotprod: false,
i8mm: false,
};
fn active_flags(&self) -> Vec<&'static str> {
let mut flags: Vec<&str> = Vec::new();
for (on, name) in [
(self.avx2, "avx2"),
(self.fma, "fma"),
(self.avx512f, "avx512f"),
(self.avx512bw, "avx512bw"),
(self.avx512vnni, "avx512vnni"),
(self.neon, "neon"),
(self.dotprod, "dotprod"),
(self.i8mm, "i8mm"),
] {
if on {
flags.push(name);
}
}
flags
}
pub fn report(&self) -> String {
format!(
"cpu: tier={} [{}]",
self.tier.label(),
self.active_flags().join(" ")
)
}
pub fn descriptor(&self) -> String {
let flags = self.active_flags();
if flags.is_empty() {
self.tier.label().to_string()
} else {
flags.join(",")
}
}
pub fn ensure_supported(&self) -> Result<(), String> {
let _ = self;
Ok(())
}
}
pub fn detect() -> CpuFeatures {
#[cfg_attr(
not(any(target_arch = "x86_64", target_arch = "aarch64")),
allow(unused_mut)
)]
let mut f = CpuFeatures::NONE;
#[cfg(target_arch = "x86_64")]
{
f.avx2 = is_x86_feature_detected!("avx2");
f.fma = is_x86_feature_detected!("fma");
f.avx512f = is_x86_feature_detected!("avx512f");
f.avx512bw = is_x86_feature_detected!("avx512bw");
f.avx512vnni = is_x86_feature_detected!("avx512vnni");
f.tier = if f.avx512f && f.avx2 && f.fma && cfg!(feature = "avx512") {
CpuTier::Avx512
} else if f.avx2 && f.fma {
CpuTier::Avx2
} else {
CpuTier::Scalar
};
}
#[cfg(target_arch = "aarch64")]
{
f.neon = std::arch::is_aarch64_feature_detected!("neon");
f.dotprod = std::arch::is_aarch64_feature_detected!("dotprod");
f.i8mm = std::arch::is_aarch64_feature_detected!("i8mm");
f.tier = if f.neon && f.dotprod && f.i8mm {
CpuTier::NeonI8mm
} else if f.neon && f.dotprod {
CpuTier::NeonDotprod
} else if f.neon {
CpuTier::Neon
} else {
CpuTier::Scalar
};
}
apply_env_override(f)
}
fn apply_env_override(f: CpuFeatures) -> CpuFeatures {
match std::env::var("CERA_CPU_TIER") {
Ok(val) => with_tier_override(f, CpuTier::parse(&val)),
Err(_) => f,
}
}
fn with_tier_override(mut f: CpuFeatures, forced: Option<CpuTier>) -> CpuFeatures {
if let Some(t) = forced {
if t < f.tier {
f.tier = t;
}
}
f
}
pub fn cpu_features() -> &'static CpuFeatures {
static FEATURES: OnceLock<CpuFeatures> = OnceLock::new();
FEATURES.get_or_init(detect)
}
pub fn cpu_tier() -> CpuTier {
cpu_features().tier
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CoreTopology {
pub perf_core_count: usize,
pub pin_cores: Vec<usize>,
}
#[cfg(any(target_os = "linux", target_os = "android"))]
const MAX_AUTO_THREADS: usize = 6;
#[cfg(any(target_os = "linux", target_os = "android"))]
const MAX_CPUS: usize = 512;
#[cfg(any(target_os = "linux", target_os = "android"))]
const CAP_MID: u32 = 400;
pub fn core_topology() -> &'static CoreTopology {
static TOPOLOGY: OnceLock<CoreTopology> = OnceLock::new();
TOPOLOGY.get_or_init(detect_topology)
}
pub(crate) fn env_usize(name: &str) -> Option<usize> {
std::env::var(name)
.ok()
.and_then(|v| v.trim().parse::<usize>().ok())
.filter(|&n| n >= 1)
}
pub fn performance_core_count() -> usize {
core_topology().perf_core_count
}
pub fn detect_topology() -> CoreTopology {
let forced = env_usize("CERA_THREADS");
#[cfg(any(target_os = "linux", target_os = "android"))]
if let Some(topo) = detect_topology_sysfs() {
return apply_thread_override(topo, forced);
}
#[cfg(any(target_os = "macos", target_os = "ios"))]
if let Some(count) = macos_perf_core_count() {
return apply_thread_override(
CoreTopology {
perf_core_count: count,
pin_cores: Vec::new(),
},
forced,
);
}
let n = std::thread::available_parallelism()
.map(|p| p.get())
.unwrap_or(1);
apply_thread_override(
CoreTopology {
perf_core_count: n,
pin_cores: Vec::new(),
},
forced,
)
}
fn apply_thread_override(mut topo: CoreTopology, forced: Option<usize>) -> CoreTopology {
if let Some(n) = forced {
topo.perf_core_count = n;
topo.pin_cores.truncate(n);
}
topo
}
#[cfg(any(target_os = "linux", target_os = "android"))]
fn detect_topology_sysfs() -> Option<CoreTopology> {
let caps = read_per_cpu_u32("cpu_capacity");
let mut cores: Vec<(usize, u32)>;
if !caps.is_empty() {
if caps.iter().all(|&(_, c)| c == caps[0].1) {
return None;
}
cores = caps.into_iter().filter(|&(_, c)| c >= CAP_MID).collect();
} else {
let freqs = read_per_cpu_u32("cpufreq/cpuinfo_max_freq");
let max = freqs.iter().map(|&(_, f)| f).max()?;
let cutoff = (max / 100) * 85;
if freqs.iter().all(|&(_, f)| f >= cutoff) {
return None;
}
cores = freqs.into_iter().filter(|&(_, f)| f >= cutoff).collect();
}
if cores.is_empty() {
return None;
}
let sibling_sets: std::collections::HashMap<usize, String> =
read_per_cpu_trimmed("topology/thread_siblings_list")
.into_iter()
.collect();
let mut seen_sets = std::collections::HashSet::new();
cores.retain(|&(cpu, _)| match sibling_sets.get(&cpu) {
Some(set) => seen_sets.insert(set.clone()),
None => true,
});
cores.sort_by(|a, b| b.1.cmp(&a.1).then(a.0.cmp(&b.0)));
cores.truncate(MAX_AUTO_THREADS);
let pin_cores: Vec<usize> = cores.iter().map(|&(i, _)| i).collect();
Some(CoreTopology {
perf_core_count: pin_cores.len(),
pin_cores,
})
}
#[cfg(any(target_os = "linux", target_os = "android"))]
fn read_per_cpu_trimmed(file: &str) -> Vec<(usize, String)> {
let mut values = Vec::new();
for cpu in 0..MAX_CPUS {
let dir = format!("/sys/devices/system/cpu/cpu{cpu}");
if !std::path::Path::new(&dir).is_dir() {
break;
}
if let Ok(s) = std::fs::read_to_string(format!("{dir}/{file}")) {
values.push((cpu, s.trim().to_string()));
}
}
values
}
#[cfg(any(target_os = "linux", target_os = "android"))]
fn read_per_cpu_u32(file: &str) -> Vec<(usize, u32)> {
read_per_cpu_trimmed(file)
.into_iter()
.filter_map(|(cpu, s)| s.parse().ok().map(|v| (cpu, v)))
.collect()
}
#[cfg(any(target_os = "macos", target_os = "ios"))]
fn macos_perf_core_count() -> Option<usize> {
if cfg!(miri) {
return None;
}
unsafe extern "C" {
fn sysctlbyname(
name: *const std::ffi::c_char,
oldp: *mut std::ffi::c_void,
oldlenp: *mut usize,
newp: *const std::ffi::c_void,
newlen: usize,
) -> i32;
}
let name = c"hw.perflevel0.logicalcpu";
let mut value: i32 = 0;
let mut size = std::mem::size_of::<i32>();
let ret = unsafe {
sysctlbyname(
name.as_ptr(),
&mut value as *mut _ as *mut std::ffi::c_void,
&mut size,
std::ptr::null(),
0,
)
};
if ret == 0 && value > 0 {
Some(value as usize)
} else {
None
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn topology_has_at_least_one_thread() {
let topo = detect_topology();
assert!(topo.perf_core_count >= 1);
assert_eq!(core_topology().perf_core_count, performance_core_count());
}
#[test]
fn thread_override_sets_count_and_caps_pins() {
let base = CoreTopology {
perf_core_count: 3,
pin_cores: vec![7, 6, 5],
};
let two = apply_thread_override(base.clone(), Some(2));
assert_eq!(two.perf_core_count, 2);
assert_eq!(two.pin_cores, vec![7, 6]);
let five = apply_thread_override(base.clone(), Some(5));
assert_eq!(five.perf_core_count, 5);
assert_eq!(five.pin_cores, vec![7, 6, 5]);
assert_eq!(apply_thread_override(base.clone(), None), base);
}
#[test]
fn tier_ordering_is_monotonic_per_arch() {
assert!(CpuTier::Scalar < CpuTier::Avx2);
assert!(CpuTier::Avx2 < CpuTier::Avx512);
assert!(CpuTier::Scalar < CpuTier::Neon);
assert!(CpuTier::Neon < CpuTier::NeonDotprod);
assert!(CpuTier::NeonDotprod < CpuTier::NeonI8mm);
}
#[test]
fn descriptor_is_compact_sorted_and_never_empty() {
assert_eq!(CpuFeatures::NONE.descriptor(), "scalar");
let neon = CpuFeatures {
tier: CpuTier::NeonI8mm,
neon: true,
dotprod: true,
i8mm: true,
..CpuFeatures::NONE
};
assert_eq!(neon.descriptor(), "neon,dotprod,i8mm");
let x86 = CpuFeatures {
tier: CpuTier::Avx2,
avx2: true,
fma: true,
..CpuFeatures::NONE
};
assert_eq!(x86.descriptor(), "avx2,fma");
assert!(x86.report().contains("[avx2 fma]"));
}
#[test]
fn detect_is_stable_and_cached() {
assert_eq!(*cpu_features(), detect());
assert_eq!(cpu_features().tier, cpu_tier());
}
#[test]
fn detected_tier_matches_arch() {
let t = detect().tier;
#[cfg(target_arch = "x86_64")]
assert!(matches!(
t,
CpuTier::Scalar | CpuTier::Avx2 | CpuTier::Avx512
));
#[cfg(target_arch = "aarch64")]
assert!(matches!(
t,
CpuTier::Scalar | CpuTier::Neon | CpuTier::NeonDotprod | CpuTier::NeonI8mm
));
#[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
assert_eq!(t, CpuTier::Scalar);
}
#[test]
fn env_override_only_downgrades() {
let at = |t: CpuTier| CpuFeatures {
tier: t,
..CpuFeatures::NONE
};
assert_eq!(
with_tier_override(at(CpuTier::Avx2), Some(CpuTier::Scalar)).tier,
CpuTier::Scalar
);
assert_eq!(
with_tier_override(at(CpuTier::Avx2), Some(CpuTier::Avx512)).tier,
CpuTier::Avx2
);
assert_eq!(
with_tier_override(at(CpuTier::NeonDotprod), Some(CpuTier::NeonDotprod)).tier,
CpuTier::NeonDotprod
);
assert_eq!(
with_tier_override(at(CpuTier::Avx2), None).tier,
CpuTier::Avx2
);
}
#[test]
fn tier_label_roundtrips_through_parse() {
let mut tiers = vec![CpuTier::Scalar];
#[cfg(target_arch = "x86_64")]
tiers.extend([CpuTier::Avx2, CpuTier::Avx512]);
#[cfg(target_arch = "aarch64")]
tiers.extend([CpuTier::Neon, CpuTier::NeonDotprod, CpuTier::NeonI8mm]);
for t in tiers {
assert_eq!(CpuTier::parse(t.label()), Some(t), "label {:?}", t.label());
}
}
#[test]
fn cross_arch_override_label_is_rejected() {
#[cfg(target_arch = "aarch64")]
{
assert_eq!(CpuTier::parse("avx2"), None);
assert_eq!(CpuTier::parse("avx512"), None);
}
#[cfg(target_arch = "x86_64")]
{
assert_eq!(CpuTier::parse("neon"), None);
assert_eq!(CpuTier::parse("i8mm"), None);
}
}
#[test]
fn report_includes_tier_label() {
let r = cpu_features().report();
assert!(r.contains("tier="));
assert!(r.contains(cpu_tier().label()));
}
}