fix(linalg): preserve half-kernel fallback on AVX-512 hosts

rust/lance-linalg/build.rs

@@ -112,7 +112,12 @@ fn main() -> Result<(), String> {
         // While GCC doesn't have support for _Float16 until GCC 12, clang
         // has support for __fp16 going back to at least clang 6.
         // We use haswell since it's the oldest CPUs on AWS.
-        if let Err(err) = build_f16_with_flags("avx2", &["-march=haswell"]) {
+        // Keep the AVX2 L2 reductions lane-partitioned. A single reassociated
+        // float accumulator can exceed the public 1e-6 relative-error contract
+        // for long f16 vectors.
+        if let Err(err) =
+            build_f16_with_flags("avx2", &["-march=haswell", "-DPRECISE_F16_REDUCTION"])
+        {
             return Err(format!(
                 "Unable to build AVX2 f16 kernels.  Please use Clang >= 6 or GCC >= 12 or remove the fp16kernels feature.  Received error: {}",
                 err

rust/lance-linalg/src/distance.rs

@@ -15,6 +15,7 @@ use arrow_array::cast::AsArray;
 use arrow_array::types::{Float16Type, Float32Type, Float64Type, UInt8Type};
 use arrow_array::{Array, ArrowPrimitiveType, FixedSizeListArray, Float32Array, ListArray};
 use arrow_schema::{ArrowError, DataType};
+use lance_core::utils::cpu::SimdSupport;
 
 pub mod cosine;
 pub mod cosine_u8;
@@ -51,6 +52,94 @@ fn assert_batch_layout(vector_len: usize, batch_len: usize, dimension: usize) {
     );
 }
 
+/// Runtime backend shared by the f16 and bf16 C kernels.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum HalfBackend {
+    Avx512,
+    Avx2,
+    Neon,
+    Lsx,
+    Lasx,
+    Scalar,
+}
+
+/// Half-precision element type whose fallback kernel is being selected.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum HalfType {
+    F16,
+    Bf16,
+}
+
+/// CPU features required by the x86 fallback objects.
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
+struct X86HalfFeatures {
+    has_avx2: bool,
+    has_f16c: bool,
+    has_fma: bool,
+}
+
+/// Whether this target links the optional half-precision C objects.
+const HALF_KERNELS_COMPILED: bool = cfg!(all(
+    feature = "fp16kernels",
+    not(target_os = "windows"),
+    any(
+        target_arch = "aarch64",
+        target_arch = "x86_64",
+        target_arch = "loongarch64"
+    )
+));
+
+/// Selects a half-precision backend without reading host or build state so the
+/// exclusive SIMD tier ladder can be tested on any machine.
+fn half_backend(
+    support: SimdSupport,
+    half_type: HalfType,
+    has_kernels: bool,
+    has_avx512_kernel: bool,
+    x86_features: X86HalfFeatures,
+) -> HalfBackend {
+    if !has_kernels {
+        return HalfBackend::Scalar;
+    }
+
+    match support {
+        SimdSupport::Avx512FP16 if has_avx512_kernel => HalfBackend::Avx512,
+        // SIMD_SUPPORT reports one exclusive tier. An AVX-512 host that cannot
+        // use the optional AVX-512 C object must therefore be named here to
+        // reach the always-built x86 fallback object.
+        SimdSupport::Avx512 | SimdSupport::Avx512FP16 | SimdSupport::Avx2
+            if x86_features.has_fma
+                && match half_type {
+                    HalfType::F16 => x86_features.has_f16c,
+                    HalfType::Bf16 => x86_features.has_avx2,
+                } =>
+        {
+            HalfBackend::Avx2
+        }
+        SimdSupport::Neon => HalfBackend::Neon,
+        SimdSupport::Lsx => HalfBackend::Lsx,
+        SimdSupport::Lasx => HalfBackend::Lasx,
+        _ => HalfBackend::Scalar,
+    }
+}
+
+/// Detects every feature emitted by the x86 fallback objects.
+#[inline]
+fn x86_half_features() -> X86HalfFeatures {
+    #[cfg(target_arch = "x86_64")]
+    {
+        X86HalfFeatures {
+            has_avx2: std::is_x86_feature_detected!("avx2"),
+            has_f16c: std::is_x86_feature_detected!("f16c"),
+            has_fma: std::is_x86_feature_detected!("fma"),
+        }
+    }
+    #[cfg(not(target_arch = "x86_64"))]
+    {
+        X86HalfFeatures::default()
+    }
+}
+
 /// Number of distances computed per call into a runtime-selected batch kernel.
 ///
 /// Keeping a small output buffer amortizes the `#[target_feature]` call while
@@ -502,6 +591,143 @@ mod tests {
     use arrow_schema::Field;
     use half::f16;
 
+    const NO_X86_HALF_FEATURES: X86HalfFeatures = X86HalfFeatures {
+        has_avx2: false,
+        has_f16c: false,
+        has_fma: false,
+    };
+    const F16_X86_HALF_FEATURES: X86HalfFeatures = X86HalfFeatures {
+        has_avx2: false,
+        has_f16c: true,
+        has_fma: true,
+    };
+    const BF16_X86_HALF_FEATURES: X86HalfFeatures = X86HalfFeatures {
+        has_avx2: true,
+        has_f16c: false,
+        has_fma: true,
+    };
+    const ALL_X86_HALF_FEATURES: X86HalfFeatures = X86HalfFeatures {
+        has_avx2: true,
+        has_f16c: true,
+        has_fma: true,
+    };
+    const X86_HALF_FEATURES_WITHOUT_FMA: X86HalfFeatures = X86HalfFeatures {
+        has_avx2: true,
+        has_f16c: true,
+        has_fma: false,
+    };
+
+    #[rstest::rstest]
+    #[case::kernels_disabled(
+        SimdSupport::Avx512FP16,
+        HalfType::F16,
+        false,
+        false,
+        ALL_X86_HALF_FEATURES,
+        HalfBackend::Scalar
+    )]
+    #[case::avx512_kernel_ready(
+        SimdSupport::Avx512FP16,
+        HalfType::F16,
+        true,
+        true,
+        NO_X86_HALF_FEATURES,
+        HalfBackend::Avx512
+    )]
+    #[case::f16_fallback_from_avx512fp16(
+        SimdSupport::Avx512FP16,
+        HalfType::F16,
+        true,
+        false,
+        F16_X86_HALF_FEATURES,
+        HalfBackend::Avx2
+    )]
+    #[case::bf16_fallback_from_avx512fp16(
+        SimdSupport::Avx512FP16,
+        HalfType::Bf16,
+        true,
+        false,
+        BF16_X86_HALF_FEATURES,
+        HalfBackend::Avx2
+    )]
+    #[case::fallback_from_avx512(
+        SimdSupport::Avx512,
+        HalfType::Bf16,
+        true,
+        false,
+        BF16_X86_HALF_FEATURES,
+        HalfBackend::Avx2
+    )]
+    #[case::f16_missing_f16c(
+        SimdSupport::Avx512FP16,
+        HalfType::F16,
+        true,
+        false,
+        BF16_X86_HALF_FEATURES,
+        HalfBackend::Scalar
+    )]
+    #[case::bf16_missing_avx2(
+        SimdSupport::Avx512FP16,
+        HalfType::Bf16,
+        true,
+        false,
+        F16_X86_HALF_FEATURES,
+        HalfBackend::Scalar
+    )]
+    #[case::fallback_missing_fma(
+        SimdSupport::Avx512FP16,
+        HalfType::F16,
+        true,
+        false,
+        X86_HALF_FEATURES_WITHOUT_FMA,
+        HalfBackend::Scalar
+    )]
+    #[case::avx2_tier(
+        SimdSupport::Avx2,
+        HalfType::Bf16,
+        true,
+        false,
+        BF16_X86_HALF_FEATURES,
+        HalfBackend::Avx2
+    )]
+    #[case::neon(
+        SimdSupport::Neon,
+        HalfType::F16,
+        true,
+        false,
+        NO_X86_HALF_FEATURES,
+        HalfBackend::Neon
+    )]
+    #[case::lsx(
+        SimdSupport::Lsx,
+        HalfType::Bf16,
+        true,
+        false,
+        NO_X86_HALF_FEATURES,
+        HalfBackend::Lsx
+    )]
+    #[case::lasx(
+        SimdSupport::Lasx,
+        HalfType::Bf16,
+        true,
+        false,
+        NO_X86_HALF_FEATURES,
+        HalfBackend::Lasx
+    )]
+    fn half_backend_follows_exclusive_tier_and_feature_requirements(
+        #[case] support: SimdSupport,
+        #[case] half_type: HalfType,
+        #[case] has_kernels: bool,
+        #[case] has_avx512_kernel: bool,
+        #[case] features: X86HalfFeatures,
+        #[case] expected: HalfBackend,
+    ) {
+        assert_eq!(
+            half_backend(support, half_type, has_kernels, has_avx512_kernel, features),
+            expected
+        );
+    }
+
     #[cfg(target_arch = "x86_64")]
     #[test]
     fn test_x86_runtime_feature_report() {

rust/lance-linalg/src/distance/cosine.rs

@@ -24,8 +24,12 @@ use lance_arrow::{ArrowFloatType, FixedSizeListArrayExt, FloatArray};
 #[allow(unused_imports)]
 use lance_core::utils::cpu::{SIMD_SUPPORT, SimdSupport};
 
-use super::{Dot, norm_l2::norm_l2};
-use super::{Normalize, dot::dot};
+#[cfg(feature = "fp16kernels")]
+use super::HalfBackend;
+use super::{
+    Dot, HALF_KERNELS_COMPILED, HalfType, Normalize, assert_equal_lengths, dot::dot, half_backend,
+    norm_l2::norm_l2, x86_half_features,
+};
 #[cfg(all(target_arch = "x86_64", not(target_feature = "avx2")))]
 use crate::distance::BatchKind;
 #[allow(unused_imports)]
@@ -108,29 +112,36 @@ mod bf16_kernel {
 
 impl Cosine for bf16 {
     fn cosine_fast(x: &[Self], x_norm: f32, y: &[Self]) -> f32 {
-        match *SIMD_SUPPORT {
+        assert_equal_lengths(x.len(), y.len());
+        match half_backend(
+            *SIMD_SUPPORT,
+            HalfType::Bf16,
+            HALF_KERNELS_COMPILED,
+            cfg!(all(kernel_support = "avx512_bf16", target_arch = "x86_64")),
+            x86_half_features(),
+        ) {
             #[cfg(all(feature = "fp16kernels", target_arch = "aarch64"))]
-            SimdSupport::Neon => unsafe {
+            HalfBackend::Neon => unsafe {
                 bf16_kernel::cosine_bf16_neon(x.as_ptr(), x_norm, y.as_ptr(), y.len() as u32)
             },
             #[cfg(all(
                 feature = "fp16kernels",
                 kernel_support = "avx512_bf16",
                 target_arch = "x86_64"
             ))]
-            SimdSupport::Avx512FP16 => unsafe {
+            HalfBackend::Avx512 => unsafe {
                 bf16_kernel::cosine_bf16_avx512(x.as_ptr(), x_norm, y.as_ptr(), y.len() as u32)
             },
             #[cfg(all(feature = "fp16kernels", target_arch = "x86_64"))]
-            SimdSupport::Avx2 | SimdSupport::Avx512 => unsafe {
+            HalfBackend::Avx2 => unsafe {
                 bf16_kernel::cosine_bf16_avx2(x.as_ptr(), x_norm, y.as_ptr(), y.len() as u32)
             },
             #[cfg(all(feature = "fp16kernels", target_arch = "loongarch64"))]
-            SimdSupport::Lasx => unsafe {
+            HalfBackend::Lasx => unsafe {
                 bf16_kernel::cosine_bf16_lasx(x.as_ptr(), x_norm, y.as_ptr(), y.len() as u32)
             },
             #[cfg(all(feature = "fp16kernels", target_arch = "loongarch64"))]
-            SimdSupport::Lsx => unsafe {
+            HalfBackend::Lsx => unsafe {
                 bf16_kernel::cosine_bf16_lsx(x.as_ptr(), x_norm, y.as_ptr(), y.len() as u32)
             },
             // SimdSupport::AvxFma and SimdSupport::Avx fall through here:
@@ -164,34 +175,41 @@ mod kernel {
 
 impl Cosine for f16 {
     fn cosine_fast(x: &[Self], x_norm: f32, y: &[Self]) -> f32 {
-        match *SIMD_SUPPORT {
+        assert_equal_lengths(x.len(), y.len());
+        match half_backend(
+            *SIMD_SUPPORT,
+            HalfType::F16,
+            HALF_KERNELS_COMPILED,
+            cfg!(all(kernel_support = "avx512_f16", target_arch = "x86_64")),
+            x86_half_features(),
+        ) {
             #[cfg(all(feature = "fp16kernels", target_arch = "aarch64"))]
-            SimdSupport::Neon => unsafe {
+            HalfBackend::Neon => unsafe {
                 kernel::cosine_f16_neon(x.as_ptr(), x_norm, y.as_ptr(), y.len() as u32)
             },
             #[cfg(all(
                 feature = "fp16kernels",
                 kernel_support = "avx512_f16",
                 target_arch = "x86_64"
             ))]
-            SimdSupport::Avx512FP16 => unsafe {
+            HalfBackend::Avx512 => unsafe {
                 kernel::cosine_f16_avx512(x.as_ptr(), x_norm, y.as_ptr(), y.len() as u32)
             },
             #[cfg(all(feature = "fp16kernels", target_arch = "x86_64"))]
-            SimdSupport::Avx2 | SimdSupport::Avx512 => unsafe {
+            HalfBackend::Avx2 => unsafe {
                 kernel::cosine_f16_avx2(x.as_ptr(), x_norm, y.as_ptr(), y.len() as u32)
             },
             #[cfg(all(feature = "fp16kernels", target_arch = "loongarch64"))]
-            SimdSupport::Lasx => unsafe {
+            HalfBackend::Lasx => unsafe {
                 kernel::cosine_f16_lasx(x.as_ptr(), x_norm, y.as_ptr(), y.len() as u32)
             },
             #[cfg(all(feature = "fp16kernels", target_arch = "loongarch64"))]
-            SimdSupport::Lsx => unsafe {
+            HalfBackend::Lsx => unsafe {
                 kernel::cosine_f16_lsx(x.as_ptr(), x_norm, y.as_ptr(), y.len() as u32)
             },
-            // SimdSupport::AvxFma and SimdSupport::Avx fall through here:
-            // the f16 C kernels are compiled with `-march=haswell` minimum
-            // (AVX2), so they cannot run on AVX-only or AVX+FMA hosts.
+            // SimdSupport::AvxFma and SimdSupport::Avx retain their scalar
+            // route; this fallback only extends the tiers the C kernel already
+            // served to Avx512FP16 after checking F16C and FMA.
             _ => cosine_scalar(x, x_norm, y),
         }
     }
@@ -1430,6 +1448,22 @@ mod tests {
         1.0 - xy / x_sq / y_sq
     }
 
+    #[test]
+    fn half_cosine_rejects_mismatched_lengths_before_ffi() {
+        let f16_x = [f16::from_f32(1.0)];
+        let f16_y = [f16::from_f32(1.0), f16::from_f32(2.0)];
+        assert!(
+            std::panic::catch_unwind(|| <f16 as Cosine>::cosine_fast(&f16_x, 1.0, &f16_y)).is_err()
+        );
+
+        let bf16_x = [bf16::from_f32(1.0)];
+        let bf16_y = [bf16::from_f32(1.0), bf16::from_f32(2.0)];
+        assert!(
+            std::panic::catch_unwind(|| <bf16 as Cosine>::cosine_fast(&bf16_x, 1.0, &bf16_y))
+                .is_err()
+        );
+    }
+
     #[test]
     fn test_cosine() {
         let x: Float32Array = (1..9).map(|v| v as f32).collect();