diff --git a/diskann-quantization/src/bits/slice.rs b/diskann-quantization/src/bits/slice.rs index dcfaeff7c2..49d0996a13 100644 --- a/diskann-quantization/src/bits/slice.rs +++ b/diskann-quantization/src/bits/slice.rs @@ -383,6 +383,13 @@ unsafe impl PermutationStrategy<4> for BitTranspose { // Errors // //////////// +#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)] +#[error("decoded output has {output} elements, expected {input}")] +pub(crate) struct DecodeLengthMismatch { + input: usize, + output: usize, +} + #[derive(Debug, Error, Clone, Copy)] #[error("input span has length {got} bytes but expected {expected}")] pub struct ConstructionError { @@ -840,6 +847,59 @@ pub type PolyBitSlice = pub type BoxedBitSlice = PolyBitSlice; +impl BitSlice<'_, NBITS, Unsigned, Dense> +where + Unsigned: Representation, +{ + /// Decode and convert directly into the destination, ignoring unused trailing bits. + /// + /// # Errors + /// + /// Returns [`DecodeLengthMismatch`] without writing or calling `convert` if the + /// output length differs from the logical input length. + #[inline] + pub(crate) fn decode_into( + &self, + output: &mut [T], + mut convert: impl FnMut(u8) -> T, + ) -> Result<(), DecodeLengthMismatch> { + if output.len() != self.len() { + return Err(DecodeLengthMismatch { + input: self.len(), + output: output.len(), + }); + } + let packed = self.as_slice(); + let divisor = 1usize << NBITS.trailing_zeros(); + let codes_per_block = 8 / divisor; + let bytes_per_block = NBITS / divisor; + let blocks = output.len() / codes_per_block; + let (output, tail) = output.split_at_mut(blocks * codes_per_block); + let (packed, rest) = packed.split_at(blocks * bytes_per_block); + let mask = (1u64 << NBITS) - 1; + // Use the smallest byte-aligned group: 8 / gcd(NBITS, 8) codes. + for (block, bytes) in output + .chunks_exact_mut(codes_per_block) + .zip(packed.chunks_exact(bytes_per_block)) + { + let mut word = [0u8; 8]; + word[..bytes_per_block].copy_from_slice(bytes); + let word = u64::from_le_bytes(word); + for (i, out) in block.iter_mut().enumerate() { + *out = convert(((word >> (i * NBITS)) & mask) as u8); + } + } + let word = rest + .iter() + .enumerate() + .fold(0u64, |word, (i, &byte)| word | (u64::from(byte) << (8 * i))); + for (i, out) in tail.iter_mut().enumerate() { + *out = convert(((word >> (i * NBITS)) & mask) as u8); + } + Ok(()) + } +} + /////////////////////////////// // Special Cased Conversions // /////////////////////////////// @@ -947,6 +1007,93 @@ mod tests { use super::*; use crate::{bits::Static, test_util::AlwaysFails}; + fn check_bulk_decoding() + where + Unsigned: Representation, + { + for len in [0, 1, 2, 3, 7, 8, 9, 15, 16, 17, 255, 256, 257, 513] { + for offset in [0, 1, 7] { + let byte_count = BitSlice::::bytes_for(len); + let mut bytes = vec![0xff; offset + byte_count]; + let bytes = &mut bytes[offset..]; + let mut codes = MutBitSlice::::new(bytes, len).unwrap(); + for i in 0..len { + codes + .set(i, ((i * 37 + i / 2) & ((1 << NBITS) - 1)) as i64) + .unwrap(); + } + let expected: Vec<_> = (0..len).map(|i| codes.get(i).unwrap()).collect(); + let mut guarded = vec![i64::MIN; len + 2]; + let mut calls = 0; + codes + .reborrow() + .decode_into(&mut guarded[1..=len], |value| { + assert_eq!(i64::from(value), expected[calls]); + calls += 1; + i64::from(value) + }) + .unwrap(); + assert_eq!(calls, len); + assert_eq!(&guarded[1..=len], expected); + assert_eq!(guarded[0], i64::MIN); + assert_eq!(guarded[len + 1], i64::MIN); + } + } + } + + #[test] + fn test_decode_into_preserves_order_tails_and_bounds() { + check_bulk_decoding::<1>(); + check_bulk_decoding::<2>(); + check_bulk_decoding::<3>(); + check_bulk_decoding::<4>(); + check_bulk_decoding::<5>(); + check_bulk_decoding::<6>(); + check_bulk_decoding::<7>(); + check_bulk_decoding::<8>(); + } + + #[test] + fn test_decode_into_applies_conversion() { + let codes = BitSlice::<8, Unsigned>::new(&[0, 128, 255][..], 3).unwrap(); + let mut output = [0i16; 3]; + codes + .decode_into(&mut output, |value| i16::from(value) * 3) + .unwrap(); + assert_eq!(output, [0, 384, 765]); + let codes = BoxedBitSlice::<4, Unsigned>::new_boxed(2); + let mut output = [1u8; 2]; + codes + .reborrow() + .decode_into(&mut output, |value| value) + .unwrap(); + assert_eq!(output, [0, 0]); + } + + #[test] + fn test_decode_into_length_mismatch_does_not_write_or_convert() { + let codes = BitSlice::<4, Unsigned>::new(&[0x21u8][..], 2).unwrap(); + for len in [0, 1, 3] { + let mut output = vec![99; len]; + let mut calls = 0; + let error = codes + .decode_into(&mut output, |value| { + calls += 1; + value + }) + .unwrap_err(); + assert_eq!( + error, + DecodeLengthMismatch { + input: 2, + output: len + } + ); + assert_eq!(calls, 0); + assert_eq!(output, vec![99; len]); + } + } + //////////// // Errors // //////////// diff --git a/diskann-quantization/src/matrix_kernels/maxsim/mod.rs b/diskann-quantization/src/matrix_kernels/maxsim/mod.rs index 412f9d98c5..e320afcfd5 100644 --- a/diskann-quantization/src/matrix_kernels/maxsim/mod.rs +++ b/diskann-quantization/src/matrix_kernels/maxsim/mod.rs @@ -19,6 +19,11 @@ pub(crate) mod packed_f32_x_unpacked_f16; pub(crate) mod packed_f32_x_unpacked_f32; pub(crate) mod packed_i8_x_unpacked_i8; +#[cfg_attr( + not(test), + expect(dead_code, reason = "TODO: will integrate in the next PR") +)] +pub(crate) mod packed_minmax_x_unpacked_minmax; #[cfg(test)] mod test; diff --git a/diskann-quantization/src/matrix_kernels/maxsim/packed_minmax_x_unpacked_minmax.rs b/diskann-quantization/src/matrix_kernels/maxsim/packed_minmax_x_unpacked_minmax.rs new file mode 100644 index 0000000000..d0173c5e41 --- /dev/null +++ b/diskann-quantization/src/matrix_kernels/maxsim/packed_minmax_x_unpacked_minmax.rs @@ -0,0 +1,2369 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! MaxSim for packed MinMax A and tiled MinMax B. +//! +//! B tiles are decoded once and reused across A panels. Compensation precedes +//! reduction because coefficients differ per B vector. Codes are zero-padded to PACK; +//! compensation and accumulator limits use logical K. +//! +//! Outputs are unnegated maximum inner products. NaN candidates are ignored; +//! empty B or all-NaN comparisons leave negative infinity. + +use std::{ + num::NonZeroUsize, + ops::{Add, Mul}, +}; + +use diskann_wide::{ + Const, SIMDCast, SIMDDotProduct, SIMDMinMax, SIMDReinterpret, SIMDVector, SupportedLaneCount, + arch::{Architecture, Scalar, Target}, +}; + +use crate::{ + bits::{Representation, Unsigned}, + matrix_kernels::{ + Cache, + blocks::{packed, unpacked}, + bounds, driver, + num::{Bytes, DimK, Elements}, + ptr::Slice, + util::LoadStore, + }, + minmax::MinMaxMeta, + multi_vector::{BlockTransposed, MatRef, Overflow}, +}; + +use super::packed_f32_x_unpacked_f32::Params; + +diskann_wide::alias!(i8x32 = i8x32); +diskann_wide::alias!(i8x64 = i8x64); +diskann_wide::alias!(i16x16 = i16x16); +diskann_wide::alias!(i16x32 = i16x32); +diskann_wide::alias!(i32x16 = i32x16); +diskann_wide::alias!(i32x8 = i32x8); +diskann_wide::alias!(u8x32 = u8x32); +diskann_wide::alias!(u8x64 = u8x64); +diskann_wide::alias!(u32x4 = u32x4); +diskann_wide::alias!(u32x8 = u32x8); +diskann_wide::alias!(u32x16 = u32x16); + +// Signed backends determine the shared accumulation limit. +const fn max_k() -> usize { + assert!(A_BITS > 0 && A_BITS <= 8 && B_BITS > 0 && B_BITS <= 8); + (i32::MAX as usize) / (((1 << A_BITS) - 1) * ((1 << B_BITS) - 1)) +} + +#[derive(Debug, Clone, Copy, Default)] +pub(crate) struct Coefficients { + a: T, + b: T, + n: T, +} + +#[derive(Debug)] +pub(crate) struct PackedA +where + Unsigned: Representation, +{ + codes: BlockTransposed, + meta: Vec>, +} + +impl PackedA +where + T: Copy + Default + From, + Unsigned: Representation, +{ + /// Zero-pad partial rows and K packs. + /// + /// # Errors + /// + /// Returns [`Overflow`] if the packed data or coefficient allocation size overflows. + #[expect(clippy::expect_used)] + pub(crate) fn new(a: MatRef<'_, MinMaxMeta>) -> Result { + let k = a.repr().intrinsic_dim(); + let mut codes = BlockTransposed::try_new(a.num_vectors(), k)?; + let panels = codes.num_blocks(); + Overflow::check::>(panels, 1)?; + let mut meta = vec![ + Coefficients { + a: [0.0; MR], + b: [0.0; MR], + n: [0.0; MR], + }; + panels + ]; + for (i, row) in a.rows().enumerate() { + let m = row.meta(); + let coefficients = &mut meta[i / MR]; + coefficients.a[i % MR] = m.a; + coefficients.b[i % MR] = m.b; + coefficients.n[i % MR] = m.n; + let v = row.vector(); + let mut dst = codes + .get_row_mut(i) + .expect("A and its packing have equal row counts"); + for j in 0..k { + // SAFETY: MinMaxMeta exposes exactly k codes for every row. + dst.set(j, T::from(unsafe { v.get_unchecked(j) } as u8)); + } + } + Ok(Self { codes, meta }) + } + + fn as_view(&self) -> Option> { + Some(AView { + codes: packed::View::from_block_transposed(self.codes.as_view())?, + meta: &self.meta, + k: DimK::new(NonZeroUsize::new(self.codes.ncols())?), + }) + } + + fn panel_bytes(&self) -> Bytes { + let codes = Elements::::new(MR * self.codes.padded_ncols()).bytes(); + Bytes::new(codes.value() + size_of::>()) + } +} + +#[derive(Clone, Copy)] +struct AView<'a, T, const MR: usize, const PACK: usize> { + codes: packed::View<'a, T, MR, PACK>, + meta: &'a [Coefficients<[f32; MR]>], + k: DimK, +} + +impl AView<'_, T, MR, PACK> { + #[inline] + fn visit_panels(&self, mut visitor: impl FnMut(APanel<'_, T, MR, PACK>, usize)) { + // SAFETY: PackedA constructs codes and coefficients with the same panel count and K. + unsafe { + self.codes.visit_panels(self.k, |codes, i| { + visitor( + APanel { + codes, + meta: &self.meta[i], + }, + i, + ); + }); + } + } +} + +#[derive(Clone, Copy)] +struct APanel<'a, T, const MR: usize, const PACK: usize> { + codes: packed::Panel<'a, T, MR, PACK>, + meta: &'a Coefficients<[f32; MR]>, +} + +//--------// +// Driver // +//--------// + +/// Each drive overwrites C; results do not accumulate across calls. +pub(crate) struct Driver< + 'a, + W: ExtraWide, + const B_BITS: usize, + const MR: usize, + const NR: usize, + const PACK: usize, +> where + Unsigned: Representation, +{ + arch: W, + a: AView<'a, W::AElem, MR, PACK>, + b: MatRef<'a, MinMaxMeta>, + c: &'a mut [f32], + b_rows_per_tile: NonZeroUsize, +} + +impl< + 'a, + W: ExtraWide, + const B_BITS: usize, + const MR: usize, + const NR: usize, + const PACK: usize, +> Driver<'a, W, B_BITS, MR, NR, PACK> +where + Unsigned: Representation, +{ + /// # Panics + /// + /// Panics if K exceeds the accumulator limit for these bit widths. + /// Valid integer dot products may still round when converted to f32. + /// + /// # Safety + /// + /// A must have nonzero row count and dimension, B must have the same dimension, + /// and C must have exactly A's logical row count. + #[expect(clippy::expect_used)] + pub(crate) unsafe fn new( + arch: W, + a: &'a PackedA, + b: MatRef<'a, MinMaxMeta>, + c: &'a mut [f32], + cache: Cache, + ) -> Self + where + Unsigned: Representation, + Self: driver::Drive, + { + const { + assert!( + B_BITS <= W::MAX_B_BITS, + "B bit width exceeds the backend's supported range" + ) + }; + let rows = a.codes.nrows(); + let a_bytes = a.panel_bytes(); + let a = a.as_view().expect("A and K must be nonempty"); + let dim = a.k.value().get(); + bounds::check_eq!(bounds::Bound::new(b.repr().intrinsic_dim()), a.k); + bounds::check_eq!(bounds::Bound::new(c.len()), rows); + let max = const { max_k::() }; + assert!( + dim <= max, + "contraction dimension {dim} exceeds the MinMax{A_BITS} x MinMax{B_BITS} limit of {max}", + ); + let b_rows_per_tile = Params::new( + cache, + a_bytes, + BScratch::::row_bytes::(a.k), + NR, + ) + .b_cols_in_l1; + Self { + arch, + a, + b, + c, + b_rows_per_tile, + } + } +} + +impl driver::Drive + for Driver<'_, W, B_BITS, MR, NR, PACK> +where + W: ExtraWide + LoadStore, + Unsigned: Representation, + for<'a> PanelKernel<'a, W, MR, NR, PACK>: driver::PanelKernel, +{ + fn drive(&mut self) { + self.arch.run( + #[inline] + || { + self.c.fill(f32::NEG_INFINITY); + let mut scratch = BScratch::::default(); + let mut start = 0; + while let Some(remaining) = NonZeroUsize::new(self.b.num_vectors() - start) { + let count = remaining.min(self.b_rows_per_tile); + // SAFETY: The driver contract fixes B's K; remaining/count bound this tile. + let b = unsafe { + scratch.prepare::( + self.b, + start, + count, + self.a.k, + W::convert_b, + ) + }; + self.a.visit_panels(|a, panel| { + let row = panel * MR; + let end = (row + MR).min(self.c.len()); + let c = self.arch.load(&self.c[row..end]); + // SAFETY: A and B share K by the driver contract. + let kernel = + unsafe { PanelKernel::::new(self.arch, a, b, c) }; + let c = self.arch.run(kernel); + self.arch.store(c, &mut self.c[row..end]); + }); + start += count.get(); + } + }, + ); + } +} + +#[derive(Default)] +struct BScratch { + codes: Vec, + meta: Vec>, +} + +/// Codes and coefficients share row order; code rows are zero-padded to PACK. +#[derive(Clone, Copy)] +struct BTile<'a, T, const PACK: usize> { + codes: unpacked::View<'a, T>, + meta: &'a [Coefficients], + k: DimK, + physical_k: DimK, +} + +impl BTile<'_, T, PACK> { + /// # Safety + /// + /// `codes` must be the NR-row panel of this tile starting at `start`. + #[inline(always)] + unsafe fn bind_panel<'a, const NR: usize>( + &'a self, + codes: unpacked::Panel<'a, T, NR>, + start: usize, + ) -> BPanel<'a, T, NR, PACK> { + // SAFETY: Preparation gives each code row one coefficient record in the same order. + let meta = unsafe { + &*self + .meta + .as_ptr() + .add(start) + .cast::<[Coefficients; NR]>() + }; + BPanel { + codes, + meta, + k: self.k, + physical_k: self.physical_k, + } + } +} + +struct BPanel<'a, T, const NR: usize, const PACK: usize> { + codes: unpacked::Panel<'a, T, NR>, + meta: &'a [Coefficients; NR], + k: DimK, + physical_k: DimK, +} + +#[inline] +#[expect(clippy::expect_used)] +fn padded_k(k: DimK) -> DimK { + DimK::new( + NonZeroUsize::new(k.value().get().next_multiple_of(PACK)) + .expect("padding a positive dimension remains positive"), + ) +} + +impl BScratch { + /// Padding bypasses `convert`. + /// + /// # Safety + /// + /// B's dimension must equal K, and `start..start + count` must be in bounds. + #[expect(clippy::expect_used)] + unsafe fn prepare( + &mut self, + b: MatRef<'_, MinMaxMeta>, + start: usize, + count: NonZeroUsize, + k: DimK, + mut convert: impl FnMut(u8) -> T, + ) -> BTile<'_, T, PACK> + where + Unsigned: Representation, + { + let dim = k.value().get(); + let physical_k = padded_k::(k); + let stride = physical_k.value().get(); + self.codes.resize(count.get() * stride, T::default()); + self.meta.resize(count.get(), Coefficients::default()); + for i in 0..count.get() { + // SAFETY: The caller guarantees that the tile range is in bounds. + let row = unsafe { b.get_row_unchecked(start + i) }; + let meta = row.meta(); + self.meta[i] = Coefficients { + a: meta.a, + b: meta.b, + n: meta.n, + }; + let dst = &mut self.codes[i * stride..(i + 1) * stride]; + row.vector() + .decode_into(&mut dst[..dim], &mut convert) + .expect("B output must match the logical code count"); + dst[dim..].fill(T::default()); + } + BTile { + // SAFETY: The buffers contain count rows, each with stride initialized codes. + codes: unsafe { unpacked::View::new(Slice::new(&self.codes), count, physical_k) }, + meta: &self.meta, + k, + physical_k, + } + } + + fn row_bytes(k: DimK) -> Bytes { + let codes = Elements::::new(padded_k::(k).value().get()).bytes(); + Bytes::new(codes.value() + size_of::>()) + } +} + +//-------------// +// PanelKernel // +//-------------// + +struct PanelKernel<'a, W: ExtraWide, const MR: usize, const NR: usize, const PACK: usize> +{ + arch: W, + a: APanel<'a, W::AElem, MR, PACK>, + b: BTile<'a, W::BElem, PACK>, + c: [f32; MR], +} + +impl<'a, W: ExtraWide, const MR: usize, const NR: usize, const PACK: usize> + PanelKernel<'a, W, MR, NR, PACK> +{ + /// # Safety + /// + /// A's logical dimension must equal B's. + unsafe fn new( + arch: W, + a: APanel<'a, W::AElem, MR, PACK>, + b: BTile<'a, W::BElem, PACK>, + c: [f32; MR], + ) -> Self { + bounds::check_eq!(a.codes.k(), b.k); + Self { arch, a, b, c } + } +} + +impl, const MR: usize, const NR: usize, const PACK: usize, const N: usize> + unpacked::PanelVisitor for &mut PanelKernel<'_, W, MR, NR, PACK> +{ + #[inline(always)] + fn visit(&mut self, codes: unpacked::Panel<'_, W::BElem, N>, start: usize) { + let this = &mut **self; + // SAFETY: Both full and trailing panels come from this kernel's B tile. + let b = unsafe { this.b.bind_panel(codes, start) }; + let mut kernel = MicroKernel { + arch: this.arch, + a: this.a, + b, + c: &mut this.c, + }; + driver::MicroKernel::micro_kernel(&mut kernel); + } +} + +impl, const MR: usize, const NR: usize, const PACK: usize> + Target for PanelKernel<'_, W, MR, NR, PACK> +where + Self: driver::PanelKernel, +{ + #[inline(always)] + fn run(mut self, _: W) -> [f32; MR] { + driver::PanelKernel::panel_kernel(&mut self); + self.c + } +} + +macro_rules! panel_kernel { + ($nr:literal, [$($tail:literal),*]) => { + impl, const MR: usize, const PACK: usize> + driver::PanelKernel for PanelKernel<'_, W, MR, $nr, PACK> + { + #[inline(always)] + fn panel_kernel(&mut self) { + let b = self.b; + let mut visitor = self; + // SAFETY: Preparation constructs B with physical_k elements per row. + let tail = unsafe { + b.codes.visit_panels::<$nr>(b.physical_k, &mut *visitor) + }; + if let Some(tail) = tail { + $( + if let Some(panel) = tail.try_as_panel::<$tail>() { + unpacked::PanelVisitor::visit(&mut visitor, panel, tail.start()); + return; + } + )* + unreachable!("all remainders below NR are handled"); + } + } + } + }; +} + +panel_kernel!(2, [1]); +panel_kernel!(4, [1, 2, 3]); +panel_kernel!(8, [1, 2, 3, 4, 5, 6, 7]); + +//-------------// +// MicroKernel // +//-------------// + +struct MicroKernel<'a, W: ExtraWide, const MR: usize, const NR: usize, const PACK: usize> +{ + arch: W, + a: APanel<'a, W::AElem, MR, PACK>, + b: BPanel<'a, W::BElem, NR, PACK>, + c: &'a mut [f32; MR], +} + +/// # Safety +/// +/// `a` must contain at least MR * PACK readable elements, and for every +/// `j < NR`, `b.add(bstride * j)` must contain PACK readable elements. +#[inline(always)] +unsafe fn accumulate_pack( + arch: W, + a: Slice<'_, W::AElem>, + b: Slice<'_, W::BElem>, + bstride: Elements, + acc: &mut [W::Acc; NR], +) where + W: ExtraWide, +{ + // SAFETY: The caller provides a full padded A pack. + let a = unsafe { arch.load_a(a) }; + for (j, acc) in acc.iter_mut().enumerate() { + // SAFETY: The caller guarantees a full pack in every B column. + let b = unsafe { arch.splat(b.add(bstride * j)) }; + *acc = arch.dot(*acc, a, b); + } +} + +impl, const MR: usize, const NR: usize, const PACK: usize> + driver::MicroKernel for MicroKernel<'_, W, MR, NR, PACK> +{ + #[inline(always)] + fn micro_kernel(&mut self) { + let physical_k = self.b.physical_k; + let k = self.b.k.value().get(); + let mut acc = [self.arch.zero_acc(); NR]; + let ap = self.a.codes.as_ptr(); + let bp = self.b.codes.as_ptr(); + let astride = self.a.codes.pack_stride(); + let bstride = self.b.codes.stride(physical_k); + + let packs = physical_k.value().get() / PACK; + for pack in 0..packs { + // SAFETY: A and every B row are zero-padded to a multiple of PACK. + unsafe { + accumulate_pack::( + self.arch, + ap.add(astride * pack), + bp.add(Elements::new(pack * PACK)), + bstride, + &mut acc, + ); + } + } + let a_meta = Coefficients { + a: &self.a.meta.a, + b: &self.a.meta.b, + n: &self.a.meta.n, + }; + self.arch.max_into(acc, a_meta, self.b.meta, k, self.c); + } +} +//----------// +// Backends // +//----------// + +/// Backends accept A codes in 0..=255; B is limited by `MAX_B_BITS`. +pub(crate) trait ExtraWide: Architecture { + /// Largest unsigned B width valid for every A code. + const MAX_B_BITS: usize; + + type AElem: Copy + Default + From; + type BElem: Copy + Default; + type Wide: Copy; + type Splat: Copy; + type Acc: Copy; + + /// `code` fits in MAX_B_BITS unsigned bits. + fn convert_b(code: u8) -> Self::BElem; + + /// # Safety + /// + /// `a` must contain at least MR * PACK initialized, readable elements. + unsafe fn load_a(self, a: Slice<'_, Self::AElem>) -> Self::Wide; + + fn zero_acc(self) -> Self::Acc; + + /// Accumulate MR dot products of PACK elements without mixing rows. + fn dot(self, acc: Self::Acc, a: Self::Wide, b: Self::Splat) -> Self::Acc; + + /// Compensate with logical K and update maxima, ignoring NaNs. + fn max_into( + self, + acc: [Self::Acc; NR], + query: Coefficients<&[f32; MR]>, + docs: &[Coefficients; NR], + k: usize, + scores: &mut [f32; MR], + ); + + /// # Safety + /// + /// B must contain PACK readable elements, including any zero padding. + unsafe fn splat(self, b: Slice<'_, Self::BElem>) -> Self::Splat; +} + +#[inline(always)] +fn compensated_max_into( + arch: F::Arch, + acc: impl IntoIterator, + query: Coefficients<&[f32; MR]>, + docs: &[Coefficients; NR], + k: usize, + scores: &mut [f32; MR], +) where + I: SIMDCast, + F: SIMDVector> + + SIMDMinMax + + Add + + Mul, + Const: SupportedLaneCount, +{ + let a = Coefficients { + a: F::from_array(arch, *query.a), + b: F::from_array(arch, *query.b), + n: F::from_array(arch, *query.n), + }; + let mut best = F::from_array(arch, *scores); + for (acc, b) in acc.into_iter().zip(docs) { + let dot = acc.simd_cast(); + let a_b = F::splat(arch, b.a); + let b_b = F::splat(arch, b.b); + let n_b = F::splat(arch, b.n); + let dim = F::splat(arch, k as f32); + let ip = ((a.a * a_b) * dot + a.n * b_b) + n_b * a.b + (a.b * b_b) * dim; + best = best.max_simd_standard(ip); + } + *scores = best.to_array(); +} + +/// # Safety +/// +/// `b` must contain at least PACK initialized, readable elements. +#[inline(always)] +unsafe fn load_pack(b: Slice<'_, T>) -> [T; PACK] { + bounds::check_ge!(b.len(), PACK); + // SAFETY: The caller provides a full pack; no array alignment is required. + unsafe { b.as_ptr().cast::<[T; PACK]>().read_unaligned() } +} + +/// # Safety +/// +/// B must contain two readable elements. +#[inline(always)] +unsafe fn pack_i16_pair(b: Slice<'_, i16>) -> u32 { + // SAFETY: The caller provides two readable elements. + let [lo, hi] = unsafe { load_pack::<_, 2>(b) }; + u32::from(lo as u16) | (u32::from(hi as u16) << 16) +} + +impl ExtraWide<16, 4> for Scalar { + const MAX_B_BITS: usize = 7; + + type AElem = u8; + type BElem = i8; + type Wide = u8x64; + type Splat = i8x64; + type Acc = i32x16; + + #[inline(always)] + unsafe fn load_a(self, a: Slice<'_, Self::AElem>) -> Self::Wide { + // SAFETY: The caller provides 64 initialized elements. + Self::Wide::from_array(self, unsafe { load_pack::<_, 64>(a) }) + } + + #[inline(always)] + fn zero_acc(self) -> Self::Acc { + SIMDVector::default(self) + } + + #[inline(always)] + fn max_into( + self, + acc: [Self::Acc; NR], + query: Coefficients<&[f32; 16]>, + docs: &[Coefficients; NR], + k: usize, + scores: &mut [f32; 16], + ) { + compensated_max_into(self, acc, query, docs, k, scores); + } + + #[inline(always)] + fn dot(self, acc: Self::Acc, a: Self::Wide, b: Self::Splat) -> Self::Acc { + acc.dot_simd(a, b) + } + + fn convert_b(code: u8) -> i8 { + code as i8 + } + + #[inline(always)] + unsafe fn splat(self, b: Slice<'_, i8>) -> Self::Splat { + // SAFETY: The caller provides four readable elements. + let values = unsafe { load_pack::<_, 4>(b) }; + let quad = u32::from_le_bytes(values.map(|x| x as u8)); + u32x16::::splat(self, quad).reinterpret_simd() + } +} + +impl ExtraWide<16, 2> for Scalar { + const MAX_B_BITS: usize = 8; + + type AElem = i16; + type BElem = i16; + type Wide = i16x32; + type Splat = Self::Wide; + type Acc = i32x16; + + #[inline(always)] + unsafe fn load_a(self, a: Slice<'_, Self::AElem>) -> Self::Wide { + // SAFETY: The caller provides 32 initialized elements. + Self::Wide::from_array(self, unsafe { load_pack::<_, 32>(a) }) + } + + #[inline(always)] + fn zero_acc(self) -> Self::Acc { + SIMDVector::default(self) + } + + #[inline(always)] + fn max_into( + self, + acc: [Self::Acc; NR], + query: Coefficients<&[f32; 16]>, + docs: &[Coefficients; NR], + k: usize, + scores: &mut [f32; 16], + ) { + compensated_max_into(self, acc, query, docs, k, scores); + } + + #[inline(always)] + fn dot(self, acc: Self::Acc, a: Self::Wide, b: Self::Splat) -> Self::Acc { + acc.dot_simd(a, b) + } + + fn convert_b(code: u8) -> i16 { + i16::from(code) + } + + #[inline(always)] + unsafe fn splat(self, b: Slice<'_, i16>) -> Self::Splat { + // SAFETY: Inherited from the caller. + u32x16::::splat(self, unsafe { pack_i16_pair(b) }).reinterpret_simd() + } +} + +impl ExtraWide<8, 2> for Scalar { + const MAX_B_BITS: usize = 8; + type AElem = i16; + type BElem = i16; + type Wide = i16x16; + type Splat = Self::Wide; + type Acc = i32x8; + + #[inline(always)] + unsafe fn load_a(self, a: Slice<'_, Self::AElem>) -> Self::Wide { + // SAFETY: The caller provides 16 initialized elements. + Self::Wide::from_array(self, unsafe { load_pack::<_, 16>(a) }) + } + + #[inline(always)] + fn zero_acc(self) -> Self::Acc { + SIMDVector::default(self) + } + + #[inline(always)] + fn max_into( + self, + acc: [Self::Acc; NR], + query: Coefficients<&[f32; 8]>, + docs: &[Coefficients; NR], + k: usize, + scores: &mut [f32; 8], + ) { + compensated_max_into(self, acc, query, docs, k, scores); + } + + #[inline(always)] + fn dot(self, acc: Self::Acc, a: Self::Wide, b: Self::Splat) -> Self::Acc { + acc.dot_simd(a, b) + } + + fn convert_b(code: u8) -> i16 { + i16::from(code) + } + #[inline(always)] + unsafe fn splat(self, b: Slice<'_, i16>) -> Self::Splat { + // SAFETY: The caller provides two readable elements. + u32x8::::splat(self, unsafe { pack_i16_pair(b) }).reinterpret_simd() + } +} + +#[cfg(target_arch = "x86_64")] +mod x86_64 { + use super::*; + use diskann_wide::SplitJoin; + use diskann_wide::arch::x86_64::{V3, V4}; + + #[inline(always)] + fn compensated_max_into_v3( + arch: V3, + acc: [i32x16; NR], + query: Coefficients<&[f32; 16]>, + docs: &[Coefficients; NR], + k: usize, + scores: &mut [f32; 16], + ) { + let a = query.a.as_chunks::<8>().0; + let b = query.b.as_chunks::<8>().0; + let n = query.n.as_chunks::<8>().0; + let scores = scores.as_chunks_mut::<8>().0; + // Finish one native-width row group before loading the next group's coefficients. + compensated_max_into( + arch, + acc.iter().map(|v| v.split().lo), + Coefficients { + a: &a[0], + b: &b[0], + n: &n[0], + }, + docs, + k, + &mut scores[0], + ); + compensated_max_into( + arch, + acc.iter().map(|v| v.split().hi), + Coefficients { + a: &a[1], + b: &b[1], + n: &n[1], + }, + docs, + k, + &mut scores[1], + ); + } + + impl ExtraWide<8, 4> for V3 { + const MAX_B_BITS: usize = 6; + type AElem = u8; + type BElem = i8; + type Wide = u8x32; + type Splat = i8x32; + type Acc = i32x8; + + #[inline(always)] + unsafe fn load_a(self, a: Slice<'_, Self::AElem>) -> Self::Wide { + // SAFETY: The caller provides 32 initialized elements. + Self::Wide::from_array(self, unsafe { load_pack::<_, 32>(a) }) + } + + #[inline(always)] + fn zero_acc(self) -> Self::Acc { + SIMDVector::default(self) + } + + #[inline(always)] + fn max_into( + self, + acc: [Self::Acc; NR], + query: Coefficients<&[f32; 8]>, + docs: &[Coefficients; NR], + k: usize, + scores: &mut [f32; 8], + ) { + compensated_max_into(self, acc, query, docs, k, scores); + } + + fn convert_b(code: u8) -> i8 { + code as i8 + } + #[inline(always)] + fn dot(self, acc: Self::Acc, a: Self::Wide, b: Self::Splat) -> Self::Acc { + // The largest pair is 2 * 255 * 63 = 32130, below i16 saturation. + acc.dot_simd(self.mm256_maddubs_epi16(a, b), i16x16::::splat(self, 1)) + } + #[inline(always)] + unsafe fn splat(self, b: Slice<'_, i8>) -> Self::Splat { + // SAFETY: The caller provides four readable elements. + let values = unsafe { load_pack::<_, 4>(b) }; + let quad = u32::from_le_bytes(values.map(|v| v as u8)); + u32x8::::splat(self, quad).reinterpret_simd() + } + } + + impl ExtraWide<8, 2> for V3 { + const MAX_B_BITS: usize = 8; + type AElem = i16; + type BElem = i16; + type Wide = i16x16; + type Splat = Self::Wide; + type Acc = i32x8; + + #[inline(always)] + unsafe fn load_a(self, a: Slice<'_, Self::AElem>) -> Self::Wide { + // SAFETY: The caller provides 16 initialized elements. + Self::Wide::from_array(self, unsafe { load_pack::<_, 16>(a) }) + } + + #[inline(always)] + fn zero_acc(self) -> Self::Acc { + SIMDVector::default(self) + } + + #[inline(always)] + fn max_into( + self, + acc: [Self::Acc; NR], + query: Coefficients<&[f32; 8]>, + docs: &[Coefficients; NR], + k: usize, + scores: &mut [f32; 8], + ) { + compensated_max_into(self, acc, query, docs, k, scores); + } + + #[inline(always)] + fn dot(self, acc: Self::Acc, a: Self::Wide, b: Self::Splat) -> Self::Acc { + acc.dot_simd(a, b) + } + + fn convert_b(code: u8) -> i16 { + i16::from(code) + } + #[inline(always)] + unsafe fn splat(self, b: Slice<'_, i16>) -> Self::Splat { + // SAFETY: The caller provides two readable elements. + u32x8::::splat(self, unsafe { pack_i16_pair(b) }).reinterpret_simd() + } + } + + impl ExtraWide<16, 4> for V3 { + const MAX_B_BITS: usize = 6; + + type AElem = u8; + type BElem = i8; + type Wide = u8x64; + type Splat = i8x64; + type Acc = i32x16; + + #[inline(always)] + unsafe fn load_a(self, a: Slice<'_, Self::AElem>) -> Self::Wide { + // SAFETY: The caller provides 64 initialized elements. + Self::Wide::from_array(self, unsafe { load_pack::<_, 64>(a) }) + } + + #[inline(always)] + fn zero_acc(self) -> Self::Acc { + SIMDVector::default(self) + } + + #[inline(always)] + fn max_into( + self, + acc: [Self::Acc; NR], + query: Coefficients<&[f32; 16]>, + docs: &[Coefficients; NR], + k: usize, + scores: &mut [f32; 16], + ) { + compensated_max_into_v3(self, acc, query, docs, k, scores); + } + + fn convert_b(code: u8) -> i8 { + code as i8 + } + + #[inline(always)] + fn dot(self, acc: Self::Acc, a: Self::Wide, b: Self::Splat) -> Self::Acc { + let a = a.split(); + let b = b.split(); + let acc = acc.split(); + let ones = i16x16::::splat(self, 1); + // Each saturated i16 pair is at most 2 * 255 * 63 = 32130, so it is exact. + i32x16::::new( + acc.lo.dot_simd(self.mm256_maddubs_epi16(a.lo, b.lo), ones), + acc.hi.dot_simd(self.mm256_maddubs_epi16(a.hi, b.hi), ones), + ) + } + + #[inline(always)] + unsafe fn splat(self, b: Slice<'_, i8>) -> Self::Splat { + // SAFETY: The caller provides four readable elements. + let values = unsafe { load_pack::<_, 4>(b) }; + let quad = u32::from_le_bytes(values.map(|x| x as u8)); + let half: i8x32 = u32x8::::splat(self, quad).reinterpret_simd(); + i8x64::::new(half, half) + } + } + + impl ExtraWide<16, 2> for V3 { + const MAX_B_BITS: usize = 8; + + type AElem = i16; + type BElem = i16; + type Wide = i16x32; + type Splat = Self::Wide; + type Acc = i32x16; + + #[inline(always)] + unsafe fn load_a(self, a: Slice<'_, Self::AElem>) -> Self::Wide { + // SAFETY: The caller provides 32 initialized elements. + Self::Wide::from_array(self, unsafe { load_pack::<_, 32>(a) }) + } + + #[inline(always)] + fn zero_acc(self) -> Self::Acc { + SIMDVector::default(self) + } + + #[inline(always)] + fn max_into( + self, + acc: [Self::Acc; NR], + query: Coefficients<&[f32; 16]>, + docs: &[Coefficients; NR], + k: usize, + scores: &mut [f32; 16], + ) { + compensated_max_into_v3(self, acc, query, docs, k, scores); + } + + #[inline(always)] + fn dot(self, acc: Self::Acc, a: Self::Wide, b: Self::Splat) -> Self::Acc { + acc.dot_simd(a, b) + } + + fn convert_b(code: u8) -> i16 { + i16::from(code) + } + + #[inline(always)] + unsafe fn splat(self, b: Slice<'_, i16>) -> Self::Splat { + // SAFETY: Inherited from the caller. + let pair = unsafe { pack_i16_pair(b) }; + let half: i16x16 = u32x8::::splat(self, pair).reinterpret_simd(); + i16x32::::new(half, half) + } + } + + impl ExtraWide<16, 4> for V4 { + const MAX_B_BITS: usize = 7; + + type AElem = u8; + type BElem = i8; + type Wide = u8x64; + type Splat = i8x64; + type Acc = i32x16; + + #[inline(always)] + unsafe fn load_a(self, a: Slice<'_, Self::AElem>) -> Self::Wide { + // SAFETY: The caller provides 64 initialized elements. + Self::Wide::from_array(self, unsafe { load_pack::<_, 64>(a) }) + } + + #[inline(always)] + fn zero_acc(self) -> Self::Acc { + SIMDVector::default(self) + } + + #[inline(always)] + fn max_into( + self, + acc: [Self::Acc; NR], + query: Coefficients<&[f32; 16]>, + docs: &[Coefficients; NR], + k: usize, + scores: &mut [f32; 16], + ) { + compensated_max_into(self, acc, query, docs, k, scores); + } + + #[inline(always)] + fn dot(self, acc: Self::Acc, a: Self::Wide, b: Self::Splat) -> Self::Acc { + acc.dot_simd(a, b) + } + + fn convert_b(code: u8) -> i8 { + code as i8 + } + + #[inline(always)] + unsafe fn splat(self, b: Slice<'_, i8>) -> Self::Splat { + // SAFETY: The caller provides four readable elements. + let values = unsafe { load_pack::<_, 4>(b) }; + let quad = u32::from_le_bytes(values.map(|x| x as u8)); + u32x16::::splat(self, quad).reinterpret_simd() + } + } + + impl ExtraWide<16, 2> for V4 { + const MAX_B_BITS: usize = 8; + + type AElem = i16; + type BElem = i16; + type Wide = i16x32; + type Splat = Self::Wide; + type Acc = i32x16; + + #[inline(always)] + unsafe fn load_a(self, a: Slice<'_, Self::AElem>) -> Self::Wide { + // SAFETY: The caller provides 32 initialized elements. + Self::Wide::from_array(self, unsafe { load_pack::<_, 32>(a) }) + } + + #[inline(always)] + fn zero_acc(self) -> Self::Acc { + SIMDVector::default(self) + } + + #[inline(always)] + fn max_into( + self, + acc: [Self::Acc; NR], + query: Coefficients<&[f32; 16]>, + docs: &[Coefficients; NR], + k: usize, + scores: &mut [f32; 16], + ) { + compensated_max_into(self, acc, query, docs, k, scores); + } + + #[inline(always)] + fn dot(self, acc: Self::Acc, a: Self::Wide, b: Self::Splat) -> Self::Acc { + acc.dot_simd(a, b) + } + + fn convert_b(code: u8) -> i16 { + i16::from(code) + } + + #[inline(always)] + unsafe fn splat(self, b: Slice<'_, i16>) -> Self::Splat { + // SAFETY: Inherited from the caller. + u32x16::::splat(self, unsafe { pack_i16_pair(b) }).reinterpret_simd() + } + } +} + +#[cfg(target_arch = "aarch64")] +mod aarch64 { + use super::*; + use diskann_wide::arch::aarch64::Neon; + + impl ExtraWide<8, 4> for Neon { + const MAX_B_BITS: usize = 8; + type AElem = u8; + type BElem = u8; + type Wide = u8x32; + type Splat = Self::Wide; + type Acc = u32x8; + + #[inline(always)] + unsafe fn load_a(self, a: Slice<'_, Self::AElem>) -> Self::Wide { + // SAFETY: The caller provides 32 initialized elements. + Self::Wide::from_array(self, unsafe { load_pack::<_, 32>(a) }) + } + + #[inline(always)] + fn zero_acc(self) -> Self::Acc { + SIMDVector::default(self) + } + + #[inline(always)] + fn max_into( + self, + acc: [Self::Acc; NR], + query: Coefficients<&[f32; 8]>, + docs: &[Coefficients; NR], + k: usize, + scores: &mut [f32; 8], + ) { + compensated_max_into(self, acc, query, docs, k, scores); + } + + #[inline(always)] + fn dot(self, acc: Self::Acc, a: Self::Wide, b: Self::Splat) -> Self::Acc { + acc.dot_simd(a, b) + } + + fn convert_b(code: u8) -> u8 { + code + } + #[inline(always)] + unsafe fn splat(self, b: Slice<'_, u8>) -> Self::Splat { + // SAFETY: The caller provides four readable elements. + let values = unsafe { load_pack::<_, 4>(b) }; + let quarter = u32x4::::splat(self, u32::from_le_bytes(values)).reinterpret_simd(); + u8x32::::new(quarter, quarter) + } + } + + impl ExtraWide<16, 4> for Neon { + const MAX_B_BITS: usize = 8; + + type AElem = u8; + type BElem = u8; + type Wide = u8x64; + type Splat = Self::Wide; + type Acc = u32x16; + + #[inline(always)] + unsafe fn load_a(self, a: Slice<'_, Self::AElem>) -> Self::Wide { + // SAFETY: The caller provides 64 initialized elements. + Self::Wide::from_array(self, unsafe { load_pack::<_, 64>(a) }) + } + + #[inline(always)] + fn zero_acc(self) -> Self::Acc { + SIMDVector::default(self) + } + + #[inline(always)] + fn max_into( + self, + acc: [Self::Acc; NR], + query: Coefficients<&[f32; 16]>, + docs: &[Coefficients; NR], + k: usize, + scores: &mut [f32; 16], + ) { + compensated_max_into(self, acc, query, docs, k, scores); + } + + #[inline(always)] + fn dot(self, acc: Self::Acc, a: Self::Wide, b: Self::Splat) -> Self::Acc { + acc.dot_simd(a, b) + } + + fn convert_b(code: u8) -> u8 { + code + } + + #[inline(always)] + unsafe fn splat(self, b: Slice<'_, u8>) -> Self::Splat { + // SAFETY: The caller provides four readable elements. + let values = unsafe { load_pack::<_, 4>(b) }; + let quad = u32::from_le_bytes(values); + let quarter = u32x4::::splat(self, quad).reinterpret_simd(); + let half = u8x32::::new(quarter, quarter); + u8x64::::new(half, half) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use diskann_vector::PureDistanceFunction; + use rand::{Rng, SeedableRng, rngs::StdRng}; + + use crate::{ + bits::{Representation, Unsigned}, + distances, + matrix_kernels::maxsim::test::packed_x_unpacked_test_dims, + minmax::{DataRef, MinMaxCompensation, MinMaxIP}, + multi_vector::{Defaulted, Mat}, + test_util::Check, + }; + + #[test] + fn i16_pair_preserves_lane_order_and_bits() { + for (values, expected) in [ + ([1, 255], 0x00ff_0001), + ([255, 1], 0x0001_00ff), + ([-1, i16::MIN], 0x8000_ffff), + ] { + // SAFETY: values contains both initialized elements of the pair. + assert_eq!(unsafe { pack_i16_pair(Slice::new(&values)) }, expected); + } + } + + fn random_minmax_matrix( + rows: usize, + k: usize, + rng: &mut StdRng, + ) -> Mat> + where + Unsigned: Representation, + { + let mut mat = Mat::new(MinMaxMeta::::new(rows, k), Defaulted).unwrap(); + for mut row in mat.rows_mut() { + let a = rng.random_range(0.001..2.0); + let b = rng.random_range(-3.0..3.0); + let mut sum = 0; + for j in 0..k { + let code = rng.random_range(0..(1 << BITS)); + row.vector_mut().set(j, code).unwrap(); + sum += code; + } + row.set_meta(MinMaxCompensation { + dim: k as u32, + a, + b, + n: a * sum as f32, + norm_squared: 0.0, + }); + } + mat + } + + fn reference_max_inner_products( + a: MatRef<'_, MinMaxMeta>, + b: MatRef<'_, MinMaxMeta>, + ) -> Vec + where + Unsigned: Representation + Representation, + { + let k = a.repr().intrinsic_dim(); + assert_eq!(k, b.repr().intrinsic_dim()); + a.rows() + .map(|a| { + let av = a.vector(); + let am = a.meta(); + b.rows() + .map(|b| { + let bv = b.vector(); + let bm = b.meta(); + let dot: u64 = (0..k) + .map(|j| av.get(j).unwrap() as u64 * bv.get(j).unwrap() as u64) + .sum(); + ((am.a * bm.a) * dot as f32 + am.n * bm.b) + + bm.n * am.b + + (am.b * bm.b) * k as f32 + }) + .fold(f32::NEG_INFINITY, f32::max) + }) + .collect() + } + + #[track_caller] + fn assert_scores_close(actual: &[f32], expected: &[f32], context: std::fmt::Arguments<'_>) { + assert_eq!(actual.len(), expected.len(), "{context}"); + for (i, (&got, &want)) in actual.iter().zip(expected).enumerate() { + if got == want { + continue; + } + assert!( + got.is_finite() && want.is_finite(), + "{context}: row {i}: got {got}, expected {want}" + ); + Check::absrel(1e-5, 1e-5) + .check(got, want) + .unwrap_or_else(|err| panic!("{context}: row {i}: {err}")); + } + } + + fn checked_driver< + 'a, + W, + const A_BITS: usize, + const B_BITS: usize, + const MR: usize, + const NR: usize, + const PACK: usize, + >( + arch: W, + a: &'a PackedA, + b: MatRef<'a, MinMaxMeta>, + c: &'a mut [f32], + ) -> Driver<'a, W, B_BITS, MR, NR, PACK> + where + W: ExtraWide + LoadStore, + Unsigned: Representation + Representation, + for<'b> PanelKernel<'b, W, MR, NR, PACK>: driver::PanelKernel, + { + assert!(a.codes.nrows() > 0 && a.codes.ncols() > 0); + assert_eq!(a.codes.nrows(), c.len()); + assert_eq!(a.codes.ncols(), b.repr().intrinsic_dim()); + // SAFETY: Nonempty A/K, matching dimensions and output length were checked above. + unsafe { Driver::::new(arch, a, b, c, Cache::detect()) } + } + + fn assert_driver_contract< + W, + const B_BITS: usize, + const MR: usize, + const NR: usize, + const PACK: usize, + >( + arch: W, + ) where + W: ExtraWide + LoadStore, + Unsigned: Representation, + for<'a> PanelKernel<'a, W, MR, NR, PACK>: driver::PanelKernel, + { + let mut rng = StdRng::seed_from_u64(0x63de91a0ba5c8742); + for shape in packed_x_unpacked_test_dims(MR, NR) { + for k in [shape.k, shape.k.next_multiple_of(PACK)] { + let a = random_minmax_matrix::<8>(shape.total_a_rows, k, &mut rng); + let b = random_minmax_matrix::(shape.total_b_cols, k, &mut rng); + let packed = PackedA::::new(a.as_view()).unwrap(); + let expected = reference_max_inner_products::<8, B_BITS>(a.as_view(), b.as_view()); + let mut guarded = vec![654321.0; a.num_vectors() + 2]; + let mut driver = checked_driver::<_, 8, B_BITS, MR, NR, PACK>( + arch, + &packed, + b.as_view(), + &mut guarded[1..=a.num_vectors()], + ); + driver.b_rows_per_tile = NonZeroUsize::new(shape.b_cols_per_tile).unwrap(); + driver.c.fill(f32::INFINITY); + driver::Drive::drive(&mut driver); + assert_scores_close( + driver.c, + &expected, + format_args!( + "forced B tiling: {shape:?}, K={k}, MR={MR}, NR={NR}, PACK={PACK}" + ), + ); + assert_eq!(guarded[0], 654321.0); + assert_eq!(guarded[a.num_vectors() + 1], 654321.0); + } + } + assert_compensated_reduction::(arch); + + let a = random_minmax_matrix::<8>(MR + 1, 5, &mut rng); + let b = random_minmax_matrix::(0, 5, &mut rng); + let packed = PackedA::::new(a.as_view()).unwrap(); + let mut c = vec![f32::INFINITY; a.num_vectors()]; + driver::Drive::drive(&mut checked_driver::<_, 8, B_BITS, MR, NR, PACK>( + arch, + &packed, + b.as_view(), + &mut c, + )); + assert_eq!(c, vec![f32::NEG_INFINITY; a.num_vectors()]); + } + + fn constant_minmax_matrix( + rows: usize, + k: usize, + code: i64, + a: f32, + b: f32, + ) -> Mat> + where + Unsigned: Representation, + { + let mut mat = Mat::new(MinMaxMeta::::new(rows, k), Defaulted).unwrap(); + for mut row in mat.rows_mut() { + for j in 0..k { + row.vector_mut().set(j, code).unwrap(); + } + row.set_meta(MinMaxCompensation { + dim: k as u32, + a, + b, + n: a * (k as f32 * code as f32), + norm_squared: 0.0, + }); + } + mat + } + + fn assert_compensated_reduction< + W, + const B_BITS: usize, + const MR: usize, + const NR: usize, + const PACK: usize, + >( + arch: W, + ) where + W: ExtraWide + LoadStore, + Unsigned: Representation, + for<'a> PanelKernel<'a, W, MR, NR, PACK>: driver::PanelKernel, + { + let a = constant_minmax_matrix::<8>(MR + 1, 3, 1, 1.0, 0.0); + let mut b = constant_minmax_matrix::(NR + 1, 3, (1 << B_BITS) - 1, 1.0, 0.0); + let winner_scale = (1 << B_BITS) as f32; + { + let mut row = b.get_row_mut(0).unwrap(); + row.vector_mut().set(0, 1).unwrap(); + row.vector_mut().set(1, 1).unwrap(); + row.vector_mut().set(2, 1).unwrap(); + row.set_meta(MinMaxCompensation { + dim: 3, + a: winner_scale, + b: 0.0, + n: 3.0 * winner_scale, + norm_squared: 0.0, + }); + } + let packed = PackedA::::new(a.as_view()).unwrap(); + let mut c = vec![0.0; MR + 1]; + let mut driver = + checked_driver::<_, 8, B_BITS, MR, NR, PACK>(arch, &packed, b.as_view(), &mut c); + // The first tile contains competing raw dots; its compensated winner must survive the tail. + driver.b_rows_per_tile = NonZeroUsize::new(NR).unwrap(); + driver::Drive::drive(&mut driver); + assert_eq!(c, vec![3.0 * winner_scale; MR + 1]); + + let a = constant_minmax_matrix::<8>(MR + 1, 3, 1, 1.0, -2.0); + let b = constant_minmax_matrix::((NR - 1).max(1), 3, 1, 1.0, 0.0); + let packed = PackedA::::new(a.as_view()).unwrap(); + driver::Drive::drive(&mut checked_driver::<_, 8, B_BITS, MR, NR, PACK>( + arch, + &packed, + b.as_view(), + &mut c, + )); + // Invalid B lanes must not introduce a zero candidate. + assert_eq!(c, vec![-3.0; MR + 1]); + } + + #[cfg(not(miri))] + #[test] + fn accumulator_is_exact_at_dimension_limit() { + let arch = Scalar::new(); + let k = max_k::<8, 8>(); + let a = constant_minmax_matrix::<8>(1, k, 255, 1.0, 0.0); + let b = constant_minmax_matrix::<8>(1, k, 255, 1.0, 0.0); + let packed = PackedA::::new(a.as_view()).unwrap(); + let mut c = [0.0]; + driver::Drive::drive(&mut checked_driver::<_, 8, 8, 8, 2, 2>( + arch, + &packed, + b.as_view(), + &mut c, + )); + assert_eq!(c[0], (k * 255 * 255) as f32); + } + + #[test] + fn scalar_driver_obeys_maxsim_contract() { + assert_driver_contract::(Scalar::new()); + assert_driver_target_shapes::(Scalar::new()); + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn avx2_driver_obeys_maxsim_contract() { + if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { + assert_driver_contract::<_, 4, 16, 4, 4>(arch); + assert_driver_target_shapes::<_, 4, 16, 4, 4>(arch); + } else { + eprintln!("AVX2 unavailable; backend not executed"); + } + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn avx512_driver_obeys_maxsim_contract() { + if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked() { + assert_driver_contract::<_, 4, 16, 8, 4>(arch); + assert_driver_target_shapes::<_, 4, 16, 8, 4>(arch); + } else { + eprintln!("V4 unavailable; backend not executed"); + } + } + + #[cfg(target_arch = "aarch64")] + #[test] + fn neon_driver_obeys_maxsim_contract() { + if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { + assert_driver_contract::<_, 4, 16, 4, 4>(arch); + assert_driver_target_shapes::<_, 4, 16, 4, 4>(arch); + } else { + eprintln!("Neon/dotprod unavailable; backend not executed"); + } + } + + #[test] + fn smaller_panels_preserve_tiling_and_nan_reduction() { + fn check(arch: W) + where + W: ExtraWide<8, PACK> + LoadStore, + Unsigned: Representation, + for<'a> PanelKernel<'a, W, 8, NR, PACK>: driver::PanelKernel, + { + assert_driver_contract::(arch); + assert_microkernel_nan_reduction::(arch); + } + check::<_, 8, 2, 2>(Scalar::new()); + #[cfg(target_arch = "x86_64")] + if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { + check::<_, 4, 4, 4>(arch); + check::<_, 8, 4, 2>(arch); + } + #[cfg(target_arch = "aarch64")] + if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { + check::<_, 8, 4, 4>(arch); + } + } + + fn assert_a_packing() + where + T: Copy + Default + From + PartialEq + std::fmt::Debug, + Unsigned: Representation, + { + for (m, k) in [(0, 0), (0, 5), (3, 0)] { + let a = constant_minmax_matrix::(m, k, 0, 1.0, 0.0); + let packed = PackedA::::new(a.as_view()).unwrap(); + assert_eq!(packed.codes.nrows(), m); + assert_eq!(packed.codes.ncols(), k); + assert!(packed.as_view().is_none()); + } + let mut rng = StdRng::seed_from_u64(0x0ae4b97f2c186d35); + let a = random_minmax_matrix::(MR + 1, 5, &mut rng); + let packed = PackedA::::new(a.as_view()).unwrap(); + let padded_cols = 5usize.div_ceil(PACK) * PACK; + assert_eq!(packed.codes.nrows(), MR + 1); + assert_eq!(packed.codes.ncols(), 5); + assert_eq!(packed.codes.padded_nrows(), 2 * MR); + assert_eq!(packed.codes.padded_ncols(), padded_cols); + for row in 0..2 * MR { + for col in 0..padded_cols { + let index = (row / MR) * MR * padded_cols + + (col / PACK) * MR * PACK + + (row % MR) * PACK + + col % PACK; + let expected = if row <= MR && col < 5 { + T::from(a.get_row(row).unwrap().vector().get(col).unwrap() as u8) + } else { + T::default() + }; + assert_eq!(packed.codes.as_slice()[index], expected); + } + let meta = a.get_row(row).map(|v| v.meta()).unwrap_or_default(); + let panel = &packed.meta[row / MR]; + assert_eq!(panel.a[row % MR], meta.a); + assert_eq!(panel.b[row % MR], meta.b); + assert_eq!(panel.n[row % MR], meta.n); + } + let mut visited = 0; + packed.as_view().unwrap().visit_panels(|panel, i| { + assert_eq!(i, visited); + assert!(std::ptr::eq(panel.meta, &packed.meta[i])); + assert_eq!( + panel.codes.as_ptr().as_ptr(), + packed + .codes + .as_slice() + .as_ptr() + .wrapping_add(i * MR * padded_cols), + ); + visited += 1; + }); + assert_eq!(visited, packed.meta.len()); + } + + #[test] + fn a_packing_preserves_codes_metadata_and_padding() { + macro_rules! check { + ($($bits:literal),+) => {$( + assert_a_packing::(); + assert_a_packing::(); + )+}; + } + check!(1, 2, 3, 4, 5, 6, 7, 8); + assert_a_packing::(); + } + + #[test] + fn cache_footprints_match_prepared_layouts() { + fn check, const PACK: usize>( + cases: [(usize, usize, usize); 2], + ) { + for (k, a_bytes, b_bytes) in cases { + let dim = DimK::new(NonZeroUsize::new(k).unwrap()); + assert_eq!(BScratch::::row_bytes::(dim).value(), b_bytes); + for rows in [1, 16, 17] { + let a = constant_minmax_matrix::<8>(rows, k, 0, 1.0, 0.0); + let packed = PackedA::::new(a.as_view()).unwrap(); + assert_eq!(packed.panel_bytes().value(), a_bytes); + } + } + } + check::([(256, 4288, 268), (257, 4352, 272)]); + check::([(256, 8384, 524), (257, 8448, 528)]); + } + + #[test] + fn compact_metadata_preserves_scores_across_tile_sizes() { + let arch = Scalar::new(); + let mut rng = StdRng::seed_from_u64(0x9f23ea1640b187ac); + let a = random_minmax_matrix::<8>(9, 3, &mut rng); + let b = random_minmax_matrix::<8>(11, 3, &mut rng); + let packed = PackedA::::new(a.as_view()).unwrap(); + // This budget gives different rounded tiles for compact and canonical metadata. + let cache = Cache::new( + NonZeroUsize::new(packed.panel_bytes().value() + 81).unwrap(), + NonZeroUsize::new(4096).unwrap(), + ); + let expected = reference_max_inner_products(a.as_view(), b.as_view()); + let mut scores = vec![f32::NAN; a.num_vectors()]; + // SAFETY: A is nonempty; both inputs have K=3 and scores covers every A row. + let mut driver = unsafe { + Driver::<_, 8, 8, 2, 2>::new::<8>(arch, &packed, b.as_view(), &mut scores, cache) + }; + assert_eq!(driver.b_rows_per_tile.get(), 6); + driver::Drive::drive(&mut driver); + assert_scores_close(driver.c, &expected, format_args!("compact metadata tiling")); + + let previous_row_bytes = Bytes::new(4 * size_of::() + size_of::()); + driver.b_rows_per_tile = + Params::new(cache, packed.panel_bytes(), previous_row_bytes, 2).b_cols_in_l1; + assert_eq!(driver.b_rows_per_tile.get(), 4); + driver::Drive::drive(&mut driver); + assert_scores_close( + driver.c, + &expected, + format_args!("previous metadata tiling"), + ); + } + + #[test] + fn b_preparation_preserves_metadata_and_zero_padding() { + macro_rules! check { + ($($bits:literal),+) => {$( + assert_b_preparation::(-1, |code| code as i8); + assert_b_preparation::(-1, i16::from); + assert_b_preparation::(u8::MAX, |code| code); + )+}; + } + check!(1, 2, 3, 4, 5, 6, 7); + assert_b_preparation::(-1, i16::from); + assert_b_preparation::(u8::MAX, |code| code); + assert_b_preparation::(u16::MAX, |code| u16::from(code) + 1); + } + + #[test] + fn b_panels_bind_coefficients_at_full_and_tail_offsets() { + let k = DimK::new(NonZeroUsize::new(3).unwrap()); + let mut rng = StdRng::seed_from_u64(0xda7bc35e83012e29); + let source = random_minmax_matrix::<4>(5, k.value().get(), &mut rng); + let mut scratch = BScratch::::default(); + // SAFETY: The tile covers all five rows of source, whose dimension is K. + let tile = unsafe { + scratch.prepare::<4, 4>( + source.as_view(), + 0, + NonZeroUsize::new(5).unwrap(), + k, + |code| code, + ) + }; + let mut actual = Vec::new(); + let visit = |codes: unpacked::Panel<'_, u8, 2>, start| { + // SAFETY: The tile's traversal supplies this panel and its offset. + let panel = unsafe { tile.bind_panel(codes, start) }; + actual.extend(panel.meta.iter().map(|m| (m.a, m.b, m.n))); + }; + // SAFETY: Preparation records the view's padded dimension. + let tail = unsafe { tile.codes.visit_panels::<2>(tile.physical_k, visit) }.unwrap(); + // SAFETY: The remainder comes from the same tile and contains one row. + let panel = unsafe { tile.bind_panel(tail.try_as_panel::<1>().unwrap(), tail.start()) }; + actual.extend(panel.meta.iter().map(|m| (m.a, m.b, m.n))); + let expected: Vec<_> = source + .rows() + .map(|row| { + let m = row.meta(); + (m.a, m.b, m.n) + }) + .collect(); + assert_eq!(actual, expected); + } + + fn assert_b_preparation( + sentinel: T, + convert: impl Fn(u8) -> T, + ) where + T: Copy + Default + PartialEq + std::fmt::Debug, + Unsigned: Representation, + { + let mut scratch = BScratch::::default(); + for dim in [1, PACK - 1, PACK, PACK + 1, 15, 16, 17] { + let k = DimK::new(NonZeroUsize::new(dim).unwrap()); + let stride = crate::minmax::Data::::canonical_bytes(dim); + // Unaligned storage and dirty high bits exercise the packed input boundary. + let mut storage = vec![0xff; 1 + 3 * stride]; + let bytes = &mut storage[1..]; + for (i, raw) in bytes.chunks_exact_mut(stride).enumerate() { + let mut row = + crate::minmax::DataMutRef::::from_canonical_front_mut(raw, dim) + .unwrap(); + let mut sum = 0; + let max_code = (1 << B_BITS) - 1; + for j in 0..dim { + let code = (max_code - ((i + j) & max_code)) as i64; + row.vector_mut().set(j, code).unwrap(); + sum += code; + } + let scale = (i + 1) as f32 * 0.25; + row.set_meta(MinMaxCompensation { + dim: dim as u32, + a: scale, + b: i as f32 - 3.0, + n: scale * sum as f32, + norm_squared: (i * i + 1) as f32, + }); + } + let b = MatRef::new(MinMaxMeta::::new(3, dim), bytes).unwrap(); + for (start, count) in [(0, 3), (2, 1), (1, 2)] { + scratch.codes.fill(sentinel); + scratch.meta.fill(Coefficients::default()); + let mut conversions = 0; + // SAFETY: The selected range is within B, which has dimension K. + let tile = unsafe { + scratch.prepare::( + b, + start, + NonZeroUsize::new(count).unwrap(), + k, + |code| { + conversions += 1; + convert(code) + }, + ) + }; + assert_eq!(conversions, count * dim); + let padded = padded_k::(k); + let row_len = padded.value().get(); + // SAFETY: Preparation constructs the view with the recorded padded dimension. + let codes = unsafe { tile.codes.as_std_slice(tile.physical_k) }; + assert_eq!(codes.len(), count * row_len); + assert_eq!(tile.meta.len(), count); + for i in 0..count { + let row = b.get_row(start + i).unwrap(); + let expected = row.meta(); + let actual = tile.meta[i]; + assert_eq!( + (actual.a, actual.b, actual.n), + (expected.a, expected.b, expected.n), + ); + for j in 0..dim { + assert_eq!( + codes[i * row_len + j], + convert(row.vector().get(j).unwrap() as u8), + "B tile: K={dim}, start={start}, count={count}, row={i}, column={j}" + ); + } + assert!( + codes[i * row_len + dim..(i + 1) * row_len] + .iter() + .all(|&v| v == T::default()) + ); + } + } + } + } + + fn assert_driver_target_shapes< + W, + const B_BITS: usize, + const MR: usize, + const NR: usize, + const PACK: usize, + >( + arch: W, + ) where + W: ExtraWide + LoadStore, + Unsigned: Representation, + for<'a> PanelKernel<'a, W, MR, NR, PACK>: driver::PanelKernel, + { + let mut rng = StdRng::seed_from_u64(0x8f61d9c34a257be0); + let mut shapes: Vec<_> = (1..NR).map(|n| (MR + 1, n, PACK + 1)).collect(); + if !cfg!(miri) { + shapes.extend([(16, 16, 256), (17, 17, 257)]); + } + for (m, n, k) in shapes { + let a = random_minmax_matrix::<8>(m, k, &mut rng); + let b = random_minmax_matrix::(n, k, &mut rng); + let expected = reference_max_inner_products::<8, B_BITS>(a.as_view(), b.as_view()); + let packed = PackedA::::new(a.as_view()).unwrap(); + let mut c = vec![f32::INFINITY; m]; + driver::Drive::drive(&mut checked_driver::<_, 8, B_BITS, MR, NR, PACK>( + arch, + &packed, + b.as_view(), + &mut c, + )); + assert_scores_close( + &c, + &expected, + format_args!("M={m}, N={n}, K={k}, MR={MR}, NR={NR}, PACK={PACK}"), + ); + } + } + + fn assert_microkernel_nan_reduction( + arch: W, + ) where + W: ExtraWide, + { + let k = DimK::new(NonZeroUsize::new(1).unwrap()); + let a_codes = vec![W::AElem::default(); MR * PACK]; + let b_codes = vec![W::BElem::default(); NR * PACK]; + arch.run(|| { + let mut n = [0.0; MR]; + n[MR - 1] = f32::NAN; + let a_meta = Coefficients { + a: [1.0; MR], + b: [1.0; MR], + n, + }; + for (initial, score) in [ + (f32::NEG_INFINITY, f32::NAN), + (f32::NEG_INFINITY, 5.0), + (f32::NEG_INFINITY, -5.0), + (f32::NEG_INFINITY, f32::INFINITY), + (7.0, f32::NAN), + ] { + for column in [0, NR - 1] { + let mut b_meta = [Coefficients { + a: 1.0, + b: 0.0, + n: f32::NAN, + }; NR]; + b_meta[column].n = score; + let mut c = [initial; MR]; + let mut kernel = MicroKernel:: { + arch, + a: APanel { + // SAFETY: A is one full MR panel, padded to PACK for logical K=1. + codes: unsafe { packed::Panel::new(Slice::new(&a_codes), k) }, + meta: &a_meta, + }, + b: BPanel { + // SAFETY: B contains NR rows of PACK initialized codes. + codes: unsafe { + unpacked::Panel::new(Slice::new(&b_codes), padded_k::(k)) + }, + meta: &b_meta, + k, + physical_k: padded_k::(k), + }, + c: &mut c, + }; + driver::MicroKernel::micro_kernel(&mut kernel); + let mut expected = [initial.max(score); MR]; + expected[MR - 1] = initial; + assert_eq!( + c, expected, + "NaN reduction: initial={initial}, score={score}, column={column}", + ); + } + } + }); + } + + #[test] + fn scalar_microkernel_ignores_nan_candidates() { + assert_microkernel_nan_reduction::(Scalar::new()); + assert_microkernel_nan_reduction::(Scalar::new()); + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn avx2_microkernel_ignores_nan_candidates() { + if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { + assert_microkernel_nan_reduction::<_, 16, 4, 4>(arch); + assert_microkernel_nan_reduction::<_, 16, 4, 2>(arch); + } else { + eprintln!("AVX2 unavailable; NaN regression not executed"); + } + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn avx2_half_reduction_matches_whole_rows() { + use diskann_wide::arch::x86_64::V3; + + fn check(arch: V3) { + arch.run(|| { + let acc: [i32x16; NR] = core::array::from_fn(|j| { + i32x16::::from_array( + arch, + core::array::from_fn(|i| (i * 37 + j * 11) as i32), + ) + }); + let a = core::array::from_fn(|i| (i + 1) as f32 * 0.125); + let b = core::array::from_fn(|i| i as f32 * 0.25 - 2.0); + let mut n = core::array::from_fn(|i| i as f32 * 0.75); + n[9] = f32::NAN; + let query = Coefficients { + a: &a, + b: &b, + n: &n, + }; + let mut docs: [Coefficients; NR] = core::array::from_fn(|j| Coefficients { + a: (j + 1) as f32 * 0.25, + b: 1.0 - j as f32, + n: j as f32 * 0.5, + }); + for special in [0.0, f32::NAN, f32::INFINITY, f32::NEG_INFINITY] { + docs[NR - 1].n = special; + let initial = core::array::from_fn(|i| { + if i.is_multiple_of(2) { + f32::NEG_INFINITY + } else { + i as f32 + } + }); + let mut expected = initial; + compensated_max_into(arch, acc, query, &docs, 31, &mut expected); + let mut pair = initial; + let mut quad = initial; + >::max_into(arch, acc, query, &docs, 31, &mut pair); + >::max_into(arch, acc, query, &docs, 31, &mut quad); + assert_eq!( + pair.map(f32::to_bits), + expected.map(f32::to_bits), + "NR={NR}" + ); + assert_eq!( + quad.map(f32::to_bits), + expected.map(f32::to_bits), + "NR={NR}" + ); + } + }); + } + + if let Some(arch) = V3::new_checked() { + check::<1>(arch); + check::<2>(arch); + check::<4>(arch); + check::<6>(arch); + check::<8>(arch); + } + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn avx512_microkernel_ignores_nan_candidates() { + if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked() { + assert_microkernel_nan_reduction::<_, 16, 8, 4>(arch); + assert_microkernel_nan_reduction::<_, 16, 8, 2>(arch); + } else { + eprintln!("V4 unavailable; NaN regression not executed"); + } + } + + #[cfg(target_arch = "aarch64")] + #[test] + fn neon_microkernel_ignores_nan_candidates() { + if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { + assert_microkernel_nan_reduction::<_, 16, 4, 4>(arch); + } else { + eprintln!("Neon/dotprod unavailable; NaN regression not executed"); + } + } + + mod all_widths { + use super::*; + + fn assert_backend_code_range(arch: W) + where + W: ExtraWide, + { + assert!((1..=8).contains(&W::MAX_B_BITS)); + let max_b = u8::MAX >> (8 - W::MAX_B_BITS); + arch.run(|| { + let codes: Vec<_> = (0..MR * PACK) + .map(|i| { + W::AElem::from(if (i / PACK).is_multiple_of(2) { + u8::MAX + } else { + 1 + }) + }) + .collect(); + let b_codes = [W::convert_b(max_b); PACK]; + // SAFETY: The buffers contain one complete A pack and B pack. + let (a, b) = unsafe { + ( + arch.load_a(Slice::new(&codes)), + arch.splat(Slice::new(&b_codes)), + ) + }; + let acc = arch.dot(arch.zero_acc(), a, b); + let acc = arch.dot(acc, a, b); + let coefficients = Coefficients { + a: &[1.0; MR], + b: &[0.0; MR], + n: &[0.0; MR], + }; + let mut actual = [f32::NEG_INFINITY; MR]; + arch.max_into( + [acc], + coefficients, + &[Coefficients { + a: 1.0, + ..Default::default() + }], + 2 * PACK, + &mut actual, + ); + assert_eq!( + actual, + core::array::from_fn(|i| { + let a = if i.is_multiple_of(2) { 255 } else { 1 }; + (2 * PACK * a * usize::from(max_b)) as f32 + }), + ); + }); + } + + fn assert_dimension_limit() { + let max_a = (1u64 << A_BITS) - 1; + let max_b = (1u64 << B_BITS) - 1; + let limit = max_k::() as u64; + assert!(limit * max_a * max_b <= i32::MAX as u64); + assert!((limit + 1) * max_a * max_b > i32::MAX as u64); + } + + #[test] + fn dimension_limits_cover_all_input_widths() { + fn check_a() { + assert_dimension_limit::(); + assert_dimension_limit::(); + assert_dimension_limit::(); + assert_dimension_limit::(); + assert_dimension_limit::(); + assert_dimension_limit::(); + assert_dimension_limit::(); + assert_dimension_limit::(); + } + check_a::<1>(); + check_a::<2>(); + check_a::<3>(); + check_a::<4>(); + check_a::<5>(); + check_a::<6>(); + check_a::<7>(); + check_a::<8>(); + } + + #[test] + fn scalar_backends_cover_supported_code_ranges() { + assert_backend_code_range::<_, 16, 2>(Scalar::new()); + assert_backend_code_range::<_, 16, 4>(Scalar::new()); + assert_backend_code_range::<_, 8, 2>(Scalar::new()); + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn avx2_backends_cover_supported_code_ranges() { + if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { + assert_backend_code_range::<_, 16, 2>(arch); + assert_backend_code_range::<_, 16, 4>(arch); + assert_backend_code_range::<_, 8, 2>(arch); + assert_backend_code_range::<_, 8, 4>(arch); + } else { + eprintln!("AVX2 unavailable; all-width regression not executed"); + } + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn avx512_backends_cover_supported_code_ranges() { + if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked() { + assert_backend_code_range::<_, 16, 2>(arch); + assert_backend_code_range::<_, 16, 4>(arch); + } else { + eprintln!("V4 unavailable; all-width regression not executed"); + } + } + + #[cfg(target_arch = "aarch64")] + #[test] + fn neon_backends_cover_supported_code_ranges() { + if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { + assert_backend_code_range::<_, 16, 4>(arch); + assert_backend_code_range::<_, 8, 4>(arch); + } else { + eprintln!("Neon/dotprod unavailable; all-width regression not executed"); + } + } + + fn check_reference() + where + Unsigned: Representation + Representation, + MinMaxIP: for<'a, 'b> PureDistanceFunction< + DataRef<'a, A_BITS>, + DataRef<'b, B_BITS>, + distances::MathematicalResult, + >, + { + let mut rng = StdRng::seed_from_u64(0x14a7f0b2c69d538e); + for k in [3, 65] { + let a = random_minmax_matrix::(3, k, &mut rng); + let b = random_minmax_matrix::(5, k, &mut rng); + let expected: Vec<_> = a + .rows() + .map(|a| { + b.rows() + .map(|b| { + , + DataRef<'_, B_BITS>, + distances::MathematicalResult, + >>::evaluate(a, b) + .unwrap() + .into_inner() + }) + .fold(f32::NEG_INFINITY, f32::max) + }) + .collect(); + assert_scores_close( + &reference_max_inner_products::(a.as_view(), b.as_view()), + &expected, + format_args!("scalar reference: {A_BITS}x{B_BITS}, K={k}"), + ); + } + } + + #[test] + fn reference_agrees_with_existing_pairwise_combinations() { + check_reference::<1, 1>(); + check_reference::<2, 2>(); + check_reference::<3, 3>(); + check_reference::<4, 4>(); + check_reference::<5, 5>(); + check_reference::<6, 6>(); + check_reference::<7, 7>(); + check_reference::<8, 8>(); + check_reference::<8, 4>(); + check_reference::<8, 2>(); + check_reference::<8, 1>(); + } + + #[cfg(not(miri))] + #[test] + fn low_bit_inputs_can_exceed_the_eight_bit_dimension_limit() { + fn check() + where + Unsigned: Representation + Representation, + { + let k = max_k::<8, 8>() + 1; + let max_a = (1 << A_BITS) - 1; + let max_b = (1 << B_BITS) - 1; + let a = constant_minmax_matrix::(1, k, max_a, 1.0, 0.0); + let b = constant_minmax_matrix::(1, k, max_b, 1.0, 0.0); + let packed = PackedA::::new(a.as_view()).unwrap(); + let mut c = [0.0]; + driver::Drive::drive(&mut checked_driver::<_, A_BITS, B_BITS, 16, 2, 2>( + Scalar::new(), + &packed, + b.as_view(), + &mut c, + )); + assert_eq!(c[0], (k as i64 * max_a * max_b) as f32); + } + check::<1, 8>(); + check::<8, 1>(); + } + } + + mod minmax8 { + use super::*; + + fn check_backend(arch: W) + where + W: ExtraWide<16, PACK> + LoadStore, + for<'a> PanelKernel<'a, W, 16, NR, PACK>: driver::PanelKernel, + { + assert_driver_contract::(arch); + assert_driver_target_shapes::(arch); + assert_unsigned_byte_products::(arch); + } + + fn assert_unsigned_byte_products(arch: W) + where + W: ExtraWide<16, PACK> + LoadStore, + for<'a> PanelKernel<'a, W, 16, NR, PACK>: driver::PanelKernel, + { + if cfg!(miri) { + return; + } + for k in [PACK - 1, PACK, PACK + 1] { + let boundary_codes = [0, 1, 127, 128, 255]; + let mut a = constant_minmax_matrix::<8>(17, k, 0, 1.0, 0.0); + for (i, mut row) in a.rows_mut().enumerate() { + let code = boundary_codes[i % boundary_codes.len()]; + for j in 0..k { + row.vector_mut().set(j, code).unwrap(); + } + let mut meta = row.meta(); + meta.n = (k as i64 * code) as f32; + row.set_meta(meta); + } + let a = PackedA::::new(a.as_view()).unwrap(); + let mut c = vec![f32::NAN; 17]; + for code in boundary_codes { + let b = constant_minmax_matrix::<8>(1, k, code, 1.0, 0.0); + driver::Drive::drive(&mut checked_driver::( + arch, + &a, + b.as_view(), + &mut c, + )); + for (i, &actual) in c.iter().enumerate() { + assert_eq!( + actual, + (boundary_codes[i % boundary_codes.len()] as usize * code as usize * k) + as f32, + "A={}, B={code}, K={k}", + boundary_codes[i % boundary_codes.len()], + ); + } + } + } + } + + #[test] + fn scalar_driver_obeys_maxsim_contract() { + check_backend::<_, 2, 2>(Scalar::new()); + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn avx2_driver_obeys_maxsim_contract() { + if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { + check_backend::<_, 4, 2>(arch); + } else { + eprintln!("AVX2 unavailable; MinMax8 x MinMax8 backend not executed"); + } + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn avx512_driver_obeys_maxsim_contract() { + if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked() { + check_backend::<_, 8, 2>(arch); + } else { + eprintln!("V4 unavailable; MinMax8 x MinMax8 backend not executed"); + } + } + + #[cfg(target_arch = "aarch64")] + #[test] + fn neon_driver_obeys_maxsim_contract() { + if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { + check_backend::<_, 4, 4>(arch); + } else { + eprintln!("Neon/dotprod unavailable; MinMax8 x MinMax8 backend not executed"); + } + } + + #[cfg(not(miri))] + #[test] + #[should_panic(expected = "exceeds the MinMax8 x MinMax8 limit")] + fn driver_rejects_k_above_accumulator_limit() { + let k = max_k::<8, 8>() + 1; + let a = constant_minmax_matrix::<8>(1, k, 255, 1.0, 0.0); + let b = constant_minmax_matrix::<8>(1, k, 255, 1.0, 0.0); + let a = PackedA::::new(a.as_view()).unwrap(); + let mut c = [0.0]; + let _ = checked_driver::<_, 8, 8, 16, 2, 2>(Scalar::new(), &a, b.as_view(), &mut c); + } + } +} diff --git a/diskann-wide/src/arch/aarch64/double.rs b/diskann-wide/src/arch/aarch64/double.rs index 6e0f0e3683..f6d8ac5020 100644 --- a/diskann-wide/src/arch/aarch64/double.rs +++ b/diskann-wide/src/arch/aarch64/double.rs @@ -221,6 +221,22 @@ impl crate::SIMDCast for i32x8 { } } +impl crate::SIMDCast for u32x8 { + type Cast = f32x8; + #[inline(always)] + fn simd_cast(self) -> f32x8 { + f32x8::new(self.0.simd_cast(), self.1.simd_cast()) + } +} + +impl crate::SIMDCast for u32x16 { + type Cast = f32x16; + #[inline(always)] + fn simd_cast(self) -> f32x16 { + f32x16::new(self.0.simd_cast(), self.1.simd_cast()) + } +} + /////////// // Tests // /////////// @@ -492,6 +508,8 @@ mod tests { test_utils::ops::test_cast!(f32x16 => f16x16, 0xba8fe343fc9dbeff, test_neon()); test_utils::ops::test_cast!(i32x8 => f32x8, 0xba8fe343fc9dbeff, test_neon()); + test_utils::ops::test_cast!(u32x8 => f32x8, 0x1e9e37b58fb3f1a8, test_neon()); + test_utils::ops::test_cast!(u32x16 => f32x16, 0xd2b068a9bf3f9d24, test_neon()); // SplitJoin + ZipUnzip for f16x16 test_utils::ops::test_splitjoin!(f16x16 => f16x8, 0x2e301b7e12090d5c, test_neon()); diff --git a/diskann-wide/src/arch/aarch64/u32x4_.rs b/diskann-wide/src/arch/aarch64/u32x4_.rs index 3230901098..e1b0298dc0 100644 --- a/diskann-wide/src/arch/aarch64/u32x4_.rs +++ b/diskann-wide/src/arch/aarch64/u32x4_.rs @@ -139,6 +139,14 @@ impl SIMDDotProduct for u32x4 { // Reinterprets // ////////////////// +impl SIMDReinterpret for u32x4 { + #[inline(always)] + fn reinterpret_simd(self) -> u8x16 { + // SAFETY: Allowed by the Neon architecture. + u8x16(unsafe { vreinterpretq_u8_u32(self.0) }) + } +} + impl SIMDReinterpret for u32x4 { #[inline(always)] fn reinterpret_simd(self) -> i8x16 { @@ -204,6 +212,7 @@ mod tests { // Reinterprets test_utils::ops::test_reinterpret!(u32x4 => i8x16, 0xbfad755f32d25e5c, test_neon()); + test_utils::ops::test_reinterpret!(u32x4 => u8x16, 0xd56e9c21a7183b4f, test_neon()); // Conversions test_utils::ops::test_lossless_convert!(u16x4 => u32x4, 0x2e07f4c89da53bb1, test_neon()); diff --git a/diskann-wide/src/arch/x86_64/v3/conversion.rs b/diskann-wide/src/arch/x86_64/v3/conversion.rs index a03ac92d66..a2cdd19dab 100644 --- a/diskann-wide/src/arch/x86_64/v3/conversion.rs +++ b/diskann-wide/src/arch/x86_64/v3/conversion.rs @@ -7,7 +7,7 @@ use std::arch::x86_64::*; use super::{ - f16x8, f16x16, f32x8, f32x16, i8x16, i8x32, i16x8, i16x16, i32x8, u8x16, u8x32, u32x8, + f16x8, f16x16, f32x8, f32x16, i8x16, i8x32, i16x8, i16x16, i32x8, i32x16, u8x16, u8x32, u32x8, }; use crate::{LoHi, SIMDCast, SIMDReinterpret, SplitJoin, helpers}; @@ -83,6 +83,14 @@ impl SIMDCast for f32x16 { } } +impl SIMDCast for i32x16 { + type Cast = f32x16; + #[inline(always)] + fn simd_cast(self) -> f32x16 { + self.split().map(|x| -> f32x8 { x.simd_cast() }).join() + } +} + // i8 to i16 helpers::unsafe_map_conversion!(i8x16, i16x16, _mm256_cvtepi8_epi16, "avx2"); @@ -177,4 +185,5 @@ mod test_x86_conversions { test_utils::ops::test_cast!(f32x16 => f16x16, 0x7b9c9afee7e6ac63, V3::new_checked_uncached()); test_utils::ops::test_cast!(i32x8 => f32x8, 0xde4fbf25c554b29e, V3::new_checked_uncached()); + test_utils::ops::test_cast!(i32x16 => f32x16, 0x83e571a946db20cf, V3::new_checked_uncached()); } diff --git a/diskann-wide/src/arch/x86_64/v4/conversion.rs b/diskann-wide/src/arch/x86_64/v4/conversion.rs index 16f679db69..37393ebc0f 100644 --- a/diskann-wide/src/arch/x86_64/v4/conversion.rs +++ b/diskann-wide/src/arch/x86_64/v4/conversion.rs @@ -8,8 +8,8 @@ use std::arch::x86_64::*; use super::{ f16x8_::f16x8, f16x16_::f16x16, f32x8_::f32x8, f32x16_::f32x16, i8x16_::i8x16, i8x32_::i8x32, - i8x64_::i8x64, i16x8_::i16x8, i16x16_::i16x16, i16x32_::i16x32, i32x8_::i32x8, u8x16_::u8x16, - u8x32_::u8x32, u8x64_::u8x64, u32x8_::u32x8, u32x16_::u32x16, u64x8_::u64x8, + i8x64_::i8x64, i16x8_::i16x8, i16x16_::i16x16, i16x32_::i16x32, i32x8_::i32x8, i32x16_::i32x16, + u8x16_::u8x16, u8x32_::u8x32, u8x64_::u8x64, u32x8_::u32x8, u32x16_::u32x16, u64x8_::u64x8, }; use crate::{SIMDCast, SIMDReinterpret, SIMDVector, arch::x86_64::v3, helpers}; @@ -78,6 +78,7 @@ helpers::unsafe_map_conversion!(u8x32, i16x32, _mm512_cvtepu8_epi16, "avx512bw") // i32 to f32 helpers::unsafe_map_cast!(i32x8 => (f32, f32x8), _mm256_cvtepi32_ps, "avx"); +helpers::unsafe_map_cast!(i32x16 => (f32, f32x16), _mm512_cvtepi32_ps, "avx512f"); helpers::unsafe_map_cast!(i16x16 => (u8, u8x16), _mm256_cvtepi16_epi8, "avx512bw,avx512vl"); helpers::unsafe_map_cast!(i16x16 => (i8, i8x16), _mm256_cvtepi16_epi8, "avx512bw,avx512vl"); @@ -220,6 +221,7 @@ mod test_x86_conversions { test_utils::ops::test_cast!(f32x16 => f16x16, 0xb3cbae34def475df, V4::new_checked_uncached()); test_utils::ops::test_cast!(i32x8 => f32x8, 0xde4fbf25c554b29e, V4::new_checked_uncached()); + test_utils::ops::test_cast!(i32x16 => f32x16, 0x42ec9b7d219536a1, V4::new_checked_uncached()); test_utils::ops::test_cast!(i16x16 => u8x16, 0x0f81df9e640b0269, V4::new_checked_uncached()); test_utils::ops::test_cast!(i16x16 => i8x16, 0x4ab1546b9d0e4046, V4::new_checked_uncached()); diff --git a/diskann-wide/src/emulated.rs b/diskann-wide/src/emulated.rs index be12697949..7161adee83 100644 --- a/diskann-wide/src/emulated.rs +++ b/diskann-wide/src/emulated.rs @@ -665,6 +665,7 @@ impl_cast!(f32 => f16, 8); impl_cast!(f32 => f16, 16); impl_cast!(i32 => f32, 8); +impl_cast!(i32 => f32, 16); /////////////// // SplitJoin // @@ -1175,6 +1176,7 @@ mod test_emulated { test_utils::ops::test_cast!(Emulated => Emulated, 0x2b637e21afd9ef6c, SC); test_utils::ops::test_cast!(Emulated => Emulated, 0x2b08e8ec7e49323b, SC); + test_utils::ops::test_cast!(Emulated => Emulated, 0x83e571a946db20cf, SC); ////////////// // ZipUnzip //