polars_utils/
cpuid.rs

1// So much conditional stuff going on here...
2#![allow(dead_code, unreachable_code, unused)]
3
4use std::sync::OnceLock;
5
6#[cfg(target_arch = "x86_64")]
7use raw_cpuid::CpuId;
8
9#[cfg(target_feature = "bmi2")]
10#[inline(never)]
11#[cold]
12fn detect_fast_bmi2() -> bool {
13    let cpu_id = CpuId::new();
14    let vendor = cpu_id.get_vendor_info().expect("could not read cpu vendor");
15    if vendor.as_str() == "AuthenticAMD" || vendor.as_str() == "HygonGenuine" {
16        let features = cpu_id
17            .get_feature_info()
18            .expect("could not read cpu feature info");
19        let family_id = features.family_id();
20
21        // Hardcoded blacklist of known-bad AMD families.
22        // We'll assume any future releases that support BMI2 have a
23        // proper implementation.
24        !(0x15..=0x18).contains(&family_id)
25    } else {
26        true
27    }
28}
29
30#[inline(always)]
31pub fn has_fast_bmi2() -> bool {
32    #[cfg(target_feature = "bmi2")]
33    {
34        static CACHE: OnceLock<bool> = OnceLock::new();
35        return *CACHE.get_or_init(detect_fast_bmi2);
36    }
37
38    false
39}
40
41#[inline]
42pub fn is_avx512_enabled() -> bool {
43    #[cfg(target_arch = "x86_64")]
44    {
45        static CACHE: OnceLock<bool> = OnceLock::new();
46        return *CACHE.get_or_init(|| {
47            if !std::arch::is_x86_feature_detected!("avx512f") {
48                return false;
49            }
50
51            if std::env::var("POLARS_DISABLE_AVX512")
52                .map(|var| var == "1")
53                .unwrap_or(false)
54            {
55                return false;
56            }
57
58            true
59        });
60    }
61
62    false
63}