aocl-utils 0.1.0

//! Safe wrappers around the AMD AOCL-Utils CPU identification API.
//!
//! All AOCL `au_cpuid_*` routines that take a `cpu_num` migrate the calling
//! thread to that core if `cpu_num != AU_CURRENT_CPU_NUM`. The [`Cpu`] enum
//! makes this distinction explicit so callers cannot accidentally pin
//! themselves to core 0 by passing a literal `0u32`.

use aocl_utils_sys as sys;

/// Sentinel passed to AOCL to mean "use whatever core the calling thread is
/// already running on, do not migrate". Equal to `UINT32_MAX`.
const AU_CURRENT_CPU_NUM: u32 = u32::MAX;

/// Which CPU to query.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Cpu {
    /// The current thread's CPU. Does not cause thread migration.
    Current,
    /// A specific core by 0-based index. **Calling AOCL with a specific
    /// core will migrate the current thread to that core** for the
    /// duration of the call.
    Specific(u32),
}

impl Cpu {
    fn as_raw(self) -> u32 {
        match self {
            Cpu::Current => AU_CURRENT_CPU_NUM,
            Cpu::Specific(n) => n,
        }
    }
}

/// AMD Zen sub-architectures recognized by AOCL.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum ZenArch {
    Zen,
    ZenPlus,
    Zen2,
    Zen3,
    Zen4,
    Zen5,
}

/// x86-64 microarchitecture levels (the System V psABI levels).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[non_exhaustive]
pub enum X86_64Level {
    V2,
    V3,
    V4,
}

/// CPU vendor + family/model details, as reported by `au_cpuid_get_vendor`.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct VendorInfo {
    pub vendor_id: String,
    pub family_id: String,
    pub model_id: String,
    pub stepping_id: String,
    pub uarch_id: String,
}

/// Returns `true` if the queried CPU is manufactured by AMD.
pub fn is_amd(cpu: Cpu) -> bool {
    // SAFETY: `au_cpuid_is_amd` is a pure CPUID query with no preconditions.
    unsafe { sys::au_cpuid_is_amd(cpu.as_raw()) }
}

/// Returns `true` if the queried CPU is in the Zen family (Zen, Zen+,
/// Zen2, Zen3, Zen4, Zen5; not Zen6+).
pub fn is_zen_family(cpu: Cpu) -> bool {
    unsafe { sys::au_cpuid_arch_is_zen_family(cpu.as_raw()) }
}

/// Returns the Zen sub-architecture of the queried CPU, if any.
pub fn zen_arch(cpu: Cpu) -> Option<ZenArch> {
    let raw = cpu.as_raw();
    // Order matters: AOCL's `arch_is_zenN` returns true for ZenN *and later*,
    // so we test from the newest backwards to get the most specific answer.
    unsafe {
        if sys::au_cpuid_arch_is_zen5(raw) {
            Some(ZenArch::Zen5)
        } else if sys::au_cpuid_arch_is_zen4(raw) {
            Some(ZenArch::Zen4)
        } else if sys::au_cpuid_arch_is_zen3(raw) {
            Some(ZenArch::Zen3)
        } else if sys::au_cpuid_arch_is_zen2(raw) {
            Some(ZenArch::Zen2)
        } else if sys::au_cpuid_arch_is_zenplus(raw) {
            Some(ZenArch::ZenPlus)
        } else if sys::au_cpuid_arch_is_zen(raw) {
            Some(ZenArch::Zen)
        } else {
            None
        }
    }
}

/// Returns the highest x86-64 microarchitecture level supported by the
/// queried CPU, or `None` if it does not even meet x86-64-v2.
pub fn x86_64_level(cpu: Cpu) -> Option<X86_64Level> {
    let raw = cpu.as_raw();
    unsafe {
        if sys::au_cpuid_arch_is_x86_64v4(raw) {
            Some(X86_64Level::V4)
        } else if sys::au_cpuid_arch_is_x86_64v3(raw) {
            Some(X86_64Level::V3)
        } else if sys::au_cpuid_arch_is_x86_64v2(raw) {
            Some(X86_64Level::V2)
        } else {
            None
        }
    }
}

