From f1bdd84d46dda24b43bb7f0aaec59f980ac4cc77 Mon Sep 17 00:00:00 2001 From: UtsavGoelArm Date: Tue, 6 Oct 2026 13:34:53 +0100 Subject: [PATCH 1/6] Inner Product Kernel for f16 now uses arm fmlal instruction --- .cargo/config.toml | 4 ++-- diskann-vector/src/distance/simd.rs | 17 ++++++++++++----- 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/.cargo/config.toml b/.cargo/config.toml index de3f9a9172..5d4c2263e9 100644 --- a/.cargo/config.toml +++ b/.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" diff --git a/diskann-vector/src/distance/simd.rs b/diskann-vector/src/distance/simd.rs index d49d66958a..4e62db687c 100644 --- a/diskann-vector/src/distance/simd.rs +++ b/diskann-vector/src/distance/simd.rs @@ -1804,12 +1804,19 @@ impl SIMDSchema for IP { y: Self::Right, acc: Self::Accumulator, ) -> Self::Accumulator { - diskann_wide::alias!(f32s = ::f32x4); - - let x: f32s = x.into(); - let y: f32s = y.into(); + let mut result = acc.to_underlying(); + // SAFETY: The target is assumed to support FHM + unsafe { + std::arch::asm!( + "fmlal {result:v}.4s, {x:v}.4h, {y:v}.4h", + result = inout(vreg) result, + x = in(vreg) x.to_underlying(), + y = in(vreg) y.to_underlying(), + options(nomem, nostack), + ); + } - x.mul_add_simd(y, acc) + Self::Accumulator::from_underlying(acc.arch(), result) } #[inline(always)] From df393d896f65d85e283030c08c4b41560bb5eeea Mon Sep 17 00:00:00 2001 From: UtsavGoelArm Date: Tue, 6 Oct 2026 15:08:43 +0100 Subject: [PATCH 2/6] implement fmlal as a dotproduct operation in diskann-wide and use in kernel --- diskann-vector/src/distance/simd.rs | 14 +------- diskann-wide/src/arch/aarch64/f32x4_.rs | 45 +++++++++++++++++++++++-- diskann-wide/src/arch/aarch64/mod.rs | 5 +-- 3 files changed, 47 insertions(+), 17 deletions(-) diff --git a/diskann-vector/src/distance/simd.rs b/diskann-vector/src/distance/simd.rs index 4e62db687c..6c4d22dd3d 100644 --- a/diskann-vector/src/distance/simd.rs +++ b/diskann-vector/src/distance/simd.rs @@ -1804,19 +1804,7 @@ impl SIMDSchema for IP { y: Self::Right, acc: Self::Accumulator, ) -> Self::Accumulator { - let mut result = acc.to_underlying(); - // SAFETY: The target is assumed to support FHM - unsafe { - std::arch::asm!( - "fmlal {result:v}.4s, {x:v}.4h, {y:v}.4h", - result = inout(vreg) result, - x = in(vreg) x.to_underlying(), - y = in(vreg) y.to_underlying(), - options(nomem, nostack), - ); - } - - Self::Accumulator::from_underlying(acc.arch(), result) + acc.dot_simd(x, y) } #[inline(always)] diff --git a/diskann-wide/src/arch/aarch64/f32x4_.rs b/diskann-wide/src/arch/aarch64/f32x4_.rs index 51d86a069b..63b8f4cf5d 100644 --- a/diskann-wide/src/arch/aarch64/f32x4_.rs +++ b/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,37 @@ impl SIMDSelect for mask32x4 { } } +impl SIMDDotProduct 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(nomem, nostack) + ); + } + acc + } + // SAFETY: Calling above fmlal function enabled with `fhm` feature. + Self(unsafe { fmlal(self.0, left.0, right.0) }) + } + } +} + //------------// // Conversion // //------------// @@ -203,6 +234,16 @@ mod tests { } } + #[test] + fn test_f16_dot_product() { + if let Some(arch) = test_neon() { + let left = f16x4::from_array(arch, [1.0, 2.0, -3.0, 4.0].map(f16::from_f32)); + let right = f16x4::from_array(arch, [2.0, -1.0, 0.5, 0.25].map(f16::from_f32)); + let acc = f32x4::from_array(arch, [0.5, 1.0, 2.0, -3.0]); + assert_eq!(acc.dot_simd(left, right).to_array(), [2.5, -1.0, 0.5, -2.0]); + } + } + // Ops test_utils::ops::test_add!(f32x4, 0xcd7a8fea9a3fb727, test_neon()); test_utils::ops::test_sub!(f32x4, 0x3f6562c94c923238, test_neon()); diff --git a/diskann-wide/src/arch/aarch64/mod.rs b/diskann-wide/src/arch/aarch64/mod.rs index 945fe59e15..8aff762289 100644 --- a/diskann-wide/src/arch/aarch64/mod.rs +++ b/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,7 @@ impl Neon { pub fn new_checked() -> Option { // 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() }) From 0665c98ae4e9360ad681e83aabe38fcbe19a8209 Mon Sep 17 00:00:00 2001 From: UtsavGoelArm Date: Tue, 20 Oct 2026 15:56:25 +0100 Subject: [PATCH 3/6] apply cargo fmt for consistent formatting --- diskann-wide/src/arch/aarch64/f32x4_.rs | 20 ++++++++++++++------ diskann-wide/src/arch/aarch64/mod.rs | 6 +++++- 2 files changed, 19 insertions(+), 7 deletions(-) diff --git a/diskann-wide/src/arch/aarch64/f32x4_.rs b/diskann-wide/src/arch/aarch64/f32x4_.rs index 63b8f4cf5d..2835648615 100644 --- a/diskann-wide/src/arch/aarch64/f32x4_.rs +++ b/diskann-wide/src/arch/aarch64/f32x4_.rs @@ -6,8 +6,8 @@ use half::f16; use crate::{ - AsSIMD, Emulated, SIMDAbs, SIMDCast, SIMDDotProduct, 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 @@ -118,12 +118,20 @@ impl SIMDDotProduct for f32x4 { 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]) - })) + 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 { + 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!( diff --git a/diskann-wide/src/arch/aarch64/mod.rs b/diskann-wide/src/arch/aarch64/mod.rs index 8aff762289..55b7af0d01 100644 --- a/diskann-wide/src/arch/aarch64/mod.rs +++ b/diskann-wide/src/arch/aarch64/mod.rs @@ -235,7 +235,11 @@ impl Neon { pub fn new_checked() -> Option { // 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", target_feature = "fhm")) { + 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() }) From e347629d1641b46357fd9e827d716168713c7bbb Mon Sep 17 00:00:00 2001 From: Patrick Fox Date: Thu, 29 Oct 2026 17:03:37 +0000 Subject: [PATCH 4/6] Make +fml global for neon, use test_dot_product! --- diskann-wide/compile-aarch64-on-x86.sh | 2 +- diskann-wide/src/arch/aarch64/f32x4_.rs | 16 ++++++---------- diskann-wide/src/test_utils/dot_product.rs | 21 +++++++++++++++++++-- 3 files changed, 26 insertions(+), 13 deletions(-) diff --git a/diskann-wide/compile-aarch64-on-x86.sh b/diskann-wide/compile-aarch64-on-x86.sh index aaa0104254..ab049b9ac0 100755 --- a/diskann-wide/compile-aarch64-on-x86.sh +++ b/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 diff --git a/diskann-wide/src/arch/aarch64/f32x4_.rs b/diskann-wide/src/arch/aarch64/f32x4_.rs index 2835648615..b78105e3ab 100644 --- a/diskann-wide/src/arch/aarch64/f32x4_.rs +++ b/diskann-wide/src/arch/aarch64/f32x4_.rs @@ -139,7 +139,7 @@ impl SIMDDotProduct for f32x4 { acc = inout(vreg) acc, left = in(vreg) left, right = in(vreg) right, - options(nomem, nostack) + options(pure, nomem, nostack) ); } acc @@ -242,15 +242,11 @@ mod tests { } } - #[test] - fn test_f16_dot_product() { - if let Some(arch) = test_neon() { - let left = f16x4::from_array(arch, [1.0, 2.0, -3.0, 4.0].map(f16::from_f32)); - let right = f16x4::from_array(arch, [2.0, -1.0, 0.5, 0.25].map(f16::from_f32)); - let acc = f32x4::from_array(arch, [0.5, 1.0, 2.0, -3.0]); - assert_eq!(acc.dot_simd(left, right).to_array(), [2.5, -1.0, 0.5, -2.0]); - } - } + 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()); diff --git a/diskann-wide/src/test_utils/dot_product.rs b/diskann-wide/src/test_utils/dot_product.rs index f0424a0237..4ee9f4e1f8 100644 --- a/diskann-wide/src/test_utils/dot_product.rs +++ b/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, +}; ///////////////////////// // Expected Reductions // @@ -30,6 +35,13 @@ pub(crate) trait ExpectedDot { } } +// Each half-precision pair is promoted and fused with one f32 accumulator lane. +impl ExpectedDot 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 for DotSchema { fn expected_dot_impl(accumulator: i32, left: &[i16; 2], right: &[i16; 2]) -> i32 { @@ -157,7 +169,12 @@ 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]) + // NaNs compare unequal to themselves, but both paths should produce NaN. + #[allow(clippy::eq_op)] + if expected != expected && result[i] != result[i] { + continue; + } + assert_eq!(expected, result[i]); } } From 955f2cd2c0240319b2148a086033f76edbb956f3 Mon Sep 17 00:00:00 2001 From: UtsavGoelArm Date: Fri, 30 Oct 2026 11:03:04 +0000 Subject: [PATCH 5/6] use ScalarTraits::exact_eq for NaN safety in testing --- diskann-wide/src/test_utils/dot_product.rs | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/diskann-wide/src/test_utils/dot_product.rs b/diskann-wide/src/test_utils/dot_product.rs index 4ee9f4e1f8..cab63f1017 100644 --- a/diskann-wide/src/test_utils/dot_product.rs +++ b/diskann-wide/src/test_utils/dot_product.rs @@ -8,6 +8,7 @@ use half::f16; use crate::{ Const, SIMDDotProduct, SIMDVector, SupportedLaneCount, reference::ReferenceScalarOps, reference::cast_f16_to_f32, + test_utils::common::ScalarTraits, }; ///////////////////////// @@ -147,7 +148,7 @@ pub(crate) fn test_dot_product_impl< Left: SIMDVector>, Right: SIMDVector>, DotSchema: ExpectedDot, - To::Scalar: PartialEq + Copy, + To::Scalar: ScalarTraits, Left::Scalar: Copy, Right::Scalar: Copy, { @@ -169,12 +170,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); - // NaNs compare unequal to themselves, but both paths should produce NaN. - #[allow(clippy::eq_op)] - if expected != expected && result[i] != result[i] { - continue; - } - assert_eq!(expected, result[i]); + assert!(expected.exact_eq(result[i])); } } From 816fbe96ff3f1832da523f7953dcc011c5edcb91 Mon Sep 17 00:00:00 2001 From: UtsavGoelArm Date: Mon, 2 Nov 2026 09:24:01 +0000 Subject: [PATCH 6/6] resolve formatting issues --- diskann-wide/src/test_utils/dot_product.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/diskann-wide/src/test_utils/dot_product.rs b/diskann-wide/src/test_utils/dot_product.rs index cab63f1017..716946ce48 100644 --- a/diskann-wide/src/test_utils/dot_product.rs +++ b/diskann-wide/src/test_utils/dot_product.rs @@ -7,8 +7,7 @@ use half::f16; use crate::{ Const, SIMDDotProduct, SIMDVector, SupportedLaneCount, reference::ReferenceScalarOps, - reference::cast_f16_to_f32, - test_utils::common::ScalarTraits, + reference::cast_f16_to_f32, test_utils::common::ScalarTraits, }; /////////////////////////