Skip to main content

aocl_utils/
cpuid.rs

1//! Safe wrappers around the AMD AOCL-Utils CPU identification API.
2//!
3//! All AOCL `au_cpuid_*` routines that take a `cpu_num` migrate the calling
4//! thread to that core if `cpu_num != AU_CURRENT_CPU_NUM`. The [`Cpu`] enum
5//! makes this distinction explicit so callers cannot accidentally pin
6//! themselves to core 0 by passing a literal `0u32`.
7
8use aocl_utils_sys as sys;
9
10/// Sentinel passed to AOCL to mean "use whatever core the calling thread is
11/// already running on, do not migrate". Equal to `UINT32_MAX`.
12const AU_CURRENT_CPU_NUM: u32 = u32::MAX;
13
14/// Which CPU to query.
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
16pub enum Cpu {
17    /// The current thread's CPU. Does not cause thread migration.
18    Current,
19    /// A specific core by 0-based index. **Calling AOCL with a specific
20    /// core will migrate the current thread to that core** for the
21    /// duration of the call.
22    Specific(u32),
23}
24
25impl Cpu {
26    fn as_raw(self) -> u32 {
27        match self {
28            Cpu::Current => AU_CURRENT_CPU_NUM,
29            Cpu::Specific(n) => n,
30        }
31    }
32}
33
34/// AMD Zen sub-architectures recognized by AOCL.
35#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
36#[non_exhaustive]
37pub enum ZenArch {
38    Zen,
39    ZenPlus,
40    Zen2,
41    Zen3,
42    Zen4,
43    Zen5,
44}
45
46/// x86-64 microarchitecture levels (the System V psABI levels).
47#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
48#[non_exhaustive]
49pub enum X86_64Level {
50    V2,
51    V3,
52    V4,
53}
54
55/// CPU vendor + family/model details, as reported by `au_cpuid_get_vendor`.
56#[derive(Debug, Clone, PartialEq, Eq)]
57pub struct VendorInfo {
58    pub vendor_id: String,
59    pub family_id: String,
60    pub model_id: String,
61    pub stepping_id: String,
62    pub uarch_id: String,
63}
64
65/// Returns `true` if the queried CPU is manufactured by AMD.
66pub fn is_amd(cpu: Cpu) -> bool {
67    // SAFETY: `au_cpuid_is_amd` is a pure CPUID query with no preconditions.
68    unsafe { sys::au_cpuid_is_amd(cpu.as_raw()) }
69}
70
71/// Returns `true` if the queried CPU is in the Zen family (Zen, Zen+,
72/// Zen2, Zen3, Zen4, Zen5; not Zen6+).
73pub fn is_zen_family(cpu: Cpu) -> bool {
74    unsafe { sys::au_cpuid_arch_is_zen_family(cpu.as_raw()) }
75}
76
77/// Returns the Zen sub-architecture of the queried CPU, if any.
78pub fn zen_arch(cpu: Cpu) -> Option<ZenArch> {
79    let raw = cpu.as_raw();
80    // Order matters: AOCL's `arch_is_zenN` returns true for ZenN *and later*,
81    // so we test from the newest backwards to get the most specific answer.
82    unsafe {
83        if sys::au_cpuid_arch_is_zen5(raw) {
84            Some(ZenArch::Zen5)
85        } else if sys::au_cpuid_arch_is_zen4(raw) {
86            Some(ZenArch::Zen4)
87        } else if sys::au_cpuid_arch_is_zen3(raw) {
88            Some(ZenArch::Zen3)
89        } else if sys::au_cpuid_arch_is_zen2(raw) {
90            Some(ZenArch::Zen2)
91        } else if sys::au_cpuid_arch_is_zenplus(raw) {
92            Some(ZenArch::ZenPlus)
93        } else if sys::au_cpuid_arch_is_zen(raw) {
94            Some(ZenArch::Zen)
95        } else {
96            None
97        }
98    }
99}
100
101/// Returns the highest x86-64 microarchitecture level supported by the
102/// queried CPU, or `None` if it does not even meet x86-64-v2.
103pub fn x86_64_level(cpu: Cpu) -> Option<X86_64Level> {
104    let raw = cpu.as_raw();
105    unsafe {
106        if sys::au_cpuid_arch_is_x86_64v4(raw) {
107            Some(X86_64Level::V4)
108        } else if sys::au_cpuid_arch_is_x86_64v3(raw) {
109            Some(X86_64Level::V3)
110        } else if sys::au_cpuid_arch_is_x86_64v2(raw) {
111            Some(X86_64Level::V2)
112        } else {
113            None
114        }
115    }
116}
117
118/// Returns `true` if the queried CPU supports the named feature flag
119/// (e.g. `"avx2"`, `"avx512f"`, `"sha"`, `"vaes"`).
120///
121/// **Caveat:** AOCL throws a C++ exception when given an unknown flag
122/// name, which Rust cannot catch — pass only flag names AOCL recognises.
123/// See `<AOCL_ROOT>/amd-utils/include/Au/Cpuid/Enum.hh` for the full
124/// list (representative subset: `sse`, `sse2`, `sse3`, `sse4_1`,
125/// `sse4_2`, `avx`, `avx2`, `avx512f`, `avx512vl`, `fma`, `aes`, `vaes`,
126/// `sha`, `bmi1`, `bmi2`, `popcnt`).
127pub fn has_flag(cpu: Cpu, flag: &str) -> bool {
128    has_flags_all(cpu, &[flag])
129}
130
131/// Returns `true` if the queried CPU supports **all** of the named flags.
132pub fn has_flags_all(cpu: Cpu, flags: &[&str]) -> bool {
133    if flags.is_empty() {
134        return true;
135    }
136    let cstrings: Vec<std::ffi::CString> = flags
137        .iter()
138        .filter_map(|s| std::ffi::CString::new(*s).ok())
139        .collect();
140    if cstrings.len() != flags.len() {
141        // A flag with an interior NUL is malformed; treat as "not supported".
142        return false;
143    }
144    let pointers: Vec<*const std::os::raw::c_char> = cstrings.iter().map(|c| c.as_ptr()).collect();
145    // SAFETY: AOCL reads `count` C-string pointers; the cstrings stay
146    // alive for the duration of this call.
147    unsafe {
148        sys::au_cpuid_has_flags_all(
149            cpu.as_raw(),
150            pointers.as_ptr(),
151            pointers.len() as std::os::raw::c_int,
152        )
153    }
154}
155
156/// Returns `true` if the queried CPU supports **any** of the named flags.
157pub fn has_flags_any(cpu: Cpu, flags: &[&str]) -> bool {
158    if flags.is_empty() {
159        return false;
160    }
161    let cstrings: Vec<std::ffi::CString> = flags
162        .iter()
163        .filter_map(|s| std::ffi::CString::new(*s).ok())
164        .collect();
165    if cstrings.len() != flags.len() {
166        return false;
167    }
168    let pointers: Vec<*const std::os::raw::c_char> = cstrings.iter().map(|c| c.as_ptr()).collect();
169    unsafe {
170        sys::au_cpuid_has_flags_any(
171            cpu.as_raw(),
172            pointers.as_ptr(),
173            pointers.len() as std::os::raw::c_int,
174        )
175    }
176}
177
178/// Fetch vendor / family / model / stepping / uarch identifiers for the
179/// queried CPU.
180pub fn vendor_info(cpu: Cpu) -> VendorInfo {
181    let mut buf = [0u8; 256];
182    // SAFETY: We pass a valid pointer + length pair; AOCL writes a
183    // NUL-terminated string of at most `len` bytes.
184    unsafe {
185        sys::au_cpuid_get_vendor(
186            cpu.as_raw(),
187            buf.as_mut_ptr() as *mut std::os::raw::c_char,
188            buf.len(),
189        );
190    }
191    let nul = buf.iter().position(|&b| b == 0).unwrap_or(buf.len());
192    let raw = std::str::from_utf8(&buf[..nul]).unwrap_or("");
193    let mut it = raw.split('\n');
194    VendorInfo {
195        vendor_id: it.next().unwrap_or("").to_string(),
196        family_id: it.next().unwrap_or("").to_string(),
197        model_id: it.next().unwrap_or("").to_string(),
198        stepping_id: it.next().unwrap_or("").to_string(),
199        uarch_id: it.next().unwrap_or("").to_string(),
200    }
201}
202
203#[cfg(test)]
204mod tests {
205    use super::*;
206
207    #[test]
208    fn current_cpu_does_not_panic() {
209        let info = vendor_info(Cpu::Current);
210        assert!(!info.vendor_id.is_empty(), "vendor_id should be populated");
211        let _ = is_amd(Cpu::Current);
212        let _ = is_zen_family(Cpu::Current);
213        let _ = zen_arch(Cpu::Current);
214        let _ = x86_64_level(Cpu::Current);
215    }
216
217    #[test]
218    fn has_flag_queries_run() {
219        // AOCL throws a C++ exception on unknown flag names, which would
220        // abort the process under Rust's panic-from-foreign-exception
221        // handling — so use only flag names AOCL recognises.
222        // SSE2 is universal on x86-64 since 2003.
223        let _ = has_flag(Cpu::Current, "sse2");
224        let _ = has_flags_all(Cpu::Current, &["sse2", "sse4_2"]);
225        let _ = has_flags_any(Cpu::Current, &["avx2", "avx512f"]);
226        // Empty list of flags
227        assert!(has_flags_all(Cpu::Current, &[]));
228        assert!(!has_flags_any(Cpu::Current, &[]));
229    }
230}