[diskann-vector] Optimise AArch64 FP16 Inner Products using FMLAL

.cargo/config.toml

@@ -7,12 +7,12 @@ rustflags = ["-C", "control-flow-guard"]
 [target.'cfg(target_arch="x86_64")']
 rustflags = ["-C", "target-cpu=x86-64-v3"]
 
-# If running on an Aarch64 CPU, this enables the `neon` and `dotprod`.
+# If running on an Aarch64 CPU, this enables the `neon`, `dotprod` and `fhm`.
 #
 # Generally speaking, CPUs like Apple M-series and Graviton 2+ will support `dotprod`, so
 # enabling this seems acceptable.
 [target.'cfg(target_arch="aarch64")']
-rustflags = ["-C", "target-feature=+neon,+dotprod"]
+rustflags = ["-C", "target-feature=+neon,+dotprod,+fhm"]
 
 [build]
 target-dir = "target"

diskann-vector/src/distance/simd.rs

@@ -1804,12 +1804,7 @@ impl SIMDSchema<Half, Half, Neon> for IP {
         y: Self::Right,
         acc: Self::Accumulator,
     ) -> Self::Accumulator {
-        diskann_wide::alias!(f32s = <Neon>::f32x4);
-
-        let x: f32s = x.into();
-        let y: f32s = y.into();
-
-        x.mul_add_simd(y, acc)
+        acc.dot_simd(x, y)
     }
 
     #[inline(always)]

diskann-wide/compile-aarch64-on-x86.sh

@@ -1,3 +1,3 @@
-RUSTFLAGS="-Ctarget-feature=+neon,+dotprod -Clinker=aarch64-linux-gnu-gcc" cargo test \
+RUSTFLAGS="-Ctarget-feature=+neon,+dotprod,+fhm -Clinker=aarch64-linux-gnu-gcc" cargo test \
     --package "$1" \
     --target aarch64-unknown-linux-gnu

diskann-wide/src/arch/aarch64/f32x4_.rs

@@ -6,8 +6,8 @@
 use half::f16;
 
 use crate::{
-    AsSIMD, Emulated, SIMDAbs, SIMDCast, SIMDMask, SIMDMinMax, SIMDMulAdd, SIMDPartialEq,
-    SIMDPartialOrd, SIMDSelect, SIMDSumTree, SIMDVector, constant::Const, helpers,
+    AsSIMD, Emulated, SIMDAbs, SIMDCast, SIMDDotProduct, SIMDMask, SIMDMinMax, SIMDMulAdd,
+    SIMDPartialEq, SIMDPartialOrd, SIMDSelect, SIMDSumTree, SIMDVector, constant::Const, helpers,
 };
 
 // AArch64 masks
@@ -111,6 +111,45 @@ impl SIMDSelect<f32x4> for mask32x4 {
     }
 }
 
+impl SIMDDotProduct<f16x4> for f32x4 {
+    #[inline(always)]
+    fn dot_simd(self, left: f16x4, right: f16x4) -> Self {
+        if cfg!(miri) {
+            let acc = self.to_array();
+            let left = left.to_array();
+            let right = right.to_array();
+            Self::from_array(
+                self.arch(),
+                std::array::from_fn(|i| {
+                    crate::cast_f16_to_f32(left[i])
+                        .mul_add(crate::cast_f16_to_f32(right[i]), acc[i])
+                }),
+            )
+        } else {
+            #[target_feature(enable = "fhm")]
+            unsafe fn fmlal(
+                mut acc: float32x4_t,
+                left: uint16x4_t,
+                right: uint16x4_t,
+            ) -> float32x4_t {
+                // SAFETY: The instruction is available with the `fhm` target feature.
+                unsafe {
+                    asm!(
+                        "fmlal {acc:v}.4s, {left:v}.4h, {right:v}.4h",
+                        acc = inout(vreg) acc,
+                        left = in(vreg) left,
+                        right = in(vreg) right,
+                        options(pure, nomem, nostack)
+                    );
+                }
+                acc
+            }
+            // SAFETY: Calling above fmlal function enabled with `fhm` feature.
+            Self(unsafe { fmlal(self.0, left.0, right.0) })
+        }
+    }
+}
+
 //------------//
 // Conversion //
 //------------//
