rt/
cpu.rs

1// cpu -- CPU feature detection for SIMD dispatch                             [src/core/cpu.s]
2//
3//   features() -> CPU_* bits     cached; SSE2 is the baseline and has no bit
4//   force_baseline()             report 0 from now on (plat calls it at process start for
5//                                `--simd sse2`: baseline kernels only, to compare results)
6//   cpuid(leaf) -> CpuidResult
7//
8// CPU_AVX2 and CPU_FMA are only reported when the OS saves the YMM state (CPUID.1:ECX.OSXSAVE
9// and XGETBV(0) bits 1 and 2).  Thread safe: every caller computes the same cached value.
10// Lives in rt so that simd (which may not depend on plat) dispatches on the same bits.
11
12use core::arch::x86_64::{CpuidResult, __cpuid_count};
13use core::sync::atomic::{AtomicU32, Ordering};
14
15pub const CPU_SSE41: u32 = 1;
16pub const CPU_AVX2: u32 = 2; // AVX2 + OS-enabled YMM state
17pub const CPU_FMA: u32 = 4; // FMA3 + OS-enabled YMM state
18pub const CPU_POPCNT: u32 = 8;
19pub const CPU_BMI2: u32 = 16;
20
21const VALID: u32 = 0x8000_0000; // cache marker
22static FEATURES: AtomicU32 = AtomicU32::new(0);
23
24#[allow(unused_unsafe)]
25pub fn cpuid(leaf: u32) -> CpuidResult {
26    // SAFETY: CPUID exists on every x86-64 processor.
27    unsafe { __cpuid_count(leaf, 0) }
28}
29
30fn xgetbv0() -> u32 {
31    let lo: u32;
32    // SAFETY: only called when CPUID reports OSXSAVE (XGETBV is enabled).
33    unsafe { core::arch::asm!("xgetbv", in("ecx") 0, out("eax") lo, out("edx") _, options(nomem, nostack)) };
34    lo
35}
36
37fn detect() -> u32 {
38    let mut f = 0;
39    let max = cpuid(0).eax;
40    if max != 0 {
41        let c1 = cpuid(1).ecx;
42        if c1 & 1 << 19 != 0 {
43            f |= CPU_SSE41;
44        }
45        if c1 & 1 << 23 != 0 {
46            f |= CPU_POPCNT;
47        }
48        let mut ymm = false; // the OS saves XMM + YMM state
49        if c1 & 1 << 27 != 0 && c1 & 1 << 28 != 0 && xgetbv0() & 6 == 6 {
50            ymm = true;
51            if c1 & 1 << 12 != 0 {
52                f |= CPU_FMA;
53            }
54        }
55        if max >= 7 {
56            let b7 = cpuid(7).ebx;
57            if b7 & 1 << 8 != 0 {
58                f |= CPU_BMI2;
59            }
60            if ymm && b7 & 1 << 5 != 0 {
61                f |= CPU_AVX2;
62            }
63        }
64    }
65    f
66}
67
68/// CPU_* bits (cached)  [cpu_features]
69#[inline]
70pub fn features() -> u32 {
71    let v = FEATURES.load(Ordering::Relaxed);
72    if v != 0 {
73        return v & !VALID;
74    }
75    let f = detect();
76    FEATURES.store(f | VALID, Ordering::Relaxed);
77    f
78}
79
80/// report no optional features from now on (`--simd sse2`)
81pub fn force_baseline() {
82    FEATURES.store(VALID, Ordering::Relaxed);
83}
rt/cpu.rs

rt/
cpu.rs
