diff --git a/diskann-quantization/src/matrix_kernels/maxsim/mod.rs b/diskann-quantization/src/matrix_kernels/maxsim/mod.rs index e320afcfd5..d16c8ba1b5 100644 --- a/diskann-quantization/src/matrix_kernels/maxsim/mod.rs +++ b/diskann-quantization/src/matrix_kernels/maxsim/mod.rs @@ -19,10 +19,6 @@ 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)] 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 index d0173c5e41..901445ef60 100644 --- 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 @@ -17,6 +17,7 @@ use std::{ ops::{Add, Mul}, }; +use diskann_utils::views::rowmajor::{self, Matrix, MatrixMut}; use diskann_wide::{ Const, SIMDCast, SIMDDotProduct, SIMDMinMax, SIMDReinterpret, SIMDVector, SupportedLaneCount, arch::{Architecture, Scalar, Target}, @@ -32,7 +33,7 @@ use crate::{ ptr::Slice, util::LoadStore, }, - minmax::MinMaxMeta, + minmax::{Coefficients, DataRef, MinMaxMeta}, multi_vector::{BlockTransposed, MatRef, Overflow}, }; @@ -51,31 +52,21 @@ 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, +pub(crate) const fn max_k(a_bits: usize, b_bits: usize) -> 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)] -pub(crate) struct PackedA -where - Unsigned: Representation, -{ +pub(crate) struct PackedA { codes: BlockTransposed, meta: Vec>, + a_bits: usize, } -impl PackedA +impl PackedA where T: Copy + Default + From, - Unsigned: Representation, { /// Zero-pad partial rows and K packs. /// @@ -83,7 +74,12 @@ where /// /// Returns [`Overflow`] if the packed data or coefficient allocation size overflows. #[expect(clippy::expect_used)] - pub(crate) fn new(a: MatRef<'_, MinMaxMeta>) -> Result { + pub(crate) fn new( + a: MatRef<'_, MinMaxMeta>, + ) -> Result + where + Unsigned: Representation, + { let k = a.repr().intrinsic_dim(); let mut codes = BlockTransposed::try_new(a.num_vectors(), k)?; let panels = codes.num_blocks(); @@ -111,7 +107,11 @@ where dst.set(j, T::from(unsafe { v.get_unchecked(j) } as u8)); } } - Ok(Self { codes, meta }) + Ok(Self { + codes, + meta, + a_bits: A_BITS, + }) } fn as_view(&self) -> Option> { @@ -167,34 +167,25 @@ struct APanel<'a, T, const MR: usize, const PACK: usize> { 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>, + b: rowmajor::Ref<'a, u8>, + decode_tile: DecodeTileFn, 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, +impl<'a, W: ExtraWide, const MR: usize, const NR: usize, const PACK: usize> + Driver<'a, W, MR, NR, PACK> { /// # Panics /// - /// Panics if K exceeds the accumulator limit for these bit widths. + /// Panics if K exceeds the accumulator limit for these bit widths or the + /// canonical document storage size overflows. /// Valid integer dot products may still round when converted to f32. /// /// # Safety @@ -202,15 +193,15 @@ where /// 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( + pub(crate) unsafe fn new( arch: W, - a: &'a PackedA, + a: &'a PackedA, b: MatRef<'a, MinMaxMeta>, c: &'a mut [f32], cache: Cache, ) -> Self where - Unsigned: Representation, + Unsigned: Representation, Self: driver::Drive, { const { @@ -219,16 +210,42 @@ where "B bit width exceeds the backend's supported range" ) }; + let dim = a.codes.ncols(); + bounds::check_eq!(bounds::Bound::new(b.repr().intrinsic_dim()), dim); + Self::from_canonical( + arch, + a, + b.as_canonical_view() + .expect("canonical document storage size must not overflow"), + c, + cache, + B_BITS, + decode_tile::, + ) + } + + // Keep common initialization out of the bit-width-specific constructors. + #[inline(never)] + #[expect(clippy::expect_used)] + fn from_canonical( + arch: W, + a: &'a PackedA, + b: rowmajor::Ref<'a, u8>, + c: &'a mut [f32], + cache: Cache, + b_bits: usize, + decode_tile: DecodeTileFn, + ) -> Self { let rows = a.codes.nrows(); let a_bytes = a.panel_bytes(); + let a_bits = a.a_bits; 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::() }; + let max = max_k(a_bits, b_bits); assert!( dim <= max, - "contraction dimension {dim} exceeds the MinMax{A_BITS} x MinMax{B_BITS} limit of {max}", + "contraction dimension {dim} exceeds the MinMax{a_bits} x MinMax{b_bits} limit of {max}", ); let b_rows_per_tile = Params::new( cache, @@ -241,48 +258,51 @@ where arch, a, b, + decode_tile, c, b_rows_per_tile, } } } -impl driver::Drive - for Driver<'_, W, B_BITS, MR, NR, PACK> +impl driver::Drive + for Driver<'_, W, 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] + #[inline(always)] || { 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) { + while let Some(remaining) = NonZeroUsize::new(self.b.nrows() - 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]); - }); + #[expect( + clippy::expect_used, + reason = "count is bounded by the remaining rows" + )] + let source = self + .b + .subview(start..start + count.get()) + .expect("B tile in bounds"); + let b = scratch.prepare::(source, self.a.k, self.decode_tile); + self.a.visit_panels( + #[inline(always)] + |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) }; + // Already inside the driver's target-feature entry. + let c = kernel.run(self.arch); + self.arch.store(c, &mut self.c[row..end]); + }, + ); start += count.get(); } }, @@ -296,6 +316,51 @@ struct BScratch { meta: Vec>, } +// Only decoding is specialized by bit width; allocation, padding, and tile views are shared. +type DecodeTileFn = fn( + source: rowmajor::Ref<'_, u8>, + k: DimK, + codes: rowmajor::Mut<'_, T>, + coefficients: &mut [Coefficients], +); + +/// Decode canonical rows into the logical columns of an already allocated tile. +/// Shared across architectures with the same decoded element type. +#[inline(never)] +#[expect( + clippy::expect_used, + reason = "the driver supplies canonical rows and output dimensions are checked" +)] +fn decode_tile( + source: rowmajor::Ref<'_, u8>, + k: DimK, + mut codes: rowmajor::Mut<'_, T>, + coefficients: &mut [Coefficients], +) where + Unsigned: Representation, +{ + let dim = k.value().get(); + assert_eq!(source.nrows(), codes.nrows(), "one output row per document"); + assert_eq!( + source.nrows(), + coefficients.len(), + "one coefficient record per document" + ); + assert!( + dim <= codes.ncols(), + "tile rows must hold the logical dimension" + ); + + for ((bytes, dst), meta) in source.rows().zip(codes.rows_mut()).zip(coefficients) { + let row = + DataRef::::from_canonical_front(bytes, dim).expect("canonical document row"); + *meta = row.meta().into(); + row.vector() + .decode_into(&mut dst[..dim], T::from_code) + .expect("B output must match the logical code count"); + } +} + /// Codes and coefficients share row order; code rows are zero-padded to PACK. #[derive(Clone, Copy)] struct BTile<'a, T, const PACK: usize> { @@ -349,43 +414,33 @@ fn padded_k(k: DimK) -> DimK { } impl BScratch { - /// Padding bypasses `convert`. + /// Prepare a nonempty tile. The decoder writes logical columns; padding bypasses it. /// - /// # Safety + /// # Panics /// - /// B's dimension must equal K, and `start..start + count` must be in bounds. + /// Panics if the source tile is empty or its output storage size overflows. #[expect(clippy::expect_used)] - unsafe fn prepare( + fn prepare( &mut self, - b: MatRef<'_, MinMaxMeta>, - start: usize, - count: NonZeroUsize, + source: rowmajor::Ref<'_, u8>, k: DimK, - mut convert: impl FnMut(u8) -> T, - ) -> BTile<'_, T, PACK> - where - Unsigned: Representation, - { + decode: impl FnOnce(rowmajor::Ref<'_, u8>, DimK, rowmajor::Mut<'_, T>, &mut [Coefficients]), + ) -> BTile<'_, T, PACK> { + let count = NonZeroUsize::new(source.nrows()).expect("B tile must be nonempty"); 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()); + let mut codes = rowmajor::Mut::try_from_data(&mut self.codes, count.get(), stride) + .expect("resized tile storage matches its shape"); + // Empty per-row fills can still lower to memset calls. + if dim < stride { + for dst in codes.rows_mut() { + dst[dim..].fill(T::default()); + } } + decode(source, k, codes, &mut self.meta); BTile { // SAFETY: The buffers contain count rows, each with stride initialized codes. codes: unsafe { unpacked::View::new(Slice::new(&self.codes), count, physical_k) }, @@ -564,20 +619,40 @@ impl, const MR: usize, const NR: usize, const PACK: usize // Backends // //----------// +/// Convert an unsigned code to the backend's decoded storage type. +pub(crate) trait UnpackedCode: Copy + Default { + fn from_code(code: u8) -> Self; +} + +impl UnpackedCode for i8 { + fn from_code(code: u8) -> Self { + code as i8 + } +} + +impl UnpackedCode for i16 { + fn from_code(code: u8) -> Self { + i16::from(code) + } +} + +impl UnpackedCode for u8 { + fn from_code(code: u8) -> Self { + code + } +} + /// 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 BElem: UnpackedCode; 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. @@ -628,11 +703,13 @@ fn compensated_max_into( 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 b = Coefficients { + a: F::splat(arch, b.a), + b: F::splat(arch, b.b), + n: 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; + let ip = a.compensate_inner_product(b, dot, dim); best = best.max_simd_standard(ip); } *scores = best.to_array(); @@ -695,10 +772,6 @@ impl ExtraWide<16, 4> for Scalar { 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. @@ -745,10 +818,6 @@ impl ExtraWide<16, 2> for Scalar { 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. @@ -792,9 +861,6 @@ impl ExtraWide<8, 2> for Scalar { 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. @@ -879,9 +945,6 @@ mod x86_64 { 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. @@ -932,9 +995,6 @@ mod x86_64 { 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. @@ -974,10 +1034,6 @@ mod x86_64 { 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(); @@ -1038,10 +1094,6 @@ mod x86_64 { 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. @@ -1088,10 +1140,6 @@ mod x86_64 { 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. @@ -1138,10 +1186,6 @@ mod x86_64 { 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. @@ -1191,9 +1235,6 @@ mod aarch64 { 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. @@ -1240,10 +1281,6 @@ mod aarch64 { 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. @@ -1362,27 +1399,26 @@ mod tests { 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, + a: &'a PackedA, b: MatRef<'a, MinMaxMeta>, c: &'a mut [f32], - ) -> Driver<'a, W, B_BITS, MR, NR, PACK> + ) -> Driver<'a, W, MR, NR, PACK> where W: ExtraWide + LoadStore, - Unsigned: Representation + Representation, + Unsigned: 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()) } + unsafe { Driver::::new(arch, a, b, c, Cache::detect()) } } fn assert_driver_contract< @@ -1403,10 +1439,10 @@ mod tests { 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 packed = PackedA::::new::<8>(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>( + let mut driver = checked_driver::<_, B_BITS, MR, NR, PACK>( arch, &packed, b.as_view(), @@ -1430,9 +1466,9 @@ mod tests { 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 packed = PackedA::::new::<8>(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>( + driver::Drive::drive(&mut checked_driver::<_, B_BITS, MR, NR, PACK>( arch, &packed, b.as_view(), @@ -1467,6 +1503,72 @@ mod tests { mat } + #[test] + fn decodes_documents_once_per_tile() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + static CALLS: AtomicUsize = AtomicUsize::new(0); + + fn counting_decode( + source: rowmajor::Ref<'_, u8>, + k: DimK, + codes: rowmajor::Mut<'_, i16>, + coefficients: &mut [Coefficients], + ) where + Unsigned: Representation, + { + CALLS.fetch_add(1, Ordering::Relaxed); + decode_tile::(source, k, codes, coefficients); + } + + fn check() + where + Unsigned: Representation, + { + let mut rng = StdRng::seed_from_u64(0xf9380b5db7cb9e41); + for rows in [1, 17] { + for docs in [0, 1, 4, 5, 11] { + let a = random_minmax_matrix::(rows, 5, &mut rng); + let b = random_minmax_matrix::(docs, 5, &mut rng); + let packed = PackedA::::new(a.as_view()).unwrap(); + let expected = reference_max_inner_products(a.as_view(), b.as_view()); + let mut scores = vec![f32::NAN; rows]; + let mut driver = checked_driver::<_, B_BITS, 8, 2, 2>( + Scalar::new(), + &packed, + b.as_view(), + &mut scores, + ); + driver.b_rows_per_tile = NonZeroUsize::new(4).unwrap(); + driver.decode_tile = counting_decode::; + CALLS.store(0, Ordering::Relaxed); + for iteration in 1..=2 { + driver::Drive::drive(&mut driver); + assert_eq!( + CALLS.load(Ordering::Relaxed), + iteration * docs.div_ceil(4), + "B_BITS={B_BITS}, rows={rows}, docs={docs}", + ); + assert_scores_close( + driver.c, + &expected, + format_args!("B_BITS={B_BITS}, rows={rows}, docs={docs}"), + ); + } + } + } + } + + check::<1>(); + check::<2>(); + check::<3>(); + check::<4>(); + check::<5>(); + check::<6>(); + check::<7>(); + check::<8>(); + } + fn assert_compensated_reduction< W, const B_BITS: usize, @@ -1496,10 +1598,10 @@ mod tests { norm_squared: 0.0, }); } - let packed = PackedA::::new(a.as_view()).unwrap(); + let packed = PackedA::::new::<8>(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); + checked_driver::<_, 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); @@ -1507,8 +1609,8 @@ mod tests { 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>( + let packed = PackedA::::new::<8>(a.as_view()).unwrap(); + driver::Drive::drive(&mut checked_driver::<_, B_BITS, MR, NR, PACK>( arch, &packed, b.as_view(), @@ -1522,12 +1624,12 @@ mod tests { #[test] fn accumulator_is_exact_at_dimension_limit() { let arch = Scalar::new(); - let k = max_k::<8, 8>(); + 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 packed = PackedA::::new(a.as_view()).unwrap(); let mut c = [0.0]; - driver::Drive::drive(&mut checked_driver::<_, 8, 8, 8, 2, 2>( + driver::Drive::drive(&mut checked_driver::<_, 8, 8, 2, 2>( arch, &packed, b.as_view(), @@ -1605,14 +1707,14 @@ mod tests { { 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(); + 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 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); @@ -1676,7 +1778,7 @@ mod tests { 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(); + let packed = PackedA::::new(a.as_view()).unwrap(); assert_eq!(packed.panel_bytes().value(), a_bytes); } } @@ -1691,7 +1793,7 @@ mod tests { 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(); + 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(), @@ -1700,9 +1802,8 @@ mod tests { 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) - }; + let mut driver = + unsafe { Driver::<_, 8, 2, 2>::new(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")); @@ -1731,7 +1832,6 @@ mod tests { 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] @@ -1740,16 +1840,11 @@ mod tests { 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 tile = scratch.prepare::<4>( + source.as_view().as_canonical_view().unwrap(), + k, + decode_tile::<4, u8>, + ); let mut actual = Vec::new(); let visit = |codes: unpacked::Panel<'_, u8, 2>, start| { // SAFETY: The tile's traversal supplies this panel and its offset. @@ -1775,7 +1870,7 @@ mod tests { sentinel: T, convert: impl Fn(u8) -> T, ) where - T: Copy + Default + PartialEq + std::fmt::Debug, + T: UnpackedCode + PartialEq + std::fmt::Debug, Unsigned: Representation, { let mut scratch = BScratch::::default(); @@ -1806,24 +1901,15 @@ mod tests { }); } let b = MatRef::new(MinMaxMeta::::new(3, dim), bytes).unwrap(); + let source = b.as_canonical_view().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 tile = scratch.prepare::( + source.subview(start..start + count).unwrap(), + k, + decode_tile::, + ); let padded = padded_k::(k); let row_len = padded.value().get(); // SAFETY: Preparation constructs the view with the recorded padded dimension. @@ -1877,9 +1963,9 @@ mod tests { 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 packed = PackedA::::new::<8>(a.as_view()).unwrap(); let mut c = vec![f32::INFINITY; m]; - driver::Drive::drive(&mut checked_driver::<_, 8, B_BITS, MR, NR, PACK>( + driver::Drive::drive(&mut checked_driver::<_, B_BITS, MR, NR, PACK>( arch, &packed, b.as_view(), @@ -2076,7 +2162,7 @@ mod tests { }) }) .collect(); - let b_codes = [W::convert_b(max_b); PACK]; + let b_codes = [W::BElem::from_code(max_b); PACK]; // SAFETY: The buffers contain one complete A pack and B pack. let (a, b) = unsafe { ( @@ -2115,7 +2201,7 @@ mod tests { fn assert_dimension_limit() { let max_a = (1u64 << A_BITS) - 1; let max_b = (1u64 << B_BITS) - 1; - let limit = max_k::() as u64; + let limit = max_k(A_BITS, B_BITS) as u64; assert!(limit * max_a * max_b <= i32::MAX as u64); assert!((limit + 1) * max_a * max_b > i32::MAX as u64); } @@ -2243,14 +2329,14 @@ mod tests { where Unsigned: Representation + Representation, { - let k = max_k::<8, 8>() + 1; + 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 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>( + driver::Drive::drive(&mut checked_driver::<_, B_BITS, 16, 2, 2>( Scalar::new(), &packed, b.as_view(), @@ -2258,7 +2344,7 @@ mod tests { )); assert_eq!(c[0], (k as i64 * max_a * max_b) as f32); } - check::<1, 8>(); + check::<1, 1>(); check::<8, 1>(); } } @@ -2296,11 +2382,11 @@ mod tests { meta.n = (k as i64 * code) as f32; row.set_meta(meta); } - let a = PackedA::::new(a.as_view()).unwrap(); + 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::( + driver::Drive::drive(&mut checked_driver::( arch, &a, b.as_view(), @@ -2358,12 +2444,12 @@ mod tests { #[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 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 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); + let _ = checked_driver::<_, 8, 16, 2, 2>(Scalar::new(), &a, b.as_view(), &mut c); } } } diff --git a/diskann-quantization/src/minmax/mod.rs b/diskann-quantization/src/minmax/mod.rs index 41c29c3d03..38d0fba5e9 100644 --- a/diskann-quantization/src/minmax/mod.rs +++ b/diskann-quantization/src/minmax/mod.rs @@ -133,9 +133,12 @@ mod vectors; // Exports // ///////////// -pub use multi::{MinMaxKernel, MinMaxMeta}; +pub use multi::{ + MinMaxKernel, MinMaxKernelFactory, MinMaxMaxSim, MinMaxMeta, build_minmax_max_sim, +}; pub use quantizer::{L2Loss, MinMaxQuantizer}; pub use recompress::{RecompressError, Recompressor}; +pub(crate) use vectors::Coefficients; pub use vectors::{ Data, DataMutRef, DataRef, DecompressError, FullQuery, FullQueryMeta, FullQueryMut, FullQueryRef, MetaParseError, MinMaxCompensation, MinMaxCosine, MinMaxCosineNormalized, diff --git a/diskann-quantization/src/minmax/multi/factory.rs b/diskann-quantization/src/minmax/multi/factory.rs new file mode 100644 index 0000000000..eced459851 --- /dev/null +++ b/diskann-quantization/src/minmax/multi/factory.rs @@ -0,0 +1,531 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use diskann_wide::arch::Scalar; +#[cfg(target_arch = "aarch64")] +use diskann_wide::arch::aarch64::Neon; +#[cfg(target_arch = "x86_64")] +use diskann_wide::arch::x86_64::{V3, V4}; + +use super::{ + kernel::{MinMaxMaxSim, MinMaxMaxSimKernel}, + prepared, + reference::ReferenceKernel, +}; +use crate::{ + bits::{Representation, Unsigned}, + matrix_kernels::maxsim::packed_minmax_x_unpacked_minmax::max_k, + minmax::MinMaxMeta, + multi_vector::{BuildMaxSimError, MatRef, MaxSimIsa, NotSupported, QueryMatRef}, +}; + +/// Build a kernel that owns a copy of the MinMax query. +/// +/// `Auto` selects a backend at construction. Explicit ISA requests never silently +/// fall back. `Reference` uses an independent scalar implementation. +/// +/// Kernels return negative maximum inner products, matching the MaxSim distance convention. +/// +/// Supports equal query/document bit widths from 1 through 8, plus 8-bit queries +/// with 1-, 2-, or 4-bit documents. Other combinations fail to compile. +/// +/// # Errors +/// +/// Returns [`BuildMaxSimError`] for an unavailable ISA, a query dimension exceeding +/// the bit-width pair's accumulator limit, or an overflowing storage size. +/// Dimension and canonical-size checks apply to every backend, including empty queries; +/// packed-size checks apply only to prepared kernels. +/// Allocation exhaustion follows the allocator's normal behavior. +/// +/// # Example +/// +/// ``` +/// use diskann_quantization::{ +/// minmax::{MinMaxMeta, build_minmax_max_sim}, +/// multi_vector::{Defaulted, Mat, MaxSimIsa, QueryMatRef}, +/// }; +/// # fn example() -> Result<(), Box> { +/// let query = Mat::new(MinMaxMeta::<8>::new(2, 16), Defaulted)?; +/// let document = Mat::new(MinMaxMeta::<4>::new(3, 16), Defaulted)?; +/// let query = QueryMatRef::from(query.as_view()); +/// let kernel = build_minmax_max_sim::<8, 4>(MaxSimIsa::Auto, query)?; +/// let mut scores = vec![0.0; kernel.nrows()]; +/// kernel.compute_max_sim(document.as_view(), &mut scores)?; +/// assert_eq!(scores, vec![0.0; 2]); +/// # Ok(()) +/// # } +/// # example().unwrap(); +/// ``` +pub fn build_minmax_max_sim( + isa: MaxSimIsa, + query: QueryMatRef<'_, MinMaxMeta>, +) -> Result, BuildMaxSimError> +where + Unsigned: MinMaxKernelFactory, +{ + >::build(isa, query) +} + +/// Precompiled query/document bit-width combinations. +#[doc(hidden)] +pub trait MinMaxKernelFactory: + Representation + Representation +{ + fn build( + isa: MaxSimIsa, + query: QueryMatRef<'_, MinMaxMeta>, + ) -> Result, BuildMaxSimError> + where + Unsigned: Representation + Representation; +} + +// Concrete, non-generic implementations keep backend compilation in this crate. +macro_rules! factory { + ($a:literal, $b:literal, $scalar:ident, $v3:ident, $v4:ident) => { + impl MinMaxKernelFactory<$a, $b> for Unsigned { + fn build( + isa: MaxSimIsa, + query: QueryMatRef<'_, MinMaxMeta<$a>>, + ) -> Result, BuildMaxSimError> { + let query = query.0; + validate_query::<$a, $b>(query)?; + let inner: Box> = match RuntimeArch::resolve(isa)? { + RuntimeArch::Scalar(a) => Box::new(prepared::$scalar::<$b>::new(a, query)?), + #[cfg(target_arch = "x86_64")] + RuntimeArch::V3(a) => Box::new(prepared::$v3::<$b>::new(a, query)?), + #[cfg(target_arch = "x86_64")] + RuntimeArch::V4(a) => Box::new(prepared::$v4::<$b>::new(a, query)?), + #[cfg(target_arch = "aarch64")] + RuntimeArch::Neon(a) => Box::new(prepared::NeonBytes::<$b>::new(a, query)?), + RuntimeArch::Reference => Box::new(ReferenceKernel::<$a, $b>::new(query)?), + }; + Ok(MinMaxMaxSim { inner }) + } + } + }; +} +factory!(1, 1, ScalarBytes, V3Bytes, V4Bytes); +factory!(2, 2, ScalarBytes, V3Bytes, V4Bytes); +factory!(3, 3, ScalarBytes, V3Bytes, V4Bytes); +factory!(4, 4, ScalarBytes, V3Bytes, V4Bytes); +factory!(5, 5, ScalarBytes, V3Bytes, V4Bytes); +factory!(6, 6, ScalarBytes, V3Bytes, V4Bytes); +factory!(7, 7, ScalarBytes, V3Words, V4Bytes); +factory!(8, 8, ScalarWords, V3Words, V4Words); +factory!(8, 1, ScalarBytes, V3Bytes, V4Bytes); +factory!(8, 2, ScalarBytes, V3Bytes, V4Bytes); +factory!(8, 4, ScalarBytes, V3Bytes, V4Bytes); + +fn validate_query( + query: MatRef<'_, MinMaxMeta>, +) -> Result<(), BuildMaxSimError> +where + Unsigned: Representation + Representation, +{ + let dim = query.repr().intrinsic_dim(); + let limit = max_k(A_BITS, B_BITS); + if dim > limit { + return Err(BuildMaxSimError::DimTooLarge(dim, limit)); + } + query.as_canonical_view()?; + Ok(()) +} + +#[derive(Debug, Clone, Copy)] +enum RuntimeArch { + Scalar(Scalar), + #[cfg(target_arch = "x86_64")] + V3(V3), + #[cfg(target_arch = "x86_64")] + V4(V4), + #[cfg(target_arch = "aarch64")] + Neon(Neon), + Reference, +} + +impl RuntimeArch { + fn resolve(isa: MaxSimIsa) -> Result { + let unavailable = || NotSupported { + isa, + reason: "requested ISA unavailable on this host", + }; + Ok(match isa { + MaxSimIsa::Auto => { + #[cfg(target_arch = "x86_64")] + { + if let Some(arch) = V4::new_checked() { + return Ok(Self::V4(arch)); + } + if let Some(arch) = V3::new_checked() { + return Ok(Self::V3(arch)); + } + } + #[cfg(target_arch = "aarch64")] + if let Some(arch) = Neon::new_checked() { + return Ok(Self::Neon(arch)); + } + Self::Scalar(Scalar::new()) + } + MaxSimIsa::Scalar => Self::Scalar(Scalar::new()), + MaxSimIsa::Reference => Self::Reference, + #[cfg(target_arch = "x86_64")] + MaxSimIsa::X86_64_V3 => Self::V3(V3::new_checked().ok_or_else(unavailable)?), + #[cfg(target_arch = "x86_64")] + MaxSimIsa::X86_64_V4 => Self::V4(V4::new_checked().ok_or_else(unavailable)?), + #[cfg(not(target_arch = "x86_64"))] + MaxSimIsa::X86_64_V3 | MaxSimIsa::X86_64_V4 => return Err(unavailable()), + #[cfg(target_arch = "aarch64")] + MaxSimIsa::Neon => Self::Neon(Neon::new_checked().ok_or_else(unavailable)?), + #[cfg(not(target_arch = "aarch64"))] + MaxSimIsa::Neon => return Err(unavailable()), + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + minmax::MinMaxCompensation, + multi_vector::{Defaulted, Mat, MaxSimError}, + }; + + fn matrix(rows: usize, dim: usize) -> Mat> + where + Unsigned: Representation, + { + let mut mat = Mat::new(MinMaxMeta::new(rows, dim), Defaulted).unwrap(); + for (i, mut row) in mat.rows_mut().enumerate() { + let mut sum = 0; + for d in 0..dim { + let code = ((i * 17 + d * 29 + 7) % (1 << BITS)) as i64; + row.vector_mut().set(d, code).unwrap(); + sum += code; + } + let a = (i + 1) as f32 / 8.0; + row.set_meta(MinMaxCompensation { + dim: dim as u32, + a, + b: i as f32 / 4.0 - 1.0, + n: a * sum as f32, + norm_squared: 0.0, + }); + } + mat + } + + fn constant_matrix( + rows: usize, + dim: usize, + code: i64, + ) -> Mat> + where + Unsigned: Representation, + { + let mut mat = Mat::new(MinMaxMeta::new(rows, dim), Defaulted).unwrap(); + for mut row in mat.rows_mut() { + for d in 0..dim { + row.vector_mut().set(d, code).unwrap(); + } + row.set_meta(MinMaxCompensation { + dim: dim as u32, + a: 1.0, + n: dim as f32 * code as f32, + ..Default::default() + }); + } + mat + } + + fn available_isas() -> impl Iterator { + [ + MaxSimIsa::Scalar, + MaxSimIsa::Auto, + MaxSimIsa::X86_64_V3, + MaxSimIsa::X86_64_V4, + MaxSimIsa::Neon, + MaxSimIsa::Reference, + ] + .into_iter() + .filter(|isa| isa.is_available()) + } + + fn assert_matches_reference() + where + Unsigned: MinMaxKernelFactory, + { + // Exercise partial query panels, document panels, and K packs on every backend. + let (rows, docs, dim) = (17, 9, 65); + let query = matrix::(rows, dim); + let document = matrix::(docs, dim); + let reference = + build_minmax_max_sim::(MaxSimIsa::Reference, query.as_view().into()) + .unwrap(); + let mut expected = vec![0.0; rows]; + reference + .compute_max_sim(document.as_view(), &mut expected) + .unwrap(); + for isa in available_isas().filter(|&isa| isa != MaxSimIsa::Reference) { + let kernel = + build_minmax_max_sim::(isa, query.as_view().into()).unwrap(); + assert_eq!(kernel.nrows(), rows); + let mut guarded = vec![123.0; rows + 2]; + kernel + .compute_max_sim(document.as_view(), &mut guarded[1..rows + 1]) + .unwrap(); + for (actual, expected) in guarded[1..rows + 1].iter().zip(&expected) { + assert!( + actual == expected + || (actual.is_finite() + && expected.is_finite() + && (actual - expected).abs() <= 1e-5 * expected.abs().max(1.0)), + "{A_BITS}x{B_BITS}, {isa}: {actual} != {expected}" + ); + } + assert_eq!((guarded[0], guarded[rows + 1]), (123.0, 123.0)); + } + } + + #[test] + fn matches_reference_across_supported_bit_widths() { + assert_matches_reference::<1, 1>(); + assert_matches_reference::<2, 2>(); + assert_matches_reference::<3, 3>(); + assert_matches_reference::<4, 4>(); + assert_matches_reference::<5, 5>(); + assert_matches_reference::<6, 6>(); + assert_matches_reference::<7, 7>(); + assert_matches_reference::<8, 8>(); + assert_matches_reference::<8, 1>(); + assert_matches_reference::<8, 2>(); + assert_matches_reference::<8, 4>(); + } + + #[test] + fn handles_empty_inputs() { + fn check() + where + Unsigned: MinMaxKernelFactory<8, B_BITS>, + { + for (rows, docs, dim, expected) in [ + (0, 5, 7, 0.0), + (3, 0, 7, f32::MAX), + (3, 5, 0, 0.0), + (3, 0, 0, f32::MAX), + ] { + let query = matrix::<8>(rows, dim); + let document = matrix::(docs, dim); + for isa in available_isas() { + let kernel = + build_minmax_max_sim::<8, B_BITS>(isa, query.as_view().into()).unwrap(); + let mut scores = vec![f32::NAN; rows]; + kernel + .compute_max_sim(document.as_view(), &mut scores) + .unwrap(); + assert_eq!( + scores, + vec![expected; rows], + "{isa}, B_BITS={B_BITS}, {rows}x{docs}x{dim}" + ); + } + } + } + check::<4>(); + check::<8>(); + } + + #[test] + fn preserves_scores_on_error() { + let query = matrix::<8>(2, 5); + for isa in available_isas() { + let kernel = build_minmax_max_sim::<8, 8>(isa, query.as_view().into()).unwrap(); + let doc = matrix::<8>(1, 6); + let mut scores = [123.0; 2]; + assert!(matches!( + kernel.compute_max_sim(doc.as_view(), &mut scores), + Err(MaxSimError::UnequalDim(6, 5)) + )); + assert_eq!(scores, [123.0; 2]); + let doc = matrix::<8>(1, 5); + let mut scores = [123.0]; + assert!(matches!( + kernel.compute_max_sim(doc.as_view(), &mut scores), + Err(MaxSimError::InvalidBufferLength(1, 2)) + )); + assert_eq!(scores, [123.0]); + } + } + + #[test] + fn rejects_unavailable_isas() { + let query = matrix::<8>(2, 5); + for isa in [MaxSimIsa::X86_64_V3, MaxSimIsa::X86_64_V4, MaxSimIsa::Neon] { + if !isa.is_available() { + assert!(matches!( + build_minmax_max_sim::<8, 8>(isa, query.as_view().into()), + Err(BuildMaxSimError::NotSupported(error)) if error.isa == isa + )); + } + } + } + + #[test] + fn rejects_excessive_dimensions() { + let limit = max_k(8, 8); + let dim = limit + 1; + let query = MatRef::new(MinMaxMeta::<8>::new(0, dim), &[] as &[u8]).unwrap(); + for isa in [MaxSimIsa::Auto, MaxSimIsa::Reference] { + let result = build_minmax_max_sim::<8, 8>(isa, query.into()); + assert!( + matches!(result, Err(BuildMaxSimError::DimTooLarge(found, max)) + if found == dim && max == limit), + "{isa}" + ); + } + } + + #[test] + fn owns_query_data() { + let document = constant_matrix::<4>(2, 3, 2); + for isa in available_isas() { + let kernel = { + let query = constant_matrix::<8>(2, 3, 1); + build_minmax_max_sim::<8, 4>(isa, query.as_view().into()).unwrap() + }; + let mut scores = [f32::NAN; 2]; + kernel + .compute_max_sim(document.as_view(), &mut scores) + .unwrap(); + assert_eq!(scores, [-6.0; 2], "{isa}"); + } + } + + #[test] + fn selects_backend_and_layout() { + fn assert_layout(isa: MaxSimIsa) + where + Unsigned: MinMaxKernelFactory, + { + let query = matrix::(1, 3); + assert_eq!( + build_minmax_max_sim::(isa, query.as_view().into()) + .unwrap() + .inner + .inner_type_id(), + std::any::TypeId::of::(), + "{isa}, A_BITS={A_BITS}, B_BITS={B_BITS}", + ); + } + + assert_layout::<7, 7, prepared::ScalarBytes<7>>(MaxSimIsa::Scalar); + assert_layout::<8, 8, prepared::ScalarWords<8>>(MaxSimIsa::Scalar); + #[cfg(target_arch = "x86_64")] + { + if V3::new_checked().is_some() { + assert_layout::<6, 6, prepared::V3Bytes<6>>(MaxSimIsa::X86_64_V3); + assert_layout::<7, 7, prepared::V3Words<7>>(MaxSimIsa::X86_64_V3); + } + if V4::new_checked().is_some() { + assert_layout::<7, 7, prepared::V4Bytes<7>>(MaxSimIsa::X86_64_V4); + assert_layout::<8, 8, prepared::V4Words<8>>(MaxSimIsa::X86_64_V4); + } + } + #[cfg(target_arch = "aarch64")] + if Neon::new_checked().is_some() { + assert_layout::<8, 8, prepared::NeonBytes<8>>(MaxSimIsa::Neon); + } + } + + #[test] + fn prepared_type_is_independent_of_query_bit_width() { + fn check() + where + Unsigned: MinMaxKernelFactory + MinMaxKernelFactory<8, B_BITS>, + { + let narrow = matrix::(1, 3); + let wide = matrix::<8>(1, 3); + for isa in available_isas().filter(|&isa| isa != MaxSimIsa::Reference) { + let narrow = + build_minmax_max_sim::(isa, narrow.as_view().into()).unwrap(); + let wide = build_minmax_max_sim::<8, B_BITS>(isa, wide.as_view().into()).unwrap(); + assert_eq!( + narrow.inner.inner_type_id(), + wide.inner.inner_type_id(), + "{isa}, B_BITS={B_BITS}", + ); + } + } + check::<1>(); + check::<2>(); + check::<4>(); + } + + #[test] + fn handles_nonfinite_scores() { + fn check() + where + Unsigned: MinMaxKernelFactory, + { + let max_code = (1i64 << B_BITS) - 1; + let max_query = (1i64 << A_BITS) - 1; + let query = constant_matrix::(1, 5, max_query); + let mut document = constant_matrix::(2, 5, max_code); + let finite_score = -(5 * max_query * max_code) as f32; + for isa in available_isas() { + let kernel = + build_minmax_max_sim::(isa, query.as_view().into()).unwrap(); + let mut scores = [0.0]; + for (scale, expected) in [ + (1.0, finite_score), + (f32::NEG_INFINITY, f32::MAX), + (f32::INFINITY, f32::NEG_INFINITY), + (f32::NAN, f32::MAX), + (1.0, finite_score), + ] { + for mut row in document.rows_mut() { + let mut meta = row.meta(); + meta.a = scale; + row.set_meta(meta); + } + kernel + .compute_max_sim(document.as_view(), &mut scores) + .unwrap(); + assert_eq!(scores, [expected], "{isa}, B_BITS={B_BITS}"); + } + let mut row = document.get_row_mut(0).unwrap(); + let mut meta = row.meta(); + meta.a = f32::NAN; + row.set_meta(meta); + kernel + .compute_max_sim(document.as_view(), &mut scores) + .unwrap(); + assert_eq!( + scores, + [finite_score], + "{isa}, B_BITS={B_BITS}, mixed NaN/finite" + ); + } + } + check::<6, 6>(); + check::<7, 7>(); + check::<8, 8>(); + check::<8, 4>(); + } + + #[cfg(not(debug_assertions))] + #[test] + fn rejects_query_size_overflow() { + // With release overflow behavior, rows * 20 wraps to zero for K=0. + // Reject this view before either packing or cloning can access its claimed rows. + let rows = 1usize << (usize::BITS - 2); + let query = MatRef::new(MinMaxMeta::<8>::new(rows, 0), &[] as &[u8]).unwrap(); + for isa in [MaxSimIsa::Auto, MaxSimIsa::Reference] { + assert!(matches!( + build_minmax_max_sim::<8, 4>(isa, query.into()), + Err(BuildMaxSimError::Overflow(_)) + )); + } + } +} diff --git a/diskann-quantization/src/minmax/multi/kernel.rs b/diskann-quantization/src/minmax/multi/kernel.rs new file mode 100644 index 0000000000..e5cf53549c --- /dev/null +++ b/diskann-quantization/src/minmax/multi/kernel.rs @@ -0,0 +1,96 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use crate::{ + bits::{Representation, Unsigned}, + minmax::MinMaxMeta, + multi_vector::{MatRef, MaxSimError}, +}; + +/// An owned MinMax query with a backend selected at construction. +#[derive(Debug)] +pub struct MinMaxMaxSim +where + Unsigned: Representation, +{ + pub(super) inner: Box>, +} + +impl MinMaxMaxSim +where + Unsigned: Representation, +{ + pub fn nrows(&self) -> usize { + self.inner.nrows() + } + + /// Overwrite scores, ignoring NaN inner products. Empty documents and query rows + /// with only NaN or negative-infinite inner products produce f32::MAX. + /// Nonempty zero-dimensional documents produce zero. + /// + /// # Errors + /// + /// Returns an error for a wrong output length or document dimension, + /// without modifying the output. + pub fn compute_max_sim( + &self, + document: MatRef<'_, MinMaxMeta>, + scores: &mut [f32], + ) -> Result<(), MaxSimError> { + self.inner.compute_max_sim(document, scores) + } +} + +// Closed implementation boundary: callers never instantiate the matrix backends. +pub(super) trait MinMaxMaxSimKernel: + Send + Sync + std::fmt::Debug + 'static +where + Unsigned: Representation, +{ + fn nrows(&self) -> usize; + fn compute_max_sim( + &self, + document: MatRef<'_, MinMaxMeta>, + scores: &mut [f32], + ) -> Result<(), MaxSimError>; + + #[cfg(test)] + fn inner_type_id(&self) -> std::any::TypeId { + std::any::TypeId::of::() + } +} + +/// Validate shape and handle empty inputs; return whether scoring is still needed. +pub(super) fn validate( + rows: usize, + dim: usize, + document: MatRef<'_, MinMaxMeta>, + scores: &mut [f32], +) -> Result +where + Unsigned: Representation, +{ + if scores.len() != rows { + return Err(MaxSimError::InvalidBufferLength(scores.len(), rows)); + } + if document.repr().intrinsic_dim() != dim { + return Err(MaxSimError::UnequalDim( + document.repr().intrinsic_dim(), + dim, + )); + } + if rows == 0 { + return Ok(false); + } + if document.num_vectors() == 0 { + scores.fill(f32::MAX); + return Ok(false); + } + if dim == 0 { + scores.fill(0.0); + return Ok(false); + } + Ok(true) +} diff --git a/diskann-quantization/src/minmax/multi/meta.rs b/diskann-quantization/src/minmax/multi/meta.rs index f006d4afb4..862b84f678 100644 --- a/diskann-quantization/src/minmax/multi/meta.rs +++ b/diskann-quantization/src/minmax/multi/meta.rs @@ -7,15 +7,17 @@ use std::ptr::NonNull; +use diskann_utils::views::rowmajor; + use super::super::MinMaxQuantizer; use super::super::vectors::DataMutRef; use crate::CompressInto; use crate::bits::{Representation, Unsigned}; -use crate::minmax::{self, Data}; +use crate::minmax::{self, Data, MinMaxCompensation}; use crate::multi_vector::matrix::{ Defaulted, NewMut, NewOwned, NewRef, Repr, ReprMut, ReprOwned, SliceError, }; -use crate::multi_vector::{LayoutError, Mat, MatMut, MatRef, Standard}; +use crate::multi_vector::{LayoutError, Mat, MatMut, MatRef, Overflow, Standard}; use crate::scalar::InputContainsNaN; use crate::utils; @@ -62,6 +64,33 @@ where } } +impl<'a, const NBITS: usize> MatRef<'a, MinMaxMeta> +where + Unsigned: Representation, +{ + /// Borrow canonical storage as a checked row-major byte view. + pub(crate) fn as_canonical_view(&self) -> Result, Overflow> { + let rows = self.num_vectors(); + let dim = self.repr().intrinsic_dim(); + let overflow = || Overflow::for_type::(rows, dim); + let stride = dim + .checked_mul(NBITS) + .ok_or_else(overflow)? + .div_ceil(8) + .checked_add(size_of::()) + .ok_or_else(overflow)?; + Overflow::check::(rows, stride)?; + // SAFETY: MatRef supplies the canonical initialized span for 'a. Arithmetic is + // checked first, including legacy shapes that could wrap in release builds. + let bytes = unsafe { std::slice::from_raw_parts(self.as_raw_ptr(), rows * stride) }; + #[expect( + clippy::expect_used, + reason = "the checked dimensions match the slice length" + )] + Ok(rowmajor::Ref::try_from_data(bytes, rows, stride).expect("validated canonical rows")) + } +} + // SAFETY: The implementation correctly computes row offsets and constructs valid // DataRef types from the underlying byte buffer. The ncols() method // returns the canonical byte size for each row, and get_row properly diff --git a/diskann-quantization/src/minmax/multi/mod.rs b/diskann-quantization/src/minmax/multi/mod.rs index d6ad4c8369..6def54ca75 100644 --- a/diskann-quantization/src/minmax/multi/mod.rs +++ b/diskann-quantization/src/minmax/multi/mod.rs @@ -5,8 +5,14 @@ //! Multi-vector support for MinMax quantized vectors. +mod factory; +mod kernel; mod max_sim; mod meta; +mod prepared; +mod reference; +pub use factory::{MinMaxKernelFactory, build_minmax_max_sim}; +pub use kernel::MinMaxMaxSim; pub use max_sim::MinMaxKernel; pub use meta::MinMaxMeta; diff --git a/diskann-quantization/src/minmax/multi/prepared.rs b/diskann-quantization/src/minmax/multi/prepared.rs new file mode 100644 index 0000000000..f2dec63bf8 --- /dev/null +++ b/diskann-quantization/src/minmax/multi/prepared.rs @@ -0,0 +1,98 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use diskann_wide::arch::Scalar; +#[cfg(target_arch = "aarch64")] +use diskann_wide::arch::aarch64::Neon; +#[cfg(target_arch = "x86_64")] +use diskann_wide::arch::x86_64::{V3, V4}; + +use super::kernel::{MinMaxMaxSimKernel, validate}; +use crate::{ + bits::{Representation, Unsigned}, + matrix_kernels::{ + self as mk, + maxsim::packed_minmax_x_unpacked_minmax::{Driver, PackedA}, + }, + minmax::MinMaxMeta, + multi_vector::{MatRef, MaxSimError, Overflow}, +}; + +// Fix ISA and packing layout; query bit widths are needed only during packing. +macro_rules! prepared_kernel { + ($name:ident, $arch:ty, $element:ty, $mr:literal, $nr:literal, $pack:literal) => { + #[derive(Debug)] + pub(super) struct $name { + arch: $arch, + query: PackedA<$element, $mr, $pack>, + rows: usize, + dim: usize, + } + + impl $name { + pub(super) fn new( + arch: $arch, + query: MatRef<'_, MinMaxMeta>, + ) -> Result + where + Unsigned: Representation, + { + Ok(Self { + arch, + query: PackedA::new(query)?, + rows: query.num_vectors(), + dim: query.repr().intrinsic_dim(), + }) + } + } + + impl MinMaxMaxSimKernel for $name + where + Unsigned: Representation, + { + fn nrows(&self) -> usize { + self.rows + } + + fn compute_max_sim( + &self, + document: MatRef<'_, MinMaxMeta>, + scores: &mut [f32], + ) -> Result<(), MaxSimError> { + if !validate(self.rows, self.dim, document, scores)? { + return Ok(()); + } + // SAFETY: Validation establishes nonempty A, positive K, matching dimensions + // and exactly one output per query row. Packing and ISA selection are fixed. + let mut driver = unsafe { + Driver::<$arch, $mr, $nr, $pack>::new( + self.arch, + &self.query, + document, + scores, + mk::Cache::detect(), + ) + }; + mk::Drive::drive(&mut driver); + for score in scores { + *score = (-*score).min(f32::MAX); + } + Ok(()) + } + } + }; +} +prepared_kernel!(ScalarBytes, Scalar, u8, 16, 2, 4); +prepared_kernel!(ScalarWords, Scalar, i16, 8, 2, 2); +#[cfg(target_arch = "x86_64")] +prepared_kernel!(V3Bytes, V3, u8, 16, 4, 4); +#[cfg(target_arch = "x86_64")] +prepared_kernel!(V3Words, V3, i16, 16, 4, 2); +#[cfg(target_arch = "x86_64")] +prepared_kernel!(V4Bytes, V4, u8, 16, 8, 4); +#[cfg(target_arch = "x86_64")] +prepared_kernel!(V4Words, V4, i16, 16, 8, 2); +#[cfg(target_arch = "aarch64")] +prepared_kernel!(NeonBytes, Neon, u8, 8, 4, 4); diff --git a/diskann-quantization/src/minmax/multi/reference.rs b/diskann-quantization/src/minmax/multi/reference.rs new file mode 100644 index 0000000000..96a6a673f6 --- /dev/null +++ b/diskann-quantization/src/minmax/multi/reference.rs @@ -0,0 +1,72 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use super::kernel::{MinMaxMaxSimKernel, validate}; +use crate::{ + bits::{Representation, Unsigned}, + minmax::{Coefficients, DataRef, MinMaxMeta}, + multi_vector::{MatRef, MaxSimError, Overflow}, +}; +use diskann_utils::views::rowmajor::{self, Matrix}; + +/// Scalar correctness baseline, independent of the prepared matrix kernels. +#[derive(Debug)] +pub(super) struct ReferenceKernel { + query: rowmajor::Owned, + dim: usize, +} + +impl ReferenceKernel +where + Unsigned: Representation, +{ + pub(super) fn new(query: MatRef<'_, MinMaxMeta>) -> Result { + Ok(Self { + query: query.as_canonical_view()?.to_rowmajor_owned(), + dim: query.repr().intrinsic_dim(), + }) + } +} + +impl MinMaxMaxSimKernel + for ReferenceKernel +where + Unsigned: Representation + Representation, +{ + fn nrows(&self) -> usize { + self.query.nrows() + } + + #[expect( + clippy::expect_used, + reason = "canonical query rows and dimensions validated before iteration" + )] + fn compute_max_sim( + &self, + document: MatRef<'_, MinMaxMeta>, + scores: &mut [f32], + ) -> Result<(), MaxSimError> { + if !validate(self.nrows(), self.dim, document, scores)? { + return Ok(()); + } + for (bytes, best) in self.query.rows().zip(scores) { + let query = DataRef::::from_canonical_front(bytes, self.dim) + .expect("canonical query row"); + *best = f32::MAX; + let q = Coefficients::from(query.meta()); + for doc in document.rows() { + let dot: u64 = (0..self.dim) + .map(|d| { + query.vector().get(d).expect("query code in bounds") as u64 + * doc.vector().get(d).expect("document code in bounds") as u64 + }) + .sum(); + let ip = q.compensate_inner_product(doc.meta().into(), dot as f32, self.dim as f32); + *best = best.min(-ip); + } + } + Ok(()) + } +} diff --git a/diskann-quantization/src/minmax/vectors.rs b/diskann-quantization/src/minmax/vectors.rs index 37a07b81a8..53875c6f58 100644 --- a/diskann-quantization/src/minmax/vectors.rs +++ b/diskann-quantization/src/minmax/vectors.rs @@ -3,6 +3,8 @@ * Licensed under the MIT license. */ +use std::ops::{Add, Mul}; + use diskann_vector::{MathematicalValue, PureDistanceFunction}; use thiserror::Error; @@ -50,6 +52,31 @@ pub struct MinMaxCompensation { pub norm_squared: f32, // - |ax * X' + bx|^2 } +#[derive(Debug, Clone, Copy, Default)] +pub(crate) struct Coefficients { + pub(crate) a: T, + pub(crate) b: T, + pub(crate) n: T, +} + +impl From for Coefficients { + fn from(meta: MinMaxCompensation) -> Self { + Self { + a: meta.a, + b: meta.b, + n: meta.n, + } + } +} + +impl + Mul> Coefficients { + /// Compensate a raw code dot product using the logical (unpadded) dimension. + #[inline(always)] + pub(crate) fn compensate_inner_product(self, other: Self, dot: T, dim: T) -> T { + ((self.a * other.a) * dot + self.n * other.b) + other.n * self.b + (self.b * other.b) * dim + } +} + const META_BYTES: usize = std::mem::size_of::(); // This will be 5 * 4 = 20 bytes. /// Error type for parsing a slice of bytes as a `DataRef` @@ -219,12 +246,11 @@ where { let raw_product = InnerProduct::evaluate(x.vector(), y.vector())?; let (xm, ym) = (x.meta(), y.meta()); - let term0 = xm.a * ym.a * raw_product.into_inner() as f32; - let term1_x = xm.n * ym.b; - let term1_y = ym.n * xm.b; - let term2 = xm.b * ym.b * (x.len() as f32); - - let v = term0 + term1_x + term1_y + term2; + let v = Coefficients::from(xm).compensate_inner_product( + ym.into(), + raw_product.into_inner() as f32, + x.len() as f32, + ); Ok(MV::new(f(v, &xm, &ym))) } diff --git a/diskann-quantization/src/multi_vector/distance/factory.rs b/diskann-quantization/src/multi_vector/distance/factory.rs index 17072e1204..416654c640 100644 --- a/diskann-quantization/src/multi_vector/distance/factory.rs +++ b/diskann-quantization/src/multi_vector/distance/factory.rs @@ -25,7 +25,7 @@ use super::kernel::{Erase, MaxSimKernel}; use super::max_sim::MaxSimError; use crate::matrix_kernels as mk; use crate::multi_vector::distance::QueryMatRef; -use crate::multi_vector::{BlockTransposed, Mat, MatRef, Standard}; +use crate::multi_vector::{BlockTransposed, Mat, MatRef, Overflow, Standard}; // ───────────────────────────────────────────────────────────────────────── // Prepared — concrete kernel for the arch-dispatched paths. @@ -757,7 +757,7 @@ impl MaxSimElement for i8 { // Factory entry point. // ───────────────────────────────────────────────────────────────────────── -/// Error returned by [`build_max_sim`]. +/// Failure to prepare a MaxSim query. #[derive(Debug, Clone, Copy, Error)] #[non_exhaustive] pub enum BuildMaxSimError { @@ -765,6 +765,8 @@ pub enum BuildMaxSimError { NotSupported(#[from] NotSupported), #[error("query-vector dim {0} exceeds the maximum of {1}")] DimTooLarge(usize, usize), + #[error(transparent)] + Overflow(#[from] Overflow), } /// Build a multi-vector MaxSim kernel for any [`MaxSimElement`] type.