@@ -203,6 +242,12 @@ mod tests {
         }
     }
 
+    test_utils::dot_product::test_dot_product!(
+        (f16x4, f16x4) => f32x4,
+        0x8f36d1c4b529a70e,
+        Some(test_neon().expect("FHM dot-product test requires Neon with fhm"))
+    );
+
     // Ops
     test_utils::ops::test_add!(f32x4, 0xcd7a8fea9a3fb727, test_neon());
     test_utils::ops::test_sub!(f32x4, 0x3f6562c94c923238, test_neon());

diskann-wide/src/arch/aarch64/mod.rs

@@ -115,7 +115,7 @@ pub(super) enum LevelInner {
 /////////////
 
 cfg_if::cfg_if! {
-    if #[cfg(all(target_feature = "neon", target_feature = "dotprod"))] {
+    if #[cfg(all(target_feature = "neon", target_feature = "dotprod", target_feature = "fhm"))] {
         pub type Current = Neon;
 
         pub const fn current() -> Current {
@@ -226,6 +226,7 @@ impl Neon {
     ///
     /// * `neon`
     /// * `dotprod`
+    /// * `fhm`
     pub const unsafe fn new() -> Self {
         Self(Hidden)
     }
@@ -234,7 +235,11 @@ impl Neon {
     pub fn new_checked() -> Option<Self> {
         // This check here ensures that if we ever switch to dynamically dispatching to
         // `Neon` that we do not forget to update `new_checked`.
-        if cfg!(all(target_feature = "neon", target_feature = "dotprod")) {
+        if cfg!(all(
+            target_feature = "neon",
+            target_feature = "dotprod",
+            target_feature = "fhm"
+        )) {
             // SAFETY: The compile-time feature check above ensures we do not accidentally
             // return an unsafe instance of `Self`.
             Some(unsafe { Self::new() })

diskann-wide/src/test_utils/dot_product.rs

@@ -3,7 +3,12 @@
  * Licensed under the MIT license.
  */
 
-use crate::{Const, SIMDDotProduct, SIMDVector, SupportedLaneCount, reference::ReferenceScalarOps};
+use half::f16;
+
+use crate::{
+    Const, SIMDDotProduct, SIMDVector, SupportedLaneCount, reference::ReferenceScalarOps,
+    reference::cast_f16_to_f32, test_utils::common::ScalarTraits,
+};
 
 /////////////////////////
 // Expected Reductions //
@@ -30,6 +35,13 @@ pub(crate) trait ExpectedDot<To, Left, Right, const EVERY: usize> {
     }
 }
 
+// Each half-precision pair is promoted and fused with one f32 accumulator lane.
+impl ExpectedDot<f32, f16, f16, 1> for DotSchema {
+    fn expected_dot_impl(accumulator: f32, left: &[f16; 1], right: &[f16; 1]) -> f32 {
+        cast_f16_to_f32(left[0]).expected_fma_(cast_f16_to_f32(right[0]), accumulator)
+    }
+}
+
 // Elements are promoted to `i32`, pairwise multiplied, summed, then accumulated.
 impl ExpectedDot<i32, i16, i16, 2> for DotSchema {
     fn expected_dot_impl(accumulator: i32, left: &[i16; 2], right: &[i16; 2]) -> i32 {
@@ -135,7 +147,7 @@ pub(crate) fn test_dot_product_impl<
     Left: SIMDVector<Arch = To::Arch, ConstLanes = Const<FROM_N>>,
     Right: SIMDVector<Arch = To::Arch, ConstLanes = Const<FROM_N>>,
     DotSchema: ExpectedDot<To::Scalar, Left::Scalar, Right::Scalar, EVERY>,
-    To::Scalar: PartialEq + Copy,
+    To::Scalar: ScalarTraits,
     Left::Scalar: Copy,
     Right::Scalar: Copy,
 {
@@ -157,7 +169,7 @@ pub(crate) fn test_dot_product_impl<
     let iter = std::iter::zip(left.chunks(EVERY), right.chunks(EVERY)).enumerate();
     for (i, (l, r)) in iter {
         let expected = DotSchema::expected_dot(accumulator[i], l, r);
-        assert_eq!(expected, result[i])
+        assert!(expected.exact_eq(result[i]));
     }
 }