#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SimdLevel {
Scalar,
Sse42,
Avx2,
Avx512,
Neon,
}
#[derive(Debug, Clone, Copy)]
pub struct CpuFeatures {
pub level: SimdLevel,
pub has_avx2: bool,
pub has_avx512f: bool,
pub has_avx512bw: bool,
pub has_sse42: bool,
pub has_neon: bool,
}
impl CpuFeatures {
#[inline]
pub fn detect() -> Self {
#[cfg(target_arch = "x86_64")]
{
let has_avx512f = is_x86_feature_detected!("avx512f");
let has_avx512bw = is_x86_feature_detected!("avx512bw");
let has_avx2 = is_x86_feature_detected!("avx2");
let has_sse42 = is_x86_feature_detected!("sse4.2");
let level = if has_avx512bw && has_avx512f {
SimdLevel::Avx512
} else if has_avx2 {
SimdLevel::Avx2
} else if has_sse42 {
SimdLevel::Sse42
} else {
SimdLevel::Scalar
};
Self {
level,
has_avx2,
has_avx512f,
has_avx512bw,
has_sse42,
has_neon: false,
}
}
#[cfg(target_arch = "aarch64")]
{
let has_neon = std::arch::is_aarch64_feature_detected!("neon");
Self {
level: if has_neon {
SimdLevel::Neon
} else {
SimdLevel::Scalar
},
has_avx2: false,
has_avx512f: false,
has_avx512bw: false,
has_sse42: false,
has_neon,
}
}
#[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
{
Self {
level: SimdLevel::Scalar,
has_avx2: false,
has_avx512f: false,
has_avx512bw: false,
has_sse42: false,
has_neon: false,
}
}
}
}
thread_local! {
static CPU_FEATURES: CpuFeatures = CpuFeatures::detect();
}
#[inline]
pub fn cpu_features() -> CpuFeatures {
CPU_FEATURES.with(|f| *f)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_detect_features() {
let features = CpuFeatures::detect();
assert!(matches!(
features.level,
SimdLevel::Scalar
| SimdLevel::Sse42
| SimdLevel::Avx2
| SimdLevel::Avx512
| SimdLevel::Neon
));
}
#[test]
fn test_cpu_features_cached() {
let f1 = cpu_features();
let f2 = cpu_features();
assert_eq!(f1.level, f2.level);
}
#[test]
#[cfg(target_arch = "x86_64")]
fn test_x86_features_consistent() {
let features = cpu_features();
if features.has_avx512bw {
assert!(features.has_avx2);
}
if features.has_avx2 {
assert!(features.has_sse42);
}
}
#[test]
#[cfg(target_arch = "aarch64")]
fn test_aarch64_has_neon() {
let features = cpu_features();
assert!(features.has_neon);
assert_eq!(features.level, SimdLevel::Neon);
}
}