use std::ops::{BitOr, BitOrAssign};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
pub struct CpuFeatures(u64);
impl CpuFeatures {
pub const SSE: Self = Self(1 << 0);
pub const SSE2: Self = Self(1 << 1);
pub const SSE3: Self = Self(1 << 2);
pub const SSSE3: Self = Self(1 << 3);
pub const SSE4_1: Self = Self(1 << 4);
pub const SSE4_2: Self = Self(1 << 5);
pub const AVX: Self = Self(1 << 6);
pub const AVX2: Self = Self(1 << 7);
pub const AVX512F: Self = Self(1 << 8);
pub const AES: Self = Self(1 << 9);
pub const PCLMULQDQ: Self = Self(1 << 10);
pub const NEON: Self = Self(1 << 11);
#[must_use]
pub const fn empty() -> Self {
Self(0)
}
#[must_use]
pub const fn bits(self) -> u64 {
self.0
}
#[must_use]
pub const fn is_empty(self) -> bool {
self.0 == 0
}
#[must_use]
pub const fn contains(self, other: Self) -> bool {
(self.0 & other.0) == other.0
}
}
impl BitOr for CpuFeatures {
type Output = Self;
fn bitor(self, rhs: Self) -> Self {
Self(self.0 | rhs.0)
}
}
impl BitOrAssign for CpuFeatures {
fn bitor_assign(&mut self, rhs: Self) {
self.0 |= rhs.0;
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct CpuInfo {
pub cores_logical: u32,
pub cores_physical: u32,
pub features: CpuFeatures,
pub cache_l1: usize,
pub cache_l2: usize,
pub cache_l3: usize,
}
impl Default for CpuInfo {
fn default() -> Self {
Self {
cores_logical: 1,
cores_physical: 1,
features: CpuFeatures::empty(),
cache_l1: 0,
cache_l2: 0,
cache_l3: 0,
}
}
}
#[must_use]
pub(super) fn probe() -> CpuInfo {
super::probe::platform::probe_cpu()
}
#[must_use]
pub(crate) fn runtime_features() -> CpuFeatures {
let mut f = CpuFeatures::empty();
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
{
if std::arch::is_x86_feature_detected!("sse") {
f |= CpuFeatures::SSE;
}
if std::arch::is_x86_feature_detected!("sse2") {
f |= CpuFeatures::SSE2;
}
if std::arch::is_x86_feature_detected!("sse3") {
f |= CpuFeatures::SSE3;
}
if std::arch::is_x86_feature_detected!("ssse3") {
f |= CpuFeatures::SSSE3;
}
if std::arch::is_x86_feature_detected!("sse4.1") {
f |= CpuFeatures::SSE4_1;
}
if std::arch::is_x86_feature_detected!("sse4.2") {
f |= CpuFeatures::SSE4_2;
}
if std::arch::is_x86_feature_detected!("avx") {
f |= CpuFeatures::AVX;
}
if std::arch::is_x86_feature_detected!("avx2") {
f |= CpuFeatures::AVX2;
}
if std::arch::is_x86_feature_detected!("avx512f") {
f |= CpuFeatures::AVX512F;
}
if std::arch::is_x86_feature_detected!("aes") {
f |= CpuFeatures::AES;
}
if std::arch::is_x86_feature_detected!("pclmulqdq") {
f |= CpuFeatures::PCLMULQDQ;
}
}
#[cfg(target_arch = "aarch64")]
{
if std::arch::is_aarch64_feature_detected!("neon") {
f |= CpuFeatures::NEON;
}
if std::arch::is_aarch64_feature_detected!("aes") {
f |= CpuFeatures::AES;
}
if std::arch::is_aarch64_feature_detected!("pmull") {
f |= CpuFeatures::PCLMULQDQ;
}
}
f
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_features_empty_contains_only_empty() {
let empty = CpuFeatures::empty();
assert!(empty.is_empty());
assert!(empty.contains(CpuFeatures::empty()));
assert!(!empty.contains(CpuFeatures::SSE2));
}
#[test]
fn test_features_bitor_combines_flags() {
let f = CpuFeatures::SSE | CpuFeatures::SSE2;
assert!(f.contains(CpuFeatures::SSE));
assert!(f.contains(CpuFeatures::SSE2));
assert!(!f.contains(CpuFeatures::AVX));
}
#[test]
fn test_features_bitor_assign_inserts_flag() {
let mut f = CpuFeatures::empty();
f |= CpuFeatures::AVX;
assert!(f.contains(CpuFeatures::AVX));
}
#[test]
fn test_features_bits_round_trip_through_constants() {
let combined = CpuFeatures::SSE | CpuFeatures::AVX2;
assert_eq!(
combined.bits(),
CpuFeatures::SSE.bits() | CpuFeatures::AVX2.bits()
);
}
#[test]
fn test_features_contains_subset() {
let all = CpuFeatures::SSE | CpuFeatures::SSE2 | CpuFeatures::AVX;
assert!(all.contains(CpuFeatures::SSE | CpuFeatures::AVX));
}
#[test]
fn test_default_cpu_info_has_one_core() {
let i = CpuInfo::default();
assert_eq!(i.cores_logical, 1);
assert_eq!(i.cores_physical, 1);
}
#[test]
fn test_probe_reports_at_least_one_logical_core() {
let i = probe();
assert!(i.cores_logical >= 1);
assert!(i.cores_physical >= 1);
assert!(
i.cores_physical <= i.cores_logical,
"physical cores ({}) cannot exceed logical cores ({})",
i.cores_physical,
i.cores_logical,
);
}
#[test]
fn test_probe_features_include_sse2_on_x86_64() {
if cfg!(all(target_arch = "x86_64", target_feature = "sse2")) {
assert!(probe().features.contains(CpuFeatures::SSE2));
}
}
#[test]
fn test_probe_features_include_neon_on_aarch64() {
if cfg!(all(target_arch = "aarch64", target_feature = "neon")) {
assert!(probe().features.contains(CpuFeatures::NEON));
}
}
#[test]
fn test_runtime_features_includes_sse2_on_any_x86_64_host() {
if cfg!(target_arch = "x86_64") {
let f = runtime_features();
assert!(
f.contains(CpuFeatures::SSE2),
"every x86_64 host has SSE2 in its baseline ISA; \
runtime_features returned 0x{:x}",
f.bits()
);
}
}
#[test]
fn test_runtime_features_includes_compile_time_baseline() {
let runtime = runtime_features();
let mut compile_time = CpuFeatures::empty();
if cfg!(target_feature = "sse2") {
compile_time |= CpuFeatures::SSE2;
}
if cfg!(target_feature = "sse4.2") {
compile_time |= CpuFeatures::SSE4_2;
}
if cfg!(target_feature = "aes") {
compile_time |= CpuFeatures::AES;
}
if cfg!(target_feature = "neon") {
compile_time |= CpuFeatures::NEON;
}
assert!(
runtime.contains(compile_time),
"runtime feature set 0x{:x} missing compile-time baseline 0x{:x}",
runtime.bits(),
compile_time.bits(),
);
}
}