lance_core/utils/
cpu.rs

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

lance_core/utils/
cpu.rs
