lance_core/utils/
cpu.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! CPU feature detection and runtime kernel selection.
5//!
6//! [`CpuFeatures`] is the set of features the host supports, detected once.
7//! Kernels declare the features they need with [`target_kernel!`](crate::target_kernel),
8//! and a [`Dispatch`] picks the first candidate the host supports.
9
10use std::fmt;
11use std::sync::LazyLock;
12
13mod dispatch;
14
15pub use dispatch::{Candidate, Dispatch, Kernel, Selection};
16
17/// A set of CPU features, named after rustc `target_feature`s.
18///
19/// Unlike [`SimdSupport`], this is not a ladder: a kernel can require any
20/// combination, such as AVX-512 without FMA or F16C on an AVX host.
21///
22/// ```
23/// use lance_core::utils::cpu::CpuFeatures;
24///
25/// // Implied features are included: AVX2 brings AVX along with it.
26/// let needed = CpuFeatures::from_target_features(&["avx2", "fma"]);
27/// assert!(needed.contains(CpuFeatures::AVX));
28///
29/// // A kernel can run here if the host has every feature it needs, and can
30/// // be called directly, without dispatch, if the build baseline does.
31/// let runnable = CpuFeatures::host().contains(needed);
32/// let inlinable = CpuFeatures::baseline().contains(needed);
33/// assert!(runnable || !inlinable);
34/// ```
35#[derive(Clone, Copy, PartialEq, Eq, Hash, Default)]
36pub struct CpuFeatures(u32);
37
38impl CpuFeatures {
39    pub const SSE2: Self = Self(1 << 0);
40    pub const SSE3: Self = Self(1 << 1);
41    pub const SSSE3: Self = Self(1 << 2);
42    pub const SSE4_1: Self = Self(1 << 3);
43    pub const SSE4_2: Self = Self(1 << 4);
44    pub const POPCNT: Self = Self(1 << 5);
45    pub const AVX: Self = Self(1 << 6);
46    pub const AVX2: Self = Self(1 << 7);
47    pub const FMA: Self = Self(1 << 8);
48    pub const F16C: Self = Self(1 << 9);
49    pub const BMI1: Self = Self(1 << 10);
50    pub const BMI2: Self = Self(1 << 11);
51    pub const AVX512F: Self = Self(1 << 12);
52    pub const AVX512BW: Self = Self(1 << 13);
53    pub const AVX512CD: Self = Self(1 << 14);
54    pub const AVX512DQ: Self = Self(1 << 15);
55    pub const AVX512VL: Self = Self(1 << 16);
56    pub const AVX512FP16: Self = Self(1 << 17);
57    pub const NEON: Self = Self(1 << 18);
58    pub const FP16: Self = Self(1 << 19);
59    pub const LSX: Self = Self(1 << 20);
60    pub const LASX: Self = Self(1 << 21);
61
62    /// The set with no features, which every host supports.
63    pub const fn empty() -> Self {
64        Self(0)
65    }
66
67    /// Whether the set has no features.
68    pub const fn is_empty(self) -> bool {
69        self.0 == 0
70    }
71
72    /// The features in either set.
73    pub const fn union(self, other: Self) -> Self {
74        Self(self.0 | other.0)
75    }
76
77    /// Whether every feature in `other` is also in `self`.
78    pub const fn contains(self, other: Self) -> bool {
79        self.0 & other.0 == other.0
80    }
81
82    /// Features supported by the host, detected once per process.
83    pub fn host() -> Self {
84        *HOST_FEATURES
85    }
86
87    /// Features the build enables for all code, e.g. via `-C target-cpu`.
88    ///
89    /// A kernel whose features are all in the baseline can be called directly
90    /// and inlined; see [`Dispatch`].
91    pub const fn baseline() -> Self {
92        BASELINE
93    }
94
95    /// The feature set for one rustc `target_feature` name, including the
96    /// features rustc enables along with it. `None` if the name is unknown.
97    ///
98    /// Implied features matter for soundness: `#[target_feature(enable =
99    /// "avx512f")]` also lets codegen use AVX2, FMA and F16C, so a kernel
100    /// declaring it must not run on a host that lacks any of them.
101    pub const fn from_target_feature(name: &str) -> Option<Self> {
102        let mut i = 0;
103        while i < FEATURE_TABLE.len() {
104            if str_eq(FEATURE_TABLE[i].name, name) {
105                return Some(FEATURE_TABLE[i].feature.with_implied());
106            }
107            i += 1;
108        }
109        None
110    }
111
112    /// Union of [`Self::from_target_feature`] over `names`.
113    ///
114    /// # Panics
115    ///
116    /// Panics if a name is unknown. [`target_kernel!`](crate::target_kernel)
117    /// evaluates this in a `const`, so a typo is a compile error there.
118    pub const fn from_target_features(names: &[&str]) -> Self {
119        let mut features = Self::empty();
120        let mut i = 0;
121        while i < names.len() {
122            match Self::from_target_feature(names[i]) {
123                Some(feature) => features = features.union(feature),
124                None => panic!("unknown target feature; add it to CpuFeatures"),
125            }
126            i += 1;
127        }
128        features
129    }
130
131    /// Every known feature, one per set.
132    #[cfg(test)]
133    pub fn all() -> impl Iterator<Item = Self> {
134        FEATURE_TABLE.iter().map(|info| info.feature)
135    }
136
137    /// The `target_feature` names of the features in this set.
138    pub fn names(self) -> impl Iterator<Item = &'static str> {
139        FEATURE_TABLE
140            .iter()
141            .filter(move |info| self.contains(info.feature))
142            .map(|info| info.name)
143    }
144
145    fn x86_names(self) -> impl Iterator<Item = &'static str> {
146        FEATURE_TABLE
147            .iter()
148            .filter(move |info| info.is_x86 && self.contains(info.feature))
149            .map(|info| info.name)
150    }
151
152    const fn with_implied(self) -> Self {
153        let mut features = self;
154        loop {
155            let mut expanded = features;
156            let mut i = 0;
157            while i < FEATURE_TABLE.len() {
158                if features.contains(FEATURE_TABLE[i].feature) {
159                    expanded = expanded.union(FEATURE_TABLE[i].implies);
160                }
161                i += 1;
162            }
163            if expanded.0 == features.0 {
164                return features;
165            }
166            features = expanded;
167        }
168    }
169}
170
171impl fmt::Debug for CpuFeatures {
172    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
173        f.debug_set().entries(self.names()).finish()
174    }
175}
176
177struct FeatureInfo {
178    feature: CpuFeatures,
179    name: &'static str,
180    is_x86: bool,
181    /// Features rustc enables along with this one. Transitive implications
182    /// are resolved by [`CpuFeatures::with_implied`].
183    implies: CpuFeatures,
184}
185
186const fn x86(feature: CpuFeatures, name: &'static str, implies: CpuFeatures) -> FeatureInfo {
187    FeatureInfo {
188        feature,
189        name,
190        is_x86: true,
191        implies,
192    }
193}
194
195const fn other(feature: CpuFeatures, name: &'static str, implies: CpuFeatures) -> FeatureInfo {
196    FeatureInfo {
197        feature,
198        name,
199        is_x86: false,
200        implies,
201    }
202}
203
204/// Ordered as [`SimdInfo::host_features`] reports them.
205const FEATURE_TABLE: &[FeatureInfo] = {
206    use CpuFeatures as F;
207    &[
208        x86(F::SSE2, "sse2", F::empty()),
209        x86(F::SSE3, "sse3", F::SSE2),
210        x86(F::SSSE3, "ssse3", F::SSE3),
211        x86(F::SSE4_1, "sse4.1", F::SSSE3),
212        x86(F::SSE4_2, "sse4.2", F::SSE4_1),
213        x86(F::POPCNT, "popcnt", F::empty()),
214        x86(F::AVX, "avx", F::SSE4_2),
215        x86(F::AVX2, "avx2", F::AVX),
216        x86(F::FMA, "fma", F::AVX),
217        x86(F::F16C, "f16c", F::AVX),
218        x86(F::BMI1, "bmi1", F::empty()),
219        x86(F::BMI2, "bmi2", F::empty()),
220        x86(F::AVX512F, "avx512f", F::AVX2.union(F::FMA).union(F::F16C)),
221        x86(F::AVX512BW, "avx512bw", F::AVX512F),
222        x86(F::AVX512CD, "avx512cd", F::AVX512F),
223        x86(F::AVX512DQ, "avx512dq", F::AVX512F),
224        x86(F::AVX512VL, "avx512vl", F::AVX512F),
225        x86(
226            F::AVX512FP16,
227            "avx512fp16",
228            F::AVX512BW.union(F::AVX512DQ).union(F::AVX512VL),
229        ),
230        other(F::NEON, "neon", F::empty()),
231        other(F::FP16, "fp16", F::NEON),
232        other(F::LSX, "lsx", F::empty()),
233        other(F::LASX, "lasx", F::LSX),
234    ]
235};
236
237const fn str_eq(a: &str, b: &str) -> bool {
238    let (a, b) = (a.as_bytes(), b.as_bytes());
239    if a.len() != b.len() {
240        return false;
241    }
242    let mut i = 0;
243    while i < a.len() {
244        if a[i] != b[i] {
245            return false;
246        }
247        i += 1;
248    }
249    true
250}
251
252const BASELINE: CpuFeatures = {
253    // `cfg!` needs a literal feature name, so this cannot loop over the table.
254    let enabled = [
255        (cfg!(target_feature = "sse2"), CpuFeatures::SSE2),
256        (cfg!(target_feature = "sse3"), CpuFeatures::SSE3),
257        (cfg!(target_feature = "ssse3"), CpuFeatures::SSSE3),
258        (cfg!(target_feature = "sse4.1"), CpuFeatures::SSE4_1),
259        (cfg!(target_feature = "sse4.2"), CpuFeatures::SSE4_2),
260        (cfg!(target_feature = "popcnt"), CpuFeatures::POPCNT),
261        (cfg!(target_feature = "avx"), CpuFeatures::AVX),
262        (cfg!(target_feature = "avx2"), CpuFeatures::AVX2),
263        (cfg!(target_feature = "fma"), CpuFeatures::FMA),
264        (cfg!(target_feature = "f16c"), CpuFeatures::F16C),
265        (cfg!(target_feature = "bmi1"), CpuFeatures::BMI1),
266        (cfg!(target_feature = "bmi2"), CpuFeatures::BMI2),
267        (cfg!(target_feature = "avx512f"), CpuFeatures::AVX512F),
268        (cfg!(target_feature = "avx512bw"), CpuFeatures::AVX512BW),
269        (cfg!(target_feature = "avx512cd"), CpuFeatures::AVX512CD),
270        (cfg!(target_feature = "avx512dq"), CpuFeatures::AVX512DQ),
271        (cfg!(target_feature = "avx512vl"), CpuFeatures::AVX512VL),
272        (cfg!(target_feature = "avx512fp16"), CpuFeatures::AVX512FP16),
273        (cfg!(target_feature = "neon"), CpuFeatures::NEON),
274        (cfg!(target_feature = "fp16"), CpuFeatures::FP16),
275        (cfg!(target_feature = "lsx"), CpuFeatures::LSX),
276        (cfg!(target_feature = "lasx"), CpuFeatures::LASX),
277    ];
278    let mut features = CpuFeatures::empty();
279    let mut i = 0;
280    while i < enabled.len() {
281        if enabled[i].0 {
282            features = features.union(enabled[i].1);
283        }
284        i += 1;
285    }
286    features
287};
288
289// Runtime detection checks only some features per target, and none on
290// targets without a detector, so it can miss features the build already
291// assumes. `Dispatch` relies on the host set containing the baseline, so add
292// it explicitly.
293static HOST_FEATURES: LazyLock<CpuFeatures> =
294    LazyLock::new(|| detect_host_features().union(BASELINE));
295
296#[cfg(target_arch = "x86_64")]
297fn detect_host_features() -> CpuFeatures {
298    let mut features = CpuFeatures::empty();
299    let mut add = |is_detected: bool, feature: CpuFeatures| {
300        if is_detected {
301            features = features.union(feature);
302        }
303    };
304    // Each call must be inline: `is_x86_feature_detected!` does its own custom
305    // input parsing and rejects feature names received via a `macro_rules!`
306    // `:literal` metavariable on some toolchains.
307    add(is_x86_feature_detected!("sse2"), CpuFeatures::SSE2);
308    add(is_x86_feature_detected!("sse3"), CpuFeatures::SSE3);
309    add(is_x86_feature_detected!("ssse3"), CpuFeatures::SSSE3);
310    add(is_x86_feature_detected!("sse4.1"), CpuFeatures::SSE4_1);
311    add(is_x86_feature_detected!("sse4.2"), CpuFeatures::SSE4_2);
312    add(is_x86_feature_detected!("popcnt"), CpuFeatures::POPCNT);
313    add(is_x86_feature_detected!("avx"), CpuFeatures::AVX);
314    add(is_x86_feature_detected!("avx2"), CpuFeatures::AVX2);
315    add(is_x86_feature_detected!("fma"), CpuFeatures::FMA);
316    add(is_x86_feature_detected!("f16c"), CpuFeatures::F16C);
317    add(is_x86_feature_detected!("bmi1"), CpuFeatures::BMI1);
318    add(is_x86_feature_detected!("bmi2"), CpuFeatures::BMI2);
319    add(is_x86_feature_detected!("avx512f"), CpuFeatures::AVX512F);
320    add(is_x86_feature_detected!("avx512bw"), CpuFeatures::AVX512BW);
321    add(is_x86_feature_detected!("avx512cd"), CpuFeatures::AVX512CD);
322    add(is_x86_feature_detected!("avx512dq"), CpuFeatures::AVX512DQ);
323    add(is_x86_feature_detected!("avx512vl"), CpuFeatures::AVX512VL);
324    add(
325        is_x86_feature_detected!("avx512fp16"),
326        CpuFeatures::AVX512FP16,
327    );
328    features
329}
330
331#[cfg(target_arch = "aarch64")]
332fn detect_host_features() -> CpuFeatures {
333    #[cfg(any(target_os = "ios", target_os = "tvos"))]
334    let has_fp16 = true;
335    #[cfg(not(any(target_os = "ios", target_os = "tvos")))]
336    let has_fp16 = aarch64::has_neon_f16_support();
337
338    if has_fp16 {
339        CpuFeatures::NEON.union(CpuFeatures::FP16)
340    } else {
341        CpuFeatures::NEON
342    }
343}
344
345#[cfg(target_arch = "loongarch64")]
346fn detect_host_features() -> CpuFeatures {
347    let mut features = CpuFeatures::empty();
348    if loongarch64::has_lsx_support() {
349        features = features.union(CpuFeatures::LSX);
350    }
351    if loongarch64::has_lasx_support() {
352        features = features.union(CpuFeatures::LASX);
353    }
354    features
355}
356
357#[cfg(not(any(
358    target_arch = "aarch64",
359    target_arch = "x86_64",
360    target_arch = "loongarch64"
361)))]
362fn detect_host_features() -> CpuFeatures {
363    CpuFeatures::empty()
364}
365
366/// A level of SIMD support for some feature.
367///
368/// `#[non_exhaustive]` so future tiers (e.g. AVX-512 BF16, AMX) can be added
369/// without breaking external `match` consumers.
370#[derive(Debug, Clone, Copy, PartialEq, Eq)]
371#[non_exhaustive]
372pub enum SimdSupport {
373    None,
374    Neon,
375    Sse,
376    /// AVX (256-bit float ops) without FMA.
377    ///
378    /// This tier does not imply that AVX2 is absent: selecting [`Self::Avx2`]
379    /// requires both AVX2 and FMA, so a host with AVX2 but no FMA selects this
380    /// tier. Intel Sandy Bridge / Ivy Bridge are the typical hosts.
381    Avx,
382    /// AVX + FMA but no AVX2.
383    /// AMD Piledriver / Steamroller / FX-7500.
384    AvxFma,
385    /// AVX2 + FMA. Intel Haswell / AMD Excavator and later.
386    ///
387    /// Selecting this tier asserts FMA is present: the kernels it dispatches to
388    /// are `#[target_feature(enable = "avx,fma")]`.
389    Avx2,
390    Avx512,
391    Avx512FP16,
392    Lsx,
393    Lasx,
394}
395
396impl fmt::Display for SimdSupport {
397    /// Formats the tier name in lowercase, matching pyarrow's
398    /// `runtime_info().simd_level` convention.
399    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
400        let name = match self {
401            Self::None => "none",
402            Self::Neon => "neon",
403            Self::Sse => "sse",
404            Self::Avx => "avx",
405            Self::AvxFma => "avx_fma",
406            Self::Avx2 => "avx2",
407            Self::Avx512 => "avx512",
408            Self::Avx512FP16 => "avx512_fp16",
409            Self::Lsx => "lsx",
410            Self::Lasx => "lasx",
411        };
412        f.write_str(name)
413    }
414}
415
416/// Snapshot of the SIMD tier lance dispatches to on the current host, plus the
417/// raw CPU features detected for diagnostic purposes.
418///
419/// Mirrors the role of `pyarrow.runtime_info()`: a single, cheap call users can
420/// make to verify which SIMD tier the runtime selected and what underlying
421/// features the host advertises. Obtain one with [`simd_info()`].
422#[derive(Debug, Clone)]
423pub struct SimdInfo {
424    /// The SIMD tier lance dispatches to at runtime on this host.
425    pub tier: SimdSupport,
426    /// The architecture name (e.g. "x86_64", "aarch64", "loongarch64").
427    pub target_arch: &'static str,
428    /// Raw CPU feature flags detected on this host (x86_64 only; empty on
429    /// other architectures). Each entry is a feature name like "avx2",
430    /// "fma", "avx512f", "popcnt", etc.
431    pub host_features: Vec<&'static str>,
432}
433
434/// Returns a snapshot of the SIMD tier lance is using on this host along with
435/// the raw CPU feature flags that drove the decision.
436///
437/// Useful for performance debugging and giving users a way to verify which
438/// dispatch tier they are hitting without rebuilding lance. See [`SimdInfo`]
439/// for the meaning of each field and [`SimdSupport`] for the tier values.
440///
441/// # Examples
442///
443/// ```
444/// use lance_core::utils::cpu::simd_info;
445///
446/// let info = simd_info();
447/// println!("dispatching to {} on {}", info.tier, info.target_arch);
448/// ```
449pub fn simd_info() -> SimdInfo {
450    SimdInfo {
451        tier: *SIMD_SUPPORT,
452        target_arch: std::env::consts::ARCH,
453        host_features: CpuFeatures::host().x86_names().collect(),
454    }
455}
456
457impl SimdSupport {
458    /// The tier lance reports for a feature set. Each architecture only sets
459    /// its own feature bits, so checking every architecture in turn is safe.
460    fn from_features(features: CpuFeatures) -> Self {
461        use CpuFeatures as F;
462        if features.contains(F::AVX512F) {
463            if features.contains(F::AVX512FP16) {
464                Self::Avx512FP16
465            } else {
466                Self::Avx512
467            }
468        } else if features.contains(F::AVX2.union(F::FMA)) {
469            Self::Avx2
470        } else if features.contains(F::AVX.union(F::FMA)) {
471            // AMD Piledriver / Steamroller / FX-7500: 256-bit float ops + FMA but no AVX2.
472            Self::AvxFma
473        } else if features.contains(F::AVX) {
474            // This includes a possible AVX2 host without FMA because the Avx2
475            // tier above requires both features.
476            Self::Avx
477        } else if features.contains(F::FP16) {
478            Self::Neon
479        } else if features.contains(F::LASX) {
480            Self::Lasx
481        } else if features.contains(F::LSX) {
482            Self::Lsx
483        } else {
484            Self::None
485        }
486    }
487}
488
489/// Support for SIMD operations, derived from [`CpuFeatures::host`].
490pub static SIMD_SUPPORT: LazyLock<SimdSupport> =
491    LazyLock::new(|| SimdSupport::from_features(CpuFeatures::host()));
492
493// Inspired by https://github.com/RustCrypto/utils/blob/master/cpufeatures/src/aarch64.rs
494// aarch64 doesn't have userspace feature detection built in, so we have to call
495// into OS-specific functions to check for features.
496
497#[cfg(all(target_arch = "aarch64", target_os = "macos"))]
498mod aarch64 {
499    pub fn has_neon_f16_support() -> bool {
500        // Maybe we can assume it's there?
501        true
502    }
503}
504
505#[cfg(all(target_arch = "aarch64", target_os = "linux"))]
506mod aarch64 {
507    pub fn has_neon_f16_support() -> bool {
508        // See: https://github.com/rust-lang/libc/blob/7ce81ca7aeb56aae7ca0237ef9353d58f3d7d2f1/src/unix/linux_like/linux/gnu/b64/aarch64/mod.rs#L533
509        let flags = unsafe { libc::getauxval(libc::AT_HWCAP) };
510        flags & libc::HWCAP_FPHP != 0
511    }
512}
513
514#[cfg(all(target_arch = "aarch64", target_os = "windows"))]
515mod aarch64 {
516    pub fn has_neon_f16_support() -> bool {
517        // https://github.com/lance-format/lance/issues/2411
518        false
519    }
520}
521
522#[cfg(target_arch = "loongarch64")]
523mod loongarch64 {
524    pub fn has_lsx_support() -> bool {
525        // See: https://github.com/rust-lang/libc/blob/7ce81ca7aeb56aae7ca0237ef9353d58f3d7d2f1/src/unix/linux_like/linux/gnu/b64/loongarch64/mod.rs#L263
526        let flags = unsafe { libc::getauxval(libc::AT_HWCAP) };
527        flags & libc::HWCAP_LOONGARCH_LSX != 0
528    }
529    pub fn has_lasx_support() -> bool {
530        // See: https://github.com/rust-lang/libc/blob/7ce81ca7aeb56aae7ca0237ef9353d58f3d7d2f1/src/unix/linux_like/linux/gnu/b64/loongarch64/mod.rs#L264
531        let flags = unsafe { libc::getauxval(libc::AT_HWCAP) };
532        flags & libc::HWCAP_LOONGARCH_LASX != 0
533    }
534}
535
536#[cfg(all(target_arch = "aarch64", target_os = "android"))]
537mod aarch64 {
538    pub fn has_neon_f16_support() -> bool {
539        false
540    }
541}
542
543#[cfg(test)]
544mod tests {
545    use super::*;
546    use rstest::rstest;
547
548    #[test]
549    fn simd_info_exposes_tier() {
550        let info = simd_info();
551        assert_eq!(info.target_arch, std::env::consts::ARCH);
552        // Tier should match the detected SIMD support.
553        assert_eq!(info.tier, *SIMD_SUPPORT);
554    }
555
556    #[cfg(target_arch = "x86_64")]
557    #[test]
558    fn simd_info_features_include_baseline() {
559        let info = simd_info();
560        // The x86_64 ABI mandates SSE2, so it must always be present on this
561        // architecture.
562        assert!(info.host_features.contains(&"sse2"));
563    }
564
565    #[cfg(not(target_arch = "x86_64"))]
566    #[test]
567    fn simd_info_features_empty_off_x86_64() {
568        let info = simd_info();
569        assert!(info.host_features.is_empty());
570    }
571
572    /// The `Avx2` and `AvxFma` tiers both dispatch to kernels declared
573    /// `#[target_feature(enable = "avx,fma")]`, so neither may be selected on a
574    /// host without FMA. AVX2 does not imply FMA in the ISA, so the detection
575    /// checks it explicitly.
576    #[cfg(target_arch = "x86_64")]
577    #[test]
578    fn avx_fma_tiers_are_only_selected_when_fma_is_detected() {
579        if matches!(*SIMD_SUPPORT, SimdSupport::Avx2 | SimdSupport::AvxFma) {
580            assert!(
581                is_x86_feature_detected!("fma"),
582                "tier {} dispatches to avx,fma kernels but the host has no FMA",
583                *SIMD_SUPPORT
584            );
585        }
586    }
587
588    #[test]
589    fn host_features_include_build_baseline() {
590        assert!(CpuFeatures::host().contains(CpuFeatures::baseline()));
591    }
592
593    /// The mask a kernel declares must cover every feature rustc enables for
594    /// it, including implied ones, or a host lacking an implied feature could
595    /// select a kernel that uses it.
596    #[rstest]
597    #[case::sse_chain("avx", CpuFeatures::SSE2.union(CpuFeatures::SSE4_2))]
598    #[case::fma_implies_avx("fma", CpuFeatures::AVX)]
599    #[case::avx512f(
600        "avx512f",
601        CpuFeatures::AVX2.union(CpuFeatures::FMA).union(CpuFeatures::F16C)
602    )]
603    #[case::avx512fp16("avx512fp16", CpuFeatures::AVX512BW.union(CpuFeatures::AVX512F))]
604    #[case::fp16("fp16", CpuFeatures::NEON)]
605    #[case::lasx("lasx", CpuFeatures::LSX)]
606    fn target_feature_mask_includes_implied_features(
607        #[case] name: &str,
608        #[case] implied: CpuFeatures,
609    ) {
610        let features = CpuFeatures::from_target_feature(name).unwrap();
611        assert!(
612            features.contains(implied),
613            "{name} -> {features:?} is missing {implied:?}"
614        );
615    }
616
617    #[test]
618    fn unknown_target_feature_is_rejected() {
619        assert_eq!(CpuFeatures::from_target_feature("avx9000"), None);
620    }
621
622    #[rstest]
623    #[case::none(&[], SimdSupport::None)]
624    #[case::avx(&["avx"], SimdSupport::Avx)]
625    #[case::avx2_without_fma(&["avx2"], SimdSupport::Avx)]
626    #[case::avx_fma(&["fma"], SimdSupport::AvxFma)]
627    #[case::avx2(&["avx2", "fma"], SimdSupport::Avx2)]
628    #[case::avx512(&["avx512f"], SimdSupport::Avx512)]
629    #[case::avx512_fp16(&["avx512fp16"], SimdSupport::Avx512FP16)]
630    #[case::neon_without_fp16(&["neon"], SimdSupport::None)]
631    #[case::neon(&["fp16"], SimdSupport::Neon)]
632    #[case::lsx(&["lsx"], SimdSupport::Lsx)]
633    #[case::lasx(&["lasx"], SimdSupport::Lasx)]
634    fn simd_support_is_derived_from_features(
635        #[case] names: &[&str],
636        #[case] expected: SimdSupport,
637    ) {
638        let features = CpuFeatures::from_target_features(names);
639        assert_eq!(SimdSupport::from_features(features), expected);
640    }
641
642    #[rstest]
643    #[case::none(SimdSupport::None, "none")]
644    #[case::neon(SimdSupport::Neon, "neon")]
645    #[case::sse(SimdSupport::Sse, "sse")]
646    #[case::avx(SimdSupport::Avx, "avx")]
647    #[case::avx_fma(SimdSupport::AvxFma, "avx_fma")]
648    #[case::avx2(SimdSupport::Avx2, "avx2")]
649    #[case::avx512(SimdSupport::Avx512, "avx512")]
650    #[case::avx512_fp16(SimdSupport::Avx512FP16, "avx512_fp16")]
651    #[case::lsx(SimdSupport::Lsx, "lsx")]
652    #[case::lasx(SimdSupport::Lasx, "lasx")]
653    fn simd_support_display_matches_lowercase_convention(
654        #[case] tier: SimdSupport,
655        #[case] expected: &str,
656    ) {
657        assert_eq!(tier.to_string(), expected);
658    }
659}
lance_core/utils/cpu.rs

lance_core/utils/
cpu.rs
