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..6c4d22dd3d 100644 --- a/diskann-vector/src/distance/simd.rs +++ b/diskann-vector/src/distance/simd.rs @@ -1804,12 +1804,7 @@ 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(); - - x.mul_add_simd(y, acc) + acc.dot_simd(x, y) } #[inline(always)] 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 51d86a069b..b78105e3ab 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,45 @@ 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(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()); diff --git a/diskann-wide/src/arch/aarch64/mod.rs b/diskann-wide/src/arch/aarch64/mod.rs index 945fe59e15..55b7af0d01 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,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")) { + 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() }) diff --git a/diskann-wide/src/test_utils/dot_product.rs b/diskann-wide/src/test_utils/dot_product.rs index f0424a0237..716946ce48 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, test_utils::common::ScalarTraits, +}; ///////////////////////// // 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 { @@ -135,7 +147,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, { @@ -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])); } }