From bdb71ff3800fe013a0807f86bc83d8a031d89f13 Mon Sep 17 00:00:00 2001 From: yangjie01 Date: Sat, 29 Aug 2026 23:04:49 +0800 Subject: [PATCH] fix(linalg): stop FP16 hosts falling to scalar when the AVX-512 f16 kernel is absent --- rust/lance-linalg/src/distance/cosine.rs | 20 ++++++++++++++++++++ rust/lance-linalg/src/distance/dot.rs | 20 ++++++++++++++++++++ rust/lance-linalg/src/distance/l2.rs | 20 ++++++++++++++++++++ rust/lance-linalg/src/distance/norm_l2.rs | 20 ++++++++++++++++++++ 4 files changed, 80 insertions(+) diff --git a/rust/lance-linalg/src/distance/cosine.rs b/rust/lance-linalg/src/distance/cosine.rs index 1512571f6fd..6f4e3a84257 100644 --- a/rust/lance-linalg/src/distance/cosine.rs +++ b/rust/lance-linalg/src/distance/cosine.rs @@ -121,6 +121,16 @@ impl Cosine for bf16 { SimdSupport::Avx512FP16 => unsafe { bf16_kernel::cosine_bf16_avx512(x.as_ptr(), x_norm, y.as_ptr(), y.len() as u32) }, + #[cfg(all( + feature = "fp16kernels", + not(kernel_support = "avx512_bf16"), + target_arch = "x86_64" + ))] + // `Avx512*` tier detection does not check FMA, and these kernels are + // compiled `-march=haswell`, which enables it. + SimdSupport::Avx512FP16 if std::is_x86_feature_detected!("fma") => unsafe { + bf16_kernel::cosine_bf16_avx2(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 { bf16_kernel::cosine_bf16_avx2(x.as_ptr(), x_norm, y.as_ptr(), y.len() as u32) @@ -177,6 +187,16 @@ impl Cosine for f16 { SimdSupport::Avx512FP16 => unsafe { kernel::cosine_f16_avx512(x.as_ptr(), x_norm, y.as_ptr(), y.len() as u32) }, + #[cfg(all( + feature = "fp16kernels", + not(kernel_support = "avx512_f16"), + target_arch = "x86_64" + ))] + // `Avx512*` tier detection does not check FMA, and these kernels are + // compiled `-march=haswell`, which enables it. + SimdSupport::Avx512FP16 if std::is_x86_feature_detected!("fma") => unsafe { + kernel::cosine_f16_avx2(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 { kernel::cosine_f16_avx2(x.as_ptr(), x_norm, y.as_ptr(), y.len() as u32) diff --git a/rust/lance-linalg/src/distance/dot.rs b/rust/lance-linalg/src/distance/dot.rs index 22274a074d2..ed9ac80f468 100644 --- a/rust/lance-linalg/src/distance/dot.rs +++ b/rust/lance-linalg/src/distance/dot.rs @@ -153,6 +153,16 @@ impl Dot for bf16 { SimdSupport::Avx512FP16 => unsafe { bf16_kernel::dot_bf16_avx512(x.as_ptr(), y.as_ptr(), x.len() as u32) }, + #[cfg(all( + feature = "fp16kernels", + not(kernel_support = "avx512_bf16"), + target_arch = "x86_64" + ))] + // `Avx512*` tier detection does not check FMA, and these kernels are + // compiled `-march=haswell`, which enables it. + SimdSupport::Avx512FP16 if std::is_x86_feature_detected!("fma") => unsafe { + bf16_kernel::dot_bf16_avx2(x.as_ptr(), y.as_ptr(), x.len() as u32) + }, #[cfg(all(feature = "fp16kernels", target_arch = "x86_64"))] SimdSupport::Avx2 | SimdSupport::Avx512 => unsafe { bf16_kernel::dot_bf16_avx2(x.as_ptr(), y.as_ptr(), x.len() as u32) @@ -210,6 +220,16 @@ impl Dot for f16 { SimdSupport::Avx512FP16 => unsafe { kernel::dot_f16_avx512(x.as_ptr(), y.as_ptr(), x.len() as u32) }, + #[cfg(all( + feature = "fp16kernels", + not(kernel_support = "avx512_f16"), + target_arch = "x86_64" + ))] + // `Avx512*` tier detection does not check FMA, and these kernels are + // compiled `-march=haswell`, which enables it. + SimdSupport::Avx512FP16 if std::is_x86_feature_detected!("fma") => unsafe { + kernel::dot_f16_avx2(x.as_ptr(), y.as_ptr(), x.len() as u32) + }, #[cfg(all(feature = "fp16kernels", target_arch = "x86_64"))] SimdSupport::Avx2 | SimdSupport::Avx512 => unsafe { kernel::dot_f16_avx2(x.as_ptr(), y.as_ptr(), x.len() as u32) diff --git a/rust/lance-linalg/src/distance/l2.rs b/rust/lance-linalg/src/distance/l2.rs index 3af13d58840..66c4d771ab4 100644 --- a/rust/lance-linalg/src/distance/l2.rs +++ b/rust/lance-linalg/src/distance/l2.rs @@ -185,6 +185,16 @@ impl L2 for bf16 { SimdSupport::Avx512FP16 => unsafe { bf16_kernel::l2_bf16_avx512(x.as_ptr(), y.as_ptr(), x.len() as u32) }, + #[cfg(all( + feature = "fp16kernels", + not(kernel_support = "avx512_bf16"), + target_arch = "x86_64" + ))] + // `Avx512*` tier detection does not check FMA, and these kernels are + // compiled `-march=haswell`, which enables it. + SimdSupport::Avx512FP16 if std::is_x86_feature_detected!("fma") => unsafe { + bf16_kernel::l2_bf16_avx2(x.as_ptr(), y.as_ptr(), x.len() as u32) + }, #[cfg(all(feature = "fp16kernels", target_arch = "x86_64"))] SimdSupport::Avx2 | SimdSupport::Avx512 => unsafe { bf16_kernel::l2_bf16_avx2(x.as_ptr(), y.as_ptr(), x.len() as u32) @@ -242,6 +252,16 @@ impl L2 for f16 { SimdSupport::Avx512FP16 => unsafe { kernel::l2_f16_avx512(x.as_ptr(), y.as_ptr(), x.len() as u32) }, + #[cfg(all( + feature = "fp16kernels", + not(kernel_support = "avx512_f16"), + target_arch = "x86_64" + ))] + // `Avx512*` tier detection does not check FMA, and these kernels are + // compiled `-march=haswell`, which enables it. + SimdSupport::Avx512FP16 if std::is_x86_feature_detected!("fma") => unsafe { + kernel::l2_f16_avx2(x.as_ptr(), y.as_ptr(), x.len() as u32) + }, #[cfg(all(feature = "fp16kernels", target_arch = "x86_64"))] SimdSupport::Avx2 | SimdSupport::Avx512 => unsafe { kernel::l2_f16_avx2(x.as_ptr(), y.as_ptr(), x.len() as u32) diff --git a/rust/lance-linalg/src/distance/norm_l2.rs b/rust/lance-linalg/src/distance/norm_l2.rs index ad0a9daa68f..500deb21175 100644 --- a/rust/lance-linalg/src/distance/norm_l2.rs +++ b/rust/lance-linalg/src/distance/norm_l2.rs @@ -61,6 +61,16 @@ impl Normalize for f16 { SimdSupport::Avx512FP16 => unsafe { kernel::norm_l2_f16_avx512(vector.as_ptr(), vector.len() as u32) }, + #[cfg(all( + feature = "fp16kernels", + not(kernel_support = "avx512_f16"), + target_arch = "x86_64" + ))] + // `Avx512*` tier detection does not check FMA, and these kernels are + // compiled `-march=haswell`, which enables it. + SimdSupport::Avx512FP16 if std::is_x86_feature_detected!("fma") => unsafe { + kernel::norm_l2_f16_avx2(vector.as_ptr(), vector.len() as u32) + }, #[cfg(all(feature = "fp16kernels", target_arch = "x86_64"))] SimdSupport::Avx2 | SimdSupport::Avx512 => unsafe { kernel::norm_l2_f16_avx2(vector.as_ptr(), vector.len() as u32) @@ -115,6 +125,16 @@ impl Normalize for bf16 { SimdSupport::Avx512FP16 => unsafe { bf16_kernel::norm_l2_bf16_avx512(vector.as_ptr(), vector.len() as u32) }, + #[cfg(all( + feature = "fp16kernels", + not(kernel_support = "avx512_bf16"), + target_arch = "x86_64" + ))] + // `Avx512*` tier detection does not check FMA, and these kernels are + // compiled `-march=haswell`, which enables it. + SimdSupport::Avx512FP16 if std::is_x86_feature_detected!("fma") => unsafe { + bf16_kernel::norm_l2_bf16_avx2(vector.as_ptr(), vector.len() as u32) + }, #[cfg(all(feature = "fp16kernels", target_arch = "x86_64"))] SimdSupport::Avx2 | SimdSupport::Avx512 => unsafe { bf16_kernel::norm_l2_bf16_avx2(vector.as_ptr(), vector.len() as u32)