arch

lancedb · Sep 15, 2024 · d2df667 · d2df667
1 parent 83652e1
commit d2df667
Showing 1 changed file with 8 additions and 5 deletions.
diff --git a/rust/lance-linalg/src/distance/dot.rs b/rust/lance-linalg/src/distance/dot.rs
@@ -85,7 +85,6 @@ impl Dot for bf16 {
     }
 }
 
-#[cfg(feature = "fp16kernels")]
 mod kernel {
     use super::*;
 
@@ -108,11 +107,15 @@ mod kernel {
 impl Dot for f16 {
     #[inline]
     fn dot(x: &[Self], y: &[Self]) -> f32 {
+        #[cfg(target_arch = "aarch64")]
+        {
+            // TODO: add SVE support.
+            if std::arch::is_aarch64_feature_detected!("neon") {
+                return unsafe { kernel::dot_f16_neon(x.as_ptr(), y.as_ptr(), x.len() as u32) };
+            }
+        }
+
         match *FP16_SIMD_SUPPORT {
-            // #[cfg(all(feature = "fp16kernels", target_arch = "aarch64"))]
-            // SimdSupport::Neon => unsafe {
-            //     kernel::dot_f16_neon(x.as_ptr(), y.as_ptr(), x.len() as u32)
-            // },
             #[cfg(all(
                 feature = "fp16kernels",
                 kernel_support = "avx512",