From 0928c14e5000d1ae5edd8c00627fe42597db94af Mon Sep 17 00:00:00 2001 From: junkui chen Date: Fri, 23 Oct 2026 12:51:20 +0800 Subject: [PATCH 1/3] Add MinMax MaxSim kernel factory Integrate the merged MinMax matrix kernel with an owned-query factory, runtime ISA selection, an independent scalar reference, and shared construction errors. Preserve QueryMatRef inputs and add focused factory and query-cloning coverage. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 94d30f43-fadc-4a6b-889b-406de424a1ab --- .../src/matrix_kernels/maxsim/mod.rs | 4 - .../maxsim/packed_minmax_x_unpacked_minmax.rs | 2 +- diskann-quantization/src/minmax/mod.rs | 5 +- .../src/minmax/multi/factory.rs | 576 ++++++++++++++++++ .../src/minmax/multi/kernel.rs | 74 +++ diskann-quantization/src/minmax/multi/meta.rs | 58 +- diskann-quantization/src/minmax/multi/mod.rs | 6 + .../src/minmax/multi/prepared.rs | 106 ++++ .../src/minmax/multi/reference.rs | 79 +++ .../src/multi_vector/distance/factory.rs | 6 +- 10 files changed, 907 insertions(+), 9 deletions(-) create mode 100644 diskann-quantization/src/minmax/multi/factory.rs create mode 100644 diskann-quantization/src/minmax/multi/kernel.rs create mode 100644 diskann-quantization/src/minmax/multi/prepared.rs create mode 100644 diskann-quantization/src/minmax/multi/reference.rs 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..8cab99881c 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 @@ -51,7 +51,7 @@ diskann_wide::alias!(u32x8 = u32x8); diskann_wide::alias!(u32x16 = u32x16); // Signed backends determine the shared accumulation limit. -const fn max_k() -> usize { +pub(crate) 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)) } diff --git a/diskann-quantization/src/minmax/mod.rs b/diskann-quantization/src/minmax/mod.rs index 41c29c3d03..5a6252ea5f 100644 --- a/diskann-quantization/src/minmax/mod.rs +++ b/diskann-quantization/src/minmax/mod.rs @@ -133,7 +133,10 @@ mod vectors; // Exports // ///////////// -pub use multi::{MinMaxKernel, MinMaxMeta}; +pub use multi::{ + MinMaxErase, MinMaxKernel, MinMaxKernelFactory, MinMaxMaxSimKernel, MinMaxMeta, + build_minmax_max_sim, +}; pub use quantizer::{L2Loss, MinMaxQuantizer}; pub use recompress::{RecompressError, Recompressor}; pub use vectors::{ diff --git a/diskann-quantization/src/minmax/multi/factory.rs b/diskann-quantization/src/minmax/multi/factory.rs new file mode 100644 index 0000000000..75a56becd4 --- /dev/null +++ b/diskann-quantization/src/minmax/multi/factory.rs @@ -0,0 +1,576 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use diskann_wide::Architecture; +#[cfg(target_arch = "aarch64")] +use diskann_wide::arch::aarch64::Neon; +#[cfg(target_arch = "x86_64")] +use diskann_wide::arch::x86_64::{V3, V4}; +use diskann_wide::arch::{Scalar, Target1}; + +use super::{kernel::MinMaxErase, prepared::Prepared, reference::ReferenceKernel}; +use crate::{ + bits::{Representation, Unsigned}, + matrix_kernels::maxsim::packed_minmax_x_unpacked_minmax::max_k, + minmax::{MinMaxCompensation, MinMaxMeta}, + multi_vector::{BuildMaxSimError, MatRef, MaxSimIsa, NotSupported, Overflow, QueryMatRef}, +}; + +/// Selects backend layouts for a document bit width. +/// +/// The built-in implementations keep document widths literal so unsupported +/// byte-dot layouts are never instantiated. +pub trait MinMaxKernelFactory: Representation { + #[doc(hidden)] + fn build>( + isa: MaxSimIsa, + query: QueryMatRef<'_, MinMaxMeta>, + erase: E, + ) -> Result + where + Unsigned: Representation + Representation; +} + +struct BuildAndErase(E); + +macro_rules! kernel_builder { + ($b:literal, $arch:ty, $mr:literal, $nr:literal, $pack:literal) => { + impl> + Target1<$arch, Result, MatRef<'_, MinMaxMeta>> + for BuildAndErase + where + Unsigned: Representation, + { + fn run( + self, + arch: $arch, + query: MatRef<'_, MinMaxMeta>, + ) -> Result { + let kernel = Prepared::<$arch, A_BITS, $mr, $nr, $pack>::new(arch, query)?; + Ok(self.0.erase(kernel)) + } + } + }; +} + +macro_rules! kernel_factory { + ( + b_bits: $b:literal, + scalar: { mr: $scalar_mr:literal, pack: $scalar_pack:literal }, + v3_pack: $v3_pack:literal, + v4_pack: $v4_pack:literal $(,)? + ) => { + kernel_builder!($b, Scalar, $scalar_mr, 2, $scalar_pack); + #[cfg(target_arch = "x86_64")] + kernel_builder!($b, V3, 16, 4, $v3_pack); + #[cfg(target_arch = "x86_64")] + kernel_builder!($b, V4, 16, 8, $v4_pack); + #[cfg(target_arch = "aarch64")] + kernel_builder!($b, Neon, 8, 4, 4); + + impl MinMaxKernelFactory<$b> for Unsigned { + fn build>( + isa: MaxSimIsa, + query: QueryMatRef<'_, MinMaxMeta>, + erase: E, + ) -> Result + where + Unsigned: Representation, + { + let query = query.0; + validate_query::(query)?; + let builder = BuildAndErase::(erase); + let unavailable = || NotSupported { + isa, + reason: "requested ISA unavailable on this host", + }; + let built = match isa { + MaxSimIsa::Auto => diskann_wide::arch::dispatch1_no_features(builder, query), + MaxSimIsa::Scalar => Scalar::new().run1(builder, query), + #[cfg(target_arch = "x86_64")] + MaxSimIsa::X86_64_V3 => V3::new_checked() + .ok_or_else(unavailable)? + .run1(builder, query), + #[cfg(target_arch = "x86_64")] + MaxSimIsa::X86_64_V4 => V4::new_checked() + .ok_or_else(unavailable)? + .run1(builder, query), + #[cfg(not(target_arch = "x86_64"))] + MaxSimIsa::X86_64_V3 | MaxSimIsa::X86_64_V4 => return Err(unavailable().into()), + #[cfg(target_arch = "aarch64")] + MaxSimIsa::Neon => Neon::new_checked() + .ok_or_else(unavailable)? + .run1(builder, query), + #[cfg(not(target_arch = "aarch64"))] + MaxSimIsa::Neon => return Err(unavailable().into()), + MaxSimIsa::Reference => { + return Ok(builder.0.erase(ReferenceKernel::new(query))); + } + }; + Ok(built?) + } + } + }; +} + +// Byte-dot layouts must handle A codes up to 255: V3's saturating intermediate +// limits B to 6 bits, while Scalar/V4's signed B bytes limit it to 7. +kernel_factory!( + b_bits: 1, + scalar: { mr: 16, pack: 4 }, + v3_pack: 4, + v4_pack: 4, +); +kernel_factory!( + b_bits: 2, + scalar: { mr: 16, pack: 4 }, + v3_pack: 4, + v4_pack: 4, +); +kernel_factory!( + b_bits: 3, + scalar: { mr: 16, pack: 4 }, + v3_pack: 4, + v4_pack: 4, +); +kernel_factory!( + b_bits: 4, + scalar: { mr: 16, pack: 4 }, + v3_pack: 4, + v4_pack: 4, +); +kernel_factory!( + b_bits: 5, + scalar: { mr: 16, pack: 4 }, + v3_pack: 4, + v4_pack: 4, +); +kernel_factory!( + b_bits: 6, + scalar: { mr: 16, pack: 4 }, + v3_pack: 4, + v4_pack: 4, +); +kernel_factory!( + b_bits: 7, + scalar: { mr: 16, pack: 4 }, + v3_pack: 2, + v4_pack: 4, +); +kernel_factory!( + b_bits: 8, + scalar: { mr: 8, pack: 2 }, + v3_pack: 2, + v4_pack: 2, +); + +/// 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. +/// +/// # 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::{BoxErase, 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, BoxErase)?; +/// 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>, + erase: E, +) -> Result +where + Unsigned: Representation + MinMaxKernelFactory, + E: MinMaxErase, +{ + >::build(isa, query, erase) +} + +fn validate_query( + query: MatRef<'_, MinMaxMeta>, +) -> Result<(), BuildMaxSimError> +where + Unsigned: Representation + Representation, +{ + let dim = query.repr().intrinsic_dim(); + let limit = max_k::(); + if dim > limit { + return Err(BuildMaxSimError::DimTooLarge(dim, limit)); + } + let rows = query.num_vectors(); + let overflow = || Overflow::for_type::(rows, dim); + let code_bytes = dim.checked_mul(A_BITS).ok_or_else(overflow)?.div_ceil(8); + let row_bytes = code_bytes + .checked_add(size_of::()) + .ok_or_else(overflow)?; + Overflow::check::(rows, row_bytes)?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + minmax::{MinMaxCompensation, MinMaxMaxSimKernel}, + multi_vector::{BoxErase, 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: Representation + 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(), BoxErase) + .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(), BoxErase).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}x{B}, {isa}: {actual} != {expected}" + ); + } + assert_eq!((guarded[0], guarded[rows + 1]), (123.0, 123.0)); + } + } + + #[test] + fn matches_reference_across_bit_widths() { + macro_rules! row { + ($a:literal) => { + assert_matches_reference::<$a, 1>(); + assert_matches_reference::<$a, 2>(); + assert_matches_reference::<$a, 3>(); + assert_matches_reference::<$a, 4>(); + assert_matches_reference::<$a, 5>(); + assert_matches_reference::<$a, 6>(); + assert_matches_reference::<$a, 7>(); + assert_matches_reference::<$a, 8>(); + }; + } + row!(1); + row!(2); + row!(3); + row!(4); + row!(5); + row!(6); + row!(7); + row!(8); + } + + #[test] + fn handles_empty_inputs() { + fn check() + where + Unsigned: MinMaxKernelFactory, + { + 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, _>(isa, query.as_view().into(), BoxErase) + .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={B}, {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(), BoxErase).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(), BoxErase), + 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] { + for result in [ + build_minmax_max_sim::<8, 8, _>(isa, query.into(), BoxErase), + >::build(isa, query.into(), BoxErase), + ] { + 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(), BoxErase).unwrap() + }; + let mut scores = [f32::NAN; 2]; + kernel + .compute_max_sim(document.as_view(), &mut scores) + .unwrap(); + assert_eq!(scores, [-6.0; 2], "{isa}"); + } + } + + struct Identify; + impl MinMaxErase for Identify + where + Unsigned: Representation, + { + type Output = std::any::TypeId; + fn erase + 'static>(self, _: K) -> Self::Output { + std::any::TypeId::of::() + } + } + + #[test] + fn selects_backend_and_layout() { + fn assert_layout( + isa: MaxSimIsa, + query: QueryMatRef<'_, MinMaxMeta<8>>, + ) where + Unsigned: MinMaxKernelFactory, + { + assert_eq!( + build_minmax_max_sim::<8, B, _>(isa, query, Identify).unwrap(), + std::any::TypeId::of::(), + "{isa}, B={B}", + ); + } + + let query = matrix::<8>(1, 3); + let query = QueryMatRef::from(query.as_view()); + assert_layout::<7, Prepared>(MaxSimIsa::Scalar, query); + assert_layout::<8, Prepared>(MaxSimIsa::Scalar, query); + #[cfg(target_arch = "x86_64")] + { + if V3::new_checked().is_some() { + assert_layout::<6, Prepared>(MaxSimIsa::X86_64_V3, query); + assert_layout::<7, Prepared>(MaxSimIsa::X86_64_V3, query); + } + if V4::new_checked().is_some() { + assert_layout::<7, Prepared>(MaxSimIsa::X86_64_V4, query); + assert_layout::<8, Prepared>(MaxSimIsa::X86_64_V4, query); + } + } + #[cfg(target_arch = "aarch64")] + if Neon::new_checked().is_some() { + assert_layout::<8, Prepared>(MaxSimIsa::Neon, query); + } + } + + #[test] + fn handles_nonfinite_scores() { + fn check() + where + Unsigned: MinMaxKernelFactory, + { + let max_code = (1i64 << B) - 1; + let query = constant_matrix::<8>(1, 5, 255); + let mut document = constant_matrix::(2, 5, max_code); + let finite_score = -(5 * 255 * max_code) as f32; + for isa in available_isas() { + let kernel = + build_minmax_max_sim::<8, B, _>(isa, query.as_view().into(), BoxErase).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={B}"); + } + 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={B}, mixed NaN/finite"); + } + } + check::<6>(); + check::<7>(); + check::<8>(); + } + + #[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(), BoxErase), + 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..31e2586219 --- /dev/null +++ b/diskann-quantization/src/minmax/multi/kernel.rs @@ -0,0 +1,74 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use crate::{ + bits::{Representation, Unsigned}, + minmax::MinMaxMeta, + multi_vector::{BoxErase, MatRef, MaxSimError}, +}; + +/// A prepared MinMax query producing one negated maximum inner product per query row. +/// Query storage is owned by the kernel; documents and output remain borrowed. +pub trait MinMaxMaxSimKernel: Send + Sync + std::fmt::Debug +where + Unsigned: Representation, +{ + fn nrows(&self) -> usize; + + /// 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 [`MaxSimError`] for a wrong output length or document dimension, + /// without modifying the output. + fn compute_max_sim( + &self, + document: MatRef<'_, MinMaxMeta>, + scores: &mut [f32], + ) -> Result<(), MaxSimError>; +} + +/// Customize factory output while the kernel's concrete type is still known. +pub trait MinMaxErase +where + Unsigned: Representation, +{ + type Output; + fn erase + 'static>(self, kernel: K) -> Self::Output; +} + +impl MinMaxErase for BoxErase +where + Unsigned: Representation, +{ + type Output = Box>; + + fn erase + 'static>(self, kernel: K) -> Self::Output { + Box::new(kernel) + } +} + +pub(super) fn validate( + rows: usize, + dim: usize, + document: MatRef<'_, MinMaxMeta>, + scores: &[f32], +) -> Result<(), MaxSimError> +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, + )); + } + Ok(()) +} diff --git a/diskann-quantization/src/minmax/multi/meta.rs b/diskann-quantization/src/minmax/multi/meta.rs index f006d4afb4..e8b5b787ed 100644 --- a/diskann-quantization/src/minmax/multi/meta.rs +++ b/diskann-quantization/src/minmax/multi/meta.rs @@ -13,7 +13,7 @@ use crate::CompressInto; use crate::bits::{Representation, Unsigned}; use crate::minmax::{self, Data}; use crate::multi_vector::matrix::{ - Defaulted, NewMut, NewOwned, NewRef, Repr, ReprMut, ReprOwned, SliceError, + Defaulted, NewCloned, NewMut, NewOwned, NewRef, Repr, ReprMut, ReprOwned, SliceError, }; use crate::multi_vector::{LayoutError, Mat, MatMut, MatRef, Standard}; use crate::scalar::InputContainsNaN; @@ -184,6 +184,21 @@ where } } +impl NewCloned for MinMaxMeta +where + Unsigned: Representation, +{ + fn new_cloned(v: MatRef<'_, Self>) -> Mat { + // SAFETY: The view covers repr.bytes() initialized bytes of canonical storage. + let bytes = unsafe { std::slice::from_raw_parts(v.as_raw_ptr(), v.repr().bytes()) }; + let owned: Box<[u8]> = bytes.into(); + + // SAFETY: The copy retains the representation's layout and uses the same + // Box<[u8]> allocation that ReprOwned::drop reconstructs. + unsafe { Mat::from_raw_parts(*v.repr(), utils::box_into_nonnull(owned).cast::()) } + } +} + // SAFETY: This checks that the slice has the correct length, which is all that is // required for [`Repr`] in this case. unsafe impl NewRef for MinMaxMeta @@ -374,6 +389,47 @@ mod tests { mod construction { use super::*; + fn test_new_cloned() + where + Unsigned: Representation, + { + for (rows, dim) in [(0, 7), (3, 0), (3, 7)] { + let repr = MinMaxMeta::::new(rows, dim); + let mut bytes: Vec = (0..repr.bytes()) + .map(|i| i.wrapping_mul(17) as u8) + .collect(); + let view = MatRef::new(repr, &bytes).unwrap(); + let cloned = view.to_owned(); + assert_eq!(cloned.num_vectors(), rows); + assert_eq!(cloned.repr().intrinsic_dim(), dim); + let expected: Vec<_> = view + .rows() + .map(|row| { + ( + bytemuck::bytes_of(&row.meta()).to_vec(), + (0..dim) + .map(|i| row.vector().get(i).unwrap()) + .collect::>(), + ) + }) + .collect(); + + bytes.fill(0); + drop(bytes); + for (row, (meta, codes)) in cloned.rows().zip(expected) { + assert_eq!(bytemuck::bytes_of(&row.meta()), meta); + assert_eq!( + (0..dim) + .map(|i| row.vector().get(i).unwrap()) + .collect::>(), + codes + ); + } + } + } + + expand_to_bitrates!(new_cloned, test_new_cloned); + /// Tests NewOwned (Mat::new with Defaulted) for various dimensions and sizes. fn test_new_owned() where diff --git a/diskann-quantization/src/minmax/multi/mod.rs b/diskann-quantization/src/minmax/multi/mod.rs index d6ad4c8369..e71bfffc9a 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::{MinMaxErase, MinMaxMaxSimKernel}; 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..2b6e1ca8a8 --- /dev/null +++ b/diskann-quantization/src/minmax/multi/prepared.rs @@ -0,0 +1,106 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use super::kernel::{MinMaxMaxSimKernel, validate}; +use crate::{ + bits::{Representation, Unsigned}, + matrix_kernels::{ + self as mk, + maxsim::packed_minmax_x_unpacked_minmax::{Driver, ExtraWide, PackedA}, + }, + minmax::MinMaxMeta, + multi_vector::{MatRef, MaxSimError, Overflow}, +}; + +#[derive(Debug)] +pub(super) struct Prepared< + W: ExtraWide, + const A_BITS: usize, + const MR: usize, + const NR: usize, + const PACK: usize, +> where + Unsigned: Representation, +{ + arch: W, + query: PackedA, + rows: usize, + dim: usize, +} + +impl< + W: ExtraWide, + const A_BITS: usize, + const MR: usize, + const NR: usize, + const PACK: usize, +> Prepared +where + Unsigned: Representation, +{ + pub(super) fn new(arch: W, query: MatRef<'_, MinMaxMeta>) -> Result { + Ok(Self { + arch, + query: PackedA::new(query)?, + rows: query.num_vectors(), + dim: query.repr().intrinsic_dim(), + }) + } +} + +impl< + W, + const A_BITS: usize, + const B_BITS: usize, + const MR: usize, + const NR: usize, + const PACK: usize, +> MinMaxMaxSimKernel for Prepared +where + W: ExtraWide + Send + Sync + std::fmt::Debug, + W::AElem: Send + Sync + std::fmt::Debug, + Unsigned: Representation + Representation, + for<'a> Driver<'a, W, B_BITS, MR, NR, PACK>: mk::Drive, +{ + fn nrows(&self) -> usize { + self.rows + } + + fn compute_max_sim( + &self, + document: MatRef<'_, MinMaxMeta>, + scores: &mut [f32], + ) -> Result<(), MaxSimError> { + validate(self.rows, self.dim, document, scores)?; + if self.rows == 0 { + return Ok(()); + } + if document.num_vectors() == 0 { + scores.fill(f32::MAX); + return Ok(()); + } + if self.dim == 0 { + scores.fill(0.0); + return Ok(()); + } + // SAFETY: The checks above establish nonempty A, positive K, matching dimensions, + // and exactly one output score per query row. + let mut driver = unsafe { + Driver::<_, B_BITS, MR, NR, PACK>::new( + self.arch, + &self.query, + document, + scores, + mk::Cache::detect(), + ) + }; + mk::Drive::drive(&mut driver); + // Match the pairwise distance sentinel when the best inner product is -infinity. + for score in scores { + *score = (-*score).min(f32::MAX); + } + Ok(()) + } +} diff --git a/diskann-quantization/src/minmax/multi/reference.rs b/diskann-quantization/src/minmax/multi/reference.rs new file mode 100644 index 0000000000..66c6fc9d10 --- /dev/null +++ b/diskann-quantization/src/minmax/multi/reference.rs @@ -0,0 +1,79 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use super::kernel::{MinMaxMaxSimKernel, validate}; +use crate::{ + bits::{Representation, Unsigned}, + minmax::MinMaxMeta, + multi_vector::{Mat, MatRef, MaxSimError}, +}; + +/// Scalar correctness baseline, independent of the prepared matrix kernels. +#[derive(Debug)] +pub(super) struct ReferenceKernel +where + Unsigned: Representation, +{ + query: Mat>, +} + +impl ReferenceKernel +where + Unsigned: Representation, +{ + pub(super) fn new(query: MatRef<'_, MinMaxMeta>) -> Self { + Self { + query: query.to_owned(), + } + } +} + +impl MinMaxMaxSimKernel + for ReferenceKernel +where + Unsigned: Representation + Representation, +{ + fn nrows(&self) -> usize { + self.query.num_vectors() + } + + #[expect( + clippy::expect_used, + reason = "dimensions checked before iterating canonical rows" + )] + fn compute_max_sim( + &self, + document: MatRef<'_, MinMaxMeta>, + scores: &mut [f32], + ) -> Result<(), MaxSimError> { + let dim = self.query.repr().intrinsic_dim(); + validate(self.query.num_vectors(), dim, document, scores)?; + if document.num_vectors() == 0 { + scores.fill(f32::MAX); + return Ok(()); + } + if dim == 0 { + scores.fill(0.0); + return Ok(()); + } + for (query, best) in self.query.rows().zip(scores) { + *best = f32::MAX; + let q = query.meta(); + for doc in document.rows() { + let dot: u64 = (0..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 b = doc.meta(); + let ip = + ((q.a * b.a) * dot as f32 + q.n * b.b) + b.n * q.b + (q.b * b.b) * dim as f32; + *best = best.min(-ip); + } + } + Ok(()) + } +} 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. From fd17b5154325b3e907eed273691ac4778a4064be Mon Sep 17 00:00:00 2001 From: junkui chen Date: Thu, 29 Oct 2026 16:41:30 +0800 Subject: [PATCH 2/3] Refactor MinMax factory to concrete dispatch Replace caller-provided erasure with a concrete MinMaxMaxSim object and library-owned constructors for each bit-width pair. Select fixed-layout kernels through RuntimeArch and retain the independent scalar reference with rowmajor-owned query storage. Keep the low-level matrix kernel unchanged. Centralize canonical view validation, use explicit bit-width names, and preserve the existing typed query and document inputs. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 94d30f43-fadc-4a6b-889b-406de424a1ab --- diskann-quantization/src/minmax/mod.rs | 3 +- .../src/minmax/multi/factory.rs | 392 ++++++++---------- .../src/minmax/multi/kernel.rs | 103 +++-- diskann-quantization/src/minmax/multi/meta.rs | 58 +-- diskann-quantization/src/minmax/multi/mod.rs | 2 +- .../src/minmax/multi/prepared.rs | 168 ++++---- .../src/minmax/multi/reference.rs | 52 ++- 7 files changed, 351 insertions(+), 427 deletions(-) diff --git a/diskann-quantization/src/minmax/mod.rs b/diskann-quantization/src/minmax/mod.rs index 5a6252ea5f..0c1b73042a 100644 --- a/diskann-quantization/src/minmax/mod.rs +++ b/diskann-quantization/src/minmax/mod.rs @@ -134,8 +134,7 @@ mod vectors; ///////////// pub use multi::{ - MinMaxErase, MinMaxKernel, MinMaxKernelFactory, MinMaxMaxSimKernel, MinMaxMeta, - build_minmax_max_sim, + MinMaxKernel, MinMaxKernelFactory, MinMaxMaxSim, MinMaxMeta, build_minmax_max_sim, }; pub use quantizer::{L2Loss, MinMaxQuantizer}; pub use recompress::{RecompressError, Recompressor}; diff --git a/diskann-quantization/src/minmax/multi/factory.rs b/diskann-quantization/src/minmax/multi/factory.rs index 75a56becd4..0c9420686a 100644 --- a/diskann-quantization/src/minmax/multi/factory.rs +++ b/diskann-quantization/src/minmax/multi/factory.rs @@ -3,169 +3,24 @@ * Licensed under the MIT license. */ -use diskann_wide::Architecture; +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 diskann_wide::arch::{Scalar, Target1}; -use super::{kernel::MinMaxErase, prepared::Prepared, reference::ReferenceKernel}; +use super::{ + kernel::{MinMaxMaxSim, MinMaxMaxSimKernel, canonical_view}, + prepared, + reference::ReferenceKernel, +}; use crate::{ bits::{Representation, Unsigned}, matrix_kernels::maxsim::packed_minmax_x_unpacked_minmax::max_k, - minmax::{MinMaxCompensation, MinMaxMeta}, - multi_vector::{BuildMaxSimError, MatRef, MaxSimIsa, NotSupported, Overflow, QueryMatRef}, + minmax::MinMaxMeta, + multi_vector::{BuildMaxSimError, MatRef, MaxSimIsa, NotSupported, QueryMatRef}, }; -/// Selects backend layouts for a document bit width. -/// -/// The built-in implementations keep document widths literal so unsupported -/// byte-dot layouts are never instantiated. -pub trait MinMaxKernelFactory: Representation { - #[doc(hidden)] - fn build>( - isa: MaxSimIsa, - query: QueryMatRef<'_, MinMaxMeta>, - erase: E, - ) -> Result - where - Unsigned: Representation + Representation; -} - -struct BuildAndErase(E); - -macro_rules! kernel_builder { - ($b:literal, $arch:ty, $mr:literal, $nr:literal, $pack:literal) => { - impl> - Target1<$arch, Result, MatRef<'_, MinMaxMeta>> - for BuildAndErase - where - Unsigned: Representation, - { - fn run( - self, - arch: $arch, - query: MatRef<'_, MinMaxMeta>, - ) -> Result { - let kernel = Prepared::<$arch, A_BITS, $mr, $nr, $pack>::new(arch, query)?; - Ok(self.0.erase(kernel)) - } - } - }; -} - -macro_rules! kernel_factory { - ( - b_bits: $b:literal, - scalar: { mr: $scalar_mr:literal, pack: $scalar_pack:literal }, - v3_pack: $v3_pack:literal, - v4_pack: $v4_pack:literal $(,)? - ) => { - kernel_builder!($b, Scalar, $scalar_mr, 2, $scalar_pack); - #[cfg(target_arch = "x86_64")] - kernel_builder!($b, V3, 16, 4, $v3_pack); - #[cfg(target_arch = "x86_64")] - kernel_builder!($b, V4, 16, 8, $v4_pack); - #[cfg(target_arch = "aarch64")] - kernel_builder!($b, Neon, 8, 4, 4); - - impl MinMaxKernelFactory<$b> for Unsigned { - fn build>( - isa: MaxSimIsa, - query: QueryMatRef<'_, MinMaxMeta>, - erase: E, - ) -> Result - where - Unsigned: Representation, - { - let query = query.0; - validate_query::(query)?; - let builder = BuildAndErase::(erase); - let unavailable = || NotSupported { - isa, - reason: "requested ISA unavailable on this host", - }; - let built = match isa { - MaxSimIsa::Auto => diskann_wide::arch::dispatch1_no_features(builder, query), - MaxSimIsa::Scalar => Scalar::new().run1(builder, query), - #[cfg(target_arch = "x86_64")] - MaxSimIsa::X86_64_V3 => V3::new_checked() - .ok_or_else(unavailable)? - .run1(builder, query), - #[cfg(target_arch = "x86_64")] - MaxSimIsa::X86_64_V4 => V4::new_checked() - .ok_or_else(unavailable)? - .run1(builder, query), - #[cfg(not(target_arch = "x86_64"))] - MaxSimIsa::X86_64_V3 | MaxSimIsa::X86_64_V4 => return Err(unavailable().into()), - #[cfg(target_arch = "aarch64")] - MaxSimIsa::Neon => Neon::new_checked() - .ok_or_else(unavailable)? - .run1(builder, query), - #[cfg(not(target_arch = "aarch64"))] - MaxSimIsa::Neon => return Err(unavailable().into()), - MaxSimIsa::Reference => { - return Ok(builder.0.erase(ReferenceKernel::new(query))); - } - }; - Ok(built?) - } - } - }; -} - -// Byte-dot layouts must handle A codes up to 255: V3's saturating intermediate -// limits B to 6 bits, while Scalar/V4's signed B bytes limit it to 7. -kernel_factory!( - b_bits: 1, - scalar: { mr: 16, pack: 4 }, - v3_pack: 4, - v4_pack: 4, -); -kernel_factory!( - b_bits: 2, - scalar: { mr: 16, pack: 4 }, - v3_pack: 4, - v4_pack: 4, -); -kernel_factory!( - b_bits: 3, - scalar: { mr: 16, pack: 4 }, - v3_pack: 4, - v4_pack: 4, -); -kernel_factory!( - b_bits: 4, - scalar: { mr: 16, pack: 4 }, - v3_pack: 4, - v4_pack: 4, -); -kernel_factory!( - b_bits: 5, - scalar: { mr: 16, pack: 4 }, - v3_pack: 4, - v4_pack: 4, -); -kernel_factory!( - b_bits: 6, - scalar: { mr: 16, pack: 4 }, - v3_pack: 4, - v4_pack: 4, -); -kernel_factory!( - b_bits: 7, - scalar: { mr: 16, pack: 4 }, - v3_pack: 2, - v4_pack: 4, -); -kernel_factory!( - b_bits: 8, - scalar: { mr: 8, pack: 2 }, - v3_pack: 2, - v4_pack: 2, -); - /// Build a kernel that owns a copy of the MinMax query. /// /// `Auto` selects a backend at construction. Explicit ISA requests never silently @@ -186,13 +41,13 @@ kernel_factory!( /// ``` /// use diskann_quantization::{ /// minmax::{MinMaxMeta, build_minmax_max_sim}, -/// multi_vector::{BoxErase, Defaulted, Mat, MaxSimIsa, QueryMatRef}, +/// 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, BoxErase)?; +/// 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]); @@ -200,18 +55,66 @@ kernel_factory!( /// # } /// # example().unwrap(); /// ``` -pub fn build_minmax_max_sim( +pub fn build_minmax_max_sim( isa: MaxSimIsa, query: QueryMatRef<'_, MinMaxMeta>, - erase: E, -) -> Result +) -> Result, BuildMaxSimError> where - Unsigned: Representation + MinMaxKernelFactory, - E: MinMaxErase, + Unsigned: MinMaxKernelFactory, +{ + >::build(isa, query) +} + +/// Precompiled query/document bit-width combinations. +#[doc(hidden)] +pub trait MinMaxKernelFactory: + Representation + Representation { - >::build(isa, query, erase) + 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::<$a, $b>::new(a, query)?), + #[cfg(target_arch = "x86_64")] + RuntimeArch::V3(a) => Box::new(prepared::$v3::<$a, $b>::new(a, query)?), + #[cfg(target_arch = "x86_64")] + RuntimeArch::V4(a) => Box::new(prepared::$v4::<$a, $b>::new(a, query)?), + #[cfg(target_arch = "aarch64")] + RuntimeArch::Neon(a) => Box::new(prepared::NeonBytes::<$a, $b>::new(a, query)?), + RuntimeArch::Reference => Box::new(ReferenceKernel::<$a, $b>::new(query)?), + }; + Ok(MinMaxMaxSim { inner }) + } + } + }; + ($($a:literal),+ $(,)?) => {$( + factory!($a, 1, ScalarBytes, V3Bytes, V4Bytes); + factory!($a, 2, ScalarBytes, V3Bytes, V4Bytes); + factory!($a, 3, ScalarBytes, V3Bytes, V4Bytes); + factory!($a, 4, ScalarBytes, V3Bytes, V4Bytes); + factory!($a, 5, ScalarBytes, V3Bytes, V4Bytes); + factory!($a, 6, ScalarBytes, V3Bytes, V4Bytes); + factory!($a, 7, ScalarBytes, V3Words, V4Bytes); + factory!($a, 8, ScalarWords, V3Words, V4Words); + )+}; +} +factory!(1, 2, 3, 4, 5, 6, 7, 8); + fn validate_query( query: MatRef<'_, MinMaxMeta>, ) -> Result<(), BuildMaxSimError> @@ -223,22 +126,67 @@ where if dim > limit { return Err(BuildMaxSimError::DimTooLarge(dim, limit)); } - let rows = query.num_vectors(); - let overflow = || Overflow::for_type::(rows, dim); - let code_bytes = dim.checked_mul(A_BITS).ok_or_else(overflow)?.div_ceil(8); - let row_bytes = code_bytes - .checked_add(size_of::()) - .ok_or_else(overflow)?; - Overflow::check::(rows, row_bytes)?; + canonical_view(query)?; 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, MinMaxMaxSimKernel}, - multi_vector::{BoxErase, Defaulted, Mat, MaxSimError}, + minmax::MinMaxCompensation, + multi_vector::{Defaulted, Mat, MaxSimError}, }; fn matrix(rows: usize, dim: usize) -> Mat> @@ -301,16 +249,16 @@ mod tests { .filter(|isa| isa.is_available()) } - fn assert_matches_reference() + fn assert_matches_reference() where - Unsigned: Representation + MinMaxKernelFactory, + 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 query = matrix::(rows, dim); + let document = matrix::(docs, dim); let reference = - build_minmax_max_sim::(MaxSimIsa::Reference, query.as_view().into(), BoxErase) + build_minmax_max_sim::(MaxSimIsa::Reference, query.as_view().into()) .unwrap(); let mut expected = vec![0.0; rows]; reference @@ -318,7 +266,7 @@ mod tests { .unwrap(); for isa in available_isas().filter(|&isa| isa != MaxSimIsa::Reference) { let kernel = - build_minmax_max_sim::(isa, query.as_view().into(), BoxErase).unwrap(); + build_minmax_max_sim::(isa, query.as_view().into()).unwrap(); assert_eq!(kernel.nrows(), rows); let mut guarded = vec![123.0; rows + 2]; kernel @@ -330,7 +278,7 @@ mod tests { || (actual.is_finite() && expected.is_finite() && (actual - expected).abs() <= 1e-5 * expected.abs().max(1.0)), - "{A}x{B}, {isa}: {actual} != {expected}" + "{A_BITS}x{B_BITS}, {isa}: {actual} != {expected}" ); } assert_eq!((guarded[0], guarded[rows + 1]), (123.0, 123.0)); @@ -363,9 +311,9 @@ mod tests { #[test] fn handles_empty_inputs() { - fn check() + fn check() where - Unsigned: MinMaxKernelFactory, + Unsigned: MinMaxKernelFactory<8, B_BITS>, { for (rows, docs, dim, expected) in [ (0, 5, 7, 0.0), @@ -374,11 +322,10 @@ mod tests { (3, 0, 0, f32::MAX), ] { let query = matrix::<8>(rows, dim); - let document = matrix::(docs, dim); + let document = matrix::(docs, dim); for isa in available_isas() { let kernel = - build_minmax_max_sim::<8, B, _>(isa, query.as_view().into(), BoxErase) - .unwrap(); + 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) @@ -386,7 +333,7 @@ mod tests { assert_eq!( scores, vec![expected; rows], - "{isa}, B={B}, {rows}x{docs}x{dim}" + "{isa}, B_BITS={B_BITS}, {rows}x{docs}x{dim}" ); } } @@ -399,8 +346,7 @@ mod tests { 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(), BoxErase).unwrap(); + 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!( @@ -424,7 +370,7 @@ mod tests { 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(), BoxErase), + build_minmax_max_sim::<8, 8>(isa, query.as_view().into()), Err(BuildMaxSimError::NotSupported(error)) if error.isa == isa )); } @@ -437,16 +383,12 @@ mod tests { let dim = limit + 1; let query = MatRef::new(MinMaxMeta::<8>::new(0, dim), &[] as &[u8]).unwrap(); for isa in [MaxSimIsa::Auto, MaxSimIsa::Reference] { - for result in [ - build_minmax_max_sim::<8, 8, _>(isa, query.into(), BoxErase), - >::build(isa, query.into(), BoxErase), - ] { - assert!( - matches!(result, Err(BuildMaxSimError::DimTooLarge(found, max)) - if found == dim && max == limit), - "{isa}", - ); - } + 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}" + ); } } @@ -456,7 +398,7 @@ mod tests { 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(), BoxErase).unwrap() + build_minmax_max_sim::<8, 4>(isa, query.as_view().into()).unwrap() }; let mut scores = [f32::NAN; 2]; kernel @@ -466,66 +408,58 @@ mod tests { } } - struct Identify; - impl MinMaxErase for Identify - where - Unsigned: Representation, - { - type Output = std::any::TypeId; - fn erase + 'static>(self, _: K) -> Self::Output { - std::any::TypeId::of::() - } - } - #[test] fn selects_backend_and_layout() { - fn assert_layout( + fn assert_layout( isa: MaxSimIsa, query: QueryMatRef<'_, MinMaxMeta<8>>, ) where - Unsigned: MinMaxKernelFactory, + Unsigned: MinMaxKernelFactory<8, B_BITS>, { assert_eq!( - build_minmax_max_sim::<8, B, _>(isa, query, Identify).unwrap(), + build_minmax_max_sim::<8, B_BITS>(isa, query) + .unwrap() + .inner + .type_id(), std::any::TypeId::of::(), - "{isa}, B={B}", + "{isa}, B_BITS={B_BITS}", ); } let query = matrix::<8>(1, 3); let query = QueryMatRef::from(query.as_view()); - assert_layout::<7, Prepared>(MaxSimIsa::Scalar, query); - assert_layout::<8, Prepared>(MaxSimIsa::Scalar, query); + assert_layout::<7, prepared::ScalarBytes<8, 7>>(MaxSimIsa::Scalar, query); + assert_layout::<8, prepared::ScalarWords<8, 8>>(MaxSimIsa::Scalar, query); #[cfg(target_arch = "x86_64")] { if V3::new_checked().is_some() { - assert_layout::<6, Prepared>(MaxSimIsa::X86_64_V3, query); - assert_layout::<7, Prepared>(MaxSimIsa::X86_64_V3, query); + assert_layout::<6, prepared::V3Bytes<8, 6>>(MaxSimIsa::X86_64_V3, query); + assert_layout::<7, prepared::V3Words<8, 7>>(MaxSimIsa::X86_64_V3, query); } if V4::new_checked().is_some() { - assert_layout::<7, Prepared>(MaxSimIsa::X86_64_V4, query); - assert_layout::<8, Prepared>(MaxSimIsa::X86_64_V4, query); + assert_layout::<7, prepared::V4Bytes<8, 7>>(MaxSimIsa::X86_64_V4, query); + assert_layout::<8, prepared::V4Words<8, 8>>(MaxSimIsa::X86_64_V4, query); } } #[cfg(target_arch = "aarch64")] if Neon::new_checked().is_some() { - assert_layout::<8, Prepared>(MaxSimIsa::Neon, query); + assert_layout::<8, prepared::NeonBytes<8, 8>>(MaxSimIsa::Neon, query); } } #[test] fn handles_nonfinite_scores() { - fn check() + fn check() where - Unsigned: MinMaxKernelFactory, + Unsigned: MinMaxKernelFactory<8, B_BITS>, { - let max_code = (1i64 << B) - 1; + let max_code = (1i64 << B_BITS) - 1; let query = constant_matrix::<8>(1, 5, 255); - let mut document = constant_matrix::(2, 5, max_code); + let mut document = constant_matrix::(2, 5, max_code); let finite_score = -(5 * 255 * max_code) as f32; for isa in available_isas() { let kernel = - build_minmax_max_sim::<8, B, _>(isa, query.as_view().into(), BoxErase).unwrap(); + build_minmax_max_sim::<8, B_BITS>(isa, query.as_view().into()).unwrap(); let mut scores = [0.0]; for (scale, expected) in [ (1.0, finite_score), @@ -542,7 +476,7 @@ mod tests { kernel .compute_max_sim(document.as_view(), &mut scores) .unwrap(); - assert_eq!(scores, [expected], "{isa}, B={B}"); + assert_eq!(scores, [expected], "{isa}, B_BITS={B_BITS}"); } let mut row = document.get_row_mut(0).unwrap(); let mut meta = row.meta(); @@ -551,7 +485,11 @@ mod tests { kernel .compute_max_sim(document.as_view(), &mut scores) .unwrap(); - assert_eq!(scores, [finite_score], "{isa}, B={B}, mixed NaN/finite"); + assert_eq!( + scores, + [finite_score], + "{isa}, B_BITS={B_BITS}, mixed NaN/finite" + ); } } check::<6>(); @@ -568,7 +506,7 @@ mod tests { 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(), BoxErase), + 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 index 31e2586219..5e763dc178 100644 --- a/diskann-quantization/src/minmax/multi/kernel.rs +++ b/diskann-quantization/src/minmax/multi/kernel.rs @@ -3,61 +3,74 @@ * Licensed under the MIT license. */ +use diskann_utils::views::rowmajor; + use crate::{ bits::{Representation, Unsigned}, - minmax::MinMaxMeta, - multi_vector::{BoxErase, MatRef, MaxSimError}, + minmax::{MinMaxCompensation, MinMaxMeta}, + multi_vector::{MatRef, MaxSimError, Overflow}, }; -/// A prepared MinMax query producing one negated maximum inner product per query row. -/// Query storage is owned by the kernel; documents and output remain borrowed. -pub trait MinMaxMaxSimKernel: Send + Sync + std::fmt::Debug +/// An owned MinMax query with a backend selected at construction. +#[derive(Debug)] +pub struct MinMaxMaxSim where Unsigned: Representation, { - fn nrows(&self) -> usize; + 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`. + /// with only NaN or negative-infinite inner products produce f32::MAX. /// Nonempty zero-dimensional documents produce zero. /// /// # Errors /// - /// Returns [`MaxSimError`] for a wrong output length or document dimension, + /// Returns an error for a wrong output length or document dimension, /// without modifying the output. - fn compute_max_sim( + pub fn compute_max_sim( &self, document: MatRef<'_, MinMaxMeta>, scores: &mut [f32], - ) -> Result<(), MaxSimError>; -} - -/// Customize factory output while the kernel's concrete type is still known. -pub trait MinMaxErase -where - Unsigned: Representation, -{ - type Output; - fn erase + 'static>(self, kernel: K) -> Self::Output; + ) -> Result<(), MaxSimError> { + self.inner.compute_max_sim(document, scores) + } } -impl MinMaxErase for BoxErase +// Closed implementation boundary: callers never instantiate the matrix backends. +pub(super) trait MinMaxMaxSimKernel: + Send + Sync + std::fmt::Debug + 'static where Unsigned: Representation, { - type Output = Box>; + fn nrows(&self) -> usize; + fn compute_max_sim( + &self, + document: MatRef<'_, MinMaxMeta>, + scores: &mut [f32], + ) -> Result<(), MaxSimError>; - fn erase + 'static>(self, kernel: K) -> Self::Output { - Box::new(kernel) + #[cfg(test)] + fn 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: &[f32], -) -> Result<(), MaxSimError> + scores: &mut [f32], +) -> Result where Unsigned: Representation, { @@ -70,5 +83,43 @@ where dim, )); } - Ok(()) + 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) +} + +/// Borrow canonical MinMax storage as a checked row-major byte view. +pub(super) fn canonical_view( + matrix: MatRef<'_, MinMaxMeta>, +) -> Result, Overflow> +where + Unsigned: Representation, +{ + let rows = matrix.num_vectors(); + let dim = matrix.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. Arithmetic is checked before + // reading the pointer, including legacy shapes that could wrap in release builds. + let bytes = unsafe { std::slice::from_raw_parts(matrix.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")) } diff --git a/diskann-quantization/src/minmax/multi/meta.rs b/diskann-quantization/src/minmax/multi/meta.rs index e8b5b787ed..f006d4afb4 100644 --- a/diskann-quantization/src/minmax/multi/meta.rs +++ b/diskann-quantization/src/minmax/multi/meta.rs @@ -13,7 +13,7 @@ use crate::CompressInto; use crate::bits::{Representation, Unsigned}; use crate::minmax::{self, Data}; use crate::multi_vector::matrix::{ - Defaulted, NewCloned, NewMut, NewOwned, NewRef, Repr, ReprMut, ReprOwned, SliceError, + Defaulted, NewMut, NewOwned, NewRef, Repr, ReprMut, ReprOwned, SliceError, }; use crate::multi_vector::{LayoutError, Mat, MatMut, MatRef, Standard}; use crate::scalar::InputContainsNaN; @@ -184,21 +184,6 @@ where } } -impl NewCloned for MinMaxMeta -where - Unsigned: Representation, -{ - fn new_cloned(v: MatRef<'_, Self>) -> Mat { - // SAFETY: The view covers repr.bytes() initialized bytes of canonical storage. - let bytes = unsafe { std::slice::from_raw_parts(v.as_raw_ptr(), v.repr().bytes()) }; - let owned: Box<[u8]> = bytes.into(); - - // SAFETY: The copy retains the representation's layout and uses the same - // Box<[u8]> allocation that ReprOwned::drop reconstructs. - unsafe { Mat::from_raw_parts(*v.repr(), utils::box_into_nonnull(owned).cast::()) } - } -} - // SAFETY: This checks that the slice has the correct length, which is all that is // required for [`Repr`] in this case. unsafe impl NewRef for MinMaxMeta @@ -389,47 +374,6 @@ mod tests { mod construction { use super::*; - fn test_new_cloned() - where - Unsigned: Representation, - { - for (rows, dim) in [(0, 7), (3, 0), (3, 7)] { - let repr = MinMaxMeta::::new(rows, dim); - let mut bytes: Vec = (0..repr.bytes()) - .map(|i| i.wrapping_mul(17) as u8) - .collect(); - let view = MatRef::new(repr, &bytes).unwrap(); - let cloned = view.to_owned(); - assert_eq!(cloned.num_vectors(), rows); - assert_eq!(cloned.repr().intrinsic_dim(), dim); - let expected: Vec<_> = view - .rows() - .map(|row| { - ( - bytemuck::bytes_of(&row.meta()).to_vec(), - (0..dim) - .map(|i| row.vector().get(i).unwrap()) - .collect::>(), - ) - }) - .collect(); - - bytes.fill(0); - drop(bytes); - for (row, (meta, codes)) in cloned.rows().zip(expected) { - assert_eq!(bytemuck::bytes_of(&row.meta()), meta); - assert_eq!( - (0..dim) - .map(|i| row.vector().get(i).unwrap()) - .collect::>(), - codes - ); - } - } - } - - expand_to_bitrates!(new_cloned, test_new_cloned); - /// Tests NewOwned (Mat::new with Defaulted) for various dimensions and sizes. fn test_new_owned() where diff --git a/diskann-quantization/src/minmax/multi/mod.rs b/diskann-quantization/src/minmax/multi/mod.rs index e71bfffc9a..6def54ca75 100644 --- a/diskann-quantization/src/minmax/multi/mod.rs +++ b/diskann-quantization/src/minmax/multi/mod.rs @@ -13,6 +13,6 @@ mod prepared; mod reference; pub use factory::{MinMaxKernelFactory, build_minmax_max_sim}; -pub use kernel::{MinMaxErase, MinMaxMaxSimKernel}; +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 index 2b6e1ca8a8..1609e3ed5c 100644 --- a/diskann-quantization/src/minmax/multi/prepared.rs +++ b/diskann-quantization/src/minmax/multi/prepared.rs @@ -3,104 +3,100 @@ * 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, ExtraWide, PackedA}, + maxsim::packed_minmax_x_unpacked_minmax::{Driver, PackedA}, }, minmax::MinMaxMeta, multi_vector::{MatRef, MaxSimError, Overflow}, }; -#[derive(Debug)] -pub(super) struct Prepared< - W: ExtraWide, - const A_BITS: usize, - const MR: usize, - const NR: usize, - const PACK: usize, -> where - Unsigned: Representation, -{ - arch: W, - query: PackedA, - rows: usize, - dim: usize, -} +// Fix ISA and packing layout, leaving only the input bit widths generic. +macro_rules! prepared_kernel { + ($name:ident, $arch:ty, $element:ty, $mr:literal, $nr:literal, $pack:literal) => { + #[derive(Debug)] + pub(super) struct $name + where + Unsigned: Representation, + { + arch: $arch, + query: PackedA<$element, A_BITS, $mr, $pack>, + rows: usize, + dim: usize, + } -impl< - W: ExtraWide, - const A_BITS: usize, - const MR: usize, - const NR: usize, - const PACK: usize, -> Prepared -where - Unsigned: Representation, -{ - pub(super) fn new(arch: W, query: MatRef<'_, MinMaxMeta>) -> Result { - Ok(Self { - arch, - query: PackedA::new(query)?, - rows: query.num_vectors(), - dim: query.repr().intrinsic_dim(), - }) - } -} + impl $name + where + Unsigned: Representation + Representation, + { + pub(super) fn new( + arch: $arch, + query: MatRef<'_, MinMaxMeta>, + ) -> Result { + Ok(Self { + arch, + query: PackedA::new(query)?, + rows: query.num_vectors(), + dim: query.repr().intrinsic_dim(), + }) + } + } -impl< - W, - const A_BITS: usize, - const B_BITS: usize, - const MR: usize, - const NR: usize, - const PACK: usize, -> MinMaxMaxSimKernel for Prepared -where - W: ExtraWide + Send + Sync + std::fmt::Debug, - W::AElem: Send + Sync + std::fmt::Debug, - Unsigned: Representation + Representation, - for<'a> Driver<'a, W, B_BITS, MR, NR, PACK>: mk::Drive, -{ - fn nrows(&self) -> usize { - self.rows - } + impl MinMaxMaxSimKernel + for $name + where + Unsigned: Representation + Representation, + { + fn nrows(&self) -> usize { + self.rows + } - fn compute_max_sim( - &self, - document: MatRef<'_, MinMaxMeta>, - scores: &mut [f32], - ) -> Result<(), MaxSimError> { - validate(self.rows, self.dim, document, scores)?; - if self.rows == 0 { - return Ok(()); - } - if document.num_vectors() == 0 { - scores.fill(f32::MAX); - return Ok(()); - } - if self.dim == 0 { - scores.fill(0.0); - return Ok(()); - } - // SAFETY: The checks above establish nonempty A, positive K, matching dimensions, - // and exactly one output score per query row. - let mut driver = unsafe { - Driver::<_, B_BITS, MR, NR, PACK>::new( - self.arch, - &self.query, - document, - scores, - mk::Cache::detect(), - ) - }; - mk::Drive::drive(&mut driver); - // Match the pairwise distance sentinel when the best inner product is -infinity. - for score in scores { - *score = (-*score).min(f32::MAX); + 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, B_BITS, $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(()) + } } - 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 index 66c6fc9d10..ec1135a518 100644 --- a/diskann-quantization/src/minmax/multi/reference.rs +++ b/diskann-quantization/src/minmax/multi/reference.rs @@ -3,74 +3,70 @@ * Licensed under the MIT license. */ -use super::kernel::{MinMaxMaxSimKernel, validate}; +use super::kernel::{MinMaxMaxSimKernel, canonical_view, validate}; use crate::{ bits::{Representation, Unsigned}, - minmax::MinMaxMeta, - multi_vector::{Mat, MatRef, MaxSimError}, + minmax::{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 -where - Unsigned: Representation, -{ - query: Mat>, +pub(super) struct ReferenceKernel { + query: rowmajor::Owned, + dim: usize, } -impl ReferenceKernel +impl ReferenceKernel where Unsigned: Representation, { - pub(super) fn new(query: MatRef<'_, MinMaxMeta>) -> Self { - Self { - query: query.to_owned(), - } + pub(super) fn new(query: MatRef<'_, MinMaxMeta>) -> Result { + Ok(Self { + query: canonical_view(query)?.to_rowmajor_owned(), + dim: query.repr().intrinsic_dim(), + }) } } impl MinMaxMaxSimKernel - for ReferenceKernel + for ReferenceKernel where Unsigned: Representation + Representation, { fn nrows(&self) -> usize { - self.query.num_vectors() + self.query.nrows() } #[expect( clippy::expect_used, - reason = "dimensions checked before iterating canonical rows" + reason = "canonical query rows and dimensions validated before iteration" )] fn compute_max_sim( &self, document: MatRef<'_, MinMaxMeta>, scores: &mut [f32], ) -> Result<(), MaxSimError> { - let dim = self.query.repr().intrinsic_dim(); - validate(self.query.num_vectors(), dim, document, scores)?; - if document.num_vectors() == 0 { - scores.fill(f32::MAX); - return Ok(()); - } - if dim == 0 { - scores.fill(0.0); + if !validate(self.nrows(), self.dim, document, scores)? { return Ok(()); } - for (query, best) in self.query.rows().zip(scores) { + 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 = query.meta(); for doc in document.rows() { - let dot: u64 = (0..dim) + 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 b = doc.meta(); - let ip = - ((q.a * b.a) * dot as f32 + q.n * b.b) + b.n * q.b + (q.b * b.b) * dim as f32; + let ip = ((q.a * b.a) * dot as f32 + q.n * b.b) + + b.n * q.b + + (q.b * b.b) * self.dim as f32; *best = best.min(-ip); } } From 01d092eb5d413ba05ffaf21afe862d9654c8ea30 Mon Sep 17 00:00:00 2001 From: Junkui Chen Date: Fri, 30 Oct 2026 23:27:23 +0800 Subject: [PATCH 3/3] [multi-vector] Reduce MinMax MaxSim specialization Limit factory support to 11 query/document bit-width pairs. Share packed query state and driver loops with one decode dispatch per tile. Preserve ISA-specific decoding and avoid empty padding fills. Reuse compensation arithmetic and clarify the kernel type-ID helper. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../maxsim/packed_minmax_x_unpacked_minmax.rs | 384 +++++++++++------- diskann-quantization/src/minmax/mod.rs | 1 + .../src/minmax/multi/factory.rs | 147 ++++--- .../src/minmax/multi/kernel.rs | 35 +- diskann-quantization/src/minmax/multi/meta.rs | 33 +- .../src/minmax/multi/prepared.rs | 28 +- .../src/minmax/multi/reference.rs | 13 +- diskann-quantization/src/minmax/vectors.rs | 38 +- 8 files changed, 410 insertions(+), 269 deletions(-) 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 8cab99881c..c4e71aa1e3 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. -pub(crate) 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 { @@ -221,15 +212,19 @@ where }; 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 = b + .as_canonical_view() + .expect("canonical document storage size must not overflow"); let b_rows_per_tile = Params::new( cache, a_bytes, @@ -241,17 +236,22 @@ where arch, a, b, + decode_tile: |arch: W, source, k, codes, coefficients| { + arch.run( + #[inline(always)] + || decode_tile::(source, k, codes, coefficients, W::convert_b), + ) + }, 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) { @@ -261,18 +261,23 @@ where 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, - ) - }; + #[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, + |source, k, codes, coefficients| { + (self.decode_tile)(self.arch, source, k, codes, coefficients) + }, + ); self.a.visit_panels(|a, panel| { let row = panel * MR; let end = (row + MR).min(self.c.len()); @@ -296,6 +301,53 @@ struct BScratch { meta: Vec>, } +// Only decoding is specialized by bit width; allocation, padding, and tile views are shared. +type DecodeTileFn = fn( + arch: W, + 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. +/// Inlined into the selected architecture's target-feature entry. +#[inline(always)] +#[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], + mut convert: impl FnMut(u8) -> T, +) 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], &mut convert) + .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 +401,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) }, @@ -628,11 +670,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(); @@ -1362,27 +1406,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 +1446,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 +1473,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 +1510,73 @@ 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( + _: Scalar, + 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, i16::from); + } + + 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 +1606,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 +1617,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 +1632,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 +1715,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 +1786,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 +1801,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 +1810,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")); @@ -1740,16 +1849,13 @@ 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, + |source, k, codes, coefficients| { + decode_tile::<4, _>(source, k, codes, coefficients, |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. @@ -1806,23 +1912,21 @@ 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| { + let tile = scratch.prepare::( + source.subview(start..start + count).unwrap(), + k, + |source, k, codes, coefficients| { + decode_tile::(source, k, codes, coefficients, |code| { conversions += 1; convert(code) - }, - ) - }; + }); + }, + ); assert_eq!(conversions, count * dim); let padded = padded_k::(k); let row_len = padded.value().get(); @@ -1877,9 +1981,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(), @@ -2115,7 +2219,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 +2347,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 +2362,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 +2400,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 +2462,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 0c1b73042a..38d0fba5e9 100644 --- a/diskann-quantization/src/minmax/mod.rs +++ b/diskann-quantization/src/minmax/mod.rs @@ -138,6 +138,7 @@ pub use multi::{ }; 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 index 0c9420686a..eced459851 100644 --- a/diskann-quantization/src/minmax/multi/factory.rs +++ b/diskann-quantization/src/minmax/multi/factory.rs @@ -10,7 +10,7 @@ use diskann_wide::arch::aarch64::Neon; use diskann_wide::arch::x86_64::{V3, V4}; use super::{ - kernel::{MinMaxMaxSim, MinMaxMaxSimKernel, canonical_view}, + kernel::{MinMaxMaxSim, MinMaxMaxSimKernel}, prepared, reference::ReferenceKernel, }; @@ -28,6 +28,9 @@ use crate::{ /// /// 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 @@ -89,31 +92,31 @@ macro_rules! factory { let query = query.0; validate_query::<$a, $b>(query)?; let inner: Box> = match RuntimeArch::resolve(isa)? { - RuntimeArch::Scalar(a) => Box::new(prepared::$scalar::<$a, $b>::new(a, query)?), + RuntimeArch::Scalar(a) => Box::new(prepared::$scalar::<$b>::new(a, query)?), #[cfg(target_arch = "x86_64")] - RuntimeArch::V3(a) => Box::new(prepared::$v3::<$a, $b>::new(a, query)?), + RuntimeArch::V3(a) => Box::new(prepared::$v3::<$b>::new(a, query)?), #[cfg(target_arch = "x86_64")] - RuntimeArch::V4(a) => Box::new(prepared::$v4::<$a, $b>::new(a, query)?), + RuntimeArch::V4(a) => Box::new(prepared::$v4::<$b>::new(a, query)?), #[cfg(target_arch = "aarch64")] - RuntimeArch::Neon(a) => Box::new(prepared::NeonBytes::<$a, $b>::new(a, query)?), + RuntimeArch::Neon(a) => Box::new(prepared::NeonBytes::<$b>::new(a, query)?), RuntimeArch::Reference => Box::new(ReferenceKernel::<$a, $b>::new(query)?), }; Ok(MinMaxMaxSim { inner }) } } }; - ($($a:literal),+ $(,)?) => {$( - factory!($a, 1, ScalarBytes, V3Bytes, V4Bytes); - factory!($a, 2, ScalarBytes, V3Bytes, V4Bytes); - factory!($a, 3, ScalarBytes, V3Bytes, V4Bytes); - factory!($a, 4, ScalarBytes, V3Bytes, V4Bytes); - factory!($a, 5, ScalarBytes, V3Bytes, V4Bytes); - factory!($a, 6, ScalarBytes, V3Bytes, V4Bytes); - factory!($a, 7, ScalarBytes, V3Words, V4Bytes); - factory!($a, 8, ScalarWords, V3Words, V4Words); - )+}; } -factory!(1, 2, 3, 4, 5, 6, 7, 8); +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>, @@ -122,11 +125,11 @@ where Unsigned: Representation + Representation, { let dim = query.repr().intrinsic_dim(); - let limit = max_k::(); + let limit = max_k(A_BITS, B_BITS); if dim > limit { return Err(BuildMaxSimError::DimTooLarge(dim, limit)); } - canonical_view(query)?; + query.as_canonical_view()?; Ok(()) } @@ -286,27 +289,18 @@ mod tests { } #[test] - fn matches_reference_across_bit_widths() { - macro_rules! row { - ($a:literal) => { - assert_matches_reference::<$a, 1>(); - assert_matches_reference::<$a, 2>(); - assert_matches_reference::<$a, 3>(); - assert_matches_reference::<$a, 4>(); - assert_matches_reference::<$a, 5>(); - assert_matches_reference::<$a, 6>(); - assert_matches_reference::<$a, 7>(); - assert_matches_reference::<$a, 8>(); - }; - } - row!(1); - row!(2); - row!(3); - row!(4); - row!(5); - row!(6); - row!(7); - row!(8); + 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] @@ -379,7 +373,7 @@ mod tests { #[test] fn rejects_excessive_dimensions() { - let limit = max_k::<8, 8>(); + 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] { @@ -410,56 +404,78 @@ mod tests { #[test] fn selects_backend_and_layout() { - fn assert_layout( - isa: MaxSimIsa, - query: QueryMatRef<'_, MinMaxMeta<8>>, - ) where - Unsigned: MinMaxKernelFactory<8, B_BITS>, + fn assert_layout(isa: MaxSimIsa) + where + Unsigned: MinMaxKernelFactory, { + let query = matrix::(1, 3); assert_eq!( - build_minmax_max_sim::<8, B_BITS>(isa, query) + build_minmax_max_sim::(isa, query.as_view().into()) .unwrap() .inner - .type_id(), + .inner_type_id(), std::any::TypeId::of::(), - "{isa}, B_BITS={B_BITS}", + "{isa}, A_BITS={A_BITS}, B_BITS={B_BITS}", ); } - let query = matrix::<8>(1, 3); - let query = QueryMatRef::from(query.as_view()); - assert_layout::<7, prepared::ScalarBytes<8, 7>>(MaxSimIsa::Scalar, query); - assert_layout::<8, prepared::ScalarWords<8, 8>>(MaxSimIsa::Scalar, query); + 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, prepared::V3Bytes<8, 6>>(MaxSimIsa::X86_64_V3, query); - assert_layout::<7, prepared::V3Words<8, 7>>(MaxSimIsa::X86_64_V3, query); + 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, prepared::V4Bytes<8, 7>>(MaxSimIsa::X86_64_V4, query); - assert_layout::<8, prepared::V4Words<8, 8>>(MaxSimIsa::X86_64_V4, query); + 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, prepared::NeonBytes<8, 8>>(MaxSimIsa::Neon, query); + assert_layout::<8, 8, prepared::NeonBytes<8>>(MaxSimIsa::Neon); } } #[test] - fn handles_nonfinite_scores() { + fn prepared_type_is_independent_of_query_bit_width() { fn check() where - Unsigned: MinMaxKernelFactory<8, B_BITS>, + 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 query = constant_matrix::<8>(1, 5, 255); + 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 * 255 * max_code) as f32; + let finite_score = -(5 * max_query * max_code) as f32; for isa in available_isas() { let kernel = - build_minmax_max_sim::<8, B_BITS>(isa, query.as_view().into()).unwrap(); + build_minmax_max_sim::(isa, query.as_view().into()).unwrap(); let mut scores = [0.0]; for (scale, expected) in [ (1.0, finite_score), @@ -492,9 +508,10 @@ mod tests { ); } } - check::<6>(); - check::<7>(); - check::<8>(); + check::<6, 6>(); + check::<7, 7>(); + check::<8, 8>(); + check::<8, 4>(); } #[cfg(not(debug_assertions))] diff --git a/diskann-quantization/src/minmax/multi/kernel.rs b/diskann-quantization/src/minmax/multi/kernel.rs index 5e763dc178..e5cf53549c 100644 --- a/diskann-quantization/src/minmax/multi/kernel.rs +++ b/diskann-quantization/src/minmax/multi/kernel.rs @@ -3,12 +3,10 @@ * Licensed under the MIT license. */ -use diskann_utils::views::rowmajor; - use crate::{ bits::{Representation, Unsigned}, - minmax::{MinMaxCompensation, MinMaxMeta}, - multi_vector::{MatRef, MaxSimError, Overflow}, + minmax::MinMaxMeta, + multi_vector::{MatRef, MaxSimError}, }; /// An owned MinMax query with a backend selected at construction. @@ -59,7 +57,7 @@ where ) -> Result<(), MaxSimError>; #[cfg(test)] - fn type_id(&self) -> std::any::TypeId { + fn inner_type_id(&self) -> std::any::TypeId { std::any::TypeId::of::() } } @@ -96,30 +94,3 @@ where } Ok(true) } - -/// Borrow canonical MinMax storage as a checked row-major byte view. -pub(super) fn canonical_view( - matrix: MatRef<'_, MinMaxMeta>, -) -> Result, Overflow> -where - Unsigned: Representation, -{ - let rows = matrix.num_vectors(); - let dim = matrix.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. Arithmetic is checked before - // reading the pointer, including legacy shapes that could wrap in release builds. - let bytes = unsafe { std::slice::from_raw_parts(matrix.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")) -} 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/prepared.rs b/diskann-quantization/src/minmax/multi/prepared.rs index 1609e3ed5c..f2dec63bf8 100644 --- a/diskann-quantization/src/minmax/multi/prepared.rs +++ b/diskann-quantization/src/minmax/multi/prepared.rs @@ -20,28 +20,25 @@ use crate::{ multi_vector::{MatRef, MaxSimError, Overflow}, }; -// Fix ISA and packing layout, leaving only the input bit widths generic. +// 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 - where - Unsigned: Representation, - { + pub(super) struct $name { arch: $arch, - query: PackedA<$element, A_BITS, $mr, $pack>, + query: PackedA<$element, $mr, $pack>, rows: usize, dim: usize, } - impl $name - where - Unsigned: Representation + Representation, - { - pub(super) fn new( + impl $name { + pub(super) fn new( arch: $arch, query: MatRef<'_, MinMaxMeta>, - ) -> Result { + ) -> Result + where + Unsigned: Representation, + { Ok(Self { arch, query: PackedA::new(query)?, @@ -51,10 +48,9 @@ macro_rules! prepared_kernel { } } - impl MinMaxMaxSimKernel - for $name + impl MinMaxMaxSimKernel for $name where - Unsigned: Representation + Representation, + Unsigned: Representation, { fn nrows(&self) -> usize { self.rows @@ -71,7 +67,7 @@ macro_rules! prepared_kernel { // 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, B_BITS, $mr, $nr, $pack>::new( + Driver::<$arch, $mr, $nr, $pack>::new( self.arch, &self.query, document, diff --git a/diskann-quantization/src/minmax/multi/reference.rs b/diskann-quantization/src/minmax/multi/reference.rs index ec1135a518..96a6a673f6 100644 --- a/diskann-quantization/src/minmax/multi/reference.rs +++ b/diskann-quantization/src/minmax/multi/reference.rs @@ -3,10 +3,10 @@ * Licensed under the MIT license. */ -use super::kernel::{MinMaxMaxSimKernel, canonical_view, validate}; +use super::kernel::{MinMaxMaxSimKernel, validate}; use crate::{ bits::{Representation, Unsigned}, - minmax::{DataRef, MinMaxMeta}, + minmax::{Coefficients, DataRef, MinMaxMeta}, multi_vector::{MatRef, MaxSimError, Overflow}, }; use diskann_utils::views::rowmajor::{self, Matrix}; @@ -24,7 +24,7 @@ where { pub(super) fn new(query: MatRef<'_, MinMaxMeta>) -> Result { Ok(Self { - query: canonical_view(query)?.to_rowmajor_owned(), + query: query.as_canonical_view()?.to_rowmajor_owned(), dim: query.repr().intrinsic_dim(), }) } @@ -55,7 +55,7 @@ where let query = DataRef::::from_canonical_front(bytes, self.dim) .expect("canonical query row"); *best = f32::MAX; - let q = query.meta(); + let q = Coefficients::from(query.meta()); for doc in document.rows() { let dot: u64 = (0..self.dim) .map(|d| { @@ -63,10 +63,7 @@ where * doc.vector().get(d).expect("document code in bounds") as u64 }) .sum(); - let b = doc.meta(); - let ip = ((q.a * b.a) * dot as f32 + q.n * b.b) - + b.n * q.b - + (q.b * b.b) * self.dim as f32; + let ip = q.compensate_inner_product(doc.meta().into(), dot as f32, self.dim as f32); *best = best.min(-ip); } } 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))) }