/// Returns `true` if the queried CPU supports the named feature flag
/// (e.g. `"avx2"`, `"avx512f"`, `"sha"`, `"vaes"`).
///
/// **Caveat:** AOCL throws a C++ exception when given an unknown flag
/// name, which Rust cannot catch — pass only flag names AOCL recognises.
/// See `<AOCL_ROOT>/amd-utils/include/Au/Cpuid/Enum.hh` for the full
/// list (representative subset: `sse`, `sse2`, `sse3`, `sse4_1`,
/// `sse4_2`, `avx`, `avx2`, `avx512f`, `avx512vl`, `fma`, `aes`, `vaes`,
/// `sha`, `bmi1`, `bmi2`, `popcnt`).
pub fn has_flag(cpu: Cpu, flag: &str) -> bool {
    has_flags_all(cpu, &[flag])
}

/// Returns `true` if the queried CPU supports **all** of the named flags.
pub fn has_flags_all(cpu: Cpu, flags: &[&str]) -> bool {
    if flags.is_empty() {
        return true;
    }
    let cstrings: Vec<std::ffi::CString> = flags
        .iter()
        .filter_map(|s| std::ffi::CString::new(*s).ok())
        .collect();
    if cstrings.len() != flags.len() {
        // A flag with an interior NUL is malformed; treat as "not supported".
        return false;
    }
    let pointers: Vec<*const std::os::raw::c_char> = cstrings.iter().map(|c| c.as_ptr()).collect();
    // SAFETY: AOCL reads `count` C-string pointers; the cstrings stay
    // alive for the duration of this call.
    unsafe {
        sys::au_cpuid_has_flags_all(
            cpu.as_raw(),
            pointers.as_ptr(),
            pointers.len() as std::os::raw::c_int,
        )
    }
}

/// Returns `true` if the queried CPU supports **any** of the named flags.
pub fn has_flags_any(cpu: Cpu, flags: &[&str]) -> bool {
    if flags.is_empty() {
        return false;
    }
    let cstrings: Vec<std::ffi::CString> = flags
        .iter()
        .filter_map(|s| std::ffi::CString::new(*s).ok())
        .collect();
    if cstrings.len() != flags.len() {
        return false;
    }
    let pointers: Vec<*const std::os::raw::c_char> = cstrings.iter().map(|c| c.as_ptr()).collect();
    unsafe {
        sys::au_cpuid_has_flags_any(
            cpu.as_raw(),
            pointers.as_ptr(),
            pointers.len() as std::os::raw::c_int,
        )
    }
}

/// Fetch vendor / family / model / stepping / uarch identifiers for the
/// queried CPU.
pub fn vendor_info(cpu: Cpu) -> VendorInfo {
    let mut buf = [0u8; 256];
    // SAFETY: We pass a valid pointer + length pair; AOCL writes a
    // NUL-terminated string of at most `len` bytes.
    unsafe {
        sys::au_cpuid_get_vendor(
            cpu.as_raw(),
            buf.as_mut_ptr() as *mut std::os::raw::c_char,
            buf.len(),
        );
    }
    let nul = buf.iter().position(|&b| b == 0).unwrap_or(buf.len());
    let raw = std::str::from_utf8(&buf[..nul]).unwrap_or("");
    let mut it = raw.split('\n');
    VendorInfo {
        vendor_id: it.next().unwrap_or("").to_string(),
        family_id: it.next().unwrap_or("").to_string(),
        model_id: it.next().unwrap_or("").to_string(),
        stepping_id: it.next().unwrap_or("").to_string(),
        uarch_id: it.next().unwrap_or("").to_string(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn current_cpu_does_not_panic() {
        let info = vendor_info(Cpu::Current);
        assert!(!info.vendor_id.is_empty(), "vendor_id should be populated");
        let _ = is_amd(Cpu::Current);
        let _ = is_zen_family(Cpu::Current);
        let _ = zen_arch(Cpu::Current);
        let _ = x86_64_level(Cpu::Current);
    }

    #[test]
    fn has_flag_queries_run() {
        // AOCL throws a C++ exception on unknown flag names, which would
        // abort the process under Rust's panic-from-foreign-exception
        // handling — so use only flag names AOCL recognises.
        // SSE2 is universal on x86-64 since 2003.
        let _ = has_flag(Cpu::Current, "sse2");
        let _ = has_flags_all(Cpu::Current, &["sse2", "sse4_2"]);
        let _ = has_flags_any(Cpu::Current, &["avx2", "avx512f"]);
        // Empty list of flags
        assert!(has_flags_all(Cpu::Current, &[]));
        assert!(!has_flags_any(Cpu::Current, &[]));
    }
}