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        // __cpuid is safe in nightly but unsafe in stable, allow both
82        #[allow(unused_unsafe)]
83        let ext_cpuid_result = unsafe { __cpuid(7) };
84        check_flag(ext_cpuid_result.edx as usize, 23)
85    }
86
87    pub fn has_avx512() -> bool {
88        is_x86_feature_detected!("avx512f")
89    }
90}
91
92// Inspired by https://github.com/RustCrypto/utils/blob/master/cpufeatures/src/aarch64.rs
93// aarch64 doesn't have userspace feature detection built in, so we have to call
94// into OS-specific functions to check for features.
95
96#[cfg(all(target_arch = "aarch64", target_os = "macos"))]
97mod aarch64 {
98    pub fn has_neon_f16_support() -> bool {
99        // Maybe we can assume it's there?
100        true
101    }
102}
103
104#[cfg(all(target_arch = "aarch64", target_os = "linux"))]
105mod aarch64 {
106    pub fn has_neon_f16_support() -> bool {
107        // See: https://github.com/rust-lang/libc/blob/7ce81ca7aeb56aae7ca0237ef9353d58f3d7d2f1/src/unix/linux_like/linux/gnu/b64/aarch64/mod.rs#L533
108        let flags = unsafe { libc::getauxval(libc::AT_HWCAP) };
109        flags & libc::HWCAP_FPHP != 0
110    }
111}
112
113#[cfg(all(target_arch = "aarch64", target_os = "windows"))]
114mod aarch64 {
115    pub fn has_neon_f16_support() -> bool {
116        // https://github.com/lance-format/lance/issues/2411
117        false
118    }
119}
120
121#[cfg(target_arch = "loongarch64")]
122mod loongarch64 {
123    pub fn has_lsx_support() -> bool {
124        // See: https://github.com/rust-lang/libc/blob/7ce81ca7aeb56aae7ca0237ef9353d58f3d7d2f1/src/unix/linux_like/linux/gnu/b64/loongarch64/mod.rs#L263
125        let flags = unsafe { libc::getauxval(libc::AT_HWCAP) };
126        flags & libc::HWCAP_LOONGARCH_LSX != 0
127    }
128    pub fn has_lasx_support() -> bool {
129        // See: https://github.com/rust-lang/libc/blob/7ce81ca7aeb56aae7ca0237ef9353d58f3d7d2f1/src/unix/linux_like/linux/gnu/b64/loongarch64/mod.rs#L264
130        let flags = unsafe { libc::getauxval(libc::AT_HWCAP) };
131        flags & libc::HWCAP_LOONGARCH_LASX != 0
132    }
133}
134
135#[cfg(all(target_arch = "aarch64", target_os = "android"))]
136mod aarch64 {
137    pub fn has_neon_f16_support() -> bool {
138        false
139    }
140}
lance_core/utils/cpu.rs

lance_core/utils/
cpu.rs
