diff --git a/diskann-quantization/src/bits/distances.rs b/diskann-quantization/src/bits/distances.rs index c466fff390..38f338d129 100644 --- a/diskann-quantization/src/bits/distances.rs +++ b/diskann-quantization/src/bits/distances.rs @@ -114,6 +114,9 @@ use diskann_wide::{ SIMDCast, SIMDDotProduct, SIMDMulAdd, SIMDReinterpret, SIMDSumTree, SIMDVector, }; +#[cfg(target_arch = "aarch64")] +use diskann_wide::{SIMDDotProduct, SIMDSumTree, SIMDVector}; + use super::{Binary, BitSlice, BitTranspose, Dense, Representation, Unsigned}; use crate::distances::{Hamming, InnerProduct, MV, MathematicalResult, SquaredL2, check_lengths}; @@ -1565,6 +1568,103 @@ impl Target2, USlice<'_, } } +#[cfg(target_arch = "aarch64")] +impl + Target2< + diskann_wide::arch::aarch64::Neon, + MathematicalResult, + USlice<'_, 4>, + USlice<'_, 4>, + > for InnerProduct +{ + #[inline(always)] + fn run( + self, + arch: diskann_wide::arch::aarch64::Neon, + x: USlice<'_, 4>, + y: USlice<'_, 4>, + ) -> MathematicalResult { + let len = check_lengths!(x, y)?; + + diskann_wide::alias!(u8s = ::u8x16); + diskann_wide::alias!(u32s = ::u32x4); + + let px_u8: *const u8 = x.as_ptr().cast(); + let py_u8: *const u8 = y.as_ptr().cast(); + + let mut i = 0; + let mut s: u32 = 0; + + // number of bytes over the underlying slice + let bytes = len / 2; + if i < bytes { + let mut s0 = u32s::default(arch); + let mut s1 = u32s::default(arch); + let mask = u8s::splat(arch, 0x0f); + while i + 16 <= bytes { + // SAFETY: load simd loads 16 bytes from offset `i` + // we have already verified that i + 16 <= bytes + let x_vec = unsafe { u8s::load_simd(arch, px_u8.add(i)) }; + // SAFETY: same logic applies for y and same conditions hold + // since the lengths are element sizes of x and y are equal. + let y_vec = unsafe { u8s::load_simd(arch, py_u8.add(i)) }; + + // compute dot product for lower 4 bits + // each set of 4 results is reduced to one lane + let first_x: u8s = x_vec & mask; + let first_y: u8s = y_vec & mask; + s0 = s0.dot_simd(first_x, first_y); + + // repeat for upper 4 bits + let second_x: u8s = (x_vec >> 4) & mask; + let second_y: u8s = (y_vec >> 4) & mask; + s1 = s1.dot_simd(second_x, second_y); + // repeat for next block + i += 16; + } + + let remaining_bytes = len / 2 - i; + + if remaining_bytes > 0 { + let remaining_vec1 = remaining_bytes.min(16); + // SAFETY: up to `remaining_bytes` can be loaded from the offset `i` + // since the floor division ensures that we only read bytes that are fully + // packed with elements from the slice. + let x_vec = unsafe { u8s::load_simd_first(arch, px_u8.add(i), remaining_vec1) }; + // SAFETY: same logic applies for y and same conditions hold + // since the lengths are element sizes of x and y are equal. + let y_vec = unsafe { u8s::load_simd_first(arch, py_u8.add(i), remaining_vec1) }; + + let first_x: u8s = x_vec & mask; + let first_y: u8s = y_vec & mask; + s0 = s0.dot_simd(first_x, first_y); + + // compute dot product for upper 4 bits, result is stored as 32x4 + let second_x: u8s = (x_vec >> 4) & mask; + let second_y: u8s = (y_vec >> 4) & mask; + s1 = s1.dot_simd(second_x, second_y); + i += remaining_bytes; + } + s = (s0 + s1).sum_tree(); + } + // Convert bytes to nibble indexes. + i *= 2; + + // Deal with the remainder the slow way (at most 1 element). + debug_assert!(len - i <= 1); + + if i != len { + // SAFETY: `i` is guaranteed to be less than `x.len()`. + let ix = unsafe { x.get_unchecked(i) } as i32; + // SAFETY: `i` is guaranteed to be less than `y.len()`. + let iy = unsafe { y.get_unchecked(i) } as i32; + s += (ix * iy) as u32; + } + + Ok(MV::new(s)) + } +} + /// Compute the inner product between bitvectors `x` and `y`. /// /// Returns an error if the arguments have different lengths. @@ -2108,7 +2208,6 @@ retarget!( 7, 6, 5, - 4, 3, 2, (8, 4), diff --git a/diskann-quantization/src/spherical/__codegen/aarch64.rs b/diskann-quantization/src/spherical/__codegen/aarch64.rs new file mode 100644 index 0000000000..5fe5599223 --- /dev/null +++ b/diskann-quantization/src/spherical/__codegen/aarch64.rs @@ -0,0 +1,51 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Explicitly instantiate the AArch64 Neon spherical inner-product paths. +use diskann_wide::arch::aarch64::Neon; + +use crate::{ + alloc::{AllocatorError, GlobalAllocator}, + spherical::{ + iface::{AsData, AsQuery, DistanceComputer, Reify}, + vectors, + }, +}; + +/// Instantiate the Neon inner-product implementation for +/// `USlice<'_, 4> × USlice<'_, 4>` in the data-to-data path. +#[inline(never)] +pub fn fourbit_neon_ip_data_data( + arch: Neon, + shift: &[f32], + dim: usize, +) -> Result { + let reify = Reify::<_, _, AsData<4>, AsData<4>>::new( + vectors::CompensatedIP::new(shift, dim), + dim, + arch, + ); + + DistanceComputer::new(reify, GlobalAllocator) +} + +/// Instantiate the Neon inner-product implementation for the four-bit +/// query-to-data path. +/// +/// `dispatch_map!(4, AsQuery<4>, Neon);` +#[inline(never)] +pub fn fourbit_neon_ip_query_data( + arch: Neon, + shift: &[f32], + dim: usize, +) -> Result { + let reify = Reify::<_, _, AsQuery<4>, AsData<4>>::new( + vectors::CompensatedIP::new(shift, dim), + dim, + arch, + ); + + DistanceComputer::new(reify, GlobalAllocator) +} diff --git a/diskann-quantization/src/spherical/__codegen/mod.rs b/diskann-quantization/src/spherical/__codegen/mod.rs index 5878a4a6d3..3c45d7c898 100644 --- a/diskann-quantization/src/spherical/__codegen/mod.rs +++ b/diskann-quantization/src/spherical/__codegen/mod.rs @@ -9,3 +9,6 @@ #[cfg(target_arch = "x86_64")] pub mod x86_64; + +#[cfg(target_arch = "aarch64")] +pub mod aarch64; diff --git a/diskann-quantization/src/spherical/iface.rs b/diskann-quantization/src/spherical/iface.rs index a97fbbd9b5..1f6a153ee6 100644 --- a/diskann-quantization/src/spherical/iface.rs +++ b/diskann-quantization/src/spherical/iface.rs @@ -1435,12 +1435,12 @@ cfg_if::cfg_if! { dispatch_map!(1, AsData<1>, Neon, downcast); dispatch_map!(2, AsData<2>, Neon, downcast); - dispatch_map!(4, AsData<4>, Neon, downcast); + dispatch_map!(4, AsData<4>, Neon); dispatch_map!(8, AsData<8>, Neon, downcast); dispatch_map!(1, AsQuery<4, bits::BitTranspose>, Neon, downcast); dispatch_map!(2, AsQuery<2>, Neon, downcast); - dispatch_map!(4, AsQuery<4>, Neon, downcast); + dispatch_map!(4, AsQuery<4>, Neon); dispatch_map!(8, AsQuery<8>, Neon, downcast); } }