lance_core/utils/
cpu.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use lazy_static::lazy_static;
5
6/// A level of SIMD support for some feature
7pub enum SimdSupport {
8    None,
9    Neon,
10    Sse,
11    Avx2,
12    Avx512,
13    Lsx,
14    Lasx,
15}
16
17lazy_static! {
18    /// Support for FP16 SIMD operations
19    pub static ref FP16_SIMD_SUPPORT: SimdSupport = {
20        #[cfg(target_arch = "aarch64")]
21        {
22            if aarch64::has_neon_f16_support() {
23                SimdSupport::Neon
24            } else {
25                SimdSupport::None
26            }
27        }
28        #[cfg(target_arch = "x86_64")]
29        {
30            if x86::has_avx512_f16_support() {
31                SimdSupport::Avx512
32            } else if is_x86_feature_detected!("avx2") {
33                SimdSupport::Avx2
34            } else {
35                SimdSupport::None
36            }
37        }
38        #[cfg(target_arch = "loongarch64")]
39        {
40            if loongarch64::has_lasx_support() {
41                SimdSupport::Lasx
42            } else if loongarch64::has_lsx_support() {
43                SimdSupport::Lsx
44            } else {
45                SimdSupport::None
46            }
47        }
48    };
49}
50
51#[cfg(target_arch = "x86_64")]
52mod x86 {
53    use core::arch::x86_64::__cpuid;
54
55    #[inline]
56    fn check_flag(x: usize, position: u32) -> bool {
57        x & (1 << position) != 0
58    }
59
60    pub fn has_avx512_f16_support() -> bool {
61        // this macro does many OS checks/etc. to determine if allowed to use AVX512
62        if !is_x86_feature_detected!("avx512f") {
63            return false;
64        }
65
66        // EAX=7, ECX=0: Extended Features (includes AVX512)
67        // More info on calling CPUID can be found here (section 1.4)
68        // https://www.intel.com/content/dam/develop/external/us/en/documents/architecture-instruction-set-extensions-programming-reference.pdf
69        let ext_cpuid_result = unsafe { __cpuid(7) };
70        check_flag(ext_cpuid_result.edx as usize, 23)
71    }
72}
73
74// Inspired by https://github.com/RustCrypto/utils/blob/master/cpufeatures/src/aarch64.rs
75// aarch64 doesn't have userspace feature detection built in, so we have to call
76// into OS-specific functions to check for features.
77
78#[cfg(all(target_arch = "aarch64", target_os = "macos"))]
79mod aarch64 {
80    pub fn has_neon_f16_support() -> bool {
81        // Maybe we can assume it's there?
82        true
83    }
84}
85
86#[cfg(all(target_arch = "aarch64", target_os = "linux"))]
87mod aarch64 {
88    pub fn has_neon_f16_support() -> bool {
89        // See: https://github.com/rust-lang/libc/blob/7ce81ca7aeb56aae7ca0237ef9353d58f3d7d2f1/src/unix/linux_like/linux/gnu/b64/aarch64/mod.rs#L533
90        let flags = unsafe { libc::getauxval(libc::AT_HWCAP) };
91        flags & libc::HWCAP_FPHP != 0
92    }
93}
94
95#[cfg(all(target_arch = "aarch64", target_os = "windows"))]
96mod aarch64 {
97    pub fn has_neon_f16_support() -> bool {
98        // https://github.com/lancedb/lance/issues/2411
99        false
100    }
101}
102
103#[cfg(target_arch = "loongarch64")]
104mod loongarch64 {
105    pub fn has_lsx_support() -> bool {
106        // See: https://github.com/rust-lang/libc/blob/7ce81ca7aeb56aae7ca0237ef9353d58f3d7d2f1/src/unix/linux_like/linux/gnu/b64/loongarch64/mod.rs#L263
107        let flags = unsafe { libc::getauxval(libc::AT_HWCAP) };
108        flags & libc::HWCAP_LOONGARCH_LSX != 0
109    }
110    pub fn has_lasx_support() -> bool {
111        // See: https://github.com/rust-lang/libc/blob/7ce81ca7aeb56aae7ca0237ef9353d58f3d7d2f1/src/unix/linux_like/linux/gnu/b64/loongarch64/mod.rs#L264
112        let flags = unsafe { libc::getauxval(libc::AT_HWCAP) };
113        flags & libc::HWCAP_LOONGARCH_LASX != 0
114    }
115}
lance_core/utils/cpu.rs

lance_core/utils/
cpu.rs
