From e465898016176c9ecf6c2ae56a71878901148eba Mon Sep 17 00:00:00 2001 From: junkui chen Date: Fri, 16 Oct 2026 19:01:56 +0800 Subject: [PATCH 1/4] Add standalone MinMax matrix MaxSim kernel Support all 1-8 bit A/B combinations through shared packing, tiled execution, compensation and maximum-inner-product reduction. Retain the tuned byte backends alongside full-range i16 paths, keep quantization widths out of the numeric microkernels, and add the required SIMD conversions and correctness coverage. Leave factory, fallback and public distance integration to follow-up PRs. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 74814a43-c576-493a-8783-585ef48c3f1c --- .../src/matrix_kernels/maxsim/mod.rs | 5 + .../maxsim/packed_minmax_x_unpacked_minmax.rs | 2145 +++++++++++++++++ diskann-wide/src/arch/aarch64/double.rs | 18 + diskann-wide/src/arch/aarch64/u32x4_.rs | 24 +- diskann-wide/src/arch/x86_64/v3/conversion.rs | 11 +- diskann-wide/src/arch/x86_64/v4/conversion.rs | 6 +- diskann-wide/src/emulated.rs | 2 + diskann-wide/src/reference.rs | 6 + 8 files changed, 2211 insertions(+), 6 deletions(-) create mode 100644 diskann-quantization/src/matrix_kernels/maxsim/packed_minmax_x_unpacked_minmax.rs diff --git a/diskann-quantization/src/matrix_kernels/maxsim/mod.rs b/diskann-quantization/src/matrix_kernels/maxsim/mod.rs index 412f9d98c5..e320afcfd5 100644 --- a/diskann-quantization/src/matrix_kernels/maxsim/mod.rs +++ b/diskann-quantization/src/matrix_kernels/maxsim/mod.rs @@ -19,6 +19,11 @@ pub(crate) mod packed_f32_x_unpacked_f16; pub(crate) mod packed_f32_x_unpacked_f32; pub(crate) mod packed_i8_x_unpacked_i8; +#[cfg_attr( + not(test), + expect(dead_code, reason = "TODO: will integrate in the next PR") +)] +pub(crate) mod packed_minmax_x_unpacked_minmax; #[cfg(test)] mod test; diff --git a/diskann-quantization/src/matrix_kernels/maxsim/packed_minmax_x_unpacked_minmax.rs b/diskann-quantization/src/matrix_kernels/maxsim/packed_minmax_x_unpacked_minmax.rs new file mode 100644 index 0000000000..a9556b38f7 --- /dev/null +++ b/diskann-quantization/src/matrix_kernels/maxsim/packed_minmax_x_unpacked_minmax.rs @@ -0,0 +1,2145 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! MaxSim for prepacked MinMax A and tiled MinMax B, with 1-8 bits per input code. +//! +//! Each B tile is decoded once and reused across A panels. Integer inner products +//! are compensated before reduction because B's coefficients vary by vector. +//! +//! Outputs are unnegated maximum inner products. NaN candidates are ignored; +//! empty B or all-NaN comparisons leave negative infinity. +//! +//! Scalar and x86 provide a general i16 path (PACK=2) and a byte path (PACK=4). +//! The byte path supports B up to six bits on V3 and seven on Scalar/V4. +//! Neon uses unsigned byte dots (PACK=4) for every bit width. +//! A packing depends only on A's bit width and the selected numeric layout, so it +//! can be reused across B formats supported by that layout. + +use std::{ + num::NonZeroUsize, + ops::{Add, Mul}, +}; + +use diskann_wide::{ + SIMDCast, SIMDDotProduct, SIMDMinMax, SIMDReinterpret, SIMDVector, + arch::{Architecture, Scalar, Target}, +}; + +use crate::{ + bits::{BitSlice, Representation, Unsigned}, + matrix_kernels::{ + Cache, + blocks::{packed, unpacked}, + bounds, driver, + num::{Bytes, DimK, Elements}, + ptr::Slice, + util::LoadStore, + }, + minmax::{MinMaxCompensation, MinMaxMeta}, + multi_vector::{BlockTransposed, MatRef, Overflow}, +}; + +use super::packed_f32_x_unpacked_f32::Params; + +diskann_wide::alias!(f32x16 = f32x16); +diskann_wide::alias!(i8x32 = i8x32); +diskann_wide::alias!(i8x64 = i8x64); +diskann_wide::alias!(i16x16 = i16x16); +diskann_wide::alias!(i16x32 = i16x32); +diskann_wide::alias!(i32x16 = i32x16); +diskann_wide::alias!(u8x32 = u8x32); +diskann_wide::alias!(u8x64 = u8x64); +diskann_wide::alias!(u32x4 = u32x4); +diskann_wide::alias!(u32x8 = u32x8); +diskann_wide::alias!(u32x16 = u32x16); + +// Signed backends determine the shared accumulation limit. +const fn max_k() -> usize { + assert!(A_BITS > 0 && A_BITS <= 8 && B_BITS > 0 && B_BITS <= 8); + (i32::MAX as usize) / (((1 << A_BITS) - 1) * ((1 << B_BITS) - 1)) +} + +#[derive(Debug, Clone, Copy)] +struct Coefficients { + a: T, + b: T, + n: T, +} + +#[derive(Debug)] +pub(crate) struct PackedA +where + Unsigned: Representation, +{ + codes: BlockTransposed, + meta: Coefficients>, +} + +impl PackedA +where + T: Copy + Default, + Unsigned: Representation, +{ + /// Pack A and copy its coefficients, zero-filling partial rows and K packs. + /// + /// # Errors + /// + /// Returns [`Overflow`] if the packed data or coefficient allocation size overflows. + #[expect(clippy::expect_used)] + pub(crate) fn new>( + _: W, + a: MatRef<'_, MinMaxMeta>, + ) -> Result { + let k = a.repr().intrinsic_dim(); + let mut codes = BlockTransposed::try_new(a.num_vectors(), k)?; + let padded_rows = codes.padded_nrows(); + Overflow::check::(padded_rows, 1)?; + let mut meta = Coefficients { + a: vec![0.0; padded_rows], + b: vec![0.0; padded_rows], + n: vec![0.0; padded_rows], + }; + for (i, row) in a.rows().enumerate() { + let m = row.meta(); + meta.a[i] = m.a; + meta.b[i] = m.b; + meta.n[i] = m.n; + let v = row.vector(); + let mut dst = codes + .get_row_mut(i) + .expect("A and its packing have equal row counts"); + for j in 0..k { + // SAFETY: MinMaxMeta exposes exactly k codes for every row. + dst.set(j, W::convert_a(unsafe { v.get_unchecked(j) } as u8)); + } + } + Ok(Self { codes, meta }) + } + + /// Cache-planning bytes for one padded A panel, including coefficients. + fn panel_bytes(&self) -> Bytes { + let codes = Elements::::new(MR * self.codes.padded_ncols()).bytes(); + Bytes::new(codes.value() + size_of::>()) + } +} + +//--------// +// Driver // +//--------// + +/// Each drive overwrites C; results do not accumulate across calls. +pub(crate) struct Driver< + 'a, + W: ExtraWide, + const B_BITS: usize, + const MR: usize, + const NR: usize, + const PACK: usize, +> where + Unsigned: Representation, +{ + arch: W, + a: packed::View<'a, W::AElem, MR, PACK>, + a_meta: Coefficients<&'a [f32]>, + b: MatRef<'a, MinMaxMeta>, + c: &'a mut [f32], + k: DimK, + b_rows_per_tile: NonZeroUsize, +} + +impl< + 'a, + W: ExtraWide, + const B_BITS: usize, + const MR: usize, + const NR: usize, + const PACK: usize, +> Driver<'a, W, B_BITS, MR, NR, PACK> +where + Unsigned: Representation, +{ + /// # Panics + /// + /// Panics if K exceeds the integer accumulator limit for these bit widths. + /// Conversion of a valid integer dot product to f32 may still round. + /// + /// B's bit width must fit the selected backend's [`ExtraWide::MAX_B_BITS`]; + /// incompatible combinations are rejected during monomorphization. + /// + /// # Safety + /// + /// A must be nonempty, both inputs must have logical dimension k, and C must have + /// exactly A's logical row count. + #[expect(clippy::expect_used)] + pub(crate) unsafe fn new( + arch: W, + a: &'a PackedA, + b: MatRef<'a, MinMaxMeta>, + c: &'a mut [f32], + k: DimK, + cache: Cache, + ) -> Self + where + Unsigned: Representation, + { + const { + assert!( + B_BITS <= W::MAX_B_BITS, + "B bit width exceeds the backend's supported range" + ) + }; + let max_k = const { max_k::() }; + assert!( + k.value().get() <= max_k, + "contraction dimension {} exceeds the MinMax{A_BITS} x MinMax{B_BITS} limit of {max_k}", + k.value().get(), + ); + bounds::check_eq!(bounds::Bound::new(a.codes.ncols()), k); + bounds::check_eq!(bounds::Bound::new(b.repr().intrinsic_dim()), k); + bounds::check_eq!(bounds::Bound::new(c.len()), a.codes.nrows()); + let a_meta = Coefficients { + a: a.meta.a.as_slice(), + b: a.meta.b.as_slice(), + n: a.meta.n.as_slice(), + }; + let a_bytes = a.panel_bytes(); + let a = packed::View::from_block_transposed(a.codes.as_view()) + .expect("A and K must be nonempty"); + for meta in [a_meta.a, a_meta.b, a_meta.n] { + bounds::check_eq!(bounds::Bound::new(meta.len()), MR * a.blocks().get()); + } + let b_rows_per_tile = + Params::new(cache, a_bytes, BScratch::::row_bytes(k), NR).b_cols_in_l1; + Self { + arch, + a, + a_meta, + b, + c, + k, + b_rows_per_tile, + } + } +} + +impl driver::Drive + for Driver<'_, W, B_BITS, MR, NR, PACK> +where + W: ExtraWide + LoadStore, + Unsigned: Representation, +{ + fn drive(&mut self) { + self.arch.run( + #[inline] + || { + self.c.fill(f32::NEG_INFINITY); + let mut scratch = BScratch::::default(); + let mut start = 0; + while let Some(remaining) = NonZeroUsize::new(self.b.num_vectors() - start) { + let count = remaining.min(self.b_rows_per_tile); + let b = scratch.prepare::(self.b, start, count, self.k); + // SAFETY: Driver's construction establishes all dimensions, metadata + // extents, and C's coverage of A. Visitors preserve these bounds. + unsafe { + self.a.visit_panels(self.k, |a, panel| { + let row = panel * MR; + let end = (row + MR).min(self.c.len()); + let c = self.arch.load(&self.c[row..end]); + let a_meta = Coefficients { + a: W::Float::load_simd(self.arch, self.a_meta.a.as_ptr().add(row)), + b: W::Float::load_simd(self.arch, self.a_meta.b.as_ptr().add(row)), + n: W::Float::load_simd(self.arch, self.a_meta.n.as_ptr().add(row)), + }; + let kernel = PanelKernel:: { + arch: self.arch, + a, + a_meta, + b: b.codes, + b_meta: b.meta, + c, + k: self.k, + }; + let c = self.arch.run(kernel); + self.arch.store(c, &mut self.c[row..end]); + }); + } + start += count.get(); + } + }, + ); + } +} + +#[derive(Default)] +struct BScratch { + codes: Vec, + meta: Vec, +} + +struct BTile<'a, T> { + codes: unpacked::View<'a, T>, + meta: &'a [MinMaxCompensation], +} + +impl BScratch { + #[expect(clippy::expect_used)] + fn prepare( + &mut self, + b: MatRef<'_, MinMaxMeta>, + start: usize, + count: NonZeroUsize, + k: DimK, + ) -> BTile<'_, T> + where + W: ExtraWide, + Unsigned: Representation, + { + let dim = k.value().get(); + assert_eq!(b.repr().intrinsic_dim(), dim); + assert!(start <= b.num_vectors() && count.get() <= b.num_vectors() - start); + self.codes.resize(count.get() * dim, T::default()); + self.meta.resize(count.get(), MinMaxCompensation::default()); + for i in 0..count.get() { + let row = b.get_row(start + i).expect("B tile range was checked"); + self.meta[i] = row.meta(); + decode_b::(row.vector(), &mut self.codes[i * dim..(i + 1) * dim]); + } + // SAFETY: The initialized codes contain count rows of exactly dim elements. + let codes = unsafe { unpacked::View::new(Slice::new(&self.codes), count, k) }; + BTile { + codes, + meta: &self.meta, + } + } + + /// Cache-planning bytes for one decoded B row, including metadata. + fn row_bytes(k: DimK) -> Bytes { + let codes = Elements::::new(k.value().get()).bytes(); + Bytes::new(codes.value() + size_of::()) + } +} + +/// Expand a B row, ignoring unused padding bits. +/// +/// # Panics +/// +/// Panics if the logical input and output lengths differ. +#[inline(always)] +#[expect(clippy::expect_used)] +fn decode_b( + codes: BitSlice<'_, B_BITS, Unsigned>, + dst: &mut [W::BElem], +) where + W: ExtraWide, + Unsigned: Representation, +{ + const { + assert!( + B_BITS <= W::MAX_B_BITS, + "B bit width exceeds the backend's supported range" + ) + }; + assert_eq!(codes.len(), dst.len()); + let packed = codes.as_slice(); + if B_BITS == 8 { + for (out, &code) in dst.iter_mut().zip(packed) { + *out = W::convert_b(code); + } + } else if B_BITS == 4 { + let (pairs, tail) = dst.as_chunks_mut::<2>(); + for (pair, &byte) in pairs.iter_mut().zip(packed) { + *pair = [W::convert_b(byte & 15), W::convert_b(byte >> 4)]; + } + if let Some(last) = tail.first_mut() { + *last = W::convert_b(packed[pairs.len()] & 15); + } + } else { + for (i, out) in dst.iter_mut().enumerate() { + *out = W::convert_b(codes.get(i).expect("input and output lengths match") as u8); + } + } +} + +//-------------// +// PanelKernel // +//-------------// + +struct PanelKernel<'a, W: ExtraWide, const MR: usize, const NR: usize, const PACK: usize> +{ + arch: W, + a: packed::Panel<'a, W::AElem, MR, PACK>, + a_meta: Coefficients, + b: unpacked::View<'a, W::BElem>, + b_meta: &'a [MinMaxCompensation], + c: [f32; MR], + k: DimK, +} + +impl, const MR: usize, const NR: usize, const PACK: usize> + Target for PanelKernel<'_, W, MR, NR, PACK> +{ + #[inline(always)] + fn run(mut self, _: W) -> [f32; MR] { + driver::PanelKernel::panel_kernel(&mut self); + self.c + } +} + +struct Visitor<'a, W: ExtraWide, const MR: usize, const PACK: usize> { + arch: W, + a: packed::Panel<'a, W::AElem, MR, PACK>, + a_meta: Coefficients, + b_meta: &'a [MinMaxCompensation], + c: &'a mut [f32; MR], + k: DimK, +} + +impl, const MR: usize, const NR: usize, const PACK: usize> + unpacked::PanelVisitor for Visitor<'_, W, MR, PACK> +{ + #[inline(always)] + fn visit(&mut self, b: unpacked::Panel<'_, W::BElem, NR>, start: usize) { + let mut kernel = MicroKernel { + arch: self.arch, + a: self.a, + a_meta: self.a_meta, + b, + b_meta: &self.b_meta[start..start + NR], + c: self.c, + k: self.k, + }; + driver::MicroKernel::micro_kernel(&mut kernel); + } +} + +impl, const MR: usize, const NR: usize, const PACK: usize> + driver::PanelKernel for PanelKernel<'_, W, MR, NR, PACK> +{ + #[inline(always)] + fn panel_kernel(&mut self) { + const { assert!(NR > 0 && NR <= 8) }; + let mut visitor = Visitor { + arch: self.arch, + a: self.a, + a_meta: self.a_meta, + b_meta: self.b_meta, + c: &mut self.c, + k: self.k, + }; + // SAFETY: PanelKernel is constructed with equal dimensions and matching B metadata. + let tail = unsafe { + self.b.visit_panels::( + self.k, + Visitor { + c: &mut *visitor.c, + ..visitor + }, + ) + }; + if let Some(tail) = tail { + macro_rules! remainder { + ($($n:literal),+) => {$( + if let Some(b) = tail.try_as_panel::<$n>() { + unpacked::PanelVisitor::visit(&mut visitor, b, tail.start()); + return; + } + )+}; + } + remainder!(1, 2, 3, 4, 5, 6, 7); + unreachable!("NR <= 8 bounds the trailing panel to at most seven B vectors"); + } + } +} + +//-------------// +// MicroKernel // +//-------------// + +struct MicroKernel<'a, W: ExtraWide, const MR: usize, const NR: usize, const PACK: usize> +{ + arch: W, + a: packed::Panel<'a, W::AElem, MR, PACK>, + a_meta: Coefficients, + b: unpacked::Panel<'a, W::BElem, NR>, + b_meta: &'a [MinMaxCompensation], + c: &'a mut [f32; MR], + k: DimK, +} + +/// Keep the hot contraction step in the caller's ISA context. +/// +/// # Safety +/// +/// `a` must contain at least `W::Wide::LANES` readable elements. `valid <= PACK`, +/// and for every `j < NR`, `b.add(bstride * j)` must contain `valid` readable elements. +#[inline(always)] +unsafe fn accumulate_pack( + arch: W, + a: Slice<'_, W::AElem>, + b: Slice<'_, W::BElem>, + bstride: Elements, + valid: usize, + acc: &mut [W::Acc; NR], +) where + W: ExtraWide, +{ + // SAFETY: The caller provides a full padded A pack. + let a = unsafe { W::Wide::load_simd(arch, a.as_ptr()) }; + for (j, acc) in acc.iter_mut().enumerate() { + // SAFETY: The caller guarantees valid elements in every B column. + let b = unsafe { arch.splat(b.add(bstride * j), valid) }; + *acc = arch.dot(*acc, a, b); + } +} + +impl, const MR: usize, const NR: usize, const PACK: usize> + driver::MicroKernel for MicroKernel<'_, W, MR, NR, PACK> +{ + #[inline(always)] + fn micro_kernel(&mut self) { + const { + assert!(W::Float::LANES == MR); + assert!(W::Wide::LANES == MR * PACK); + } + bounds::check_eq!(self.a.k(), self.k); + bounds::check_eq!(self.b.k(), self.k); + let k = self.k.value().get(); + let mut acc = [W::Acc::default(self.arch); NR]; + let ap = self.a.as_ptr(); + let bp = self.b.as_ptr(); + let astride = self.a.pack_stride(); + let bstride = self.b.stride(self.k); + let full = k / PACK; + + for pack in 0..full { + // SAFETY: A contains a full pack; every B column has PACK logical elements. + unsafe { + accumulate_pack::( + self.arch, + ap.add(astride * pack), + bp.add(Elements::new(pack * PACK)), + bstride, + PACK, + &mut acc, + ); + } + } + if !k.is_multiple_of(PACK) { + // SAFETY: A's last pack is padded. B has k % PACK remaining logical elements; + // splat zero-fills the rest, without reading a padded B extent. + unsafe { + accumulate_pack::( + self.arch, + ap.add(astride * full), + bp.add(Elements::new(full * PACK)), + bstride, + k % PACK, + &mut acc, + ); + } + } + + // SAFETY: Float has exactly MR lanes, and C is an MR-element array. + let mut best = unsafe { W::Float::load_simd(self.arch, self.c.as_ptr()) }; + for (acc, b) in acc.into_iter().zip(self.b_meta) { + let dot = acc.simd_cast(); + let a_b = W::Float::splat(self.arch, b.a); + let b_b = W::Float::splat(self.arch, b.b); + let n_b = W::Float::splat(self.arch, b.n); + let dim = W::Float::splat(self.arch, k as f32); + let a = self.a_meta; + let ip = ((a.a * a_b) * dot + a.n * b_b) + n_b * a.b + (a.b * b_b) * dim; + best = best.max_simd_standard(ip); + } + // SAFETY: C has exactly Float::LANES writable elements. + unsafe { best.store_simd(self.c.as_mut_ptr()) }; + } +} + +//----------// +// Backends // +//----------// + +/// SIMD types for MR A rows and PACK contraction elements. NR counts accumulators, +/// not lanes, so it does not change these types. +/// +/// All implementations accept A codes in 0..=255. B's range is limited by +/// both its element type and any intermediate arithmetic in the dot product. +pub(crate) trait ExtraWide: Architecture { + /// Largest complete unsigned B width supported for every A code in 0..=255. + /// This is a backend capability, not the actual A or B input width. + const MAX_B_BITS: usize; + + type AElem: Copy + Default; + type BElem: Copy + Default; + type Wide: SIMDVector; + type Splat: SIMDVector; + type Acc: SIMDVector + + SIMDDotProduct + + SIMDCast; + type Float: SIMDVector + + SIMDMinMax + + Add + + Mul; + + fn convert_a(code: u8) -> Self::AElem; + /// `code` fits in MAX_B_BITS unsigned bits. + fn convert_b(code: u8) -> Self::BElem; + + /// Accumulate codes produced by `convert_a` and `convert_b`. + #[inline(always)] + fn dot(self, acc: Self::Acc, a: Self::Wide, b: Self::Splat) -> Self::Acc { + acc.dot_simd(a, b) + } + + /// # Safety + /// + /// `valid <= PACK` and the first `valid` elements of B must be readable. + unsafe fn splat(self, b: Slice<'_, Self::BElem>, valid: usize) -> Self::Splat; +} + +/// # Safety +/// +/// `valid <= PACK` and the first `valid` elements of B must be readable. +#[inline(always)] +unsafe fn group(b: Slice<'_, T>, valid: usize) -> [T; PACK] { + core::array::from_fn(|i| { + if i < valid { + // SAFETY: i < valid, so the element is readable. + unsafe { *b.add(Elements::new(i)).as_unit().as_ref() } + } else { + T::default() + } + }) +} + +/// # Safety +/// +/// `valid <= 2` and the first `valid` elements of B must be readable. +#[inline(always)] +unsafe fn pair(b: Slice<'_, i16>, valid: usize) -> u32 { + // SAFETY: Inherited from the caller; group zero-fills a partial pair. + let [lo, hi] = unsafe { group::<_, 2>(b, valid) }; + u32::from(lo as u16) | (u32::from(hi as u16) << 16) +} + +impl ExtraWide<16, 4> for Scalar { + const MAX_B_BITS: usize = 7; + + type AElem = u8; + type BElem = i8; + type Wide = u8x64; + type Splat = i8x64; + type Acc = i32x16; + type Float = f32x16; + + fn convert_a(code: u8) -> u8 { + code + } + + fn convert_b(code: u8) -> i8 { + code as i8 + } + + #[inline(always)] + unsafe fn splat(self, b: Slice<'_, i8>, valid: usize) -> Self::Splat { + // SAFETY: Inherited from the caller. + let values = unsafe { group::<_, 4>(b, valid) }; + let quad = u32::from_le_bytes(values.map(|x| x as u8)); + u32x16::::splat(self, quad).reinterpret_simd() + } +} + +impl ExtraWide<16, 2> for Scalar { + const MAX_B_BITS: usize = 8; + + type AElem = i16; + type BElem = i16; + type Wide = i16x32; + type Splat = Self::Wide; + type Acc = i32x16; + type Float = f32x16; + + fn convert_a(code: u8) -> i16 { + i16::from(code) + } + + fn convert_b(code: u8) -> i16 { + i16::from(code) + } + + #[inline(always)] + unsafe fn splat(self, b: Slice<'_, i16>, valid: usize) -> Self::Splat { + // SAFETY: Inherited from the caller. + u32x16::::splat(self, unsafe { pair(b, valid) }).reinterpret_simd() + } +} + +#[cfg(target_arch = "x86_64")] +mod x86_64 { + use super::*; + use diskann_wide::SplitJoin; + use diskann_wide::arch::x86_64::{V3, V4}; + + impl ExtraWide<16, 4> for V3 { + const MAX_B_BITS: usize = 6; + + type AElem = u8; + type BElem = i8; + type Wide = u8x64; + type Splat = i8x64; + type Acc = i32x16; + type Float = f32x16; + + fn convert_a(code: u8) -> u8 { + code + } + + fn convert_b(code: u8) -> i8 { + code as i8 + } + + #[inline(always)] + fn dot(self, acc: Self::Acc, a: Self::Wide, b: Self::Splat) -> Self::Acc { + let a = a.split(); + let b = b.split(); + let acc = acc.split(); + let ones = i16x16::::splat(self, 1); + // Each saturated i16 pair is at most 2 * 255 * 63 = 32130, so it is exact. + i32x16::::new( + acc.lo.dot_simd(self.mm256_maddubs_epi16(a.lo, b.lo), ones), + acc.hi.dot_simd(self.mm256_maddubs_epi16(a.hi, b.hi), ones), + ) + } + + #[inline(always)] + unsafe fn splat(self, b: Slice<'_, i8>, valid: usize) -> Self::Splat { + // SAFETY: Inherited from the caller. + let values = unsafe { group::<_, 4>(b, valid) }; + let quad = u32::from_le_bytes(values.map(|x| x as u8)); + let half: i8x32 = u32x8::::splat(self, quad).reinterpret_simd(); + i8x64::::new(half, half) + } + } + + impl ExtraWide<16, 2> for V3 { + const MAX_B_BITS: usize = 8; + + type AElem = i16; + type BElem = i16; + type Wide = i16x32; + type Splat = Self::Wide; + type Acc = i32x16; + type Float = f32x16; + + fn convert_a(code: u8) -> i16 { + i16::from(code) + } + + fn convert_b(code: u8) -> i16 { + i16::from(code) + } + + #[inline(always)] + unsafe fn splat(self, b: Slice<'_, i16>, valid: usize) -> Self::Splat { + // SAFETY: Inherited from the caller. + let pair = unsafe { pair(b, valid) }; + let half: i16x16 = u32x8::::splat(self, pair).reinterpret_simd(); + i16x32::::new(half, half) + } + } + + impl ExtraWide<16, 4> for V4 { + const MAX_B_BITS: usize = 7; + + type AElem = u8; + type BElem = i8; + type Wide = u8x64; + type Splat = i8x64; + type Acc = i32x16; + type Float = f32x16; + + fn convert_a(code: u8) -> u8 { + code + } + + fn convert_b(code: u8) -> i8 { + code as i8 + } + + #[inline(always)] + unsafe fn splat(self, b: Slice<'_, i8>, valid: usize) -> Self::Splat { + // SAFETY: Inherited from the caller. + let values = unsafe { group::<_, 4>(b, valid) }; + let quad = u32::from_le_bytes(values.map(|x| x as u8)); + u32x16::::splat(self, quad).reinterpret_simd() + } + } + + impl ExtraWide<16, 2> for V4 { + const MAX_B_BITS: usize = 8; + + type AElem = i16; + type BElem = i16; + type Wide = i16x32; + type Splat = Self::Wide; + type Acc = i32x16; + type Float = f32x16; + + fn convert_a(code: u8) -> i16 { + i16::from(code) + } + + fn convert_b(code: u8) -> i16 { + i16::from(code) + } + + #[inline(always)] + unsafe fn splat(self, b: Slice<'_, i16>, valid: usize) -> Self::Splat { + // SAFETY: Inherited from the caller. + u32x16::::splat(self, unsafe { pair(b, valid) }).reinterpret_simd() + } + } +} + +#[cfg(target_arch = "aarch64")] +mod aarch64 { + use super::*; + use diskann_wide::arch::aarch64::Neon; + + impl ExtraWide<16, 4> for Neon { + const MAX_B_BITS: usize = 8; + + type AElem = u8; + type BElem = u8; + type Wide = u8x64; + type Splat = Self::Wide; + type Acc = u32x16; + type Float = f32x16; + + fn convert_a(code: u8) -> u8 { + code + } + + fn convert_b(code: u8) -> u8 { + code + } + + #[inline(always)] + unsafe fn splat(self, b: Slice<'_, u8>, valid: usize) -> Self::Splat { + // SAFETY: Inherited from the caller. + let quad = u32::from_le_bytes(unsafe { group::<_, 4>(b, valid) }); + let quarter = u32x4::::splat(self, quad).reinterpret_simd(); + let half = u8x32::::new(quarter, quarter); + u8x64::::new(half, half) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use diskann_vector::PureDistanceFunction; + use rand::{Rng, SeedableRng, rngs::StdRng}; + + use crate::{ + bits::{Representation, Unsigned}, + distances, + matrix_kernels::maxsim::test::packed_x_unpacked_test_dims, + minmax::{DataRef, MinMaxIP}, + multi_vector::{Defaulted, Mat}, + }; + + const MAX_K_8X4: usize = max_k::<8, 4>(); + + fn random_minmax_matrix( + rows: usize, + k: usize, + rng: &mut StdRng, + ) -> Mat> + where + Unsigned: Representation, + { + let mut mat = Mat::new(MinMaxMeta::::new(rows, k), Defaulted).unwrap(); + for mut row in mat.rows_mut() { + let a = rng.random_range(0.001..2.0); + let b = rng.random_range(-3.0..3.0); + let mut sum = 0; + for j in 0..k { + let code = rng.random_range(0..(1 << BITS)); + row.vector_mut().set(j, code).unwrap(); + sum += code; + } + row.set_meta(MinMaxCompensation { + dim: k as u32, + a, + b, + n: a * sum as f32, + norm_squared: 0.0, + }); + } + mat + } + + fn reference_max_inner_products( + a: MatRef<'_, MinMaxMeta>, + b: MatRef<'_, MinMaxMeta>, + ) -> Vec + where + Unsigned: Representation + Representation, + { + let k = a.repr().intrinsic_dim(); + assert_eq!(k, b.repr().intrinsic_dim()); + a.rows() + .map(|a| { + let av = a.vector(); + let am = a.meta(); + b.rows() + .map(|b| { + let bv = b.vector(); + let bm = b.meta(); + let dot: u64 = (0..k) + .map(|j| av.get(j).unwrap() as u64 * bv.get(j).unwrap() as u64) + .sum(); + ((am.a * bm.a) * dot as f32 + am.n * bm.b) + + bm.n * am.b + + (am.b * bm.b) * k as f32 + }) + .fold(f32::NEG_INFINITY, f32::max) + }) + .collect() + } + + #[track_caller] + fn assert_scores_close(actual: &[f32], expected: &[f32], context: std::fmt::Arguments<'_>) { + assert_eq!(actual.len(), expected.len(), "{context}"); + for (i, (&got, &want)) in actual.iter().zip(expected).enumerate() { + assert!( + got == want + || (got.is_finite() + && want.is_finite() + && (got - want).abs() <= 1e-5 * want.abs().max(1.0)), + "{context}: row {i}: got {got}, expected {want}" + ); + } + } + + fn checked_driver< + 'a, + W, + const A_BITS: usize, + const B_BITS: usize, + const MR: usize, + const NR: usize, + const PACK: usize, + >( + arch: W, + a: &'a PackedA, + b: MatRef<'a, MinMaxMeta>, + c: &'a mut [f32], + ) -> Driver<'a, W, B_BITS, MR, NR, PACK> + where + W: ExtraWide + LoadStore, + Unsigned: Representation + Representation, + { + assert!(a.codes.nrows() > 0); + assert_eq!(a.codes.nrows(), c.len()); + assert_eq!(a.codes.ncols(), b.repr().intrinsic_dim()); + let k = DimK::new(NonZeroUsize::new(a.codes.ncols()).unwrap()); + // SAFETY: Nonzero A/K and matching dimensions/output were checked above; + // PackedA's constructor establishes metadata extents and padding. + unsafe { Driver::::new(arch, a, b, c, k, Cache::detect()) } + } + + fn assert_driver_contract< + W, + const B_BITS: usize, + const MR: usize, + const NR: usize, + const PACK: usize, + >( + arch: W, + ) where + W: ExtraWide + LoadStore, + Unsigned: Representation, + { + let mut rng = StdRng::seed_from_u64(0x63de91a0ba5c8742); + for shape in packed_x_unpacked_test_dims(MR, NR) { + for k in [1, 2, 3, 4, 5, 7, 8, 15, 16, 17, 31, 32, 33, 127, 129] { + 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::<_, 8, MR, PACK>::new(arch, a.as_view()).unwrap(); + let expected = reference_max_inner_products::<8, B_BITS>(a.as_view(), b.as_view()); + let mut guarded = vec![654321.0; a.num_vectors() + 2]; + let mut driver = checked_driver::<_, 8, B_BITS, MR, NR, PACK>( + arch, + &packed, + b.as_view(), + &mut guarded[1..=a.num_vectors()], + ); + driver::Drive::drive(&mut driver); + assert_scores_close( + driver.c, + &expected, + format_args!("default tiling: {shape:?}, K={k}, MR={MR}, NR={NR}, PACK={PACK}"), + ); + + driver.b_rows_per_tile = NonZeroUsize::new(shape.b_cols_per_tile).unwrap(); + driver.c.fill(f32::INFINITY); + driver::Drive::drive(&mut driver); + assert_scores_close( + driver.c, + &expected, + format_args!( + "forced B tiling: {shape:?}, K={k}, MR={MR}, NR={NR}, PACK={PACK}" + ), + ); + // Reset stale scores once per invocation, never once per B tile. + driver.c.fill(f32::INFINITY); + driver::Drive::drive(&mut driver); + assert_scores_close( + driver.c, + &expected, + format_args!("repeated drive: {shape:?}, K={k}, MR={MR}, NR={NR}, PACK={PACK}"), + ); + assert_eq!(guarded[0], 654321.0); + assert_eq!(guarded[a.num_vectors() + 1], 654321.0); + } + } + assert_compensated_reduction::(arch); + assert_accumulator_limit_is_exact::(arch); + + let a = random_minmax_matrix::<8>(MR + 1, 5, &mut rng); + let b = random_minmax_matrix::(0, 5, &mut rng); + let packed = PackedA::<_, 8, MR, PACK>::new(arch, a.as_view()).unwrap(); + let mut c = vec![f32::INFINITY; a.num_vectors()]; + driver::Drive::drive(&mut checked_driver::<_, 8, B_BITS, MR, NR, PACK>( + arch, + &packed, + b.as_view(), + &mut c, + )); + assert_eq!(c, vec![f32::NEG_INFINITY; a.num_vectors()]); + } + + fn constant_minmax_matrix( + rows: usize, + k: usize, + code: i64, + a: f32, + b: f32, + ) -> Mat> + where + Unsigned: Representation, + { + let mut mat = Mat::new(MinMaxMeta::::new(rows, k), Defaulted).unwrap(); + for mut row in mat.rows_mut() { + for j in 0..k { + row.vector_mut().set(j, code).unwrap(); + } + row.set_meta(MinMaxCompensation { + dim: k as u32, + a, + b, + n: a * (k as f32 * code as f32), + norm_squared: 0.0, + }); + } + mat + } + + fn assert_compensated_reduction< + W, + const B_BITS: usize, + const MR: usize, + const NR: usize, + const PACK: usize, + >( + arch: W, + ) where + W: ExtraWide + LoadStore, + Unsigned: Representation, + { + let a = constant_minmax_matrix::<8>(MR + 1, 3, 1, 1.0, 0.0); + let mut b = constant_minmax_matrix::(NR + 1, 3, 15, 1.0, 0.0); + { + let mut row = b.get_row_mut(0).unwrap(); + row.vector_mut().set(0, 1).unwrap(); + row.vector_mut().set(1, 1).unwrap(); + row.vector_mut().set(2, 1).unwrap(); + row.set_meta(MinMaxCompensation { + dim: 3, + a: 100.0, + b: 0.0, + n: 300.0, + norm_squared: 0.0, + }); + } + let packed = PackedA::<_, 8, MR, PACK>::new(arch, a.as_view()).unwrap(); + let mut c = vec![0.0; MR + 1]; + let mut driver = + checked_driver::<_, 8, B_BITS, MR, NR, PACK>(arch, &packed, b.as_view(), &mut c); + // The compensated winner is in the first tile, not the last. + driver.b_rows_per_tile = NonZeroUsize::new(1).unwrap(); + driver::Drive::drive(&mut driver); + assert_eq!(c, vec![300.0; MR + 1]); + + let a = constant_minmax_matrix::<8>(MR + 1, 3, 1, 1.0, -2.0); + let b = constant_minmax_matrix::((NR - 1).max(1), 3, 1, 1.0, 0.0); + let packed = PackedA::<_, 8, MR, PACK>::new(arch, a.as_view()).unwrap(); + driver::Drive::drive(&mut checked_driver::<_, 8, B_BITS, MR, NR, PACK>( + arch, + &packed, + b.as_view(), + &mut c, + )); + // Invalid B lanes must not introduce a zero candidate. + assert_eq!(c, vec![-3.0; MR + 1]); + } + + fn assert_accumulator_limit_is_exact< + W, + const B_BITS: usize, + const MR: usize, + const NR: usize, + const PACK: usize, + >( + arch: W, + ) where + W: ExtraWide + LoadStore, + Unsigned: Representation, + { + if cfg!(miri) { + return; + } + let k = max_k::<8, B_BITS>(); + let max_b = (1 << B_BITS) - 1; + let a = constant_minmax_matrix::<8>(1, k, 255, 1.0, 0.0); + let b = constant_minmax_matrix::(1, k, max_b as i64, 1.0, 0.0); + let packed = PackedA::<_, 8, MR, PACK>::new(arch, a.as_view()).unwrap(); + let mut c = [0.0]; + driver::Drive::drive(&mut checked_driver::<_, 8, B_BITS, MR, NR, PACK>( + arch, + &packed, + b.as_view(), + &mut c, + )); + assert_eq!(c[0], (k * 255 * max_b) as f32); + } + + #[test] + fn scalar_driver_obeys_maxsim_contract() { + assert_driver_contract::(Scalar::new()); + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn avx2_driver_obeys_maxsim_contract() { + if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { + assert_driver_contract::<_, 4, 16, 4, 4>(arch); + } else { + eprintln!("AVX2 unavailable; backend not executed"); + } + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn avx512_driver_obeys_maxsim_contract() { + if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked() { + assert_driver_contract::<_, 4, 16, 8, 4>(arch); + } else { + eprintln!("V4 unavailable; backend not executed"); + } + } + + #[cfg(target_arch = "aarch64")] + #[test] + fn neon_driver_obeys_maxsim_contract() { + if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { + assert_driver_contract::<_, 4, 16, 4, 4>(arch); + } else { + eprintln!("Neon/dotprod unavailable; backend not executed"); + } + } + + fn assert_a_packing(arch: W) + where + W: ExtraWide, + W::AElem: PartialEq + std::fmt::Debug, + Unsigned: Representation, + { + for (m, k) in [(0, 0), (0, 5), (3, 0), (0, MAX_K_8X4), (0, MAX_K_8X4 + 1)] { + let a = constant_minmax_matrix::(m, k, 0, 1.0, 0.0); + let packed = PackedA::<_, A_BITS, MR, PACK>::new(arch, a.as_view()).unwrap(); + assert_eq!(packed.codes.nrows(), m); + assert_eq!(packed.codes.ncols(), k); + assert!(packed::View::from_block_transposed(packed.codes.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::<_, A_BITS, MR, PACK>::new(arch, a.as_view()).unwrap(); + let padded_cols = 5usize.div_ceil(PACK) * PACK; + assert_eq!(packed.codes.nrows(), MR + 1); + assert_eq!(packed.codes.ncols(), 5); + assert_eq!(packed.codes.padded_nrows(), 2 * MR); + assert_eq!(packed.codes.padded_ncols(), padded_cols); + for row in 0..2 * MR { + for col in 0..padded_cols { + let index = (row / MR) * MR * padded_cols + + (col / PACK) * MR * PACK + + (row % MR) * PACK + + col % PACK; + let expected = if row <= MR && col < 5 { + W::convert_a(a.get_row(row).unwrap().vector().get(col).unwrap() as u8) + } else { + W::AElem::default() + }; + assert_eq!(packed.codes.as_slice()[index], expected); + } + let meta = a.get_row(row).map(|v| v.meta()).unwrap_or_default(); + assert_eq!(packed.meta.a[row], meta.a); + assert_eq!(packed.meta.b[row], meta.b); + assert_eq!(packed.meta.n[row], meta.n); + } + } + + #[test] + fn packing_preserves_codes_coefficients_and_zero_padding() { + assert_a_packing::<_, 8, 16, 4>(Scalar::new()); + #[cfg(target_arch = "x86_64")] + if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { + assert_a_packing::<_, 8, 16, 4>(arch); + } + #[cfg(target_arch = "x86_64")] + if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked() { + assert_a_packing::<_, 8, 16, 4>(arch); + } + #[cfg(target_arch = "aarch64")] + if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { + assert_a_packing::<_, 8, 16, 4>(arch); + } + } + + #[test] + fn cache_footprints_match_prepared_layouts() { + for (k, a_bytes, b_bytes) in [(256, 4288, 276), (257, 4352, 277)] { + let dim = DimK::new(NonZeroUsize::new(k).unwrap()); + assert_eq!(BScratch::::row_bytes(dim).value(), b_bytes); + 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(Scalar::new(), a.as_view()).unwrap(); + assert_eq!(packed.panel_bytes().value(), a_bytes); + + // Preserve the previous view-based cache estimate. + let view = packed::View::from_block_transposed(packed.codes.as_view()).unwrap(); + assert_eq!( + packed.panel_bytes().value(), + view.block_stride(dim).bytes().value() + 3 * 16 * size_of::() + ); + } + } + } + + #[test] + fn widened_cache_footprints_match_prepared_layouts() { + for (k, a_bytes, b_bytes) in [(256, 8384, 532), (257, 8448, 534)] { + let dim = DimK::new(NonZeroUsize::new(k).unwrap()); + assert_eq!(BScratch::::row_bytes(dim).value(), b_bytes); + for rows in [1, 16, 17] { + let a = constant_minmax_matrix::<8>(rows, k, 255, 1.0, 0.0); + let packed = PackedA::::new(Scalar::new(), a.as_view()).unwrap(); + assert_eq!(packed.panel_bytes().value(), a_bytes); + let view = packed::View::from_block_transposed(packed.codes.as_view()).unwrap(); + assert_eq!( + packed.panel_bytes().value(), + view.block_stride(dim).bytes().value() + 3 * 16 * size_of::(), + ); + } + } + } + + #[cfg(not(miri))] + #[test] + #[should_panic(expected = "exceeds the MinMax8 x MinMax4 limit")] + fn driver_rejects_k_above_accumulator_limit() { + let k = MAX_K_8X4 + 1; + let a = constant_minmax_matrix::<8>(1, k, 255, 1.0, 0.0); + let b = constant_minmax_matrix::<4>(1, k, 15, 1.0, 0.0); + let packed = PackedA::::new(Scalar::new(), a.as_view()).unwrap(); + let mut c = [123.0]; + let _ = checked_driver::<_, 8, 4, 16, 2, 4>(Scalar::new(), &packed, b.as_view(), &mut c); + } + + #[test] + fn b_preparation_masks_padding_and_reuses_storage() { + let k = 5; + let stride = crate::minmax::Data::<4>::canonical_bytes(k); + let mut bytes = vec![0xff; stride * 3]; + for (row_index, raw) in bytes.chunks_exact_mut(stride).enumerate() { + let mut row = crate::minmax::DataMutRef::<4>::from_canonical_front_mut(raw, k).unwrap(); + let a = (row_index + 1) as f32; + row.set_meta(MinMaxCompensation { + dim: k as u32, + a, + b: -(row_index as f32), + n: a * (10 + 5 * row_index) as f32, + ..Default::default() + }); + for i in 0..k { + row.vector_mut().set(i, (i + row_index) as i64).unwrap(); + } + } + let b = MatRef::new(MinMaxMeta::<4>::new(3, k), &bytes).unwrap(); + let mut scratch = BScratch::::default(); + let dim = DimK::new(NonZeroUsize::new(k).unwrap()); + let tile = scratch.prepare::(b, 0, NonZeroUsize::new(3).unwrap(), dim); + assert_eq!( + // SAFETY: dim is the dimension passed to prepare. + unsafe { tile.codes.as_std_slice(dim) }, + &[0, 1, 2, 3, 4, 1, 2, 3, 4, 5, 2, 3, 4, 5, 6] + ); + let ptr = scratch.codes.as_ptr(); + let tile = scratch.prepare::(b, 2, NonZeroUsize::new(1).unwrap(), dim); + assert_eq!(tile.meta, &[b.get_row(2).unwrap().meta()]); + assert_eq!( + // SAFETY: dim is the dimension passed to prepare. + unsafe { tile.codes.as_std_slice(dim) }, + &[2, 3, 4, 5, 6] + ); + assert_eq!(scratch.codes.as_ptr(), ptr); + } + + fn assert_b_decoding_boundaries< + W, + const B_BITS: usize, + const MR: usize, + const NR: usize, + const PACK: usize, + >( + arch: W, + sentinel: W::BElem, + ) where + W: ExtraWide + LoadStore, + W::BElem: PartialEq + std::fmt::Debug, + Unsigned: Representation, + { + decode_b::(BitSlice::new(&[] as &[u8], 0).unwrap(), &mut []); + let mut rng = StdRng::seed_from_u64(0x8af450de13c769b2); + let mut scratch = BScratch::::default(); + let dimensions: &[usize] = if cfg!(miri) { + &[1, 2, 3, 4, 5, 15, 16, 17] + } else { + &[ + 1, 2, 15, 16, 17, 31, 32, 33, 63, 64, 65, 127, 128, 129, 250, 255, 256, 257, 511, + 512, 513, + ] + }; + for &dim in dimensions { + let k = DimK::new(NonZeroUsize::new(dim).unwrap()); + let a = random_minmax_matrix::<8>(MR + 1, dim, &mut rng); + let packed = PackedA::<_, 8, MR, PACK>::new(arch, a.as_view()).unwrap(); + let stride = crate::minmax::Data::::canonical_bytes(dim); + for offset in [0, 1, 3, 15] { + let mut storage = vec![0xff; offset + 5 * stride]; + let bytes = &mut storage[offset..]; + for (i, raw) in bytes.chunks_exact_mut(stride).enumerate() { + let mut row = + crate::minmax::DataMutRef::::from_canonical_front_mut(raw, dim) + .unwrap(); + let mut sum = 0; + for j in 0..dim { + // Cover every code value, and every packed byte for MinMax4. + let code = if B_BITS == 4 { + let byte = (i * 37 + j / 2) & 255; + (byte >> (4 * (j % 2))) & 15 + } else { + (i * 37 + j) & ((1 << B_BITS) - 1) + } as i64; + row.vector_mut().set(j, code).unwrap(); + sum += code; + } + let scale = (i + 1) as f32 * 0.25; + row.set_meta(MinMaxCompensation { + dim: dim as u32, + a: scale, + b: i as f32 - 3.0, + n: scale * sum as f32, + norm_squared: (i * i + 1) as f32, + }); + } + let b = MatRef::new(MinMaxMeta::::new(5, dim), bytes).unwrap(); + for row in b.rows() { + let codes = row.vector(); + if B_BITS == 4 && dim % 2 != 0 { + assert_eq!(codes.as_slice().last().unwrap() >> 4, 15); + } + let mut guarded = vec![sentinel; dim + 2]; + decode_b::(codes, &mut guarded[1..=dim]); + for j in 0..dim { + assert_eq!( + guarded[j + 1], + W::convert_b(codes.get(j).unwrap() as u8), + "row decode: K={dim}, offset={offset}, column={j}, MR={MR}, PACK={PACK}" + ); + } + assert_eq!(guarded[0], sentinel); + assert_eq!(guarded[dim + 1], sentinel); + } + for (start, count) in [(0, 1), (1, 3), (4, 1), (0, 5)] { + scratch.codes.fill(sentinel); + scratch.meta.fill(MinMaxCompensation::default()); + let tile = scratch.prepare::( + b, + start, + NonZeroUsize::new(count).unwrap(), + k, + ); + // SAFETY: k is the dimension used to prepare this tile. + let codes = unsafe { tile.codes.as_std_slice(k) }; + assert_eq!(codes.len(), count * dim); + assert_eq!(tile.meta.len(), count); + for i in 0..count { + let row = b.get_row(start + i).unwrap(); + assert_eq!(tile.meta[i], row.meta()); + for j in 0..dim { + assert_eq!( + codes[i * dim + j], + W::convert_b(row.vector().get(j).unwrap() as u8), + "B tile: K={dim}, offset={offset}, start={start}, \ + count={count}, row={i}, column={j}, MR={MR}, PACK={PACK}" + ); + } + } + } + let expected = reference_max_inner_products::<8, B_BITS>(a.as_view(), b); + let mut guarded = vec![654321.0; MR + 3]; + for tile_rows in [3, 4, 5] { + let mut driver = checked_driver::( + arch, + &packed, + b, + &mut guarded[1..=MR + 1], + ); + driver.b_rows_per_tile = NonZeroUsize::new(tile_rows).unwrap(); + driver.c.fill(f32::INFINITY); + driver::Drive::drive(&mut driver); + assert_scores_close( + driver.c, + &expected, + format_args!( + "B decoding: K={dim}, offset={offset}, B tile={tile_rows}, \ + MR={MR}, NR={NR}, PACK={PACK}" + ), + ); + assert_eq!(guarded[0], 654321.0); + assert_eq!(guarded[MR + 2], 654321.0); + } + } + } + } + + #[test] + fn scalar_b_decoding_preserves_rows_tails_and_scores() { + assert_b_decoding_boundaries::<_, 4, 16, 2, 4>(Scalar::new(), -1); + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn avx2_b_decoding_preserves_rows_tails_and_scores() { + if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { + assert_b_decoding_boundaries::<_, 4, 16, 4, 4>(arch, -1); + } else { + eprintln!("AVX2 unavailable; B decoding regression not executed"); + } + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn avx512_b_decoding_preserves_rows_tails_and_scores() { + if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked() { + assert_b_decoding_boundaries::<_, 4, 16, 8, 4>(arch, -1); + } else { + eprintln!("V4 unavailable; B decoding regression not executed"); + } + } + + #[cfg(target_arch = "aarch64")] + #[test] + fn neon_b_decoding_preserves_rows_tails_and_scores() { + if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { + assert_b_decoding_boundaries::<_, 4, 16, 4, 4>(arch, u8::MAX); + } else { + eprintln!("Neon/dotprod unavailable; B decoding regression not executed"); + } + } + + #[test] + #[should_panic] + fn b_decoding_rejects_a_short_output() { + let codes = BitSlice::new(&[0x21u8][..], 2).unwrap(); + decode_b::(codes, &mut [0]); + } + + #[test] + #[should_panic] + fn b_decoding_rejects_a_long_output() { + let codes = BitSlice::new(&[0xf1u8][..], 1).unwrap(); + decode_b::(codes, &mut [0; 2]); + } + + fn assert_packed_a_reuse< + W, + const B_BITS: usize, + const MR: usize, + const NR: usize, + const PACK: usize, + >( + arch: W, + ) where + W: ExtraWide + LoadStore, + Unsigned: Representation, + { + let k = 17; + let mut rng = StdRng::seed_from_u64(0xa619e2c5f387b04d); + let a = random_minmax_matrix::<8>(11, k, &mut rng); + let b = random_minmax_matrix::(9, k, &mut rng); + let expected = reference_max_inner_products::<8, B_BITS>(a.as_view(), b.as_view()); + let packed = PackedA::<_, 8, MR, PACK>::new(arch, a.as_view()).unwrap(); + drop(a); + let mut actual = vec![0.0; 11]; + for _ in 0..2 { + actual.fill(f32::INFINITY); + driver::Drive::drive(&mut checked_driver::<_, 8, B_BITS, MR, NR, PACK>( + arch, + &packed, + b.as_view(), + &mut actual, + )); + assert_scores_close( + &actual, + &expected, + format_args!("packed-A reuse: MR={MR}, NR={NR}, PACK={PACK}"), + ); + } + } + + #[test] + fn packed_a_outlives_source_and_supports_repeated_drives() { + assert_packed_a_reuse::<_, 4, 16, 2, 4>(Scalar::new()); + #[cfg(target_arch = "x86_64")] + if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { + assert_packed_a_reuse::<_, 4, 16, 4, 4>(arch); + } + #[cfg(target_arch = "x86_64")] + if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked() { + assert_packed_a_reuse::<_, 4, 16, 8, 4>(arch); + } + #[cfg(target_arch = "aarch64")] + if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { + assert_packed_a_reuse::<_, 4, 16, 4, 4>(arch); + } + } + + fn assert_driver_target_shapes< + W, + const B_BITS: usize, + const MR: usize, + const NR: usize, + const PACK: usize, + >( + arch: W, + ) where + W: ExtraWide + LoadStore, + Unsigned: Representation, + { + let mut rng = StdRng::seed_from_u64(0x8f61d9c34a257be0); + let dimensions: &[usize] = if cfg!(miri) { &[3] } else { &[3, 256, 257] }; + for m in [1, 16, 17] { + for n in (1..=NR + 1).chain([16, 17]) { + for &k in dimensions { + 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::<_, 8, MR, PACK>::new(arch, a.as_view()).unwrap(); + let mut c = vec![f32::INFINITY; m]; + driver::Drive::drive(&mut checked_driver::<_, 8, B_BITS, MR, NR, PACK>( + arch, + &packed, + b.as_view(), + &mut c, + )); + assert_scores_close( + &c, + &expected, + format_args!("M={m}, N={n}, K={k}, MR={MR}, NR={NR}, PACK={PACK}"), + ); + } + } + } + let a = constant_minmax_matrix::<8>(16, 257, 255, 1.0, 0.0); + let max_b = (1 << B_BITS) - 1; + let b = constant_minmax_matrix::(16, 257, max_b, 1.0, 0.0); + let packed = PackedA::<_, 8, MR, PACK>::new(arch, a.as_view()).unwrap(); + let mut c = [0.0; 16]; + driver::Drive::drive(&mut checked_driver::<_, 8, B_BITS, MR, NR, PACK>( + arch, + &packed, + b.as_view(), + &mut c, + )); + assert_eq!(c, [(257 * 255 * max_b) as f32; 16]); + } + + #[test] + fn scalar_driver_matches_reference_at_target_shapes() { + assert_driver_target_shapes::<_, 4, 16, 2, 4>(Scalar::new()); + } + + #[cfg(target_arch = "aarch64")] + #[test] + fn neon_driver_matches_reference_at_target_shapes() { + if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { + assert_driver_target_shapes::<_, 4, 16, 4, 4>(arch); + } else { + eprintln!("Neon/dotprod unavailable; target-shape regression not executed"); + } + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn avx2_driver_matches_reference_at_target_shapes() { + if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { + assert_driver_target_shapes::<_, 4, 16, 4, 4>(arch); + } else { + eprintln!("AVX2 unavailable; target-shape regression not executed"); + } + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn avx512_driver_matches_reference_at_target_shapes() { + if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked() { + assert_driver_target_shapes::<_, 4, 16, 8, 4>(arch); + } else { + eprintln!("V4 unavailable; target-shape regression not executed"); + } + } + + fn assert_nan_candidates_preserve_maxima< + W, + const B_BITS: usize, + const MR: usize, + const NR: usize, + const PACK: usize, + >( + arch: W, + ) where + W: ExtraWide + LoadStore, + Unsigned: Representation, + { + let mut a = constant_minmax_matrix::<8>(MR + 1, 1, 1, 1.0, 0.0); + // Exercise mixed SIMD lanes as well as the partial A panel. + for row in [1, MR] { + a.get_row_mut(row).unwrap().set_meta(MinMaxCompensation { + dim: 1, + a: 1.0, + b: 0.0, + n: f32::NAN, + norm_squared: 0.0, + }); + } + let packed = PackedA::<_, 8, MR, PACK>::new(arch, a.as_view()).unwrap(); + let mut cases = vec![ + vec![], + vec![10.0, f32::NAN, 5.0], + vec![f32::NAN, 10.0, 5.0], + vec![10.0, 5.0, f32::NAN], + vec![10.0, f32::NAN, f32::NAN, 5.0], + vec![-5.0, f32::NAN, -10.0], + vec![f32::NAN, f32::NAN], + ]; + let n = 2 * NR + 1; + for nan_index in [0, NR, n - 1] { + let mut scores = vec![5.0; n]; + scores[if nan_index == 0 { 1 } else { 0 }] = 10.0; + scores[nan_index] = f32::NAN; + cases.push(scores); + } + + for scores in cases { + let mut b = constant_minmax_matrix::(scores.len(), 1, 0, 1.0, 0.0); + for (mut row, offset) in b.rows_mut().zip(scores.iter().copied()) { + row.set_meta(MinMaxCompensation { + dim: 1, + a: 1.0, + b: offset, + n: 0.0, + norm_squared: 0.0, + }); + } + let expected = reference_max_inner_products::<8, B_BITS>(a.as_view(), b.as_view()); + assert_eq!( + expected[0], + scores.iter().copied().fold(f32::NEG_INFINITY, f32::max) + ); + assert_eq!(expected[1], f32::NEG_INFINITY); + assert_eq!(expected[MR], f32::NEG_INFINITY); + for tile_rows in [1, 2, NR, NR + 1, scores.len().max(1)] { + let mut c = vec![f32::NAN; MR + 1]; + let mut driver = checked_driver::<_, 8, B_BITS, MR, NR, PACK>( + arch, + &packed, + b.as_view(), + &mut c, + ); + driver.b_rows_per_tile = NonZeroUsize::new(tile_rows).unwrap(); + driver::Drive::drive(&mut driver); + assert_scores_close( + &c, + &expected, + format_args!( + "NaN reduction: scores={scores:?}, B tile={tile_rows}, \ + MR={MR}, NR={NR}, PACK={PACK}" + ), + ); + } + } + } + + #[test] + fn scalar_driver_ignores_nan_candidates() { + assert_nan_candidates_preserve_maxima::(Scalar::new()); + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn avx2_driver_ignores_nan_candidates() { + if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { + assert_nan_candidates_preserve_maxima::<_, 4, 16, 4, 4>(arch); + } else { + eprintln!("AVX2 unavailable; NaN regression not executed"); + } + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn avx512_driver_ignores_nan_candidates() { + if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked() { + assert_nan_candidates_preserve_maxima::<_, 4, 16, 8, 4>(arch); + } else { + eprintln!("V4 unavailable; NaN regression not executed"); + } + } + + #[cfg(target_arch = "aarch64")] + #[test] + fn neon_driver_ignores_nan_candidates() { + if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { + assert_nan_candidates_preserve_maxima::<_, 4, 16, 4, 4>(arch); + } else { + eprintln!("Neon/dotprod unavailable; NaN regression not executed"); + } + } + + mod all_widths { + use super::*; + + fn assert_backend_code_range(arch: W) + where + W: ExtraWide, + { + assert!((1..=8).contains(&W::MAX_B_BITS)); + let max_b = u8::MAX >> (8 - W::MAX_B_BITS); + arch.run(|| { + let a = W::Wide::splat(arch, W::convert_a(u8::MAX)); + let b = W::Splat::splat(arch, W::convert_b(max_b)); + let dot = arch.dot(W::Acc::default(arch), a, b).simd_cast(); + let mut actual = [0.0; MR]; + assert_eq!(W::Float::LANES, MR); + // SAFETY: actual has exactly Float::LANES writable elements. + unsafe { dot.store_simd(actual.as_mut_ptr()) }; + assert_eq!( + actual, + [(PACK * usize::from(u8::MAX) * usize::from(max_b)) as f32; MR], + ); + }); + } + + fn assert_bit_pair< + W, + const A_BITS: usize, + const B_BITS: usize, + const MR: usize, + const NR: usize, + const PACK: usize, + >( + arch: W, + rng: &mut StdRng, + ) where + W: ExtraWide + LoadStore, + W::AElem: PartialEq + std::fmt::Debug, + Unsigned: Representation + Representation, + { + let max_a = (1u64 << A_BITS) - 1; + let max_b = (1u64 << B_BITS) - 1; + let limit = max_k::() as u64; + assert!(limit * max_a * max_b <= i32::MAX as u64); + assert!((limit + 1) * max_a * max_b > i32::MAX as u64); + assert_a_packing::(arch); + + let dimensions: &[usize] = if cfg!(miri) { + &[1, 3] + } else { + &[1, 2, 3, 4, 5, 7, 8, 15, 16, 17, 31, 33, 255, 256, 257] + }; + for (m, n) in [ + (1, 0), + (1, 1), + (MR - 1, (NR - 1).max(1)), + (MR, NR), + (MR + 1, 2 * NR + 1), + ] { + for &k in dimensions { + let mut a = random_minmax_matrix::(m, k, rng); + let mut b = random_minmax_matrix::(n, k, rng); + if m > 1 { + let mut row = a.get_row_mut(m - 1).unwrap(); + let mut meta = row.meta(); + meta.n = f32::NAN; + row.set_meta(meta); + } + if n > 1 { + let mut row = b.get_row_mut(n - 1).unwrap(); + let mut meta = row.meta(); + meta.n = f32::NAN; + row.set_meta(meta); + } + let expected = + reference_max_inner_products::(a.as_view(), b.as_view()); + let packed = PackedA::<_, A_BITS, MR, PACK>::new(arch, a.as_view()).unwrap(); + let mut guarded = vec![654321.0; m + 2]; + let mut kernel = checked_driver::( + arch, + &packed, + b.as_view(), + &mut guarded[1..=m], + ); + for tile_rows in [kernel.b_rows_per_tile, NonZeroUsize::new(NR + 1).unwrap()] { + kernel.b_rows_per_tile = tile_rows; + kernel.c.fill(f32::INFINITY); + driver::Drive::drive(&mut kernel); + assert_scores_close( + kernel.c, + &expected, + format_args!( + "{A_BITS}x{B_BITS}: M={m}, N={n}, K={k}, B tile={tile_rows}, PACK={PACK}" + ), + ); + } + assert_eq!(guarded[0], 654321.0); + assert_eq!(guarded[m + 1], 654321.0); + } + } + + let dimensions: &[usize] = if cfg!(miri) { + &[PACK, PACK + 1] + } else { + &[PACK, PACK + 1, 257] + }; + for &k in dimensions { + let a = constant_minmax_matrix::(MR + 1, k, max_a as i64, 1.0, 0.0); + let b = constant_minmax_matrix::(NR + 1, k, max_b as i64, 1.0, 0.0); + let packed = PackedA::<_, A_BITS, MR, PACK>::new(arch, a.as_view()).unwrap(); + let mut c = vec![0.0; MR + 1]; + driver::Drive::drive(&mut checked_driver::( + arch, + &packed, + b.as_view(), + &mut c, + )); + assert_eq!(c, vec![(k as u64 * max_a * max_b) as f32; MR + 1]); + } + } + + macro_rules! check_widths { + ($arch:expr, $sentinel:expr, $nr:literal, $pack:literal; $($b:literal),+ $(,)?) => {{ + let arch = $arch; + let mut rng = StdRng::seed_from_u64(0x948cd3e02a671fb5); + assert_backend_code_range::<_, 16, $pack>(arch); + $( + assert_b_decoding_boundaries::<_, $b, 16, $nr, $pack>(arch, $sentinel); + assert_bit_pair::<_, 1, $b, 16, $nr, $pack>(arch, &mut rng); + assert_bit_pair::<_, 2, $b, 16, $nr, $pack>(arch, &mut rng); + assert_bit_pair::<_, 3, $b, 16, $nr, $pack>(arch, &mut rng); + assert_bit_pair::<_, 4, $b, 16, $nr, $pack>(arch, &mut rng); + assert_bit_pair::<_, 5, $b, 16, $nr, $pack>(arch, &mut rng); + assert_bit_pair::<_, 6, $b, 16, $nr, $pack>(arch, &mut rng); + assert_bit_pair::<_, 7, $b, 16, $nr, $pack>(arch, &mut rng); + assert_bit_pair::<_, 8, $b, 16, $nr, $pack>(arch, &mut rng); + )+ + }}; + } + + #[test] + fn scalar_supports_all_bit_pairs_and_byte_layouts() { + check_widths!(Scalar::new(), -1, 2, 2; 1, 2, 3, 4, 5, 6, 7, 8); + check_widths!(Scalar::new(), -1, 2, 4; 1, 2, 3, 4, 5, 6, 7); + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn avx2_supports_all_bit_pairs_and_byte_layouts() { + if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { + check_widths!(arch, -1, 4, 2; 1, 2, 3, 4, 5, 6, 7, 8); + check_widths!(arch, -1, 4, 4; 1, 2, 3, 4, 5, 6); + } else { + eprintln!("AVX2 unavailable; all-width regression not executed"); + } + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn avx512_supports_all_bit_pairs_and_byte_layouts() { + if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked() { + check_widths!(arch, -1, 8, 2; 1, 2, 3, 4, 5, 6, 7, 8); + check_widths!(arch, -1, 8, 4; 1, 2, 3, 4, 5, 6, 7); + } else { + eprintln!("V4 unavailable; all-width regression not executed"); + } + } + + #[cfg(target_arch = "aarch64")] + #[test] + fn neon_supports_all_bit_pairs() { + if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { + check_widths!(arch, u8::MAX, 4, 4; 1, 2, 3, 4, 5, 6, 7, 8); + } else { + eprintln!("Neon/dotprod unavailable; all-width regression not executed"); + } + } + + fn check_reference() + where + Unsigned: Representation + Representation, + MinMaxIP: for<'a, 'b> PureDistanceFunction< + DataRef<'a, A_BITS>, + DataRef<'b, B_BITS>, + distances::MathematicalResult, + >, + { + let mut rng = StdRng::seed_from_u64(0x14a7f0b2c69d538e); + for k in [1, 3, 17, 257] { + let a = random_minmax_matrix::(17, k, &mut rng); + let b = random_minmax_matrix::(19, k, &mut rng); + let expected: Vec<_> = a + .rows() + .map(|a| { + b.rows() + .map(|b| { + , + DataRef<'_, B_BITS>, + distances::MathematicalResult, + >>::evaluate(a, b) + .unwrap() + .into_inner() + }) + .fold(f32::NEG_INFINITY, f32::max) + }) + .collect(); + assert_scores_close( + &reference_max_inner_products::(a.as_view(), b.as_view()), + &expected, + format_args!("scalar reference: {A_BITS}x{B_BITS}, K={k}"), + ); + } + } + + #[test] + fn reference_agrees_with_existing_pairwise_combinations() { + check_reference::<1, 1>(); + check_reference::<2, 2>(); + check_reference::<3, 3>(); + check_reference::<4, 4>(); + check_reference::<5, 5>(); + check_reference::<6, 6>(); + check_reference::<7, 7>(); + check_reference::<8, 8>(); + check_reference::<8, 4>(); + check_reference::<8, 2>(); + check_reference::<8, 1>(); + } + + #[test] + fn packed_a_is_reused_across_b_widths() { + let arch = Scalar::new(); + let mut rng = StdRng::seed_from_u64(0xac21f49d70b365e8); + let a = random_minmax_matrix::<3>(17, 33, &mut rng); + let packed = PackedA::<_, 3, 16, 2>::new(arch, a.as_view()).unwrap(); + let mut c = vec![f32::NAN; 17]; + macro_rules! check { + ($($b:literal),+) => {$( + let b = random_minmax_matrix::<$b>(9, 33, &mut rng); + let expected = reference_max_inner_products::<3, $b>(a.as_view(), b.as_view()); + driver::Drive::drive(&mut checked_driver::<_, 3, $b, 16, 2, 2>( + arch, &packed, b.as_view(), &mut c, + )); + assert_scores_close(&c, &expected, format_args!("reused A3 with B{}", $b)); + )+}; + } + check!(1, 2, 3, 4, 5, 6, 7, 8); + } + + #[cfg(not(miri))] + #[test] + fn low_bit_inputs_can_exceed_the_eight_bit_dimension_limit() { + fn check() + where + Unsigned: Representation + Representation, + { + let k = max_k::<8, 8>() + 1; + let max_a = (1 << A_BITS) - 1; + let max_b = (1 << B_BITS) - 1; + let a = constant_minmax_matrix::(1, k, max_a, 1.0, 0.0); + let b = constant_minmax_matrix::(1, k, max_b, 1.0, 0.0); + let packed = PackedA::<_, A_BITS, 16, 2>::new(Scalar::new(), a.as_view()).unwrap(); + let mut c = [0.0]; + driver::Drive::drive(&mut checked_driver::<_, A_BITS, B_BITS, 16, 2, 2>( + Scalar::new(), + &packed, + b.as_view(), + &mut c, + )); + assert_eq!(c[0], (k as i64 * max_a * max_b) as f32); + } + check::<1, 8>(); + check::<4, 8>(); + check::<8, 1>(); + } + } + + mod minmax8 { + use super::*; + + fn check_backend(arch: W, sentinel: W::BElem) + where + W: ExtraWide<16, PACK> + LoadStore, + W::AElem: PartialEq + std::fmt::Debug, + W::BElem: PartialEq + std::fmt::Debug, + { + assert_a_packing::(arch); + assert_driver_contract::(arch); + assert_driver_target_shapes::(arch); + assert_nan_candidates_preserve_maxima::(arch); + assert_packed_a_reuse::(arch); + assert_b_decoding_boundaries::(arch, sentinel); + assert_unsigned_byte_products::(arch); + } + + fn assert_unsigned_byte_products(arch: W) + where + W: ExtraWide<16, PACK> + LoadStore, + { + if cfg!(miri) { + return; + } + for k in [1, 2, 3, 4, 5, 257, 259] { + let mut a = constant_minmax_matrix::<8>(257, k, 0, 1.0, 0.0); + for (i, mut row) in a.rows_mut().enumerate() { + let code = (i % 256) as i64; + for j in 0..k { + row.vector_mut().set(j, code).unwrap(); + } + let mut meta = row.meta(); + meta.n = (k as i64 * code) as f32; + row.set_meta(meta); + } + let a = PackedA::<_, 8, 16, PACK>::new(arch, a.as_view()).unwrap(); + let mut c = vec![f32::NAN; 257]; + for code in 0..=255 { + let b = constant_minmax_matrix::<8>(1, k, code, 1.0, 0.0); + driver::Drive::drive(&mut checked_driver::( + arch, + &a, + b.as_view(), + &mut c, + )); + for (i, &actual) in c.iter().enumerate() { + assert_eq!( + actual, + ((i % 256) * code as usize * k) as f32, + "A={}, B={code}, K={k}", + i % 256, + ); + } + } + } + } + + #[test] + fn scalar_driver_obeys_maxsim_contract() { + check_backend::<_, 2, 2>(Scalar::new(), -1); + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn avx2_driver_obeys_maxsim_contract() { + if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { + check_backend::<_, 4, 2>(arch, -1); + } else { + eprintln!("AVX2 unavailable; MinMax8 x MinMax8 backend not executed"); + } + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn avx512_driver_obeys_maxsim_contract() { + if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked() { + check_backend::<_, 8, 2>(arch, -1); + } else { + eprintln!("V4 unavailable; MinMax8 x MinMax8 backend not executed"); + } + } + + #[cfg(target_arch = "aarch64")] + #[test] + fn neon_driver_obeys_maxsim_contract() { + if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { + check_backend::<_, 4, 4>(arch, u8::MAX); + } else { + eprintln!("Neon/dotprod unavailable; MinMax8 x MinMax8 backend not executed"); + } + } + + #[cfg(not(miri))] + #[test] + #[should_panic(expected = "exceeds the MinMax8 x MinMax8 limit")] + fn driver_rejects_k_above_accumulator_limit() { + let k = 33026; + 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::<_, 8, 16, 2>::new(Scalar::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); + } + + #[test] + #[should_panic] + fn b_decoding_rejects_a_short_output() { + let codes = BitSlice::new(&[128u8, 255][..], 2).unwrap(); + decode_b::(codes, &mut [0]); + } + + #[test] + #[should_panic] + fn b_decoding_rejects_a_long_output() { + let codes = BitSlice::new(&[255u8][..], 1).unwrap(); + decode_b::(codes, &mut [0; 2]); + } + } +} diff --git a/diskann-wide/src/arch/aarch64/double.rs b/diskann-wide/src/arch/aarch64/double.rs index 7087793b0e..e2d874b13d 100644 --- a/diskann-wide/src/arch/aarch64/double.rs +++ b/diskann-wide/src/arch/aarch64/double.rs @@ -197,6 +197,22 @@ impl crate::SIMDCast for i32x8 { } } +impl crate::SIMDCast for u32x8 { + type Cast = f32x8; + #[inline(always)] + fn simd_cast(self) -> f32x8 { + f32x8::new(self.0.simd_cast(), self.1.simd_cast()) + } +} + +impl crate::SIMDCast for u32x16 { + type Cast = f32x16; + #[inline(always)] + fn simd_cast(self) -> f32x16 { + f32x16::new(self.0.simd_cast(), self.1.simd_cast()) + } +} + /////////// // Tests // /////////// @@ -460,6 +476,8 @@ mod tests { test_utils::ops::test_cast!(f32x16 => f16x16, 0xba8fe343fc9dbeff, test_neon()); test_utils::ops::test_cast!(i32x8 => f32x8, 0xba8fe343fc9dbeff, test_neon()); + test_utils::ops::test_cast!(u32x8 => f32x8, 0x1e9e37b58fb3f1a8, test_neon()); + test_utils::ops::test_cast!(u32x16 => f32x16, 0xd2b068a9bf3f9d24, test_neon()); // SplitJoin + ZipUnzip for f16x16 test_utils::ops::test_splitjoin!(f16x16 => f16x8, 0x2e301b7e12090d5c, test_neon()); diff --git a/diskann-wide/src/arch/aarch64/u32x4_.rs b/diskann-wide/src/arch/aarch64/u32x4_.rs index cefa6393c6..526e7471b6 100644 --- a/diskann-wide/src/arch/aarch64/u32x4_.rs +++ b/diskann-wide/src/arch/aarch64/u32x4_.rs @@ -4,13 +4,13 @@ */ use crate::{ - Emulated, SIMDDotProduct, SIMDMask, SIMDMulAdd, SIMDPartialEq, SIMDPartialOrd, SIMDReinterpret, - SIMDSelect, SIMDSumTree, SIMDVector, constant::Const, helpers, + Emulated, SIMDCast, SIMDDotProduct, SIMDMask, SIMDMulAdd, SIMDPartialEq, SIMDPartialOrd, + SIMDReinterpret, SIMDSelect, SIMDSumTree, SIMDVector, constant::Const, helpers, }; // AArch64 masks use super::{ - Neon, i8x16, internal, + Neon, f32x4, i8x16, internal, macros::{self, AArchLoadStore, AArchSplat}, masks::mask32x4, u8x16, u32x2, @@ -116,10 +116,24 @@ impl SIMDDotProduct for u32x4 { } } +//-------------// +// Conversions // +//-------------// + +helpers::unsafe_map_cast!(u32x4 => (f32, f32x4), vcvtq_f32_u32, "neon"); + ////////////////// // Reinterprets // ////////////////// +impl SIMDReinterpret for u32x4 { + #[inline(always)] + fn reinterpret_simd(self) -> u8x16 { + // SAFETY: Allowed by the Neon architecture. + u8x16(unsafe { vreinterpretq_u8_u32(self.0) }) + } +} + impl SIMDReinterpret for u32x4 { #[inline(always)] fn reinterpret_simd(self) -> i8x16 { @@ -183,6 +197,10 @@ mod tests { // Reductions test_utils::ops::test_sumtree!(u32x4, 0xb9ac82ab23a855da, test_neon()); + // Conversions + test_utils::ops::test_cast!(u32x4 => f32x4, 0x9a8d420371b6c5ef, test_neon()); + // Reinterprets test_utils::ops::test_reinterpret!(u32x4 => i8x16, 0xbfad755f32d25e5c, test_neon()); + test_utils::ops::test_reinterpret!(u32x4 => u8x16, 0xd56e9c21a7183b4f, test_neon()); } diff --git a/diskann-wide/src/arch/x86_64/v3/conversion.rs b/diskann-wide/src/arch/x86_64/v3/conversion.rs index a03ac92d66..a2cdd19dab 100644 --- a/diskann-wide/src/arch/x86_64/v3/conversion.rs +++ b/diskann-wide/src/arch/x86_64/v3/conversion.rs @@ -7,7 +7,7 @@ use std::arch::x86_64::*; use super::{ - f16x8, f16x16, f32x8, f32x16, i8x16, i8x32, i16x8, i16x16, i32x8, u8x16, u8x32, u32x8, + f16x8, f16x16, f32x8, f32x16, i8x16, i8x32, i16x8, i16x16, i32x8, i32x16, u8x16, u8x32, u32x8, }; use crate::{LoHi, SIMDCast, SIMDReinterpret, SplitJoin, helpers}; @@ -83,6 +83,14 @@ impl SIMDCast for f32x16 { } } +impl SIMDCast for i32x16 { + type Cast = f32x16; + #[inline(always)] + fn simd_cast(self) -> f32x16 { + self.split().map(|x| -> f32x8 { x.simd_cast() }).join() + } +} + // i8 to i16 helpers::unsafe_map_conversion!(i8x16, i16x16, _mm256_cvtepi8_epi16, "avx2"); @@ -177,4 +185,5 @@ mod test_x86_conversions { test_utils::ops::test_cast!(f32x16 => f16x16, 0x7b9c9afee7e6ac63, V3::new_checked_uncached()); test_utils::ops::test_cast!(i32x8 => f32x8, 0xde4fbf25c554b29e, V3::new_checked_uncached()); + test_utils::ops::test_cast!(i32x16 => f32x16, 0x83e571a946db20cf, V3::new_checked_uncached()); } diff --git a/diskann-wide/src/arch/x86_64/v4/conversion.rs b/diskann-wide/src/arch/x86_64/v4/conversion.rs index 16f679db69..37393ebc0f 100644 --- a/diskann-wide/src/arch/x86_64/v4/conversion.rs +++ b/diskann-wide/src/arch/x86_64/v4/conversion.rs @@ -8,8 +8,8 @@ use std::arch::x86_64::*; use super::{ f16x8_::f16x8, f16x16_::f16x16, f32x8_::f32x8, f32x16_::f32x16, i8x16_::i8x16, i8x32_::i8x32, - i8x64_::i8x64, i16x8_::i16x8, i16x16_::i16x16, i16x32_::i16x32, i32x8_::i32x8, u8x16_::u8x16, - u8x32_::u8x32, u8x64_::u8x64, u32x8_::u32x8, u32x16_::u32x16, u64x8_::u64x8, + i8x64_::i8x64, i16x8_::i16x8, i16x16_::i16x16, i16x32_::i16x32, i32x8_::i32x8, i32x16_::i32x16, + u8x16_::u8x16, u8x32_::u8x32, u8x64_::u8x64, u32x8_::u32x8, u32x16_::u32x16, u64x8_::u64x8, }; use crate::{SIMDCast, SIMDReinterpret, SIMDVector, arch::x86_64::v3, helpers}; @@ -78,6 +78,7 @@ helpers::unsafe_map_conversion!(u8x32, i16x32, _mm512_cvtepu8_epi16, "avx512bw") // i32 to f32 helpers::unsafe_map_cast!(i32x8 => (f32, f32x8), _mm256_cvtepi32_ps, "avx"); +helpers::unsafe_map_cast!(i32x16 => (f32, f32x16), _mm512_cvtepi32_ps, "avx512f"); helpers::unsafe_map_cast!(i16x16 => (u8, u8x16), _mm256_cvtepi16_epi8, "avx512bw,avx512vl"); helpers::unsafe_map_cast!(i16x16 => (i8, i8x16), _mm256_cvtepi16_epi8, "avx512bw,avx512vl"); @@ -220,6 +221,7 @@ mod test_x86_conversions { test_utils::ops::test_cast!(f32x16 => f16x16, 0xb3cbae34def475df, V4::new_checked_uncached()); test_utils::ops::test_cast!(i32x8 => f32x8, 0xde4fbf25c554b29e, V4::new_checked_uncached()); + test_utils::ops::test_cast!(i32x16 => f32x16, 0x42ec9b7d219536a1, V4::new_checked_uncached()); test_utils::ops::test_cast!(i16x16 => u8x16, 0x0f81df9e640b0269, V4::new_checked_uncached()); test_utils::ops::test_cast!(i16x16 => i8x16, 0x4ab1546b9d0e4046, V4::new_checked_uncached()); diff --git a/diskann-wide/src/emulated.rs b/diskann-wide/src/emulated.rs index be12697949..7161adee83 100644 --- a/diskann-wide/src/emulated.rs +++ b/diskann-wide/src/emulated.rs @@ -665,6 +665,7 @@ impl_cast!(f32 => f16, 8); impl_cast!(f32 => f16, 16); impl_cast!(i32 => f32, 8); +impl_cast!(i32 => f32, 16); /////////////// // SplitJoin // @@ -1175,6 +1176,7 @@ mod test_emulated { test_utils::ops::test_cast!(Emulated => Emulated, 0x2b637e21afd9ef6c, SC); test_utils::ops::test_cast!(Emulated => Emulated, 0x2b08e8ec7e49323b, SC); + test_utils::ops::test_cast!(Emulated => Emulated, 0x83e571a946db20cf, SC); ////////////// // ZipUnzip // diff --git a/diskann-wide/src/reference.rs b/diskann-wide/src/reference.rs index f260dd073a..c5ffe63d84 100644 --- a/diskann-wide/src/reference.rs +++ b/diskann-wide/src/reference.rs @@ -251,6 +251,12 @@ impl ReferenceCast for i32 { } } +impl ReferenceCast for u32 { + fn reference_cast(self) -> f32 { + self as f32 + } +} + impl ReferenceCast for i16 { fn reference_cast(self) -> u8 { self as u8 From 41710a54f05b740b0f6406b57306646cb28773e7 Mon Sep 17 00:00:00 2001 From: partychen <1669295+partychen@users.noreply.github.com> Date: Sat, 17 Oct 2026 12:44:34 +0800 Subject: [PATCH 2/4] Refine MinMax kernel preparation, backend operations, and tests Move bulk decoding into BitSlice and zero-pad prepared B rows for a unified pack loop. Decouple backend representations behind fixed-size operations, retain smaller row-panel choices, and replace redundant test sweeps with directed coverage. Make B preparation independent of backend and row-panel types by accepting a code converter. Keep backend capability validation at the Driver construction boundary. Give each backend the full microkernel compensation and reduction stage through max_into with borrowed coefficients. Keep MR16 contraction intact while V3 reduces native eight-row halves. Define backend operations explicitly and share the compensation formula; preserve floating-point order and NaN semantics. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- diskann-quantization/src/bits/mod.rs | 6 +- diskann-quantization/src/bits/slice.rs | 150 ++ .../maxsim/packed_minmax_x_unpacked_minmax.rs | 1764 +++++++++-------- 3 files changed, 1134 insertions(+), 786 deletions(-) diff --git a/diskann-quantization/src/bits/mod.rs b/diskann-quantization/src/bits/mod.rs index 6b5b06734b..62763a5923 100644 --- a/diskann-quantization/src/bits/mod.rs +++ b/diskann-quantization/src/bits/mod.rs @@ -15,7 +15,7 @@ pub mod distances; pub use length::{Dynamic, Length, Static}; pub use ptr::{AsMutPtr, AsPtr, MutSlicePtr, SlicePtr}; pub use slice::{ - Binary, BitSlice, BitSliceBase, BitTranspose, BoxedBitSlice, ConstructionError, Dense, - EncodingError, GetError, IndexOutOfBounds, MutBitSlice, PermutationStrategy, Representation, - SetError, Unsigned, + Binary, BitSlice, BitSliceBase, BitTranspose, BoxedBitSlice, ConstructionError, + DecodeLengthMismatch, Dense, EncodingError, GetError, IndexOutOfBounds, MutBitSlice, + PermutationStrategy, Representation, SetError, Unsigned, }; diff --git a/diskann-quantization/src/bits/slice.rs b/diskann-quantization/src/bits/slice.rs index dcfaeff7c2..98f8fbc430 100644 --- a/diskann-quantization/src/bits/slice.rs +++ b/diskann-quantization/src/bits/slice.rs @@ -383,6 +383,14 @@ unsafe impl PermutationStrategy<4> for BitTranspose { // Errors // //////////// +/// The destination cannot hold exactly the logical decoded values. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)] +#[error("decoded output has {output} elements, expected {input}")] +pub struct DecodeLengthMismatch { + pub input: usize, + pub output: usize, +} + #[derive(Debug, Error, Clone, Copy)] #[error("input span has length {got} bytes but expected {expected}")] pub struct ConstructionError { @@ -732,6 +740,64 @@ where } } +impl BitSliceBase +where + Repr: Representation, + Ptr: AsPtr, + Len: Length, +{ + /// Decode in logical order, applying `convert` once to each value. + /// + /// Unused bits in the last byte are ignored. Unlike repeated `get` calls, bulk + /// decoding can reuse each packed byte across adjacent values. + /// + /// # Errors + /// + /// Returns [`DecodeLengthMismatch`] without writing or calling `convert` if the + /// output length differs from the logical input length. + #[inline] + pub fn decode_into( + &self, + output: &mut [T], + mut convert: impl FnMut(i64) -> T, + ) -> Result<(), DecodeLengthMismatch> { + const { assert!(NBITS > 0 && NBITS <= 8) }; + if output.len() != self.len() { + return Err(DecodeLengthMismatch { + input: self.len(), + output: output.len(), + }); + } + let packed = self.as_slice(); + let divisor = 1usize << NBITS.trailing_zeros(); + let codes_per_block = 8 / divisor; + let bytes_per_block = NBITS / divisor; + let blocks = output.len() / codes_per_block; + let (output, tail) = output.split_at_mut(blocks * codes_per_block); + let (packed, rest) = packed.split_at(blocks * bytes_per_block); + let mask = (1u64 << NBITS) - 1; + // Use the smallest byte-aligned group: 8 / gcd(NBITS, 8) codes. + for (block, bytes) in output + .chunks_exact_mut(codes_per_block) + .zip(packed.chunks_exact(bytes_per_block)) + { + let mut word = [0u8; 8]; + word[..bytes_per_block].copy_from_slice(bytes); + let word = u64::from_le_bytes(word); + for (i, out) in block.iter_mut().enumerate() { + *out = convert(Repr::decode(((word >> (i * NBITS)) & mask) as u8)); + } + } + let mut word = [0u8; 8]; + word[..rest.len()].copy_from_slice(rest); + let word = u64::from_le_bytes(word); + for (i, out) in tail.iter_mut().enumerate() { + *out = convert(Repr::decode(((word >> (i * NBITS)) & mask) as u8)); + } + Ok(()) + } +} + impl BitSliceBase, Perm, Len> where @@ -947,6 +1013,90 @@ mod tests { use super::*; use crate::{bits::Static, test_util::AlwaysFails}; + fn check_bulk_decoding() + where + Unsigned: Representation, + { + for len in [0, 1, 2, 3, 7, 8, 9, 15, 16, 17, 255, 256, 257, 513] { + for offset in [0, 1, 7] { + let byte_count = BitSlice::::bytes_for(len); + let mut bytes = vec![0xff; offset + byte_count]; + let bytes = &mut bytes[offset..]; + let mut codes = MutBitSlice::::new(bytes, len).unwrap(); + for i in 0..len { + codes + .set(i, ((i * 37 + i / 2) & ((1 << NBITS) - 1)) as i64) + .unwrap(); + } + let expected: Vec<_> = (0..len).map(|i| codes.get(i).unwrap()).collect(); + let mut guarded = vec![i64::MIN; len + 2]; + let mut calls = 0; + codes + .decode_into(&mut guarded[1..=len], |value| { + assert_eq!(value, expected[calls]); + calls += 1; + value + }) + .unwrap(); + assert_eq!(calls, len); + assert_eq!(&guarded[1..=len], expected); + assert_eq!(guarded[0], i64::MIN); + assert_eq!(guarded[len + 1], i64::MIN); + } + } + } + + #[test] + fn bulk_decoding_preserves_order_tails_and_output_bounds() { + check_bulk_decoding::<1>(); + check_bulk_decoding::<2>(); + check_bulk_decoding::<3>(); + check_bulk_decoding::<4>(); + check_bulk_decoding::<5>(); + check_bulk_decoding::<6>(); + check_bulk_decoding::<7>(); + check_bulk_decoding::<8>(); + } + + #[test] + fn bulk_decoding_applies_representation_and_conversion() { + let codes = + BitSlice::<1, Binary, Dense, Static<3>>::new(&[0b101u8][..], Static::<3>).unwrap(); + let mut output = [0i16; 3]; + codes + .decode_into(&mut output, |value| (value * 3) as i16) + .unwrap(); + assert_eq!(output, [3, -3, 3]); + let codes = BoxedBitSlice::<4, Unsigned>::new_boxed(2); + let mut output = [1u8; 2]; + codes.decode_into(&mut output, |value| value as u8).unwrap(); + assert_eq!(output, [0, 0]); + } + + #[test] + fn bulk_decoding_length_mismatch_does_not_write_or_convert() { + let codes = BitSlice::<4, Unsigned>::new(&[0x21u8][..], 2).unwrap(); + for len in [0, 1, 3] { + let mut output = vec![99; len]; + let mut calls = 0; + let error = codes + .decode_into(&mut output, |value| { + calls += 1; + value + }) + .unwrap_err(); + assert_eq!( + error, + DecodeLengthMismatch { + input: 2, + output: len + } + ); + assert_eq!(calls, 0); + assert_eq!(output, vec![99; len]); + } + } + //////////// // Errors // //////////// diff --git a/diskann-quantization/src/matrix_kernels/maxsim/packed_minmax_x_unpacked_minmax.rs b/diskann-quantization/src/matrix_kernels/maxsim/packed_minmax_x_unpacked_minmax.rs index a9556b38f7..e712e55251 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 @@ -7,6 +7,8 @@ //! //! Each B tile is decoded once and reused across A panels. Integer inner products //! are compensated before reduction because B's coefficients vary by vector. +//! B rows are zero-padded to PACK; kernel `k` remains the logical dimension used +//! for compensation and accumulator limits. //! //! Outputs are unnegated maximum inner products. NaN candidates are ignored; //! empty B or all-NaN comparisons leave negative infinity. @@ -23,12 +25,12 @@ use std::{ }; use diskann_wide::{ - SIMDCast, SIMDDotProduct, SIMDMinMax, SIMDReinterpret, SIMDVector, + Const, SIMDCast, SIMDDotProduct, SIMDMinMax, SIMDReinterpret, SIMDVector, SupportedLaneCount, arch::{Architecture, Scalar, Target}, }; use crate::{ - bits::{BitSlice, Representation, Unsigned}, + bits::{Representation, Unsigned}, matrix_kernels::{ Cache, blocks::{packed, unpacked}, @@ -43,12 +45,12 @@ use crate::{ use super::packed_f32_x_unpacked_f32::Params; -diskann_wide::alias!(f32x16 = f32x16); diskann_wide::alias!(i8x32 = i8x32); diskann_wide::alias!(i8x64 = i8x64); diskann_wide::alias!(i16x16 = i16x16); diskann_wide::alias!(i16x32 = i16x32); diskann_wide::alias!(i32x16 = i32x16); +diskann_wide::alias!(i32x8 = i32x8); diskann_wide::alias!(u8x32 = u8x32); diskann_wide::alias!(u8x64 = u8x64); diskann_wide::alias!(u32x4 = u32x4); @@ -62,7 +64,7 @@ const fn max_k() -> usize { } #[derive(Debug, Clone, Copy)] -struct Coefficients { +pub(crate) struct Coefficients { a: T, b: T, n: T, @@ -142,7 +144,7 @@ pub(crate) struct Driver< { arch: W, a: packed::View<'a, W::AElem, MR, PACK>, - a_meta: Coefficients<&'a [f32]>, + a_meta: Coefficients<&'a [[f32; MR]]>, b: MatRef<'a, MinMaxMeta>, c: &'a mut [f32], k: DimK, @@ -185,6 +187,7 @@ where Unsigned: Representation, { const { + assert!(NR > 0 && NR <= 8, "NR must be in 1..=8"); assert!( B_BITS <= W::MAX_B_BITS, "B bit width exceeds the backend's supported range" @@ -200,18 +203,23 @@ where bounds::check_eq!(bounds::Bound::new(b.repr().intrinsic_dim()), k); bounds::check_eq!(bounds::Bound::new(c.len()), a.codes.nrows()); let a_meta = Coefficients { - a: a.meta.a.as_slice(), - b: a.meta.b.as_slice(), - n: a.meta.n.as_slice(), + a: a.meta.a.as_chunks::().0, + b: a.meta.b.as_chunks::().0, + n: a.meta.n.as_chunks::().0, }; let a_bytes = a.panel_bytes(); let a = packed::View::from_block_transposed(a.codes.as_view()) .expect("A and K must be nonempty"); for meta in [a_meta.a, a_meta.b, a_meta.n] { - bounds::check_eq!(bounds::Bound::new(meta.len()), MR * a.blocks().get()); + bounds::check_eq!(bounds::Bound::new(meta.len()), a.blocks().get()); } - let b_rows_per_tile = - Params::new(cache, a_bytes, BScratch::::row_bytes(k), NR).b_cols_in_l1; + let b_rows_per_tile = Params::new( + cache, + a_bytes, + BScratch::::row_bytes::(k), + NR, + ) + .b_cols_in_l1; Self { arch, a, @@ -239,7 +247,8 @@ where let mut start = 0; while let Some(remaining) = NonZeroUsize::new(self.b.num_vectors() - start) { let count = remaining.min(self.b_rows_per_tile); - let b = scratch.prepare::(self.b, start, count, self.k); + let b = + scratch.prepare::(self.b, start, count, self.k, W::convert_b); // SAFETY: Driver's construction establishes all dimensions, metadata // extents, and C's coverage of A. Visitors preserve these bounds. unsafe { @@ -248,9 +257,9 @@ where let end = (row + MR).min(self.c.len()); let c = self.arch.load(&self.c[row..end]); let a_meta = Coefficients { - a: W::Float::load_simd(self.arch, self.a_meta.a.as_ptr().add(row)), - b: W::Float::load_simd(self.arch, self.a_meta.b.as_ptr().add(row)), - n: W::Float::load_simd(self.arch, self.a_meta.n.as_ptr().add(row)), + a: &self.a_meta.a[panel], + b: &self.a_meta.b[panel], + n: &self.a_meta.n[panel], }; let kernel = PanelKernel:: { arch: self.arch, @@ -283,31 +292,47 @@ struct BTile<'a, T> { meta: &'a [MinMaxCompensation], } +#[inline] +#[expect(clippy::expect_used)] +fn padded_k(k: DimK) -> DimK { + DimK::new( + NonZeroUsize::new(k.value().get().next_multiple_of(PACK)) + .expect("padding a positive dimension remains positive"), + ) +} + impl BScratch { + /// Convert logical codes once each; padding uses `T::default()` without conversion. #[expect(clippy::expect_used)] - fn prepare( + fn prepare( &mut self, b: MatRef<'_, MinMaxMeta>, start: usize, count: NonZeroUsize, k: DimK, + mut convert: impl FnMut(u8) -> T, ) -> BTile<'_, T> where - W: ExtraWide, Unsigned: Representation, { let dim = k.value().get(); assert_eq!(b.repr().intrinsic_dim(), dim); assert!(start <= b.num_vectors() && count.get() <= b.num_vectors() - start); - self.codes.resize(count.get() * dim, T::default()); + 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(), MinMaxCompensation::default()); for i in 0..count.get() { let row = b.get_row(start + i).expect("B tile range was checked"); self.meta[i] = row.meta(); - decode_b::(row.vector(), &mut self.codes[i * dim..(i + 1) * dim]); - } - // SAFETY: The initialized codes contain count rows of exactly dim elements. - let codes = unsafe { unpacked::View::new(Slice::new(&self.codes), count, k) }; + let dst = &mut self.codes[i * stride..(i + 1) * stride]; + row.vector() + .decode_into(&mut dst[..dim], |code| convert(code as u8)) + .expect("B output must match the logical code count"); + dst[dim..].fill(T::default()); + } + // SAFETY: Each row has physical_k initialized elements, including zero padding. + let codes = unsafe { unpacked::View::new(Slice::new(&self.codes), count, physical_k) }; BTile { codes, meta: &self.meta, @@ -315,53 +340,12 @@ impl BScratch { } /// Cache-planning bytes for one decoded B row, including metadata. - fn row_bytes(k: DimK) -> Bytes { - let codes = Elements::::new(k.value().get()).bytes(); + fn row_bytes(k: DimK) -> Bytes { + let codes = Elements::::new(padded_k::(k).value().get()).bytes(); Bytes::new(codes.value() + size_of::()) } } -/// Expand a B row, ignoring unused padding bits. -/// -/// # Panics -/// -/// Panics if the logical input and output lengths differ. -#[inline(always)] -#[expect(clippy::expect_used)] -fn decode_b( - codes: BitSlice<'_, B_BITS, Unsigned>, - dst: &mut [W::BElem], -) where - W: ExtraWide, - Unsigned: Representation, -{ - const { - assert!( - B_BITS <= W::MAX_B_BITS, - "B bit width exceeds the backend's supported range" - ) - }; - assert_eq!(codes.len(), dst.len()); - let packed = codes.as_slice(); - if B_BITS == 8 { - for (out, &code) in dst.iter_mut().zip(packed) { - *out = W::convert_b(code); - } - } else if B_BITS == 4 { - let (pairs, tail) = dst.as_chunks_mut::<2>(); - for (pair, &byte) in pairs.iter_mut().zip(packed) { - *pair = [W::convert_b(byte & 15), W::convert_b(byte >> 4)]; - } - if let Some(last) = tail.first_mut() { - *last = W::convert_b(packed[pairs.len()] & 15); - } - } else { - for (i, out) in dst.iter_mut().enumerate() { - *out = W::convert_b(codes.get(i).expect("input and output lengths match") as u8); - } - } -} - //-------------// // PanelKernel // //-------------// @@ -370,7 +354,7 @@ struct PanelKernel<'a, W: ExtraWide, const MR: usize, const NR: usize, { arch: W, a: packed::Panel<'a, W::AElem, MR, PACK>, - a_meta: Coefficients, + a_meta: Coefficients<&'a [f32; MR]>, b: unpacked::View<'a, W::BElem>, b_meta: &'a [MinMaxCompensation], c: [f32; MR], @@ -390,7 +374,7 @@ impl, const MR: usize, const NR: usize, const PACK: usize struct Visitor<'a, W: ExtraWide, const MR: usize, const PACK: usize> { arch: W, a: packed::Panel<'a, W::AElem, MR, PACK>, - a_meta: Coefficients, + a_meta: Coefficients<&'a [f32; MR]>, b_meta: &'a [MinMaxCompensation], c: &'a mut [f32; MR], k: DimK, @@ -400,13 +384,16 @@ impl, const MR: usize, const NR: usize, const PACK: usize unpacked::PanelVisitor for Visitor<'_, W, MR, PACK> { #[inline(always)] + #[expect(clippy::expect_used)] fn visit(&mut self, b: unpacked::Panel<'_, W::BElem, NR>, start: usize) { let mut kernel = MicroKernel { arch: self.arch, a: self.a, a_meta: self.a_meta, b, - b_meta: &self.b_meta[start..start + NR], + b_meta: self.b_meta[start..] + .first_chunk::() + .expect("B panel metadata is in bounds"), c: self.c, k: self.k, }; @@ -419,7 +406,6 @@ impl, const MR: usize, const NR: usize, const PACK: usize { #[inline(always)] fn panel_kernel(&mut self) { - const { assert!(NR > 0 && NR <= 8) }; let mut visitor = Visitor { arch: self.arch, a: self.a, @@ -428,10 +414,10 @@ impl, const MR: usize, const NR: usize, const PACK: usize c: &mut self.c, k: self.k, }; - // SAFETY: PanelKernel is constructed with equal dimensions and matching B metadata. + // SAFETY: B's view is constructed with PACK-padded rows and matching metadata. let tail = unsafe { self.b.visit_panels::( - self.k, + padded_k::(self.k), Visitor { c: &mut *visitor.c, ..visitor @@ -441,7 +427,7 @@ impl, const MR: usize, const NR: usize, const PACK: usize if let Some(tail) = tail { macro_rules! remainder { ($($n:literal),+) => {$( - if let Some(b) = tail.try_as_panel::<$n>() { + if $n < NR && let Some(b) = tail.try_as_panel::<$n>() { unpacked::PanelVisitor::visit(&mut visitor, b, tail.start()); return; } @@ -461,9 +447,9 @@ struct MicroKernel<'a, W: ExtraWide, const MR: usize, const NR: usize, { arch: W, a: packed::Panel<'a, W::AElem, MR, PACK>, - a_meta: Coefficients, + a_meta: Coefficients<&'a [f32; MR]>, b: unpacked::Panel<'a, W::BElem, NR>, - b_meta: &'a [MinMaxCompensation], + b_meta: &'a [MinMaxCompensation; NR], c: &'a mut [f32; MR], k: DimK, } @@ -472,24 +458,23 @@ struct MicroKernel<'a, W: ExtraWide, const MR: usize, const NR: usize, /// /// # Safety /// -/// `a` must contain at least `W::Wide::LANES` readable elements. `valid <= PACK`, -/// and for every `j < NR`, `b.add(bstride * j)` must contain `valid` readable elements. +/// `a` must contain at least MR * PACK readable elements, and for every +/// `j < NR`, `b.add(bstride * j)` must contain PACK readable elements. #[inline(always)] unsafe fn accumulate_pack( arch: W, a: Slice<'_, W::AElem>, b: Slice<'_, W::BElem>, bstride: Elements, - valid: usize, acc: &mut [W::Acc; NR], ) where W: ExtraWide, { // SAFETY: The caller provides a full padded A pack. - let a = unsafe { W::Wide::load_simd(arch, a.as_ptr()) }; + let a = unsafe { arch.load_a(a) }; for (j, acc) in acc.iter_mut().enumerate() { - // SAFETY: The caller guarantees valid elements in every B column. - let b = unsafe { arch.splat(b.add(bstride * j), valid) }; + // SAFETY: The caller guarantees a full pack in every B column. + let b = unsafe { arch.splat(b.add(bstride * j)) }; *acc = arch.dot(*acc, a, b); } } @@ -499,71 +484,38 @@ impl, const MR: usize, const NR: usize, const PACK: usize { #[inline(always)] fn micro_kernel(&mut self) { - const { - assert!(W::Float::LANES == MR); - assert!(W::Wide::LANES == MR * PACK); - } bounds::check_eq!(self.a.k(), self.k); - bounds::check_eq!(self.b.k(), self.k); + let physical_k = padded_k::(self.k); + bounds::check_eq!(self.b.k(), physical_k); let k = self.k.value().get(); - let mut acc = [W::Acc::default(self.arch); NR]; + let mut acc = [self.arch.zero_acc(); NR]; let ap = self.a.as_ptr(); let bp = self.b.as_ptr(); let astride = self.a.pack_stride(); - let bstride = self.b.stride(self.k); - let full = k / PACK; + let bstride = self.b.stride(physical_k); - for pack in 0..full { - // SAFETY: A contains a full pack; every B column has PACK logical elements. + let packs = physical_k.value().get() / PACK; + for pack in 0..packs { + // SAFETY: A and every B row are zero-padded to a multiple of PACK. unsafe { accumulate_pack::( self.arch, ap.add(astride * pack), bp.add(Elements::new(pack * PACK)), bstride, - PACK, &mut acc, ); } } - if !k.is_multiple_of(PACK) { - // SAFETY: A's last pack is padded. B has k % PACK remaining logical elements; - // splat zero-fills the rest, without reading a padded B extent. - unsafe { - accumulate_pack::( - self.arch, - ap.add(astride * full), - bp.add(Elements::new(full * PACK)), - bstride, - k % PACK, - &mut acc, - ); - } - } - - // SAFETY: Float has exactly MR lanes, and C is an MR-element array. - let mut best = unsafe { W::Float::load_simd(self.arch, self.c.as_ptr()) }; - for (acc, b) in acc.into_iter().zip(self.b_meta) { - let dot = acc.simd_cast(); - let a_b = W::Float::splat(self.arch, b.a); - let b_b = W::Float::splat(self.arch, b.b); - let n_b = W::Float::splat(self.arch, b.n); - let dim = W::Float::splat(self.arch, k as f32); - let a = self.a_meta; - let ip = ((a.a * a_b) * dot + a.n * b_b) + n_b * a.b + (a.b * b_b) * dim; - best = best.max_simd_standard(ip); - } - // SAFETY: C has exactly Float::LANES writable elements. - unsafe { best.store_simd(self.c.as_mut_ptr()) }; + self.arch.max_into(acc, self.a_meta, self.b_meta, k, self.c); } } - //----------// // Backends // //----------// -/// SIMD types for MR A rows and PACK contraction elements. NR counts accumulators, -/// not lanes, so it does not change these types. +/// Backend operations for MR A rows and PACK contraction elements. Register +/// representation is private to each backend; NR only controls the number of accumulators. /// /// All implementations accept A codes in 0..=255. B's range is limited by /// both its element type and any intermediate arithmetic in the dot product. @@ -574,54 +526,94 @@ pub(crate) trait ExtraWide: Architecture { type AElem: Copy + Default; type BElem: Copy + Default; - type Wide: SIMDVector; - type Splat: SIMDVector; - type Acc: SIMDVector - + SIMDDotProduct - + SIMDCast; - type Float: SIMDVector - + SIMDMinMax - + Add - + Mul; + type Wide: Copy; + type Splat: Copy; + type Acc: Copy; fn convert_a(code: u8) -> Self::AElem; /// `code` fits in MAX_B_BITS unsigned bits. fn convert_b(code: u8) -> Self::BElem; - /// Accumulate codes produced by `convert_a` and `convert_b`. - #[inline(always)] - fn dot(self, acc: Self::Acc, a: Self::Wide, b: Self::Splat) -> Self::Acc { - acc.dot_simd(a, b) - } + /// # Safety + /// + /// `a` must contain at least MR * PACK initialized, readable elements. + unsafe fn load_a(self, a: Slice<'_, Self::AElem>) -> Self::Wide; + + fn zero_acc(self) -> Self::Acc; + + /// Accumulate MR dot products of PACK elements without mixing rows. + fn dot(self, acc: Self::Acc, a: Self::Wide, b: Self::Splat) -> Self::Acc; + + /// Compensate in B order using logical `k`, then update scores without lowering existing + /// maxima. NaN candidates are ignored. + fn max_into( + self, + acc: [Self::Acc; NR], + query: Coefficients<&[f32; MR]>, + docs: &[MinMaxCompensation; NR], + k: usize, + scores: &mut [f32; MR], + ); /// # Safety /// - /// `valid <= PACK` and the first `valid` elements of B must be readable. - unsafe fn splat(self, b: Slice<'_, Self::BElem>, valid: usize) -> Self::Splat; + /// B must contain PACK readable elements, including any zero padding. + unsafe fn splat(self, b: Slice<'_, Self::BElem>) -> Self::Splat; +} + +#[inline(always)] +fn compensated_max_into( + arch: F::Arch, + acc: impl IntoIterator, + query: Coefficients<&[f32; MR]>, + docs: &[MinMaxCompensation; NR], + k: usize, + scores: &mut [f32; MR], +) where + I: SIMDCast, + F: SIMDVector> + + SIMDMinMax + + Add + + Mul, + Const: SupportedLaneCount, +{ + let a = Coefficients { + a: F::from_array(arch, *query.a), + b: F::from_array(arch, *query.b), + n: F::from_array(arch, *query.n), + }; + let mut best = F::from_array(arch, *scores); + for (acc, b) in acc.into_iter().zip(docs) { + let dot = acc.simd_cast(); + let a_b = F::splat(arch, b.a); + let b_b = F::splat(arch, b.b); + let n_b = F::splat(arch, b.n); + let dim = F::splat(arch, k as f32); + let ip = ((a.a * a_b) * dot + a.n * b_b) + n_b * a.b + (a.b * b_b) * dim; + best = best.max_simd_standard(ip); + } + *scores = best.to_array(); } /// # Safety /// -/// `valid <= PACK` and the first `valid` elements of B must be readable. +/// `b` must contain at least PACK initialized, readable elements. #[inline(always)] -unsafe fn group(b: Slice<'_, T>, valid: usize) -> [T; PACK] { - core::array::from_fn(|i| { - if i < valid { - // SAFETY: i < valid, so the element is readable. - unsafe { *b.add(Elements::new(i)).as_unit().as_ref() } - } else { - T::default() - } - }) +unsafe fn load_pack(b: Slice<'_, T>) -> [T; PACK] { + bounds::check_ge!(b.len(), PACK); + // SAFETY: The caller provides a full pack; no array alignment is required. + unsafe { b.as_ptr().cast::<[T; PACK]>().read_unaligned() } } +/// Pack two i16 bit patterns into the low and high halves of a u32. +/// /// # Safety /// -/// `valid <= 2` and the first `valid` elements of B must be readable. +/// B must contain two readable elements. #[inline(always)] -unsafe fn pair(b: Slice<'_, i16>, valid: usize) -> u32 { - // SAFETY: Inherited from the caller; group zero-fills a partial pair. - let [lo, hi] = unsafe { group::<_, 2>(b, valid) }; +unsafe fn pack_i16_pair(b: Slice<'_, i16>) -> u32 { + // SAFETY: The caller provides two readable elements, including any zero padding. + let [lo, hi] = unsafe { load_pack::<_, 2>(b) }; u32::from(lo as u16) | (u32::from(hi as u16) << 16) } @@ -633,7 +625,34 @@ impl ExtraWide<16, 4> for Scalar { type Wide = u8x64; type Splat = i8x64; type Acc = i32x16; - type Float = f32x16; + + #[inline(always)] + unsafe fn load_a(self, a: Slice<'_, Self::AElem>) -> Self::Wide { + // SAFETY: The caller provides 64 initialized elements. + Self::Wide::from_array(self, unsafe { load_pack::<_, 64>(a) }) + } + + #[inline(always)] + fn zero_acc(self) -> Self::Acc { + SIMDVector::default(self) + } + + #[inline(always)] + fn max_into( + self, + acc: [Self::Acc; NR], + query: Coefficients<&[f32; 16]>, + docs: &[MinMaxCompensation; NR], + k: usize, + scores: &mut [f32; 16], + ) { + compensated_max_into(self, acc, query, docs, k, scores); + } + + #[inline(always)] + fn dot(self, acc: Self::Acc, a: Self::Wide, b: Self::Splat) -> Self::Acc { + acc.dot_simd(a, b) + } fn convert_a(code: u8) -> u8 { code @@ -644,9 +663,9 @@ impl ExtraWide<16, 4> for Scalar { } #[inline(always)] - unsafe fn splat(self, b: Slice<'_, i8>, valid: usize) -> Self::Splat { - // SAFETY: Inherited from the caller. - let values = unsafe { group::<_, 4>(b, valid) }; + unsafe fn splat(self, b: Slice<'_, i8>) -> Self::Splat { + // SAFETY: The caller provides four readable elements, including any zero padding. + let values = unsafe { load_pack::<_, 4>(b) }; let quad = u32::from_le_bytes(values.map(|x| x as u8)); u32x16::::splat(self, quad).reinterpret_simd() } @@ -660,7 +679,34 @@ impl ExtraWide<16, 2> for Scalar { type Wide = i16x32; type Splat = Self::Wide; type Acc = i32x16; - type Float = f32x16; + + #[inline(always)] + unsafe fn load_a(self, a: Slice<'_, Self::AElem>) -> Self::Wide { + // SAFETY: The caller provides 32 initialized elements. + Self::Wide::from_array(self, unsafe { load_pack::<_, 32>(a) }) + } + + #[inline(always)] + fn zero_acc(self) -> Self::Acc { + SIMDVector::default(self) + } + + #[inline(always)] + fn max_into( + self, + acc: [Self::Acc; NR], + query: Coefficients<&[f32; 16]>, + docs: &[MinMaxCompensation; NR], + k: usize, + scores: &mut [f32; 16], + ) { + compensated_max_into(self, acc, query, docs, k, scores); + } + + #[inline(always)] + fn dot(self, acc: Self::Acc, a: Self::Wide, b: Self::Splat) -> Self::Acc { + acc.dot_simd(a, b) + } fn convert_a(code: u8) -> i16 { i16::from(code) @@ -671,9 +717,58 @@ impl ExtraWide<16, 2> for Scalar { } #[inline(always)] - unsafe fn splat(self, b: Slice<'_, i16>, valid: usize) -> Self::Splat { + unsafe fn splat(self, b: Slice<'_, i16>) -> Self::Splat { // SAFETY: Inherited from the caller. - u32x16::::splat(self, unsafe { pair(b, valid) }).reinterpret_simd() + u32x16::::splat(self, unsafe { pack_i16_pair(b) }).reinterpret_simd() + } +} + +impl ExtraWide<8, 2> for Scalar { + const MAX_B_BITS: usize = 8; + type AElem = i16; + type BElem = i16; + type Wide = i16x16; + type Splat = Self::Wide; + type Acc = i32x8; + + #[inline(always)] + unsafe fn load_a(self, a: Slice<'_, Self::AElem>) -> Self::Wide { + // SAFETY: The caller provides 16 initialized elements. + Self::Wide::from_array(self, unsafe { load_pack::<_, 16>(a) }) + } + + #[inline(always)] + fn zero_acc(self) -> Self::Acc { + SIMDVector::default(self) + } + + #[inline(always)] + fn max_into( + self, + acc: [Self::Acc; NR], + query: Coefficients<&[f32; 8]>, + docs: &[MinMaxCompensation; NR], + k: usize, + scores: &mut [f32; 8], + ) { + compensated_max_into(self, acc, query, docs, k, scores); + } + + #[inline(always)] + fn dot(self, acc: Self::Acc, a: Self::Wide, b: Self::Splat) -> Self::Acc { + acc.dot_simd(a, b) + } + + fn convert_a(code: u8) -> i16 { + i16::from(code) + } + fn convert_b(code: u8) -> i16 { + i16::from(code) + } + #[inline(always)] + unsafe fn splat(self, b: Slice<'_, i16>) -> Self::Splat { + // SAFETY: The caller provides two readable elements. + u32x8::::splat(self, unsafe { pack_i16_pair(b) }).reinterpret_simd() } } @@ -683,6 +778,146 @@ mod x86_64 { use diskann_wide::SplitJoin; use diskann_wide::arch::x86_64::{V3, V4}; + #[inline(always)] + fn compensated_max_into_v3( + arch: V3, + acc: [i32x16; NR], + query: Coefficients<&[f32; 16]>, + docs: &[MinMaxCompensation; NR], + k: usize, + scores: &mut [f32; 16], + ) { + let a = query.a.as_chunks::<8>().0; + let b = query.b.as_chunks::<8>().0; + let n = query.n.as_chunks::<8>().0; + let scores = scores.as_chunks_mut::<8>().0; + // Finish one native-width row group before loading the next group's coefficients. + compensated_max_into( + arch, + acc.iter().map(|v| v.split().lo), + Coefficients { + a: &a[0], + b: &b[0], + n: &n[0], + }, + docs, + k, + &mut scores[0], + ); + compensated_max_into( + arch, + acc.iter().map(|v| v.split().hi), + Coefficients { + a: &a[1], + b: &b[1], + n: &n[1], + }, + docs, + k, + &mut scores[1], + ); + } + + impl ExtraWide<8, 4> for V3 { + const MAX_B_BITS: usize = 6; + type AElem = u8; + type BElem = i8; + type Wide = u8x32; + type Splat = i8x32; + type Acc = i32x8; + + #[inline(always)] + unsafe fn load_a(self, a: Slice<'_, Self::AElem>) -> Self::Wide { + // SAFETY: The caller provides 32 initialized elements. + Self::Wide::from_array(self, unsafe { load_pack::<_, 32>(a) }) + } + + #[inline(always)] + fn zero_acc(self) -> Self::Acc { + SIMDVector::default(self) + } + + #[inline(always)] + fn max_into( + self, + acc: [Self::Acc; NR], + query: Coefficients<&[f32; 8]>, + docs: &[MinMaxCompensation; NR], + k: usize, + scores: &mut [f32; 8], + ) { + compensated_max_into(self, acc, query, docs, k, scores); + } + + fn convert_a(code: u8) -> u8 { + code + } + fn convert_b(code: u8) -> i8 { + code as i8 + } + #[inline(always)] + fn dot(self, acc: Self::Acc, a: Self::Wide, b: Self::Splat) -> Self::Acc { + // The largest pair is 2 * 255 * 63 = 32130, below i16 saturation. + acc.dot_simd(self.mm256_maddubs_epi16(a, b), i16x16::::splat(self, 1)) + } + #[inline(always)] + unsafe fn splat(self, b: Slice<'_, i8>) -> Self::Splat { + // SAFETY: The caller provides four readable elements. + let values = unsafe { load_pack::<_, 4>(b) }; + let quad = u32::from_le_bytes(values.map(|v| v as u8)); + u32x8::::splat(self, quad).reinterpret_simd() + } + } + + impl ExtraWide<8, 2> for V3 { + const MAX_B_BITS: usize = 8; + type AElem = i16; + type BElem = i16; + type Wide = i16x16; + type Splat = Self::Wide; + type Acc = i32x8; + + #[inline(always)] + unsafe fn load_a(self, a: Slice<'_, Self::AElem>) -> Self::Wide { + // SAFETY: The caller provides 16 initialized elements. + Self::Wide::from_array(self, unsafe { load_pack::<_, 16>(a) }) + } + + #[inline(always)] + fn zero_acc(self) -> Self::Acc { + SIMDVector::default(self) + } + + #[inline(always)] + fn max_into( + self, + acc: [Self::Acc; NR], + query: Coefficients<&[f32; 8]>, + docs: &[MinMaxCompensation; NR], + k: usize, + scores: &mut [f32; 8], + ) { + compensated_max_into(self, acc, query, docs, k, scores); + } + + #[inline(always)] + fn dot(self, acc: Self::Acc, a: Self::Wide, b: Self::Splat) -> Self::Acc { + acc.dot_simd(a, b) + } + + fn convert_a(code: u8) -> i16 { + i16::from(code) + } + fn convert_b(code: u8) -> i16 { + i16::from(code) + } + #[inline(always)] + unsafe fn splat(self, b: Slice<'_, i16>) -> Self::Splat { + // SAFETY: The caller provides two readable elements. + u32x8::::splat(self, unsafe { pack_i16_pair(b) }).reinterpret_simd() + } + } + impl ExtraWide<16, 4> for V3 { const MAX_B_BITS: usize = 6; @@ -691,7 +926,29 @@ mod x86_64 { type Wide = u8x64; type Splat = i8x64; type Acc = i32x16; - type Float = f32x16; + + #[inline(always)] + unsafe fn load_a(self, a: Slice<'_, Self::AElem>) -> Self::Wide { + // SAFETY: The caller provides 64 initialized elements. + Self::Wide::from_array(self, unsafe { load_pack::<_, 64>(a) }) + } + + #[inline(always)] + fn zero_acc(self) -> Self::Acc { + SIMDVector::default(self) + } + + #[inline(always)] + fn max_into( + self, + acc: [Self::Acc; NR], + query: Coefficients<&[f32; 16]>, + docs: &[MinMaxCompensation; NR], + k: usize, + scores: &mut [f32; 16], + ) { + compensated_max_into_v3(self, acc, query, docs, k, scores); + } fn convert_a(code: u8) -> u8 { code @@ -715,9 +972,9 @@ mod x86_64 { } #[inline(always)] - unsafe fn splat(self, b: Slice<'_, i8>, valid: usize) -> Self::Splat { - // SAFETY: Inherited from the caller. - let values = unsafe { group::<_, 4>(b, valid) }; + unsafe fn splat(self, b: Slice<'_, i8>) -> Self::Splat { + // SAFETY: The caller provides four readable elements, including any zero padding. + let values = unsafe { load_pack::<_, 4>(b) }; let quad = u32::from_le_bytes(values.map(|x| x as u8)); let half: i8x32 = u32x8::::splat(self, quad).reinterpret_simd(); i8x64::::new(half, half) @@ -732,7 +989,34 @@ mod x86_64 { type Wide = i16x32; type Splat = Self::Wide; type Acc = i32x16; - type Float = f32x16; + + #[inline(always)] + unsafe fn load_a(self, a: Slice<'_, Self::AElem>) -> Self::Wide { + // SAFETY: The caller provides 32 initialized elements. + Self::Wide::from_array(self, unsafe { load_pack::<_, 32>(a) }) + } + + #[inline(always)] + fn zero_acc(self) -> Self::Acc { + SIMDVector::default(self) + } + + #[inline(always)] + fn max_into( + self, + acc: [Self::Acc; NR], + query: Coefficients<&[f32; 16]>, + docs: &[MinMaxCompensation; NR], + k: usize, + scores: &mut [f32; 16], + ) { + compensated_max_into_v3(self, acc, query, docs, k, scores); + } + + #[inline(always)] + fn dot(self, acc: Self::Acc, a: Self::Wide, b: Self::Splat) -> Self::Acc { + acc.dot_simd(a, b) + } fn convert_a(code: u8) -> i16 { i16::from(code) @@ -743,9 +1027,9 @@ mod x86_64 { } #[inline(always)] - unsafe fn splat(self, b: Slice<'_, i16>, valid: usize) -> Self::Splat { + unsafe fn splat(self, b: Slice<'_, i16>) -> Self::Splat { // SAFETY: Inherited from the caller. - let pair = unsafe { pair(b, valid) }; + let pair = unsafe { pack_i16_pair(b) }; let half: i16x16 = u32x8::::splat(self, pair).reinterpret_simd(); i16x32::::new(half, half) } @@ -759,7 +1043,34 @@ mod x86_64 { type Wide = u8x64; type Splat = i8x64; type Acc = i32x16; - type Float = f32x16; + + #[inline(always)] + unsafe fn load_a(self, a: Slice<'_, Self::AElem>) -> Self::Wide { + // SAFETY: The caller provides 64 initialized elements. + Self::Wide::from_array(self, unsafe { load_pack::<_, 64>(a) }) + } + + #[inline(always)] + fn zero_acc(self) -> Self::Acc { + SIMDVector::default(self) + } + + #[inline(always)] + fn max_into( + self, + acc: [Self::Acc; NR], + query: Coefficients<&[f32; 16]>, + docs: &[MinMaxCompensation; NR], + k: usize, + scores: &mut [f32; 16], + ) { + compensated_max_into(self, acc, query, docs, k, scores); + } + + #[inline(always)] + fn dot(self, acc: Self::Acc, a: Self::Wide, b: Self::Splat) -> Self::Acc { + acc.dot_simd(a, b) + } fn convert_a(code: u8) -> u8 { code @@ -770,9 +1081,9 @@ mod x86_64 { } #[inline(always)] - unsafe fn splat(self, b: Slice<'_, i8>, valid: usize) -> Self::Splat { - // SAFETY: Inherited from the caller. - let values = unsafe { group::<_, 4>(b, valid) }; + unsafe fn splat(self, b: Slice<'_, i8>) -> Self::Splat { + // SAFETY: The caller provides four readable elements, including any zero padding. + let values = unsafe { load_pack::<_, 4>(b) }; let quad = u32::from_le_bytes(values.map(|x| x as u8)); u32x16::::splat(self, quad).reinterpret_simd() } @@ -786,7 +1097,34 @@ mod x86_64 { type Wide = i16x32; type Splat = Self::Wide; type Acc = i32x16; - type Float = f32x16; + + #[inline(always)] + unsafe fn load_a(self, a: Slice<'_, Self::AElem>) -> Self::Wide { + // SAFETY: The caller provides 32 initialized elements. + Self::Wide::from_array(self, unsafe { load_pack::<_, 32>(a) }) + } + + #[inline(always)] + fn zero_acc(self) -> Self::Acc { + SIMDVector::default(self) + } + + #[inline(always)] + fn max_into( + self, + acc: [Self::Acc; NR], + query: Coefficients<&[f32; 16]>, + docs: &[MinMaxCompensation; NR], + k: usize, + scores: &mut [f32; 16], + ) { + compensated_max_into(self, acc, query, docs, k, scores); + } + + #[inline(always)] + fn dot(self, acc: Self::Acc, a: Self::Wide, b: Self::Splat) -> Self::Acc { + acc.dot_simd(a, b) + } fn convert_a(code: u8) -> i16 { i16::from(code) @@ -797,9 +1135,9 @@ mod x86_64 { } #[inline(always)] - unsafe fn splat(self, b: Slice<'_, i16>, valid: usize) -> Self::Splat { + unsafe fn splat(self, b: Slice<'_, i16>) -> Self::Splat { // SAFETY: Inherited from the caller. - u32x16::::splat(self, unsafe { pair(b, valid) }).reinterpret_simd() + u32x16::::splat(self, unsafe { pack_i16_pair(b) }).reinterpret_simd() } } } @@ -809,6 +1147,57 @@ mod aarch64 { use super::*; use diskann_wide::arch::aarch64::Neon; + impl ExtraWide<8, 4> for Neon { + const MAX_B_BITS: usize = 8; + type AElem = u8; + type BElem = u8; + type Wide = u8x32; + type Splat = Self::Wide; + type Acc = u32x8; + + #[inline(always)] + unsafe fn load_a(self, a: Slice<'_, Self::AElem>) -> Self::Wide { + // SAFETY: The caller provides 32 initialized elements. + Self::Wide::from_array(self, unsafe { load_pack::<_, 32>(a) }) + } + + #[inline(always)] + fn zero_acc(self) -> Self::Acc { + SIMDVector::default(self) + } + + #[inline(always)] + fn max_into( + self, + acc: [Self::Acc; NR], + query: Coefficients<&[f32; 8]>, + docs: &[MinMaxCompensation; NR], + k: usize, + scores: &mut [f32; 8], + ) { + compensated_max_into(self, acc, query, docs, k, scores); + } + + #[inline(always)] + fn dot(self, acc: Self::Acc, a: Self::Wide, b: Self::Splat) -> Self::Acc { + acc.dot_simd(a, b) + } + + fn convert_a(code: u8) -> u8 { + code + } + fn convert_b(code: u8) -> u8 { + code + } + #[inline(always)] + unsafe fn splat(self, b: Slice<'_, u8>) -> Self::Splat { + // SAFETY: The caller provides four readable elements. + let values = unsafe { load_pack::<_, 4>(b) }; + let quarter = u32x4::::splat(self, u32::from_le_bytes(values)).reinterpret_simd(); + u8x32::::new(quarter, quarter) + } + } + impl ExtraWide<16, 4> for Neon { const MAX_B_BITS: usize = 8; @@ -817,7 +1206,34 @@ mod aarch64 { type Wide = u8x64; type Splat = Self::Wide; type Acc = u32x16; - type Float = f32x16; + + #[inline(always)] + unsafe fn load_a(self, a: Slice<'_, Self::AElem>) -> Self::Wide { + // SAFETY: The caller provides 64 initialized elements. + Self::Wide::from_array(self, unsafe { load_pack::<_, 64>(a) }) + } + + #[inline(always)] + fn zero_acc(self) -> Self::Acc { + SIMDVector::default(self) + } + + #[inline(always)] + fn max_into( + self, + acc: [Self::Acc; NR], + query: Coefficients<&[f32; 16]>, + docs: &[MinMaxCompensation; NR], + k: usize, + scores: &mut [f32; 16], + ) { + compensated_max_into(self, acc, query, docs, k, scores); + } + + #[inline(always)] + fn dot(self, acc: Self::Acc, a: Self::Wide, b: Self::Splat) -> Self::Acc { + acc.dot_simd(a, b) + } fn convert_a(code: u8) -> u8 { code @@ -828,9 +1244,10 @@ mod aarch64 { } #[inline(always)] - unsafe fn splat(self, b: Slice<'_, u8>, valid: usize) -> Self::Splat { - // SAFETY: Inherited from the caller. - let quad = u32::from_le_bytes(unsafe { group::<_, 4>(b, valid) }); + unsafe fn splat(self, b: Slice<'_, u8>) -> Self::Splat { + // SAFETY: The caller provides four readable elements, including any zero padding. + let values = unsafe { load_pack::<_, 4>(b) }; + let quad = u32::from_le_bytes(values); let quarter = u32x4::::splat(self, quad).reinterpret_simd(); let half = u8x32::::new(quarter, quarter); u8x64::::new(half, half) @@ -853,7 +1270,17 @@ mod tests { multi_vector::{Defaulted, Mat}, }; - const MAX_K_8X4: usize = max_k::<8, 4>(); + #[test] + fn i16_pair_preserves_lane_order_and_bits() { + for (values, expected) in [ + ([1, 255], 0x00ff_0001), + ([255, 1], 0x0001_00ff), + ([-1, i16::MIN], 0x8000_ffff), + ] { + // SAFETY: values contains both initialized elements of the pair. + assert_eq!(unsafe { pack_i16_pair(Slice::new(&values)) }, expected); + } + } fn random_minmax_matrix( rows: usize, @@ -968,7 +1395,7 @@ mod tests { { let mut rng = StdRng::seed_from_u64(0x63de91a0ba5c8742); for shape in packed_x_unpacked_test_dims(MR, NR) { - for k in [1, 2, 3, 4, 5, 7, 8, 15, 16, 17, 31, 32, 33, 127, 129] { + 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::<_, 8, MR, PACK>::new(arch, a.as_view()).unwrap(); @@ -980,13 +1407,6 @@ mod tests { b.as_view(), &mut guarded[1..=a.num_vectors()], ); - driver::Drive::drive(&mut driver); - assert_scores_close( - driver.c, - &expected, - format_args!("default tiling: {shape:?}, K={k}, MR={MR}, NR={NR}, PACK={PACK}"), - ); - driver.b_rows_per_tile = NonZeroUsize::new(shape.b_cols_per_tile).unwrap(); driver.c.fill(f32::INFINITY); driver::Drive::drive(&mut driver); @@ -997,20 +1417,11 @@ mod tests { "forced B tiling: {shape:?}, K={k}, MR={MR}, NR={NR}, PACK={PACK}" ), ); - // Reset stale scores once per invocation, never once per B tile. - driver.c.fill(f32::INFINITY); - driver::Drive::drive(&mut driver); - assert_scores_close( - driver.c, - &expected, - format_args!("repeated drive: {shape:?}, K={k}, MR={MR}, NR={NR}, PACK={PACK}"), - ); assert_eq!(guarded[0], 654321.0); assert_eq!(guarded[a.num_vectors() + 1], 654321.0); } } assert_compensated_reduction::(arch); - assert_accumulator_limit_is_exact::(arch); let a = random_minmax_matrix::<8>(MR + 1, 5, &mut rng); let b = random_minmax_matrix::(0, 5, &mut rng); @@ -1064,7 +1475,8 @@ mod tests { Unsigned: Representation, { let a = constant_minmax_matrix::<8>(MR + 1, 3, 1, 1.0, 0.0); - let mut b = constant_minmax_matrix::(NR + 1, 3, 15, 1.0, 0.0); + let mut b = constant_minmax_matrix::(NR + 1, 3, (1 << B_BITS) - 1, 1.0, 0.0); + let winner_scale = (1 << B_BITS) as f32; { let mut row = b.get_row_mut(0).unwrap(); row.vector_mut().set(0, 1).unwrap(); @@ -1072,9 +1484,9 @@ mod tests { row.vector_mut().set(2, 1).unwrap(); row.set_meta(MinMaxCompensation { dim: 3, - a: 100.0, + a: winner_scale, b: 0.0, - n: 300.0, + n: 3.0 * winner_scale, norm_squared: 0.0, }); } @@ -1082,10 +1494,10 @@ mod tests { let mut c = vec![0.0; MR + 1]; let mut driver = checked_driver::<_, 8, B_BITS, MR, NR, PACK>(arch, &packed, b.as_view(), &mut c); - // The compensated winner is in the first tile, not the last. - driver.b_rows_per_tile = NonZeroUsize::new(1).unwrap(); + // The first tile contains competing raw dots; its compensated winner must survive the tail. + driver.b_rows_per_tile = NonZeroUsize::new(NR).unwrap(); driver::Drive::drive(&mut driver); - assert_eq!(c, vec![300.0; MR + 1]); + assert_eq!(c, vec![3.0 * winner_scale; MR + 1]); let a = constant_minmax_matrix::<8>(MR + 1, 3, 1, 1.0, -2.0); let b = constant_minmax_matrix::((NR - 1).max(1), 3, 1, 1.0, 0.0); @@ -1100,34 +1512,22 @@ mod tests { assert_eq!(c, vec![-3.0; MR + 1]); } - fn assert_accumulator_limit_is_exact< - W, - const B_BITS: usize, - const MR: usize, - const NR: usize, - const PACK: usize, - >( - arch: W, - ) where - W: ExtraWide + LoadStore, - Unsigned: Representation, - { - if cfg!(miri) { - return; - } - let k = max_k::<8, B_BITS>(); - let max_b = (1 << B_BITS) - 1; + #[cfg(not(miri))] + #[test] + fn accumulator_is_exact_at_dimension_limit() { + let arch = Scalar::new(); + let k = max_k::<8, 8>(); let a = constant_minmax_matrix::<8>(1, k, 255, 1.0, 0.0); - let b = constant_minmax_matrix::(1, k, max_b as i64, 1.0, 0.0); - let packed = PackedA::<_, 8, MR, PACK>::new(arch, a.as_view()).unwrap(); + let b = constant_minmax_matrix::<8>(1, k, 255, 1.0, 0.0); + let packed = PackedA::<_, 8, 8, 2>::new(arch, a.as_view()).unwrap(); let mut c = [0.0]; - driver::Drive::drive(&mut checked_driver::<_, 8, B_BITS, MR, NR, PACK>( + driver::Drive::drive(&mut checked_driver::<_, 8, 8, 8, 2, 2>( arch, &packed, b.as_view(), &mut c, )); - assert_eq!(c[0], (k * 255 * max_b) as f32); + assert_eq!(c[0], (k * 255 * 255) as f32); } #[test] @@ -1165,13 +1565,35 @@ mod tests { } } + #[test] + fn smaller_panels_preserve_tiling_and_nan_reduction() { + fn check(arch: W) + where + W: ExtraWide<8, PACK> + LoadStore, + Unsigned: Representation, + { + assert_driver_contract::(arch); + assert_microkernel_nan_reduction::(arch); + } + check::<_, 8, 2, 2>(Scalar::new()); + #[cfg(target_arch = "x86_64")] + if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { + check::<_, 4, 4, 4>(arch); + check::<_, 8, 4, 2>(arch); + } + #[cfg(target_arch = "aarch64")] + if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { + check::<_, 8, 4, 4>(arch); + } + } + fn assert_a_packing(arch: W) where W: ExtraWide, W::AElem: PartialEq + std::fmt::Debug, Unsigned: Representation, { - for (m, k) in [(0, 0), (0, 5), (3, 0), (0, MAX_K_8X4), (0, MAX_K_8X4 + 1)] { + 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::<_, A_BITS, MR, PACK>::new(arch, a.as_view()).unwrap(); assert_eq!(packed.codes.nrows(), m); @@ -1207,28 +1629,23 @@ mod tests { } #[test] - fn packing_preserves_codes_coefficients_and_zero_padding() { - assert_a_packing::<_, 8, 16, 4>(Scalar::new()); - #[cfg(target_arch = "x86_64")] - if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { - assert_a_packing::<_, 8, 16, 4>(arch); - } - #[cfg(target_arch = "x86_64")] - if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked() { - assert_a_packing::<_, 8, 16, 4>(arch); - } - #[cfg(target_arch = "aarch64")] - if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { - assert_a_packing::<_, 8, 16, 4>(arch); - } + fn a_packing_preserves_codes_metadata_and_padding() { + macro_rules! check { + ($($bits:literal),+) => {$( + assert_a_packing::<_, $bits, 16, 4>(Scalar::new()); + assert_a_packing::<_, $bits, 16, 2>(Scalar::new()); + )+}; + } + check!(1, 2, 3, 4, 5, 6, 7, 8); + assert_a_packing::<_, 3, 8, 2>(Scalar::new()); } #[test] fn cache_footprints_match_prepared_layouts() { - for (k, a_bytes, b_bytes) in [(256, 4288, 276), (257, 4352, 277)] { + for (k, a_bytes, b_bytes) in [(256, 4288, 276), (257, 4352, 280)] { let dim = DimK::new(NonZeroUsize::new(k).unwrap()); - assert_eq!(BScratch::::row_bytes(dim).value(), b_bytes); - assert_eq!(BScratch::::row_bytes(dim).value(), b_bytes); + assert_eq!(BScratch::::row_bytes::<4>(dim).value(), b_bytes); + assert_eq!(BScratch::::row_bytes::<4>(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(Scalar::new(), a.as_view()).unwrap(); @@ -1246,9 +1663,9 @@ mod tests { #[test] fn widened_cache_footprints_match_prepared_layouts() { - for (k, a_bytes, b_bytes) in [(256, 8384, 532), (257, 8448, 534)] { + for (k, a_bytes, b_bytes) in [(256, 8384, 532), (257, 8448, 536)] { let dim = DimK::new(NonZeroUsize::new(k).unwrap()); - assert_eq!(BScratch::::row_bytes(dim).value(), b_bytes); + assert_eq!(BScratch::::row_bytes::<2>(dim).value(), b_bytes); for rows in [1, 16, 17] { let a = constant_minmax_matrix::<8>(rows, k, 255, 1.0, 0.0); let packed = PackedA::::new(Scalar::new(), a.as_view()).unwrap(); @@ -1262,236 +1679,91 @@ mod tests { } } - #[cfg(not(miri))] - #[test] - #[should_panic(expected = "exceeds the MinMax8 x MinMax4 limit")] - fn driver_rejects_k_above_accumulator_limit() { - let k = MAX_K_8X4 + 1; - let a = constant_minmax_matrix::<8>(1, k, 255, 1.0, 0.0); - let b = constant_minmax_matrix::<4>(1, k, 15, 1.0, 0.0); - let packed = PackedA::::new(Scalar::new(), a.as_view()).unwrap(); - let mut c = [123.0]; - let _ = checked_driver::<_, 8, 4, 16, 2, 4>(Scalar::new(), &packed, b.as_view(), &mut c); - } - #[test] - fn b_preparation_masks_padding_and_reuses_storage() { - let k = 5; - let stride = crate::minmax::Data::<4>::canonical_bytes(k); - let mut bytes = vec![0xff; stride * 3]; - for (row_index, raw) in bytes.chunks_exact_mut(stride).enumerate() { - let mut row = crate::minmax::DataMutRef::<4>::from_canonical_front_mut(raw, k).unwrap(); - let a = (row_index + 1) as f32; - row.set_meta(MinMaxCompensation { - dim: k as u32, - a, - b: -(row_index as f32), - n: a * (10 + 5 * row_index) as f32, - ..Default::default() - }); - for i in 0..k { - row.vector_mut().set(i, (i + row_index) as i64).unwrap(); - } - } - let b = MatRef::new(MinMaxMeta::<4>::new(3, k), &bytes).unwrap(); - let mut scratch = BScratch::::default(); - let dim = DimK::new(NonZeroUsize::new(k).unwrap()); - let tile = scratch.prepare::(b, 0, NonZeroUsize::new(3).unwrap(), dim); - assert_eq!( - // SAFETY: dim is the dimension passed to prepare. - unsafe { tile.codes.as_std_slice(dim) }, - &[0, 1, 2, 3, 4, 1, 2, 3, 4, 5, 2, 3, 4, 5, 6] - ); - let ptr = scratch.codes.as_ptr(); - let tile = scratch.prepare::(b, 2, NonZeroUsize::new(1).unwrap(), dim); - assert_eq!(tile.meta, &[b.get_row(2).unwrap().meta()]); - assert_eq!( - // SAFETY: dim is the dimension passed to prepare. - unsafe { tile.codes.as_std_slice(dim) }, - &[2, 3, 4, 5, 6] - ); - assert_eq!(scratch.codes.as_ptr(), ptr); + fn b_preparation_preserves_metadata_and_zero_padding() { + assert_b_preparation::(-1, |code| code as i8); + assert_b_preparation::(-1, i16::from); + assert_b_preparation::(u8::MAX, |code| code); + assert_b_preparation::(u16::MAX, |code| u16::from(code) + 1); } - fn assert_b_decoding_boundaries< - W, - const B_BITS: usize, - const MR: usize, - const NR: usize, - const PACK: usize, - >( - arch: W, - sentinel: W::BElem, + fn assert_b_preparation( + sentinel: T, + convert: impl Fn(u8) -> T, ) where - W: ExtraWide + LoadStore, - W::BElem: PartialEq + std::fmt::Debug, + T: Copy + Default + PartialEq + std::fmt::Debug, Unsigned: Representation, { - decode_b::(BitSlice::new(&[] as &[u8], 0).unwrap(), &mut []); - let mut rng = StdRng::seed_from_u64(0x8af450de13c769b2); - let mut scratch = BScratch::::default(); - let dimensions: &[usize] = if cfg!(miri) { - &[1, 2, 3, 4, 5, 15, 16, 17] - } else { - &[ - 1, 2, 15, 16, 17, 31, 32, 33, 63, 64, 65, 127, 128, 129, 250, 255, 256, 257, 511, - 512, 513, - ] - }; - for &dim in dimensions { + let mut scratch = BScratch::::default(); + for dim in [PACK, PACK + 1] { let k = DimK::new(NonZeroUsize::new(dim).unwrap()); - let a = random_minmax_matrix::<8>(MR + 1, dim, &mut rng); - let packed = PackedA::<_, 8, MR, PACK>::new(arch, a.as_view()).unwrap(); let stride = crate::minmax::Data::::canonical_bytes(dim); - for offset in [0, 1, 3, 15] { - let mut storage = vec![0xff; offset + 5 * stride]; - let bytes = &mut storage[offset..]; - for (i, raw) in bytes.chunks_exact_mut(stride).enumerate() { - let mut row = - crate::minmax::DataMutRef::::from_canonical_front_mut(raw, dim) - .unwrap(); - let mut sum = 0; - for j in 0..dim { - // Cover every code value, and every packed byte for MinMax4. - let code = if B_BITS == 4 { - let byte = (i * 37 + j / 2) & 255; - (byte >> (4 * (j % 2))) & 15 - } else { - (i * 37 + j) & ((1 << B_BITS) - 1) - } as i64; - row.vector_mut().set(j, code).unwrap(); - sum += code; - } - let scale = (i + 1) as f32 * 0.25; - row.set_meta(MinMaxCompensation { - dim: dim as u32, - a: scale, - b: i as f32 - 3.0, - n: scale * sum as f32, - norm_squared: (i * i + 1) as f32, - }); + // Unaligned storage and dirty high bits exercise the packed input boundary. + let mut storage = vec![0xff; 1 + 3 * stride]; + let bytes = &mut storage[1..]; + for (i, raw) in bytes.chunks_exact_mut(stride).enumerate() { + let mut row = + crate::minmax::DataMutRef::::from_canonical_front_mut(raw, dim) + .unwrap(); + let mut sum = 0; + let max_code = (1 << B_BITS) - 1; + for j in 0..dim { + let code = (max_code - ((i + j) & max_code)) as i64; + row.vector_mut().set(j, code).unwrap(); + sum += code; } - let b = MatRef::new(MinMaxMeta::::new(5, dim), bytes).unwrap(); - for row in b.rows() { - let codes = row.vector(); - if B_BITS == 4 && dim % 2 != 0 { - assert_eq!(codes.as_slice().last().unwrap() >> 4, 15); - } - let mut guarded = vec![sentinel; dim + 2]; - decode_b::(codes, &mut guarded[1..=dim]); + let scale = (i + 1) as f32 * 0.25; + row.set_meta(MinMaxCompensation { + dim: dim as u32, + a: scale, + b: i as f32 - 3.0, + n: scale * sum as f32, + norm_squared: (i * i + 1) as f32, + }); + } + let b = MatRef::new(MinMaxMeta::::new(3, dim), bytes).unwrap(); + for (start, count) in [(0, 3), (2, 1), (1, 2)] { + scratch.codes.fill(sentinel); + scratch.meta.fill(MinMaxCompensation::default()); + let mut conversions = 0; + let tile = scratch.prepare::( + b, + start, + NonZeroUsize::new(count).unwrap(), + k, + |code| { + conversions += 1; + convert(code) + }, + ); + assert_eq!(conversions, count * dim); + let padded = padded_k::(k); + let row_len = padded.value().get(); + // SAFETY: prepare constructs the view with this padded dimension. + let codes = unsafe { tile.codes.as_std_slice(padded) }; + assert_eq!(codes.len(), count * row_len); + assert_eq!(tile.meta.len(), count); + for i in 0..count { + let row = b.get_row(start + i).unwrap(); + assert_eq!(tile.meta[i], row.meta()); for j in 0..dim { assert_eq!( - guarded[j + 1], - W::convert_b(codes.get(j).unwrap() as u8), - "row decode: K={dim}, offset={offset}, column={j}, MR={MR}, PACK={PACK}" + codes[i * row_len + j], + convert(row.vector().get(j).unwrap() as u8), + "B tile: K={dim}, start={start}, count={count}, row={i}, column={j}" ); } - assert_eq!(guarded[0], sentinel); - assert_eq!(guarded[dim + 1], sentinel); - } - for (start, count) in [(0, 1), (1, 3), (4, 1), (0, 5)] { - scratch.codes.fill(sentinel); - scratch.meta.fill(MinMaxCompensation::default()); - let tile = scratch.prepare::( - b, - start, - NonZeroUsize::new(count).unwrap(), - k, - ); - // SAFETY: k is the dimension used to prepare this tile. - let codes = unsafe { tile.codes.as_std_slice(k) }; - assert_eq!(codes.len(), count * dim); - assert_eq!(tile.meta.len(), count); - for i in 0..count { - let row = b.get_row(start + i).unwrap(); - assert_eq!(tile.meta[i], row.meta()); - for j in 0..dim { - assert_eq!( - codes[i * dim + j], - W::convert_b(row.vector().get(j).unwrap() as u8), - "B tile: K={dim}, offset={offset}, start={start}, \ - count={count}, row={i}, column={j}, MR={MR}, PACK={PACK}" - ); - } - } - } - let expected = reference_max_inner_products::<8, B_BITS>(a.as_view(), b); - let mut guarded = vec![654321.0; MR + 3]; - for tile_rows in [3, 4, 5] { - let mut driver = checked_driver::( - arch, - &packed, - b, - &mut guarded[1..=MR + 1], - ); - driver.b_rows_per_tile = NonZeroUsize::new(tile_rows).unwrap(); - driver.c.fill(f32::INFINITY); - driver::Drive::drive(&mut driver); - assert_scores_close( - driver.c, - &expected, - format_args!( - "B decoding: K={dim}, offset={offset}, B tile={tile_rows}, \ - MR={MR}, NR={NR}, PACK={PACK}" - ), + assert!( + codes[i * row_len + dim..(i + 1) * row_len] + .iter() + .all(|&v| v == T::default()) ); - assert_eq!(guarded[0], 654321.0); - assert_eq!(guarded[MR + 2], 654321.0); } } } } - #[test] - fn scalar_b_decoding_preserves_rows_tails_and_scores() { - assert_b_decoding_boundaries::<_, 4, 16, 2, 4>(Scalar::new(), -1); - } - - #[cfg(target_arch = "x86_64")] - #[test] - fn avx2_b_decoding_preserves_rows_tails_and_scores() { - if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { - assert_b_decoding_boundaries::<_, 4, 16, 4, 4>(arch, -1); - } else { - eprintln!("AVX2 unavailable; B decoding regression not executed"); - } - } - - #[cfg(target_arch = "x86_64")] - #[test] - fn avx512_b_decoding_preserves_rows_tails_and_scores() { - if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked() { - assert_b_decoding_boundaries::<_, 4, 16, 8, 4>(arch, -1); - } else { - eprintln!("V4 unavailable; B decoding regression not executed"); - } - } - - #[cfg(target_arch = "aarch64")] - #[test] - fn neon_b_decoding_preserves_rows_tails_and_scores() { - if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { - assert_b_decoding_boundaries::<_, 4, 16, 4, 4>(arch, u8::MAX); - } else { - eprintln!("Neon/dotprod unavailable; B decoding regression not executed"); - } - } - - #[test] - #[should_panic] - fn b_decoding_rejects_a_short_output() { - let codes = BitSlice::new(&[0x21u8][..], 2).unwrap(); - decode_b::(codes, &mut [0]); - } - - #[test] - #[should_panic] - fn b_decoding_rejects_a_long_output() { - let codes = BitSlice::new(&[0xf1u8][..], 1).unwrap(); - decode_b::(codes, &mut [0; 2]); - } - - fn assert_packed_a_reuse< + fn assert_driver_target_shapes< W, const B_BITS: usize, const MR: usize, @@ -1503,98 +1775,31 @@ mod tests { W: ExtraWide + LoadStore, Unsigned: Representation, { - let k = 17; - let mut rng = StdRng::seed_from_u64(0xa619e2c5f387b04d); - let a = random_minmax_matrix::<8>(11, k, &mut rng); - let b = random_minmax_matrix::(9, k, &mut rng); - let expected = reference_max_inner_products::<8, B_BITS>(a.as_view(), b.as_view()); - let packed = PackedA::<_, 8, MR, PACK>::new(arch, a.as_view()).unwrap(); - drop(a); - let mut actual = vec![0.0; 11]; - for _ in 0..2 { - actual.fill(f32::INFINITY); + let mut rng = StdRng::seed_from_u64(0x8f61d9c34a257be0); + let mut shapes: Vec<_> = (1..NR).map(|n| (MR + 1, n, PACK + 1)).collect(); + if !cfg!(miri) { + shapes.extend([(16, 16, 256), (17, 17, 257)]); + } + for (m, n, k) in shapes { + let a = random_minmax_matrix::<8>(m, k, &mut rng); + let b = random_minmax_matrix::(n, k, &mut rng); + let expected = reference_max_inner_products::<8, B_BITS>(a.as_view(), b.as_view()); + let packed = PackedA::<_, 8, MR, PACK>::new(arch, a.as_view()).unwrap(); + let mut c = vec![f32::INFINITY; m]; driver::Drive::drive(&mut checked_driver::<_, 8, B_BITS, MR, NR, PACK>( arch, &packed, b.as_view(), - &mut actual, + &mut c, )); assert_scores_close( - &actual, + &c, &expected, - format_args!("packed-A reuse: MR={MR}, NR={NR}, PACK={PACK}"), + format_args!("M={m}, N={n}, K={k}, MR={MR}, NR={NR}, PACK={PACK}"), ); } } - #[test] - fn packed_a_outlives_source_and_supports_repeated_drives() { - assert_packed_a_reuse::<_, 4, 16, 2, 4>(Scalar::new()); - #[cfg(target_arch = "x86_64")] - if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { - assert_packed_a_reuse::<_, 4, 16, 4, 4>(arch); - } - #[cfg(target_arch = "x86_64")] - if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked() { - assert_packed_a_reuse::<_, 4, 16, 8, 4>(arch); - } - #[cfg(target_arch = "aarch64")] - if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { - assert_packed_a_reuse::<_, 4, 16, 4, 4>(arch); - } - } - - fn assert_driver_target_shapes< - W, - const B_BITS: usize, - const MR: usize, - const NR: usize, - const PACK: usize, - >( - arch: W, - ) where - W: ExtraWide + LoadStore, - Unsigned: Representation, - { - let mut rng = StdRng::seed_from_u64(0x8f61d9c34a257be0); - let dimensions: &[usize] = if cfg!(miri) { &[3] } else { &[3, 256, 257] }; - for m in [1, 16, 17] { - for n in (1..=NR + 1).chain([16, 17]) { - for &k in dimensions { - 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::<_, 8, MR, PACK>::new(arch, a.as_view()).unwrap(); - let mut c = vec![f32::INFINITY; m]; - driver::Drive::drive(&mut checked_driver::<_, 8, B_BITS, MR, NR, PACK>( - arch, - &packed, - b.as_view(), - &mut c, - )); - assert_scores_close( - &c, - &expected, - format_args!("M={m}, N={n}, K={k}, MR={MR}, NR={NR}, PACK={PACK}"), - ); - } - } - } - let a = constant_minmax_matrix::<8>(16, 257, 255, 1.0, 0.0); - let max_b = (1 << B_BITS) - 1; - let b = constant_minmax_matrix::(16, 257, max_b, 1.0, 0.0); - let packed = PackedA::<_, 8, MR, PACK>::new(arch, a.as_view()).unwrap(); - let mut c = [0.0; 16]; - driver::Drive::drive(&mut checked_driver::<_, 8, B_BITS, MR, NR, PACK>( - arch, - &packed, - b.as_view(), - &mut c, - )); - assert_eq!(c, [(257 * 255 * max_b) as f32; 16]); - } - #[test] fn scalar_driver_matches_reference_at_target_shapes() { assert_driver_target_shapes::<_, 4, 16, 2, 4>(Scalar::new()); @@ -1630,97 +1835,76 @@ mod tests { } } - fn assert_nan_candidates_preserve_maxima< - W, - const B_BITS: usize, - const MR: usize, - const NR: usize, - const PACK: usize, - >( + fn assert_microkernel_nan_reduction( arch: W, ) where - W: ExtraWide + LoadStore, - Unsigned: Representation, + W: ExtraWide, { - let mut a = constant_minmax_matrix::<8>(MR + 1, 1, 1, 1.0, 0.0); - // Exercise mixed SIMD lanes as well as the partial A panel. - for row in [1, MR] { - a.get_row_mut(row).unwrap().set_meta(MinMaxCompensation { - dim: 1, - a: 1.0, - b: 0.0, - n: f32::NAN, - norm_squared: 0.0, - }); - } - let packed = PackedA::<_, 8, MR, PACK>::new(arch, a.as_view()).unwrap(); - let mut cases = vec![ - vec![], - vec![10.0, f32::NAN, 5.0], - vec![f32::NAN, 10.0, 5.0], - vec![10.0, 5.0, f32::NAN], - vec![10.0, f32::NAN, f32::NAN, 5.0], - vec![-5.0, f32::NAN, -10.0], - vec![f32::NAN, f32::NAN], - ]; - let n = 2 * NR + 1; - for nan_index in [0, NR, n - 1] { - let mut scores = vec![5.0; n]; - scores[if nan_index == 0 { 1 } else { 0 }] = 10.0; - scores[nan_index] = f32::NAN; - cases.push(scores); - } - - for scores in cases { - let mut b = constant_minmax_matrix::(scores.len(), 1, 0, 1.0, 0.0); - for (mut row, offset) in b.rows_mut().zip(scores.iter().copied()) { - row.set_meta(MinMaxCompensation { - dim: 1, - a: 1.0, - b: offset, - n: 0.0, - norm_squared: 0.0, - }); - } - let expected = reference_max_inner_products::<8, B_BITS>(a.as_view(), b.as_view()); - assert_eq!( - expected[0], - scores.iter().copied().fold(f32::NEG_INFINITY, f32::max) - ); - assert_eq!(expected[1], f32::NEG_INFINITY); - assert_eq!(expected[MR], f32::NEG_INFINITY); - for tile_rows in [1, 2, NR, NR + 1, scores.len().max(1)] { - let mut c = vec![f32::NAN; MR + 1]; - let mut driver = checked_driver::<_, 8, B_BITS, MR, NR, PACK>( - arch, - &packed, - b.as_view(), - &mut c, - ); - driver.b_rows_per_tile = NonZeroUsize::new(tile_rows).unwrap(); - driver::Drive::drive(&mut driver); - assert_scores_close( - &c, - &expected, - format_args!( - "NaN reduction: scores={scores:?}, B tile={tile_rows}, \ - MR={MR}, NR={NR}, PACK={PACK}" - ), - ); + let k = DimK::new(NonZeroUsize::new(1).unwrap()); + let a_codes = vec![W::AElem::default(); MR * PACK]; + let b_codes = vec![W::BElem::default(); NR * PACK]; + arch.run(|| { + let mut n = [0.0; MR]; + n[MR - 1] = f32::NAN; + let a_meta = Coefficients { + a: &[1.0; MR], + b: &[1.0; MR], + n: &n, + }; + for (initial, score) in [ + (f32::NEG_INFINITY, f32::NAN), + (f32::NEG_INFINITY, 5.0), + (f32::NEG_INFINITY, -5.0), + (f32::NEG_INFINITY, f32::INFINITY), + (7.0, f32::NAN), + ] { + for column in [0, NR - 1] { + let mut b_meta = [MinMaxCompensation { + dim: 1, + a: 1.0, + b: 0.0, + n: f32::NAN, + norm_squared: 0.0, + }; NR]; + b_meta[column].n = score; + let mut c = [initial; MR]; + let mut kernel = MicroKernel:: { + arch, + // SAFETY: A is one full MR panel, padded to PACK for logical K=1. + a: unsafe { packed::Panel::new(Slice::new(&a_codes), k) }, + a_meta, + // SAFETY: B contains NR rows of PACK initialized elements. + b: unsafe { + unpacked::Panel::new(Slice::new(&b_codes), padded_k::(k)) + }, + b_meta: &b_meta, + c: &mut c, + k, + }; + driver::MicroKernel::micro_kernel(&mut kernel); + let mut expected = [initial.max(score); MR]; + expected[MR - 1] = initial; + assert_eq!( + c, expected, + "NaN reduction: initial={initial}, score={score}, column={column}", + ); + } } - } + }); } #[test] - fn scalar_driver_ignores_nan_candidates() { - assert_nan_candidates_preserve_maxima::(Scalar::new()); + fn scalar_microkernel_ignores_nan_candidates() { + assert_microkernel_nan_reduction::(Scalar::new()); + assert_microkernel_nan_reduction::(Scalar::new()); } #[cfg(target_arch = "x86_64")] #[test] - fn avx2_driver_ignores_nan_candidates() { + fn avx2_microkernel_ignores_nan_candidates() { if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { - assert_nan_candidates_preserve_maxima::<_, 4, 16, 4, 4>(arch); + assert_microkernel_nan_reduction::<_, 16, 4, 4>(arch); + assert_microkernel_nan_reduction::<_, 16, 4, 2>(arch); } else { eprintln!("AVX2 unavailable; NaN regression not executed"); } @@ -1728,9 +1912,78 @@ mod tests { #[cfg(target_arch = "x86_64")] #[test] - fn avx512_driver_ignores_nan_candidates() { + fn avx2_half_reduction_matches_whole_rows() { + use diskann_wide::arch::x86_64::V3; + + fn check(arch: V3) { + arch.run(|| { + let acc: [i32x16; NR] = core::array::from_fn(|j| { + i32x16::::from_array( + arch, + core::array::from_fn(|i| (i * 37 + j * 11) as i32), + ) + }); + let a = core::array::from_fn(|i| (i + 1) as f32 * 0.125); + let b = core::array::from_fn(|i| i as f32 * 0.25 - 2.0); + let mut n = core::array::from_fn(|i| i as f32 * 0.75); + n[9] = f32::NAN; + let query = Coefficients { + a: &a, + b: &b, + n: &n, + }; + let mut docs: [MinMaxCompensation; NR] = + core::array::from_fn(|j| MinMaxCompensation { + dim: 31, + a: (j + 1) as f32 * 0.25, + b: 1.0 - j as f32, + n: j as f32 * 0.5, + norm_squared: 0.0, + }); + for special in [0.0, f32::NAN, f32::INFINITY, f32::NEG_INFINITY] { + docs[NR - 1].n = special; + let initial = core::array::from_fn(|i| { + if i.is_multiple_of(2) { + f32::NEG_INFINITY + } else { + i as f32 + } + }); + let mut expected = initial; + compensated_max_into(arch, acc, query, &docs, 31, &mut expected); + let mut pair = initial; + let mut quad = initial; + >::max_into(arch, acc, query, &docs, 31, &mut pair); + >::max_into(arch, acc, query, &docs, 31, &mut quad); + assert_eq!( + pair.map(f32::to_bits), + expected.map(f32::to_bits), + "NR={NR}" + ); + assert_eq!( + quad.map(f32::to_bits), + expected.map(f32::to_bits), + "NR={NR}" + ); + } + }); + } + + if let Some(arch) = V3::new_checked() { + check::<1>(arch); + check::<2>(arch); + check::<4>(arch); + check::<6>(arch); + check::<8>(arch); + } + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn avx512_microkernel_ignores_nan_candidates() { if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked() { - assert_nan_candidates_preserve_maxima::<_, 4, 16, 8, 4>(arch); + assert_microkernel_nan_reduction::<_, 16, 8, 4>(arch); + assert_microkernel_nan_reduction::<_, 16, 8, 2>(arch); } else { eprintln!("V4 unavailable; NaN regression not executed"); } @@ -1738,9 +1991,9 @@ mod tests { #[cfg(target_arch = "aarch64")] #[test] - fn neon_driver_ignores_nan_candidates() { + fn neon_microkernel_ignores_nan_candidates() { if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { - assert_nan_candidates_preserve_maxima::<_, 4, 16, 4, 4>(arch); + assert_microkernel_nan_reduction::<_, 16, 4, 4>(arch); } else { eprintln!("Neon/dotprod unavailable; NaN regression not executed"); } @@ -1756,23 +2009,53 @@ mod tests { assert!((1..=8).contains(&W::MAX_B_BITS)); let max_b = u8::MAX >> (8 - W::MAX_B_BITS); arch.run(|| { - let a = W::Wide::splat(arch, W::convert_a(u8::MAX)); - let b = W::Splat::splat(arch, W::convert_b(max_b)); - let dot = arch.dot(W::Acc::default(arch), a, b).simd_cast(); - let mut actual = [0.0; MR]; - assert_eq!(W::Float::LANES, MR); - // SAFETY: actual has exactly Float::LANES writable elements. - unsafe { dot.store_simd(actual.as_mut_ptr()) }; + let codes: Vec<_> = (0..MR * PACK) + .map(|i| { + W::convert_a(if (i / PACK).is_multiple_of(2) { + u8::MAX + } else { + 1 + }) + }) + .collect(); + let b_codes = [W::convert_b(max_b); PACK]; + // SAFETY: The buffers contain one complete A pack and B pack. + let (a, b) = unsafe { + ( + arch.load_a(Slice::new(&codes)), + arch.splat(Slice::new(&b_codes)), + ) + }; + let acc = arch.dot(arch.zero_acc(), a, b); + let acc = arch.dot(acc, a, b); + let coefficients = Coefficients { + a: &[1.0; MR], + b: &[0.0; MR], + n: &[0.0; MR], + }; + let mut actual = [f32::NEG_INFINITY; MR]; + arch.max_into( + [acc], + coefficients, + &[MinMaxCompensation { + a: 1.0, + ..Default::default() + }], + 2 * PACK, + &mut actual, + ); assert_eq!( actual, - [(PACK * usize::from(u8::MAX) * usize::from(max_b)) as f32; MR], + core::array::from_fn(|i| { + let a = if i.is_multiple_of(2) { 255 } else { 1 }; + (2 * PACK * a * usize::from(max_b)) as f32 + }), ); }); } - fn assert_bit_pair< + fn assert_b_width< W, - const A_BITS: usize, const B_BITS: usize, const MR: usize, const NR: usize, @@ -1782,121 +2065,74 @@ mod tests { rng: &mut StdRng, ) where W: ExtraWide + LoadStore, - W::AElem: PartialEq + std::fmt::Debug, - Unsigned: Representation + Representation, + Unsigned: Representation, { - let max_a = (1u64 << A_BITS) - 1; + let max_a = u64::from(u8::MAX); let max_b = (1u64 << B_BITS) - 1; - let limit = max_k::() as u64; + let limit = max_k::<8, 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); - assert_a_packing::(arch); - let dimensions: &[usize] = if cfg!(miri) { - &[1, 3] - } else { - &[1, 2, 3, 4, 5, 7, 8, 15, 16, 17, 31, 33, 255, 256, 257] - }; - for (m, n) in [ - (1, 0), - (1, 1), - (MR - 1, (NR - 1).max(1)), - (MR, NR), - (MR + 1, 2 * NR + 1), - ] { - for &k in dimensions { - let mut a = random_minmax_matrix::(m, k, rng); - let mut b = random_minmax_matrix::(n, k, rng); - if m > 1 { - let mut row = a.get_row_mut(m - 1).unwrap(); - let mut meta = row.meta(); - meta.n = f32::NAN; - row.set_meta(meta); - } - if n > 1 { - let mut row = b.get_row_mut(n - 1).unwrap(); - let mut meta = row.meta(); - meta.n = f32::NAN; - row.set_meta(meta); - } - let expected = - reference_max_inner_products::(a.as_view(), b.as_view()); - let packed = PackedA::<_, A_BITS, MR, PACK>::new(arch, a.as_view()).unwrap(); - let mut guarded = vec![654321.0; m + 2]; - let mut kernel = checked_driver::( - arch, - &packed, - b.as_view(), - &mut guarded[1..=m], - ); - for tile_rows in [kernel.b_rows_per_tile, NonZeroUsize::new(NR + 1).unwrap()] { - kernel.b_rows_per_tile = tile_rows; - kernel.c.fill(f32::INFINITY); - driver::Drive::drive(&mut kernel); - assert_scores_close( - kernel.c, - &expected, - format_args!( - "{A_BITS}x{B_BITS}: M={m}, N={n}, K={k}, B tile={tile_rows}, PACK={PACK}" - ), - ); - } - assert_eq!(guarded[0], 654321.0); - assert_eq!(guarded[m + 1], 654321.0); - } - } + let (m, n, k) = (MR + 1, NR + 1, PACK + 1); + let a = random_minmax_matrix::<8>(m, k, rng); + let b = random_minmax_matrix::(n, k, rng); + let expected = reference_max_inner_products::<8, B_BITS>(a.as_view(), b.as_view()); + let packed = PackedA::<_, 8, MR, PACK>::new(arch, a.as_view()).unwrap(); + let mut guarded = vec![654321.0; m + 2]; + let mut kernel = checked_driver::( + arch, + &packed, + b.as_view(), + &mut guarded[1..=m], + ); + driver::Drive::drive(&mut kernel); + assert_scores_close( + kernel.c, + &expected, + format_args!("8x{B_BITS}: M={m}, N={n}, K={k}, PACK={PACK}"), + ); + assert_eq!(guarded[0], 654321.0); + assert_eq!(guarded[m + 1], 654321.0); - let dimensions: &[usize] = if cfg!(miri) { - &[PACK, PACK + 1] - } else { - &[PACK, PACK + 1, 257] - }; - for &k in dimensions { - let a = constant_minmax_matrix::(MR + 1, k, max_a as i64, 1.0, 0.0); - let b = constant_minmax_matrix::(NR + 1, k, max_b as i64, 1.0, 0.0); - let packed = PackedA::<_, A_BITS, MR, PACK>::new(arch, a.as_view()).unwrap(); - let mut c = vec![0.0; MR + 1]; - driver::Drive::drive(&mut checked_driver::( - arch, - &packed, - b.as_view(), - &mut c, - )); - assert_eq!(c, vec![(k as u64 * max_a * max_b) as f32; MR + 1]); - } + let a = constant_minmax_matrix::<8>(1, k, max_a as i64, 1.0, 0.0); + let b = constant_minmax_matrix::(1, k, max_b as i64, 1.0, 0.0); + let packed = PackedA::<_, 8, MR, PACK>::new(arch, a.as_view()).unwrap(); + let mut c = [0.0]; + driver::Drive::drive(&mut checked_driver::( + arch, + &packed, + b.as_view(), + &mut c, + )); + assert_eq!(c, [(k as u64 * max_a * max_b) as f32]); } macro_rules! check_widths { - ($arch:expr, $sentinel:expr, $nr:literal, $pack:literal; $($b:literal),+ $(,)?) => {{ + ($arch:expr, $mr:literal, $nr:literal, $pack:literal; $($b:literal),+ $(,)?) => {{ let arch = $arch; let mut rng = StdRng::seed_from_u64(0x948cd3e02a671fb5); - assert_backend_code_range::<_, 16, $pack>(arch); + assert_backend_code_range::<_, $mr, $pack>(arch); $( - assert_b_decoding_boundaries::<_, $b, 16, $nr, $pack>(arch, $sentinel); - assert_bit_pair::<_, 1, $b, 16, $nr, $pack>(arch, &mut rng); - assert_bit_pair::<_, 2, $b, 16, $nr, $pack>(arch, &mut rng); - assert_bit_pair::<_, 3, $b, 16, $nr, $pack>(arch, &mut rng); - assert_bit_pair::<_, 4, $b, 16, $nr, $pack>(arch, &mut rng); - assert_bit_pair::<_, 5, $b, 16, $nr, $pack>(arch, &mut rng); - assert_bit_pair::<_, 6, $b, 16, $nr, $pack>(arch, &mut rng); - assert_bit_pair::<_, 7, $b, 16, $nr, $pack>(arch, &mut rng); - assert_bit_pair::<_, 8, $b, 16, $nr, $pack>(arch, &mut rng); + assert_b_width::<_, $b, $mr, $nr, $pack>(arch, &mut rng); )+ }}; } #[test] - fn scalar_supports_all_bit_pairs_and_byte_layouts() { - check_widths!(Scalar::new(), -1, 2, 2; 1, 2, 3, 4, 5, 6, 7, 8); - check_widths!(Scalar::new(), -1, 2, 4; 1, 2, 3, 4, 5, 6, 7); + fn scalar_supports_all_b_widths_and_layouts() { + check_widths!(Scalar::new(), 16, 2, 2; 1, 2, 3, 4, 5, 6, 7, 8); + check_widths!(Scalar::new(), 16, 2, 4; 1, 2, 3, 4, 5, 6, 7); + check_widths!(Scalar::new(), 8, 2, 2; 1, 2, 3, 4, 5, 6, 7, 8); } #[cfg(target_arch = "x86_64")] #[test] - fn avx2_supports_all_bit_pairs_and_byte_layouts() { + fn avx2_supports_all_b_widths_and_layouts() { if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { - check_widths!(arch, -1, 4, 2; 1, 2, 3, 4, 5, 6, 7, 8); - check_widths!(arch, -1, 4, 4; 1, 2, 3, 4, 5, 6); + check_widths!(arch, 16, 4, 2; 1, 2, 3, 4, 5, 6, 7, 8); + check_widths!(arch, 16, 4, 4; 1, 2, 3, 4, 5, 6); + check_widths!(arch, 8, 4, 2; 1, 2, 3, 4, 5, 6, 7, 8); + check_widths!(arch, 8, 4, 4; 1, 2, 3, 4, 5, 6); } else { eprintln!("AVX2 unavailable; all-width regression not executed"); } @@ -1904,10 +2140,10 @@ mod tests { #[cfg(target_arch = "x86_64")] #[test] - fn avx512_supports_all_bit_pairs_and_byte_layouts() { + fn avx512_supports_all_b_widths_and_layouts() { if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked() { - check_widths!(arch, -1, 8, 2; 1, 2, 3, 4, 5, 6, 7, 8); - check_widths!(arch, -1, 8, 4; 1, 2, 3, 4, 5, 6, 7); + check_widths!(arch, 16, 8, 2; 1, 2, 3, 4, 5, 6, 7, 8); + check_widths!(arch, 16, 8, 4; 1, 2, 3, 4, 5, 6, 7); } else { eprintln!("V4 unavailable; all-width regression not executed"); } @@ -1915,9 +2151,10 @@ mod tests { #[cfg(target_arch = "aarch64")] #[test] - fn neon_supports_all_bit_pairs() { + fn neon_supports_all_b_widths_and_layouts() { if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { - check_widths!(arch, u8::MAX, 4, 4; 1, 2, 3, 4, 5, 6, 7, 8); + check_widths!(arch, 16, 4, 4; 1, 2, 3, 4, 5, 6, 7, 8); + check_widths!(arch, 8, 4, 4; 1, 2, 3, 4, 5, 6, 7, 8); } else { eprintln!("Neon/dotprod unavailable; all-width regression not executed"); } @@ -1933,9 +2170,9 @@ mod tests { >, { let mut rng = StdRng::seed_from_u64(0x14a7f0b2c69d538e); - for k in [1, 3, 17, 257] { - let a = random_minmax_matrix::(17, k, &mut rng); - let b = random_minmax_matrix::(19, k, &mut rng); + for k in [3, 65] { + let a = random_minmax_matrix::(3, k, &mut rng); + let b = random_minmax_matrix::(5, k, &mut rng); let expected: Vec<_> = a .rows() .map(|a| { @@ -1975,26 +2212,6 @@ mod tests { check_reference::<8, 1>(); } - #[test] - fn packed_a_is_reused_across_b_widths() { - let arch = Scalar::new(); - let mut rng = StdRng::seed_from_u64(0xac21f49d70b365e8); - let a = random_minmax_matrix::<3>(17, 33, &mut rng); - let packed = PackedA::<_, 3, 16, 2>::new(arch, a.as_view()).unwrap(); - let mut c = vec![f32::NAN; 17]; - macro_rules! check { - ($($b:literal),+) => {$( - let b = random_minmax_matrix::<$b>(9, 33, &mut rng); - let expected = reference_max_inner_products::<3, $b>(a.as_view(), b.as_view()); - driver::Drive::drive(&mut checked_driver::<_, 3, $b, 16, 2, 2>( - arch, &packed, b.as_view(), &mut c, - )); - assert_scores_close(&c, &expected, format_args!("reused A3 with B{}", $b)); - )+}; - } - check!(1, 2, 3, 4, 5, 6, 7, 8); - } - #[cfg(not(miri))] #[test] fn low_bit_inputs_can_exceed_the_eight_bit_dimension_limit() { @@ -2018,7 +2235,6 @@ mod tests { assert_eq!(c[0], (k as i64 * max_a * max_b) as f32); } check::<1, 8>(); - check::<4, 8>(); check::<8, 1>(); } } @@ -2026,18 +2242,12 @@ mod tests { mod minmax8 { use super::*; - fn check_backend(arch: W, sentinel: W::BElem) + fn check_backend(arch: W) where W: ExtraWide<16, PACK> + LoadStore, - W::AElem: PartialEq + std::fmt::Debug, - W::BElem: PartialEq + std::fmt::Debug, { - assert_a_packing::(arch); assert_driver_contract::(arch); assert_driver_target_shapes::(arch); - assert_nan_candidates_preserve_maxima::(arch); - assert_packed_a_reuse::(arch); - assert_b_decoding_boundaries::(arch, sentinel); assert_unsigned_byte_products::(arch); } @@ -2048,10 +2258,11 @@ mod tests { if cfg!(miri) { return; } - for k in [1, 2, 3, 4, 5, 257, 259] { - let mut a = constant_minmax_matrix::<8>(257, k, 0, 1.0, 0.0); + for k in [PACK - 1, PACK, PACK + 1] { + let boundary_codes = [0, 1, 127, 128, 255]; + let mut a = constant_minmax_matrix::<8>(17, k, 0, 1.0, 0.0); for (i, mut row) in a.rows_mut().enumerate() { - let code = (i % 256) as i64; + let code = boundary_codes[i % boundary_codes.len()]; for j in 0..k { row.vector_mut().set(j, code).unwrap(); } @@ -2060,8 +2271,8 @@ mod tests { row.set_meta(meta); } let a = PackedA::<_, 8, 16, PACK>::new(arch, a.as_view()).unwrap(); - let mut c = vec![f32::NAN; 257]; - for code in 0..=255 { + let mut c = vec![f32::NAN; 17]; + for code in boundary_codes { let b = constant_minmax_matrix::<8>(1, k, code, 1.0, 0.0); driver::Drive::drive(&mut checked_driver::( arch, @@ -2072,9 +2283,10 @@ mod tests { for (i, &actual) in c.iter().enumerate() { assert_eq!( actual, - ((i % 256) * code as usize * k) as f32, + (boundary_codes[i % boundary_codes.len()] as usize * code as usize * k) + as f32, "A={}, B={code}, K={k}", - i % 256, + boundary_codes[i % boundary_codes.len()], ); } } @@ -2083,14 +2295,14 @@ mod tests { #[test] fn scalar_driver_obeys_maxsim_contract() { - check_backend::<_, 2, 2>(Scalar::new(), -1); + check_backend::<_, 2, 2>(Scalar::new()); } #[cfg(target_arch = "x86_64")] #[test] fn avx2_driver_obeys_maxsim_contract() { if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { - check_backend::<_, 4, 2>(arch, -1); + check_backend::<_, 4, 2>(arch); } else { eprintln!("AVX2 unavailable; MinMax8 x MinMax8 backend not executed"); } @@ -2100,7 +2312,7 @@ mod tests { #[test] fn avx512_driver_obeys_maxsim_contract() { if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked() { - check_backend::<_, 8, 2>(arch, -1); + check_backend::<_, 8, 2>(arch); } else { eprintln!("V4 unavailable; MinMax8 x MinMax8 backend not executed"); } @@ -2110,7 +2322,7 @@ mod tests { #[test] fn neon_driver_obeys_maxsim_contract() { if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { - check_backend::<_, 4, 4>(arch, u8::MAX); + check_backend::<_, 4, 4>(arch); } else { eprintln!("Neon/dotprod unavailable; MinMax8 x MinMax8 backend not executed"); } @@ -2120,26 +2332,12 @@ mod tests { #[test] #[should_panic(expected = "exceeds the MinMax8 x MinMax8 limit")] fn driver_rejects_k_above_accumulator_limit() { - let k = 33026; + 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::<_, 8, 16, 2>::new(Scalar::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); } - - #[test] - #[should_panic] - fn b_decoding_rejects_a_short_output() { - let codes = BitSlice::new(&[128u8, 255][..], 2).unwrap(); - decode_b::(codes, &mut [0]); - } - - #[test] - #[should_panic] - fn b_decoding_rejects_a_long_output() { - let codes = BitSlice::new(&[255u8][..], 1).unwrap(); - decode_b::(codes, &mut [0; 2]); - } } } From 2d95565a81217de8e5afc02eb3435db15e2e1e8e Mon Sep 17 00:00:00 2001 From: partychen <1669295+partychen@users.noreply.github.com> Date: Mon, 19 Oct 2026 16:58:26 +0800 Subject: [PATCH 3/4] Decouple MinMax packing and compact prepared coefficients Use From for architecture-independent A packing. Project canonical B metadata into the existing Coefficients type at preparation, preserving dimension checks and updating cache footprints and focused regression tests. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../maxsim/packed_minmax_x_unpacked_minmax.rs | 222 +++++++++--------- 1 file changed, 109 insertions(+), 113 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 e712e55251..f3cfa24278 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 @@ -7,6 +7,7 @@ //! //! Each B tile is decoded once and reused across A panels. Integer inner products //! are compensated before reduction because B's coefficients vary by vector. +//! B preparation retains only the three coefficients used by MaxSim. //! B rows are zero-padded to PACK; kernel `k` remains the logical dimension used //! for compensation and accumulator limits. //! @@ -39,7 +40,7 @@ use crate::{ ptr::Slice, util::LoadStore, }, - minmax::{MinMaxCompensation, MinMaxMeta}, + minmax::MinMaxMeta, multi_vector::{BlockTransposed, MatRef, Overflow}, }; @@ -63,7 +64,7 @@ const fn max_k() -> usize { (i32::MAX as usize) / (((1 << A_BITS) - 1) * ((1 << B_BITS) - 1)) } -#[derive(Debug, Clone, Copy)] +#[derive(Debug, Clone, Copy, Default)] pub(crate) struct Coefficients { a: T, b: T, @@ -81,7 +82,7 @@ where impl PackedA where - T: Copy + Default, + T: Copy + Default + From, Unsigned: Representation, { /// Pack A and copy its coefficients, zero-filling partial rows and K packs. @@ -90,10 +91,7 @@ where /// /// Returns [`Overflow`] if the packed data or coefficient allocation size overflows. #[expect(clippy::expect_used)] - pub(crate) fn new>( - _: W, - a: MatRef<'_, MinMaxMeta>, - ) -> Result { + pub(crate) fn new(a: MatRef<'_, MinMaxMeta>) -> Result { let k = a.repr().intrinsic_dim(); let mut codes = BlockTransposed::try_new(a.num_vectors(), k)?; let padded_rows = codes.padded_nrows(); @@ -114,7 +112,7 @@ where .expect("A and its packing have equal row counts"); for j in 0..k { // SAFETY: MinMaxMeta exposes exactly k codes for every row. - dst.set(j, W::convert_a(unsafe { v.get_unchecked(j) } as u8)); + dst.set(j, T::from(unsafe { v.get_unchecked(j) } as u8)); } } Ok(Self { codes, meta }) @@ -284,12 +282,12 @@ where #[derive(Default)] struct BScratch { codes: Vec, - meta: Vec, + meta: Vec>, } struct BTile<'a, T> { codes: unpacked::View<'a, T>, - meta: &'a [MinMaxCompensation], + meta: &'a [Coefficients], } #[inline] @@ -321,10 +319,15 @@ impl BScratch { 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(), MinMaxCompensation::default()); + self.meta.resize(count.get(), Coefficients::default()); for i in 0..count.get() { let row = b.get_row(start + i).expect("B tile range was checked"); - self.meta[i] = row.meta(); + 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], |code| convert(code as u8)) @@ -342,7 +345,7 @@ impl BScratch { /// Cache-planning bytes for one decoded B row, including metadata. fn row_bytes(k: DimK) -> Bytes { let codes = Elements::::new(padded_k::(k).value().get()).bytes(); - Bytes::new(codes.value() + size_of::()) + Bytes::new(codes.value() + size_of::>()) } } @@ -356,7 +359,7 @@ struct PanelKernel<'a, W: ExtraWide, const MR: usize, const NR: usize, a: packed::Panel<'a, W::AElem, MR, PACK>, a_meta: Coefficients<&'a [f32; MR]>, b: unpacked::View<'a, W::BElem>, - b_meta: &'a [MinMaxCompensation], + b_meta: &'a [Coefficients], c: [f32; MR], k: DimK, } @@ -375,7 +378,7 @@ struct Visitor<'a, W: ExtraWide, const MR: usize, const PACK: usize> { arch: W, a: packed::Panel<'a, W::AElem, MR, PACK>, a_meta: Coefficients<&'a [f32; MR]>, - b_meta: &'a [MinMaxCompensation], + b_meta: &'a [Coefficients], c: &'a mut [f32; MR], k: DimK, } @@ -449,7 +452,7 @@ struct MicroKernel<'a, W: ExtraWide, const MR: usize, const NR: usize, a: packed::Panel<'a, W::AElem, MR, PACK>, a_meta: Coefficients<&'a [f32; MR]>, b: unpacked::Panel<'a, W::BElem, NR>, - b_meta: &'a [MinMaxCompensation; NR], + b_meta: &'a [Coefficients; NR], c: &'a mut [f32; MR], k: DimK, } @@ -524,13 +527,12 @@ pub(crate) trait ExtraWide: Architecture { /// This is a backend capability, not the actual A or B input width. const MAX_B_BITS: usize; - type AElem: Copy + Default; + type AElem: Copy + Default + From; type BElem: Copy + Default; type Wide: Copy; type Splat: Copy; type Acc: Copy; - fn convert_a(code: u8) -> Self::AElem; /// `code` fits in MAX_B_BITS unsigned bits. fn convert_b(code: u8) -> Self::BElem; @@ -550,7 +552,7 @@ pub(crate) trait ExtraWide: Architecture { self, acc: [Self::Acc; NR], query: Coefficients<&[f32; MR]>, - docs: &[MinMaxCompensation; NR], + docs: &[Coefficients; NR], k: usize, scores: &mut [f32; MR], ); @@ -566,7 +568,7 @@ fn compensated_max_into( arch: F::Arch, acc: impl IntoIterator, query: Coefficients<&[f32; MR]>, - docs: &[MinMaxCompensation; NR], + docs: &[Coefficients; NR], k: usize, scores: &mut [f32; MR], ) where @@ -642,7 +644,7 @@ impl ExtraWide<16, 4> for Scalar { self, acc: [Self::Acc; NR], query: Coefficients<&[f32; 16]>, - docs: &[MinMaxCompensation; NR], + docs: &[Coefficients; NR], k: usize, scores: &mut [f32; 16], ) { @@ -654,10 +656,6 @@ impl ExtraWide<16, 4> for Scalar { acc.dot_simd(a, b) } - fn convert_a(code: u8) -> u8 { - code - } - fn convert_b(code: u8) -> i8 { code as i8 } @@ -696,7 +694,7 @@ impl ExtraWide<16, 2> for Scalar { self, acc: [Self::Acc; NR], query: Coefficients<&[f32; 16]>, - docs: &[MinMaxCompensation; NR], + docs: &[Coefficients; NR], k: usize, scores: &mut [f32; 16], ) { @@ -708,10 +706,6 @@ impl ExtraWide<16, 2> for Scalar { acc.dot_simd(a, b) } - fn convert_a(code: u8) -> i16 { - i16::from(code) - } - fn convert_b(code: u8) -> i16 { i16::from(code) } @@ -747,7 +741,7 @@ impl ExtraWide<8, 2> for Scalar { self, acc: [Self::Acc; NR], query: Coefficients<&[f32; 8]>, - docs: &[MinMaxCompensation; NR], + docs: &[Coefficients; NR], k: usize, scores: &mut [f32; 8], ) { @@ -759,9 +753,6 @@ impl ExtraWide<8, 2> for Scalar { acc.dot_simd(a, b) } - fn convert_a(code: u8) -> i16 { - i16::from(code) - } fn convert_b(code: u8) -> i16 { i16::from(code) } @@ -783,7 +774,7 @@ mod x86_64 { arch: V3, acc: [i32x16; NR], query: Coefficients<&[f32; 16]>, - docs: &[MinMaxCompensation; NR], + docs: &[Coefficients; NR], k: usize, scores: &mut [f32; 16], ) { @@ -842,16 +833,13 @@ mod x86_64 { self, acc: [Self::Acc; NR], query: Coefficients<&[f32; 8]>, - docs: &[MinMaxCompensation; NR], + docs: &[Coefficients; NR], k: usize, scores: &mut [f32; 8], ) { compensated_max_into(self, acc, query, docs, k, scores); } - fn convert_a(code: u8) -> u8 { - code - } fn convert_b(code: u8) -> i8 { code as i8 } @@ -893,7 +881,7 @@ mod x86_64 { self, acc: [Self::Acc; NR], query: Coefficients<&[f32; 8]>, - docs: &[MinMaxCompensation; NR], + docs: &[Coefficients; NR], k: usize, scores: &mut [f32; 8], ) { @@ -905,9 +893,6 @@ mod x86_64 { acc.dot_simd(a, b) } - fn convert_a(code: u8) -> i16 { - i16::from(code) - } fn convert_b(code: u8) -> i16 { i16::from(code) } @@ -943,17 +928,13 @@ mod x86_64 { self, acc: [Self::Acc; NR], query: Coefficients<&[f32; 16]>, - docs: &[MinMaxCompensation; NR], + docs: &[Coefficients; NR], k: usize, scores: &mut [f32; 16], ) { compensated_max_into_v3(self, acc, query, docs, k, scores); } - fn convert_a(code: u8) -> u8 { - code - } - fn convert_b(code: u8) -> i8 { code as i8 } @@ -1006,7 +987,7 @@ mod x86_64 { self, acc: [Self::Acc; NR], query: Coefficients<&[f32; 16]>, - docs: &[MinMaxCompensation; NR], + docs: &[Coefficients; NR], k: usize, scores: &mut [f32; 16], ) { @@ -1018,10 +999,6 @@ mod x86_64 { acc.dot_simd(a, b) } - fn convert_a(code: u8) -> i16 { - i16::from(code) - } - fn convert_b(code: u8) -> i16 { i16::from(code) } @@ -1060,7 +1037,7 @@ mod x86_64 { self, acc: [Self::Acc; NR], query: Coefficients<&[f32; 16]>, - docs: &[MinMaxCompensation; NR], + docs: &[Coefficients; NR], k: usize, scores: &mut [f32; 16], ) { @@ -1072,10 +1049,6 @@ mod x86_64 { acc.dot_simd(a, b) } - fn convert_a(code: u8) -> u8 { - code - } - fn convert_b(code: u8) -> i8 { code as i8 } @@ -1114,7 +1087,7 @@ mod x86_64 { self, acc: [Self::Acc; NR], query: Coefficients<&[f32; 16]>, - docs: &[MinMaxCompensation; NR], + docs: &[Coefficients; NR], k: usize, scores: &mut [f32; 16], ) { @@ -1126,10 +1099,6 @@ mod x86_64 { acc.dot_simd(a, b) } - fn convert_a(code: u8) -> i16 { - i16::from(code) - } - fn convert_b(code: u8) -> i16 { i16::from(code) } @@ -1171,7 +1140,7 @@ mod aarch64 { self, acc: [Self::Acc; NR], query: Coefficients<&[f32; 8]>, - docs: &[MinMaxCompensation; NR], + docs: &[Coefficients; NR], k: usize, scores: &mut [f32; 8], ) { @@ -1183,9 +1152,6 @@ mod aarch64 { acc.dot_simd(a, b) } - fn convert_a(code: u8) -> u8 { - code - } fn convert_b(code: u8) -> u8 { code } @@ -1223,7 +1189,7 @@ mod aarch64 { self, acc: [Self::Acc; NR], query: Coefficients<&[f32; 16]>, - docs: &[MinMaxCompensation; NR], + docs: &[Coefficients; NR], k: usize, scores: &mut [f32; 16], ) { @@ -1235,10 +1201,6 @@ mod aarch64 { acc.dot_simd(a, b) } - fn convert_a(code: u8) -> u8 { - code - } - fn convert_b(code: u8) -> u8 { code } @@ -1266,7 +1228,7 @@ mod tests { bits::{Representation, Unsigned}, distances, matrix_kernels::maxsim::test::packed_x_unpacked_test_dims, - minmax::{DataRef, MinMaxIP}, + minmax::{DataRef, MinMaxCompensation, MinMaxIP}, multi_vector::{Defaulted, Mat}, }; @@ -1398,7 +1360,7 @@ 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::<_, 8, MR, PACK>::new(arch, a.as_view()).unwrap(); + let packed = PackedA::::new(a.as_view()).unwrap(); let expected = reference_max_inner_products::<8, B_BITS>(a.as_view(), b.as_view()); let mut guarded = vec![654321.0; a.num_vectors() + 2]; let mut driver = checked_driver::<_, 8, B_BITS, MR, NR, PACK>( @@ -1425,7 +1387,7 @@ 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::<_, 8, MR, PACK>::new(arch, a.as_view()).unwrap(); + let packed = PackedA::::new(a.as_view()).unwrap(); let mut c = vec![f32::INFINITY; a.num_vectors()]; driver::Drive::drive(&mut checked_driver::<_, 8, B_BITS, MR, NR, PACK>( arch, @@ -1490,7 +1452,7 @@ mod tests { norm_squared: 0.0, }); } - let packed = PackedA::<_, 8, MR, PACK>::new(arch, a.as_view()).unwrap(); + let packed = PackedA::::new(a.as_view()).unwrap(); let mut c = vec![0.0; MR + 1]; let mut driver = checked_driver::<_, 8, B_BITS, MR, NR, PACK>(arch, &packed, b.as_view(), &mut c); @@ -1501,7 +1463,7 @@ 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::<_, 8, MR, PACK>::new(arch, a.as_view()).unwrap(); + let packed = PackedA::::new(a.as_view()).unwrap(); driver::Drive::drive(&mut checked_driver::<_, 8, B_BITS, MR, NR, PACK>( arch, &packed, @@ -1519,7 +1481,7 @@ mod tests { 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::<_, 8, 8, 2>::new(arch, 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>( arch, @@ -1587,22 +1549,21 @@ mod tests { } } - fn assert_a_packing(arch: W) + fn assert_a_packing() where - W: ExtraWide, - W::AElem: PartialEq + std::fmt::Debug, + T: Copy + Default + From + PartialEq + std::fmt::Debug, Unsigned: Representation, { for (m, k) in [(0, 0), (0, 5), (3, 0)] { let a = constant_minmax_matrix::(m, k, 0, 1.0, 0.0); - let packed = PackedA::<_, A_BITS, MR, PACK>::new(arch, 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::View::from_block_transposed(packed.codes.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::<_, A_BITS, MR, PACK>::new(arch, 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); @@ -1615,9 +1576,9 @@ mod tests { + (row % MR) * PACK + col % PACK; let expected = if row <= MR && col < 5 { - W::convert_a(a.get_row(row).unwrap().vector().get(col).unwrap() as u8) + T::from(a.get_row(row).unwrap().vector().get(col).unwrap() as u8) } else { - W::AElem::default() + T::default() }; assert_eq!(packed.codes.as_slice()[index], expected); } @@ -1632,23 +1593,23 @@ mod tests { fn a_packing_preserves_codes_metadata_and_padding() { macro_rules! check { ($($bits:literal),+) => {$( - assert_a_packing::<_, $bits, 16, 4>(Scalar::new()); - assert_a_packing::<_, $bits, 16, 2>(Scalar::new()); + assert_a_packing::(); + assert_a_packing::(); )+}; } check!(1, 2, 3, 4, 5, 6, 7, 8); - assert_a_packing::<_, 3, 8, 2>(Scalar::new()); + assert_a_packing::(); } #[test] fn cache_footprints_match_prepared_layouts() { - for (k, a_bytes, b_bytes) in [(256, 4288, 276), (257, 4352, 280)] { + for (k, a_bytes, b_bytes) in [(256, 4288, 268), (257, 4352, 272)] { let dim = DimK::new(NonZeroUsize::new(k).unwrap()); assert_eq!(BScratch::::row_bytes::<4>(dim).value(), b_bytes); assert_eq!(BScratch::::row_bytes::<4>(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(Scalar::new(), a.as_view()).unwrap(); + let packed = PackedA::::new(a.as_view()).unwrap(); assert_eq!(packed.panel_bytes().value(), a_bytes); // Preserve the previous view-based cache estimate. @@ -1663,12 +1624,12 @@ mod tests { #[test] fn widened_cache_footprints_match_prepared_layouts() { - for (k, a_bytes, b_bytes) in [(256, 8384, 532), (257, 8448, 536)] { + for (k, a_bytes, b_bytes) in [(256, 8384, 524), (257, 8448, 528)] { let dim = DimK::new(NonZeroUsize::new(k).unwrap()); assert_eq!(BScratch::::row_bytes::<2>(dim).value(), b_bytes); for rows in [1, 16, 17] { let a = constant_minmax_matrix::<8>(rows, k, 255, 1.0, 0.0); - let packed = PackedA::::new(Scalar::new(), a.as_view()).unwrap(); + let packed = PackedA::::new(a.as_view()).unwrap(); assert_eq!(packed.panel_bytes().value(), a_bytes); let view = packed::View::from_block_transposed(packed.codes.as_view()).unwrap(); assert_eq!( @@ -1679,6 +1640,41 @@ mod tests { } } + #[test] + fn compact_metadata_preserves_scores_across_tile_sizes() { + let arch = Scalar::new(); + let mut rng = StdRng::seed_from_u64(0x9f23ea1640b187ac); + let a = random_minmax_matrix::<8>(9, 3, &mut rng); + let b = random_minmax_matrix::<8>(11, 3, &mut rng); + let packed = PackedA::::new(a.as_view()).unwrap(); + let k = DimK::new(NonZeroUsize::new(3).unwrap()); + // This budget gives different rounded tiles for compact and canonical metadata. + let cache = Cache::new( + NonZeroUsize::new(packed.panel_bytes().value() + 81).unwrap(), + NonZeroUsize::new(4096).unwrap(), + ); + let expected = reference_max_inner_products(a.as_view(), b.as_view()); + let mut scores = vec![f32::NAN; a.num_vectors()]; + // SAFETY: A is nonempty; both inputs have K=3 and scores covers every A row. + let mut driver = unsafe { + Driver::<_, 8, 8, 2, 2>::new::<8>(arch, &packed, b.as_view(), &mut scores, k, cache) + }; + assert_eq!(driver.b_rows_per_tile.get(), 6); + driver::Drive::drive(&mut driver); + assert_scores_close(driver.c, &expected, format_args!("compact metadata tiling")); + + let previous_row_bytes = Bytes::new(4 * size_of::() + size_of::()); + driver.b_rows_per_tile = + Params::new(cache, packed.panel_bytes(), previous_row_bytes, 2).b_cols_in_l1; + assert_eq!(driver.b_rows_per_tile.get(), 4); + driver::Drive::drive(&mut driver); + assert_scores_close( + driver.c, + &expected, + format_args!("previous metadata tiling"), + ); + } + #[test] fn b_preparation_preserves_metadata_and_zero_padding() { assert_b_preparation::(-1, |code| code as i8); @@ -1724,7 +1720,7 @@ mod tests { let b = MatRef::new(MinMaxMeta::::new(3, dim), bytes).unwrap(); for (start, count) in [(0, 3), (2, 1), (1, 2)] { scratch.codes.fill(sentinel); - scratch.meta.fill(MinMaxCompensation::default()); + scratch.meta.fill(Coefficients::default()); let mut conversions = 0; let tile = scratch.prepare::( b, @@ -1745,7 +1741,12 @@ mod tests { assert_eq!(tile.meta.len(), count); for i in 0..count { let row = b.get_row(start + i).unwrap(); - assert_eq!(tile.meta[i], row.meta()); + let expected = row.meta(); + let actual = tile.meta[i]; + assert_eq!( + (actual.a, actual.b, actual.n), + (expected.a, expected.b, expected.n), + ); for j in 0..dim { assert_eq!( codes[i * row_len + j], @@ -1784,7 +1785,7 @@ 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::<_, 8, MR, PACK>::new(arch, a.as_view()).unwrap(); + let packed = PackedA::::new(a.as_view()).unwrap(); let mut c = vec![f32::INFINITY; m]; driver::Drive::drive(&mut checked_driver::<_, 8, B_BITS, MR, NR, PACK>( arch, @@ -1859,12 +1860,10 @@ mod tests { (7.0, f32::NAN), ] { for column in [0, NR - 1] { - let mut b_meta = [MinMaxCompensation { - dim: 1, + let mut b_meta = [Coefficients { a: 1.0, b: 0.0, n: f32::NAN, - norm_squared: 0.0, }; NR]; b_meta[column].n = score; let mut c = [initial; MR]; @@ -1932,14 +1931,11 @@ mod tests { b: &b, n: &n, }; - let mut docs: [MinMaxCompensation; NR] = - core::array::from_fn(|j| MinMaxCompensation { - dim: 31, - a: (j + 1) as f32 * 0.25, - b: 1.0 - j as f32, - n: j as f32 * 0.5, - norm_squared: 0.0, - }); + let mut docs: [Coefficients; NR] = core::array::from_fn(|j| Coefficients { + a: (j + 1) as f32 * 0.25, + b: 1.0 - j as f32, + n: j as f32 * 0.5, + }); for special in [0.0, f32::NAN, f32::INFINITY, f32::NEG_INFINITY] { docs[NR - 1].n = special; let initial = core::array::from_fn(|i| { @@ -2011,7 +2007,7 @@ mod tests { arch.run(|| { let codes: Vec<_> = (0..MR * PACK) .map(|i| { - W::convert_a(if (i / PACK).is_multiple_of(2) { + W::AElem::from(if (i / PACK).is_multiple_of(2) { u8::MAX } else { 1 @@ -2037,7 +2033,7 @@ mod tests { arch.max_into( [acc], coefficients, - &[MinMaxCompensation { + &[Coefficients { a: 1.0, ..Default::default() }], @@ -2077,7 +2073,7 @@ mod tests { let a = random_minmax_matrix::<8>(m, k, rng); let b = random_minmax_matrix::(n, k, rng); let expected = reference_max_inner_products::<8, B_BITS>(a.as_view(), b.as_view()); - let packed = PackedA::<_, 8, MR, PACK>::new(arch, a.as_view()).unwrap(); + let packed = PackedA::::new(a.as_view()).unwrap(); let mut guarded = vec![654321.0; m + 2]; let mut kernel = checked_driver::( arch, @@ -2096,7 +2092,7 @@ mod tests { let a = constant_minmax_matrix::<8>(1, k, max_a as i64, 1.0, 0.0); let b = constant_minmax_matrix::(1, k, max_b as i64, 1.0, 0.0); - let packed = PackedA::<_, 8, MR, PACK>::new(arch, a.as_view()).unwrap(); + let packed = PackedA::::new(a.as_view()).unwrap(); let mut c = [0.0]; driver::Drive::drive(&mut checked_driver::( arch, @@ -2224,7 +2220,7 @@ mod tests { 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::<_, A_BITS, 16, 2>::new(Scalar::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>( Scalar::new(), @@ -2270,7 +2266,7 @@ mod tests { meta.n = (k as i64 * code) as f32; row.set_meta(meta); } - let a = PackedA::<_, 8, 16, PACK>::new(arch, 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); @@ -2335,7 +2331,7 @@ mod tests { 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::<_, 8, 16, 2>::new(Scalar::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); } From 0977338a8508b5abac5b3b894dc95f924770bfe6 Mon Sep 17 00:00:00 2001 From: junkui chen Date: Thu, 22 Oct 2026 15:33:54 +0800 Subject: [PATCH 4/4] Simplify MinMax preparation and kernel contracts Clarify BitSlice decode_into test names without changing test behavior. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 74814a43-c576-493a-8783-585ef48c3f1c --- diskann-quantization/src/bits/mod.rs | 6 +- diskann-quantization/src/bits/slice.rs | 141 ++-- .../maxsim/packed_minmax_x_unpacked_minmax.rs | 776 +++++++++--------- 3 files changed, 475 insertions(+), 448 deletions(-) diff --git a/diskann-quantization/src/bits/mod.rs b/diskann-quantization/src/bits/mod.rs index 62763a5923..6b5b06734b 100644 --- a/diskann-quantization/src/bits/mod.rs +++ b/diskann-quantization/src/bits/mod.rs @@ -15,7 +15,7 @@ pub mod distances; pub use length::{Dynamic, Length, Static}; pub use ptr::{AsMutPtr, AsPtr, MutSlicePtr, SlicePtr}; pub use slice::{ - Binary, BitSlice, BitSliceBase, BitTranspose, BoxedBitSlice, ConstructionError, - DecodeLengthMismatch, Dense, EncodingError, GetError, IndexOutOfBounds, MutBitSlice, - PermutationStrategy, Representation, SetError, Unsigned, + Binary, BitSlice, BitSliceBase, BitTranspose, BoxedBitSlice, ConstructionError, Dense, + EncodingError, GetError, IndexOutOfBounds, MutBitSlice, PermutationStrategy, Representation, + SetError, Unsigned, }; diff --git a/diskann-quantization/src/bits/slice.rs b/diskann-quantization/src/bits/slice.rs index 98f8fbc430..49d0996a13 100644 --- a/diskann-quantization/src/bits/slice.rs +++ b/diskann-quantization/src/bits/slice.rs @@ -383,12 +383,11 @@ unsafe impl PermutationStrategy<4> for BitTranspose { // Errors // //////////// -/// The destination cannot hold exactly the logical decoded values. #[derive(Debug, Clone, Copy, PartialEq, Eq, Error)] #[error("decoded output has {output} elements, expected {input}")] -pub struct DecodeLengthMismatch { - pub input: usize, - pub output: usize, +pub(crate) struct DecodeLengthMismatch { + input: usize, + output: usize, } #[derive(Debug, Error, Clone, Copy)] @@ -740,64 +739,6 @@ where } } -impl BitSliceBase -where - Repr: Representation, - Ptr: AsPtr, - Len: Length, -{ - /// Decode in logical order, applying `convert` once to each value. - /// - /// Unused bits in the last byte are ignored. Unlike repeated `get` calls, bulk - /// decoding can reuse each packed byte across adjacent values. - /// - /// # Errors - /// - /// Returns [`DecodeLengthMismatch`] without writing or calling `convert` if the - /// output length differs from the logical input length. - #[inline] - pub fn decode_into( - &self, - output: &mut [T], - mut convert: impl FnMut(i64) -> T, - ) -> Result<(), DecodeLengthMismatch> { - const { assert!(NBITS > 0 && NBITS <= 8) }; - if output.len() != self.len() { - return Err(DecodeLengthMismatch { - input: self.len(), - output: output.len(), - }); - } - let packed = self.as_slice(); - let divisor = 1usize << NBITS.trailing_zeros(); - let codes_per_block = 8 / divisor; - let bytes_per_block = NBITS / divisor; - let blocks = output.len() / codes_per_block; - let (output, tail) = output.split_at_mut(blocks * codes_per_block); - let (packed, rest) = packed.split_at(blocks * bytes_per_block); - let mask = (1u64 << NBITS) - 1; - // Use the smallest byte-aligned group: 8 / gcd(NBITS, 8) codes. - for (block, bytes) in output - .chunks_exact_mut(codes_per_block) - .zip(packed.chunks_exact(bytes_per_block)) - { - let mut word = [0u8; 8]; - word[..bytes_per_block].copy_from_slice(bytes); - let word = u64::from_le_bytes(word); - for (i, out) in block.iter_mut().enumerate() { - *out = convert(Repr::decode(((word >> (i * NBITS)) & mask) as u8)); - } - } - let mut word = [0u8; 8]; - word[..rest.len()].copy_from_slice(rest); - let word = u64::from_le_bytes(word); - for (i, out) in tail.iter_mut().enumerate() { - *out = convert(Repr::decode(((word >> (i * NBITS)) & mask) as u8)); - } - Ok(()) - } -} - impl BitSliceBase, Perm, Len> where @@ -906,6 +847,59 @@ pub type PolyBitSlice = pub type BoxedBitSlice = PolyBitSlice; +impl BitSlice<'_, NBITS, Unsigned, Dense> +where + Unsigned: Representation, +{ + /// Decode and convert directly into the destination, ignoring unused trailing bits. + /// + /// # Errors + /// + /// Returns [`DecodeLengthMismatch`] without writing or calling `convert` if the + /// output length differs from the logical input length. + #[inline] + pub(crate) fn decode_into( + &self, + output: &mut [T], + mut convert: impl FnMut(u8) -> T, + ) -> Result<(), DecodeLengthMismatch> { + if output.len() != self.len() { + return Err(DecodeLengthMismatch { + input: self.len(), + output: output.len(), + }); + } + let packed = self.as_slice(); + let divisor = 1usize << NBITS.trailing_zeros(); + let codes_per_block = 8 / divisor; + let bytes_per_block = NBITS / divisor; + let blocks = output.len() / codes_per_block; + let (output, tail) = output.split_at_mut(blocks * codes_per_block); + let (packed, rest) = packed.split_at(blocks * bytes_per_block); + let mask = (1u64 << NBITS) - 1; + // Use the smallest byte-aligned group: 8 / gcd(NBITS, 8) codes. + for (block, bytes) in output + .chunks_exact_mut(codes_per_block) + .zip(packed.chunks_exact(bytes_per_block)) + { + let mut word = [0u8; 8]; + word[..bytes_per_block].copy_from_slice(bytes); + let word = u64::from_le_bytes(word); + for (i, out) in block.iter_mut().enumerate() { + *out = convert(((word >> (i * NBITS)) & mask) as u8); + } + } + let word = rest + .iter() + .enumerate() + .fold(0u64, |word, (i, &byte)| word | (u64::from(byte) << (8 * i))); + for (i, out) in tail.iter_mut().enumerate() { + *out = convert(((word >> (i * NBITS)) & mask) as u8); + } + Ok(()) + } +} + /////////////////////////////// // Special Cased Conversions // /////////////////////////////// @@ -1032,10 +1026,11 @@ mod tests { let mut guarded = vec![i64::MIN; len + 2]; let mut calls = 0; codes + .reborrow() .decode_into(&mut guarded[1..=len], |value| { - assert_eq!(value, expected[calls]); + assert_eq!(i64::from(value), expected[calls]); calls += 1; - value + i64::from(value) }) .unwrap(); assert_eq!(calls, len); @@ -1047,7 +1042,7 @@ mod tests { } #[test] - fn bulk_decoding_preserves_order_tails_and_output_bounds() { + fn test_decode_into_preserves_order_tails_and_bounds() { check_bulk_decoding::<1>(); check_bulk_decoding::<2>(); check_bulk_decoding::<3>(); @@ -1059,22 +1054,24 @@ mod tests { } #[test] - fn bulk_decoding_applies_representation_and_conversion() { - let codes = - BitSlice::<1, Binary, Dense, Static<3>>::new(&[0b101u8][..], Static::<3>).unwrap(); + fn test_decode_into_applies_conversion() { + let codes = BitSlice::<8, Unsigned>::new(&[0, 128, 255][..], 3).unwrap(); let mut output = [0i16; 3]; codes - .decode_into(&mut output, |value| (value * 3) as i16) + .decode_into(&mut output, |value| i16::from(value) * 3) .unwrap(); - assert_eq!(output, [3, -3, 3]); + assert_eq!(output, [0, 384, 765]); let codes = BoxedBitSlice::<4, Unsigned>::new_boxed(2); let mut output = [1u8; 2]; - codes.decode_into(&mut output, |value| value as u8).unwrap(); + codes + .reborrow() + .decode_into(&mut output, |value| value) + .unwrap(); assert_eq!(output, [0, 0]); } #[test] - fn bulk_decoding_length_mismatch_does_not_write_or_convert() { + fn test_decode_into_length_mismatch_does_not_write_or_convert() { let codes = BitSlice::<4, Unsigned>::new(&[0x21u8][..], 2).unwrap(); for len in [0, 1, 3] { let mut output = vec![99; len]; 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 f3cfa24278..d0173c5e41 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 @@ -3,22 +3,14 @@ * Licensed under the MIT license. */ -//! MaxSim for prepacked MinMax A and tiled MinMax B, with 1-8 bits per input code. +//! MaxSim for packed MinMax A and tiled MinMax B. //! -//! Each B tile is decoded once and reused across A panels. Integer inner products -//! are compensated before reduction because B's coefficients vary by vector. -//! B preparation retains only the three coefficients used by MaxSim. -//! B rows are zero-padded to PACK; kernel `k` remains the logical dimension used -//! for compensation and accumulator limits. +//! B tiles are decoded once and reused across A panels. Compensation precedes +//! reduction because coefficients differ per B vector. Codes are zero-padded to PACK; +//! compensation and accumulator limits use logical K. //! //! Outputs are unnegated maximum inner products. NaN candidates are ignored; //! empty B or all-NaN comparisons leave negative infinity. -//! -//! Scalar and x86 provide a general i16 path (PACK=2) and a byte path (PACK=4). -//! The byte path supports B up to six bits on V3 and seven on Scalar/V4. -//! Neon uses unsigned byte dots (PACK=4) for every bit width. -//! A packing depends only on A's bit width and the selected numeric layout, so it -//! can be reused across B formats supported by that layout. use std::{ num::NonZeroUsize, @@ -77,7 +69,7 @@ where Unsigned: Representation, { codes: BlockTransposed, - meta: Coefficients>, + meta: Vec>, } impl PackedA @@ -85,7 +77,7 @@ where T: Copy + Default + From, Unsigned: Representation, { - /// Pack A and copy its coefficients, zero-filling partial rows and K packs. + /// Zero-pad partial rows and K packs. /// /// # Errors /// @@ -94,18 +86,22 @@ where pub(crate) fn new(a: MatRef<'_, MinMaxMeta>) -> Result { let k = a.repr().intrinsic_dim(); let mut codes = BlockTransposed::try_new(a.num_vectors(), k)?; - let padded_rows = codes.padded_nrows(); - Overflow::check::(padded_rows, 1)?; - let mut meta = Coefficients { - a: vec![0.0; padded_rows], - b: vec![0.0; padded_rows], - n: vec![0.0; padded_rows], - }; + let panels = codes.num_blocks(); + Overflow::check::>(panels, 1)?; + let mut meta = vec![ + Coefficients { + a: [0.0; MR], + b: [0.0; MR], + n: [0.0; MR], + }; + panels + ]; for (i, row) in a.rows().enumerate() { let m = row.meta(); - meta.a[i] = m.a; - meta.b[i] = m.b; - meta.n[i] = m.n; + let coefficients = &mut meta[i / MR]; + coefficients.a[i % MR] = m.a; + coefficients.b[i % MR] = m.b; + coefficients.n[i % MR] = m.n; let v = row.vector(); let mut dst = codes .get_row_mut(i) @@ -118,13 +114,51 @@ where Ok(Self { codes, meta }) } - /// Cache-planning bytes for one padded A panel, including coefficients. + fn as_view(&self) -> Option> { + Some(AView { + codes: packed::View::from_block_transposed(self.codes.as_view())?, + meta: &self.meta, + k: DimK::new(NonZeroUsize::new(self.codes.ncols())?), + }) + } + fn panel_bytes(&self) -> Bytes { let codes = Elements::::new(MR * self.codes.padded_ncols()).bytes(); Bytes::new(codes.value() + size_of::>()) } } +#[derive(Clone, Copy)] +struct AView<'a, T, const MR: usize, const PACK: usize> { + codes: packed::View<'a, T, MR, PACK>, + meta: &'a [Coefficients<[f32; MR]>], + k: DimK, +} + +impl AView<'_, T, MR, PACK> { + #[inline] + fn visit_panels(&self, mut visitor: impl FnMut(APanel<'_, T, MR, PACK>, usize)) { + // SAFETY: PackedA constructs codes and coefficients with the same panel count and K. + unsafe { + self.codes.visit_panels(self.k, |codes, i| { + visitor( + APanel { + codes, + meta: &self.meta[i], + }, + i, + ); + }); + } + } +} + +#[derive(Clone, Copy)] +struct APanel<'a, T, const MR: usize, const PACK: usize> { + codes: packed::Panel<'a, T, MR, PACK>, + meta: &'a Coefficients<[f32; MR]>, +} + //--------// // Driver // //--------// @@ -141,11 +175,9 @@ pub(crate) struct Driver< Unsigned: Representation, { arch: W, - a: packed::View<'a, W::AElem, MR, PACK>, - a_meta: Coefficients<&'a [[f32; MR]]>, + a: AView<'a, W::AElem, MR, PACK>, b: MatRef<'a, MinMaxMeta>, c: &'a mut [f32], - k: DimK, b_rows_per_tile: NonZeroUsize, } @@ -162,69 +194,54 @@ where { /// # Panics /// - /// Panics if K exceeds the integer accumulator limit for these bit widths. - /// Conversion of a valid integer dot product to f32 may still round. - /// - /// B's bit width must fit the selected backend's [`ExtraWide::MAX_B_BITS`]; - /// incompatible combinations are rejected during monomorphization. + /// Panics if K exceeds the accumulator limit for these bit widths. + /// Valid integer dot products may still round when converted to f32. /// /// # Safety /// - /// A must be nonempty, both inputs must have logical dimension k, and C must have - /// exactly A's logical row count. + /// A must have nonzero row count and dimension, B must have the same dimension, + /// and C must have exactly A's logical row count. #[expect(clippy::expect_used)] pub(crate) unsafe fn new( arch: W, a: &'a PackedA, b: MatRef<'a, MinMaxMeta>, c: &'a mut [f32], - k: DimK, cache: Cache, ) -> Self where Unsigned: Representation, + Self: driver::Drive, { const { - assert!(NR > 0 && NR <= 8, "NR must be in 1..=8"); assert!( B_BITS <= W::MAX_B_BITS, "B bit width exceeds the backend's supported range" ) }; - let max_k = const { max_k::() }; + let rows = a.codes.nrows(); + let a_bytes = a.panel_bytes(); + let a = a.as_view().expect("A and K must be nonempty"); + let dim = a.k.value().get(); + bounds::check_eq!(bounds::Bound::new(b.repr().intrinsic_dim()), a.k); + bounds::check_eq!(bounds::Bound::new(c.len()), rows); + let max = const { max_k::() }; assert!( - k.value().get() <= max_k, - "contraction dimension {} exceeds the MinMax{A_BITS} x MinMax{B_BITS} limit of {max_k}", - k.value().get(), + dim <= max, + "contraction dimension {dim} exceeds the MinMax{A_BITS} x MinMax{B_BITS} limit of {max}", ); - bounds::check_eq!(bounds::Bound::new(a.codes.ncols()), k); - bounds::check_eq!(bounds::Bound::new(b.repr().intrinsic_dim()), k); - bounds::check_eq!(bounds::Bound::new(c.len()), a.codes.nrows()); - let a_meta = Coefficients { - a: a.meta.a.as_chunks::().0, - b: a.meta.b.as_chunks::().0, - n: a.meta.n.as_chunks::().0, - }; - let a_bytes = a.panel_bytes(); - let a = packed::View::from_block_transposed(a.codes.as_view()) - .expect("A and K must be nonempty"); - for meta in [a_meta.a, a_meta.b, a_meta.n] { - bounds::check_eq!(bounds::Bound::new(meta.len()), a.blocks().get()); - } let b_rows_per_tile = Params::new( cache, a_bytes, - BScratch::::row_bytes::(k), + BScratch::::row_bytes::(a.k), NR, ) .b_cols_in_l1; Self { arch, a, - a_meta, b, c, - k, b_rows_per_tile, } } @@ -235,6 +252,7 @@ impl + LoadStore, Unsigned: Representation, + for<'a> PanelKernel<'a, W, MR, NR, PACK>: driver::PanelKernel, { fn drive(&mut self) { self.arch.run( @@ -245,33 +263,26 @@ where let mut start = 0; while let Some(remaining) = NonZeroUsize::new(self.b.num_vectors() - start) { let count = remaining.min(self.b_rows_per_tile); - let b = - scratch.prepare::(self.b, start, count, self.k, W::convert_b); - // SAFETY: Driver's construction establishes all dimensions, metadata - // extents, and C's coverage of A. Visitors preserve these bounds. - unsafe { - self.a.visit_panels(self.k, |a, panel| { - let row = panel * MR; - let end = (row + MR).min(self.c.len()); - let c = self.arch.load(&self.c[row..end]); - let a_meta = Coefficients { - a: &self.a_meta.a[panel], - b: &self.a_meta.b[panel], - n: &self.a_meta.n[panel], - }; - let kernel = PanelKernel:: { - arch: self.arch, - a, - a_meta, - b: b.codes, - b_meta: b.meta, - c, - k: self.k, - }; - let c = self.arch.run(kernel); - self.arch.store(c, &mut self.c[row..end]); - }); - } + // SAFETY: The driver contract fixes B's K; remaining/count bound this tile. + let b = unsafe { + scratch.prepare::( + self.b, + start, + count, + self.a.k, + W::convert_b, + ) + }; + self.a.visit_panels(|a, panel| { + let row = panel * MR; + let end = (row + MR).min(self.c.len()); + let c = self.arch.load(&self.c[row..end]); + // SAFETY: A and B share K by the driver contract. + let kernel = + unsafe { PanelKernel::::new(self.arch, a, b, c) }; + let c = self.arch.run(kernel); + self.arch.store(c, &mut self.c[row..end]); + }); start += count.get(); } }, @@ -285,9 +296,47 @@ struct BScratch { meta: Vec>, } -struct BTile<'a, T> { +/// Codes and coefficients share row order; code rows are zero-padded to PACK. +#[derive(Clone, Copy)] +struct BTile<'a, T, const PACK: usize> { codes: unpacked::View<'a, T>, meta: &'a [Coefficients], + k: DimK, + physical_k: DimK, +} + +impl BTile<'_, T, PACK> { + /// # Safety + /// + /// `codes` must be the NR-row panel of this tile starting at `start`. + #[inline(always)] + unsafe fn bind_panel<'a, const NR: usize>( + &'a self, + codes: unpacked::Panel<'a, T, NR>, + start: usize, + ) -> BPanel<'a, T, NR, PACK> { + // SAFETY: Preparation gives each code row one coefficient record in the same order. + let meta = unsafe { + &*self + .meta + .as_ptr() + .add(start) + .cast::<[Coefficients; NR]>() + }; + BPanel { + codes, + meta, + k: self.k, + physical_k: self.physical_k, + } + } +} + +struct BPanel<'a, T, const NR: usize, const PACK: usize> { + codes: unpacked::Panel<'a, T, NR>, + meta: &'a [Coefficients; NR], + k: DimK, + physical_k: DimK, } #[inline] @@ -300,28 +349,31 @@ fn padded_k(k: DimK) -> DimK { } impl BScratch { - /// Convert logical codes once each; padding uses `T::default()` without conversion. + /// Padding bypasses `convert`. + /// + /// # Safety + /// + /// B's dimension must equal K, and `start..start + count` must be in bounds. #[expect(clippy::expect_used)] - fn prepare( + unsafe fn prepare( &mut self, b: MatRef<'_, MinMaxMeta>, start: usize, count: NonZeroUsize, k: DimK, mut convert: impl FnMut(u8) -> T, - ) -> BTile<'_, T> + ) -> BTile<'_, T, PACK> where Unsigned: Representation, { let dim = k.value().get(); - assert_eq!(b.repr().intrinsic_dim(), dim); - assert!(start <= b.num_vectors() && count.get() <= b.num_vectors() - start); 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() { - let row = b.get_row(start + i).expect("B tile range was checked"); + // 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, @@ -330,19 +382,19 @@ impl BScratch { }; let dst = &mut self.codes[i * stride..(i + 1) * stride]; row.vector() - .decode_into(&mut dst[..dim], |code| convert(code as u8)) + .decode_into(&mut dst[..dim], &mut convert) .expect("B output must match the logical code count"); dst[dim..].fill(T::default()); } - // SAFETY: Each row has physical_k initialized elements, including zero padding. - let codes = unsafe { unpacked::View::new(Slice::new(&self.codes), count, physical_k) }; BTile { - codes, + // SAFETY: The buffers contain count rows, each with stride initialized codes. + codes: unsafe { unpacked::View::new(Slice::new(&self.codes), count, physical_k) }, meta: &self.meta, + k, + physical_k, } } - /// Cache-planning bytes for one decoded B row, including metadata. fn row_bytes(k: DimK) -> Bytes { let codes = Elements::::new(padded_k::(k).value().get()).bytes(); Bytes::new(codes.value() + size_of::>()) @@ -356,92 +408,89 @@ impl BScratch { struct PanelKernel<'a, W: ExtraWide, const MR: usize, const NR: usize, const PACK: usize> { arch: W, - a: packed::Panel<'a, W::AElem, MR, PACK>, - a_meta: Coefficients<&'a [f32; MR]>, - b: unpacked::View<'a, W::BElem>, - b_meta: &'a [Coefficients], + a: APanel<'a, W::AElem, MR, PACK>, + b: BTile<'a, W::BElem, PACK>, c: [f32; MR], - k: DimK, } -impl, const MR: usize, const NR: usize, const PACK: usize> - Target for PanelKernel<'_, W, MR, NR, PACK> +impl<'a, W: ExtraWide, const MR: usize, const NR: usize, const PACK: usize> + PanelKernel<'a, W, MR, NR, PACK> { - #[inline(always)] - fn run(mut self, _: W) -> [f32; MR] { - driver::PanelKernel::panel_kernel(&mut self); - self.c + /// # Safety + /// + /// A's logical dimension must equal B's. + unsafe fn new( + arch: W, + a: APanel<'a, W::AElem, MR, PACK>, + b: BTile<'a, W::BElem, PACK>, + c: [f32; MR], + ) -> Self { + bounds::check_eq!(a.codes.k(), b.k); + Self { arch, a, b, c } } } -struct Visitor<'a, W: ExtraWide, const MR: usize, const PACK: usize> { - arch: W, - a: packed::Panel<'a, W::AElem, MR, PACK>, - a_meta: Coefficients<&'a [f32; MR]>, - b_meta: &'a [Coefficients], - c: &'a mut [f32; MR], - k: DimK, -} - -impl, const MR: usize, const NR: usize, const PACK: usize> - unpacked::PanelVisitor for Visitor<'_, W, MR, PACK> +impl, const MR: usize, const NR: usize, const PACK: usize, const N: usize> + unpacked::PanelVisitor for &mut PanelKernel<'_, W, MR, NR, PACK> { #[inline(always)] - #[expect(clippy::expect_used)] - fn visit(&mut self, b: unpacked::Panel<'_, W::BElem, NR>, start: usize) { + fn visit(&mut self, codes: unpacked::Panel<'_, W::BElem, N>, start: usize) { + let this = &mut **self; + // SAFETY: Both full and trailing panels come from this kernel's B tile. + let b = unsafe { this.b.bind_panel(codes, start) }; let mut kernel = MicroKernel { - arch: self.arch, - a: self.a, - a_meta: self.a_meta, + arch: this.arch, + a: this.a, b, - b_meta: self.b_meta[start..] - .first_chunk::() - .expect("B panel metadata is in bounds"), - c: self.c, - k: self.k, + c: &mut this.c, }; driver::MicroKernel::micro_kernel(&mut kernel); } } impl, const MR: usize, const NR: usize, const PACK: usize> - driver::PanelKernel for PanelKernel<'_, W, MR, NR, PACK> + Target for PanelKernel<'_, W, MR, NR, PACK> +where + Self: driver::PanelKernel, { #[inline(always)] - fn panel_kernel(&mut self) { - let mut visitor = Visitor { - arch: self.arch, - a: self.a, - a_meta: self.a_meta, - b_meta: self.b_meta, - c: &mut self.c, - k: self.k, - }; - // SAFETY: B's view is constructed with PACK-padded rows and matching metadata. - let tail = unsafe { - self.b.visit_panels::( - padded_k::(self.k), - Visitor { - c: &mut *visitor.c, - ..visitor - }, - ) - }; - if let Some(tail) = tail { - macro_rules! remainder { - ($($n:literal),+) => {$( - if $n < NR && let Some(b) = tail.try_as_panel::<$n>() { - unpacked::PanelVisitor::visit(&mut visitor, b, tail.start()); - return; - } - )+}; + fn run(mut self, _: W) -> [f32; MR] { + driver::PanelKernel::panel_kernel(&mut self); + self.c + } +} + +macro_rules! panel_kernel { + ($nr:literal, [$($tail:literal),*]) => { + impl, const MR: usize, const PACK: usize> + driver::PanelKernel for PanelKernel<'_, W, MR, $nr, PACK> + { + #[inline(always)] + fn panel_kernel(&mut self) { + let b = self.b; + let mut visitor = self; + // SAFETY: Preparation constructs B with physical_k elements per row. + let tail = unsafe { + b.codes.visit_panels::<$nr>(b.physical_k, &mut *visitor) + }; + if let Some(tail) = tail { + $( + if let Some(panel) = tail.try_as_panel::<$tail>() { + unpacked::PanelVisitor::visit(&mut visitor, panel, tail.start()); + return; + } + )* + unreachable!("all remainders below NR are handled"); + } } - remainder!(1, 2, 3, 4, 5, 6, 7); - unreachable!("NR <= 8 bounds the trailing panel to at most seven B vectors"); } - } + }; } +panel_kernel!(2, [1]); +panel_kernel!(4, [1, 2, 3]); +panel_kernel!(8, [1, 2, 3, 4, 5, 6, 7]); + //-------------// // MicroKernel // //-------------// @@ -449,16 +498,11 @@ impl, const MR: usize, const NR: usize, const PACK: usize struct MicroKernel<'a, W: ExtraWide, const MR: usize, const NR: usize, const PACK: usize> { arch: W, - a: packed::Panel<'a, W::AElem, MR, PACK>, - a_meta: Coefficients<&'a [f32; MR]>, - b: unpacked::Panel<'a, W::BElem, NR>, - b_meta: &'a [Coefficients; NR], + a: APanel<'a, W::AElem, MR, PACK>, + b: BPanel<'a, W::BElem, NR, PACK>, c: &'a mut [f32; MR], - k: DimK, } -/// Keep the hot contraction step in the caller's ISA context. -/// /// # Safety /// /// `a` must contain at least MR * PACK readable elements, and for every @@ -487,15 +531,13 @@ impl, const MR: usize, const NR: usize, const PACK: usize { #[inline(always)] fn micro_kernel(&mut self) { - bounds::check_eq!(self.a.k(), self.k); - let physical_k = padded_k::(self.k); - bounds::check_eq!(self.b.k(), physical_k); - let k = self.k.value().get(); + let physical_k = self.b.physical_k; + let k = self.b.k.value().get(); let mut acc = [self.arch.zero_acc(); NR]; - let ap = self.a.as_ptr(); - let bp = self.b.as_ptr(); - let astride = self.a.pack_stride(); - let bstride = self.b.stride(physical_k); + let ap = self.a.codes.as_ptr(); + let bp = self.b.codes.as_ptr(); + let astride = self.a.codes.pack_stride(); + let bstride = self.b.codes.stride(physical_k); let packs = physical_k.value().get() / PACK; for pack in 0..packs { @@ -510,21 +552,21 @@ impl, const MR: usize, const NR: usize, const PACK: usize ); } } - self.arch.max_into(acc, self.a_meta, self.b_meta, k, self.c); + let a_meta = Coefficients { + a: &self.a.meta.a, + b: &self.a.meta.b, + n: &self.a.meta.n, + }; + self.arch.max_into(acc, a_meta, self.b.meta, k, self.c); } } //----------// // Backends // //----------// -/// Backend operations for MR A rows and PACK contraction elements. Register -/// representation is private to each backend; NR only controls the number of accumulators. -/// -/// All implementations accept A codes in 0..=255. B's range is limited by -/// both its element type and any intermediate arithmetic in the dot product. +/// Backends accept A codes in 0..=255; B is limited by `MAX_B_BITS`. pub(crate) trait ExtraWide: Architecture { - /// Largest complete unsigned B width supported for every A code in 0..=255. - /// This is a backend capability, not the actual A or B input width. + /// Largest unsigned B width valid for every A code. const MAX_B_BITS: usize; type AElem: Copy + Default + From; @@ -546,8 +588,7 @@ pub(crate) trait ExtraWide: Architecture { /// Accumulate MR dot products of PACK elements without mixing rows. fn dot(self, acc: Self::Acc, a: Self::Wide, b: Self::Splat) -> Self::Acc; - /// Compensate in B order using logical `k`, then update scores without lowering existing - /// maxima. NaN candidates are ignored. + /// Compensate with logical K and update maxima, ignoring NaNs. fn max_into( self, acc: [Self::Acc; NR], @@ -607,14 +648,12 @@ unsafe fn load_pack(b: Slice<'_, T>) -> [T; PACK] { unsafe { b.as_ptr().cast::<[T; PACK]>().read_unaligned() } } -/// Pack two i16 bit patterns into the low and high halves of a u32. -/// /// # Safety /// /// B must contain two readable elements. #[inline(always)] unsafe fn pack_i16_pair(b: Slice<'_, i16>) -> u32 { - // SAFETY: The caller provides two readable elements, including any zero padding. + // SAFETY: The caller provides two readable elements. let [lo, hi] = unsafe { load_pack::<_, 2>(b) }; u32::from(lo as u16) | (u32::from(hi as u16) << 16) } @@ -662,7 +701,7 @@ impl ExtraWide<16, 4> for Scalar { #[inline(always)] unsafe fn splat(self, b: Slice<'_, i8>) -> Self::Splat { - // SAFETY: The caller provides four readable elements, including any zero padding. + // SAFETY: The caller provides four readable elements. let values = unsafe { load_pack::<_, 4>(b) }; let quad = u32::from_le_bytes(values.map(|x| x as u8)); u32x16::::splat(self, quad).reinterpret_simd() @@ -954,7 +993,7 @@ mod x86_64 { #[inline(always)] unsafe fn splat(self, b: Slice<'_, i8>) -> Self::Splat { - // SAFETY: The caller provides four readable elements, including any zero padding. + // SAFETY: The caller provides four readable elements. let values = unsafe { load_pack::<_, 4>(b) }; let quad = u32::from_le_bytes(values.map(|x| x as u8)); let half: i8x32 = u32x8::::splat(self, quad).reinterpret_simd(); @@ -1055,7 +1094,7 @@ mod x86_64 { #[inline(always)] unsafe fn splat(self, b: Slice<'_, i8>) -> Self::Splat { - // SAFETY: The caller provides four readable elements, including any zero padding. + // SAFETY: The caller provides four readable elements. let values = unsafe { load_pack::<_, 4>(b) }; let quad = u32::from_le_bytes(values.map(|x| x as u8)); u32x16::::splat(self, quad).reinterpret_simd() @@ -1207,7 +1246,7 @@ mod aarch64 { #[inline(always)] unsafe fn splat(self, b: Slice<'_, u8>) -> Self::Splat { - // SAFETY: The caller provides four readable elements, including any zero padding. + // SAFETY: The caller provides four readable elements. let values = unsafe { load_pack::<_, 4>(b) }; let quad = u32::from_le_bytes(values); let quarter = u32x4::::splat(self, quad).reinterpret_simd(); @@ -1230,6 +1269,7 @@ mod tests { matrix_kernels::maxsim::test::packed_x_unpacked_test_dims, minmax::{DataRef, MinMaxCompensation, MinMaxIP}, multi_vector::{Defaulted, Mat}, + test_util::Check, }; #[test] @@ -1306,13 +1346,16 @@ mod tests { fn assert_scores_close(actual: &[f32], expected: &[f32], context: std::fmt::Arguments<'_>) { assert_eq!(actual.len(), expected.len(), "{context}"); for (i, (&got, &want)) in actual.iter().zip(expected).enumerate() { + if got == want { + continue; + } assert!( - got == want - || (got.is_finite() - && want.is_finite() - && (got - want).abs() <= 1e-5 * want.abs().max(1.0)), + got.is_finite() && want.is_finite(), "{context}: row {i}: got {got}, expected {want}" ); + Check::absrel(1e-5, 1e-5) + .check(got, want) + .unwrap_or_else(|err| panic!("{context}: row {i}: {err}")); } } @@ -1333,14 +1376,13 @@ mod tests { where W: ExtraWide + LoadStore, Unsigned: Representation + Representation, + for<'b> PanelKernel<'b, W, MR, NR, PACK>: driver::PanelKernel, { - assert!(a.codes.nrows() > 0); + 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()); - let k = DimK::new(NonZeroUsize::new(a.codes.ncols()).unwrap()); - // SAFETY: Nonzero A/K and matching dimensions/output were checked above; - // PackedA's constructor establishes metadata extents and padding. - unsafe { Driver::::new(arch, a, b, c, k, Cache::detect()) } + // SAFETY: Nonempty A/K, matching dimensions and output length were checked above. + unsafe { Driver::::new(arch, a, b, c, Cache::detect()) } } fn assert_driver_contract< @@ -1354,6 +1396,7 @@ mod tests { ) where W: ExtraWide + LoadStore, Unsigned: Representation, + for<'a> PanelKernel<'a, W, MR, NR, PACK>: driver::PanelKernel, { let mut rng = StdRng::seed_from_u64(0x63de91a0ba5c8742); for shape in packed_x_unpacked_test_dims(MR, NR) { @@ -1435,6 +1478,7 @@ mod tests { ) where W: ExtraWide + LoadStore, Unsigned: Representation, + for<'a> PanelKernel<'a, W, MR, NR, PACK>: driver::PanelKernel, { let a = constant_minmax_matrix::<8>(MR + 1, 3, 1, 1.0, 0.0); let mut b = constant_minmax_matrix::(NR + 1, 3, (1 << B_BITS) - 1, 1.0, 0.0); @@ -1495,6 +1539,7 @@ mod tests { #[test] fn scalar_driver_obeys_maxsim_contract() { assert_driver_contract::(Scalar::new()); + assert_driver_target_shapes::(Scalar::new()); } #[cfg(target_arch = "x86_64")] @@ -1502,6 +1547,7 @@ mod tests { fn avx2_driver_obeys_maxsim_contract() { if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { assert_driver_contract::<_, 4, 16, 4, 4>(arch); + assert_driver_target_shapes::<_, 4, 16, 4, 4>(arch); } else { eprintln!("AVX2 unavailable; backend not executed"); } @@ -1512,6 +1558,7 @@ mod tests { fn avx512_driver_obeys_maxsim_contract() { if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked() { assert_driver_contract::<_, 4, 16, 8, 4>(arch); + assert_driver_target_shapes::<_, 4, 16, 8, 4>(arch); } else { eprintln!("V4 unavailable; backend not executed"); } @@ -1522,6 +1569,7 @@ mod tests { fn neon_driver_obeys_maxsim_contract() { if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { assert_driver_contract::<_, 4, 16, 4, 4>(arch); + assert_driver_target_shapes::<_, 4, 16, 4, 4>(arch); } else { eprintln!("Neon/dotprod unavailable; backend not executed"); } @@ -1533,6 +1581,7 @@ mod tests { where W: ExtraWide<8, PACK> + LoadStore, Unsigned: Representation, + for<'a> PanelKernel<'a, W, 8, NR, PACK>: driver::PanelKernel, { assert_driver_contract::(arch); assert_microkernel_nan_reduction::(arch); @@ -1559,7 +1608,7 @@ mod tests { let packed = PackedA::::new(a.as_view()).unwrap(); assert_eq!(packed.codes.nrows(), m); assert_eq!(packed.codes.ncols(), k); - assert!(packed::View::from_block_transposed(packed.codes.as_view()).is_none()); + assert!(packed.as_view().is_none()); } let mut rng = StdRng::seed_from_u64(0x0ae4b97f2c186d35); let a = random_minmax_matrix::(MR + 1, 5, &mut rng); @@ -1583,10 +1632,26 @@ mod tests { assert_eq!(packed.codes.as_slice()[index], expected); } let meta = a.get_row(row).map(|v| v.meta()).unwrap_or_default(); - assert_eq!(packed.meta.a[row], meta.a); - assert_eq!(packed.meta.b[row], meta.b); - assert_eq!(packed.meta.n[row], meta.n); - } + let panel = &packed.meta[row / MR]; + assert_eq!(panel.a[row % MR], meta.a); + assert_eq!(panel.b[row % MR], meta.b); + assert_eq!(panel.n[row % MR], meta.n); + } + let mut visited = 0; + packed.as_view().unwrap().visit_panels(|panel, i| { + assert_eq!(i, visited); + assert!(std::ptr::eq(panel.meta, &packed.meta[i])); + assert_eq!( + panel.codes.as_ptr().as_ptr(), + packed + .codes + .as_slice() + .as_ptr() + .wrapping_add(i * MR * padded_cols), + ); + visited += 1; + }); + assert_eq!(visited, packed.meta.len()); } #[test] @@ -1603,41 +1668,21 @@ mod tests { #[test] fn cache_footprints_match_prepared_layouts() { - for (k, a_bytes, b_bytes) in [(256, 4288, 268), (257, 4352, 272)] { - let dim = DimK::new(NonZeroUsize::new(k).unwrap()); - assert_eq!(BScratch::::row_bytes::<4>(dim).value(), b_bytes); - assert_eq!(BScratch::::row_bytes::<4>(dim).value(), b_bytes); - for rows in [1, 16, 17] { - let a = constant_minmax_matrix::<8>(rows, k, 0, 1.0, 0.0); - let packed = PackedA::::new(a.as_view()).unwrap(); - assert_eq!(packed.panel_bytes().value(), a_bytes); - - // Preserve the previous view-based cache estimate. - let view = packed::View::from_block_transposed(packed.codes.as_view()).unwrap(); - assert_eq!( - packed.panel_bytes().value(), - view.block_stride(dim).bytes().value() + 3 * 16 * size_of::() - ); - } - } - } - - #[test] - fn widened_cache_footprints_match_prepared_layouts() { - for (k, a_bytes, b_bytes) in [(256, 8384, 524), (257, 8448, 528)] { - let dim = DimK::new(NonZeroUsize::new(k).unwrap()); - assert_eq!(BScratch::::row_bytes::<2>(dim).value(), b_bytes); - for rows in [1, 16, 17] { - let a = constant_minmax_matrix::<8>(rows, k, 255, 1.0, 0.0); - let packed = PackedA::::new(a.as_view()).unwrap(); - assert_eq!(packed.panel_bytes().value(), a_bytes); - let view = packed::View::from_block_transposed(packed.codes.as_view()).unwrap(); - assert_eq!( - packed.panel_bytes().value(), - view.block_stride(dim).bytes().value() + 3 * 16 * size_of::(), - ); + fn check, const PACK: usize>( + cases: [(usize, usize, usize); 2], + ) { + for (k, a_bytes, b_bytes) in cases { + let dim = DimK::new(NonZeroUsize::new(k).unwrap()); + assert_eq!(BScratch::::row_bytes::(dim).value(), b_bytes); + for rows in [1, 16, 17] { + let a = constant_minmax_matrix::<8>(rows, k, 0, 1.0, 0.0); + let packed = PackedA::::new(a.as_view()).unwrap(); + assert_eq!(packed.panel_bytes().value(), a_bytes); + } } } + check::([(256, 4288, 268), (257, 4352, 272)]); + check::([(256, 8384, 524), (257, 8448, 528)]); } #[test] @@ -1647,7 +1692,6 @@ mod tests { 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 k = DimK::new(NonZeroUsize::new(3).unwrap()); // This budget gives different rounded tiles for compact and canonical metadata. let cache = Cache::new( NonZeroUsize::new(packed.panel_bytes().value() + 81).unwrap(), @@ -1657,7 +1701,7 @@ mod tests { 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, k, cache) + Driver::<_, 8, 8, 2, 2>::new::<8>(arch, &packed, b.as_view(), &mut scores, cache) }; assert_eq!(driver.b_rows_per_tile.get(), 6); driver::Drive::drive(&mut driver); @@ -1677,12 +1721,56 @@ mod tests { #[test] fn b_preparation_preserves_metadata_and_zero_padding() { - assert_b_preparation::(-1, |code| code as i8); + macro_rules! check { + ($($bits:literal),+) => {$( + assert_b_preparation::(-1, |code| code as i8); + assert_b_preparation::(-1, i16::from); + assert_b_preparation::(u8::MAX, |code| code); + )+}; + } + check!(1, 2, 3, 4, 5, 6, 7); assert_b_preparation::(-1, i16::from); assert_b_preparation::(u8::MAX, |code| code); assert_b_preparation::(u16::MAX, |code| u16::from(code) + 1); } + #[test] + fn b_panels_bind_coefficients_at_full_and_tail_offsets() { + let k = DimK::new(NonZeroUsize::new(3).unwrap()); + let mut rng = StdRng::seed_from_u64(0xda7bc35e83012e29); + let source = random_minmax_matrix::<4>(5, k.value().get(), &mut rng); + let mut scratch = BScratch::::default(); + // SAFETY: The tile covers all five rows of source, whose dimension is K. + let tile = unsafe { + scratch.prepare::<4, 4>( + source.as_view(), + 0, + NonZeroUsize::new(5).unwrap(), + k, + |code| code, + ) + }; + let mut actual = Vec::new(); + let visit = |codes: unpacked::Panel<'_, u8, 2>, start| { + // SAFETY: The tile's traversal supplies this panel and its offset. + let panel = unsafe { tile.bind_panel(codes, start) }; + actual.extend(panel.meta.iter().map(|m| (m.a, m.b, m.n))); + }; + // SAFETY: Preparation records the view's padded dimension. + let tail = unsafe { tile.codes.visit_panels::<2>(tile.physical_k, visit) }.unwrap(); + // SAFETY: The remainder comes from the same tile and contains one row. + let panel = unsafe { tile.bind_panel(tail.try_as_panel::<1>().unwrap(), tail.start()) }; + actual.extend(panel.meta.iter().map(|m| (m.a, m.b, m.n))); + let expected: Vec<_> = source + .rows() + .map(|row| { + let m = row.meta(); + (m.a, m.b, m.n) + }) + .collect(); + assert_eq!(actual, expected); + } + fn assert_b_preparation( sentinel: T, convert: impl Fn(u8) -> T, @@ -1691,7 +1779,7 @@ mod tests { Unsigned: Representation, { let mut scratch = BScratch::::default(); - for dim in [PACK, PACK + 1] { + for dim in [1, PACK - 1, PACK, PACK + 1, 15, 16, 17] { let k = DimK::new(NonZeroUsize::new(dim).unwrap()); let stride = crate::minmax::Data::::canonical_bytes(dim); // Unaligned storage and dirty high bits exercise the packed input boundary. @@ -1722,21 +1810,24 @@ mod tests { scratch.codes.fill(sentinel); scratch.meta.fill(Coefficients::default()); let mut conversions = 0; - let tile = scratch.prepare::( - b, - start, - NonZeroUsize::new(count).unwrap(), - k, - |code| { - conversions += 1; - convert(code) - }, - ); + // SAFETY: The selected range is within B, which has dimension K. + let tile = unsafe { + scratch.prepare::( + b, + start, + NonZeroUsize::new(count).unwrap(), + k, + |code| { + conversions += 1; + convert(code) + }, + ) + }; assert_eq!(conversions, count * dim); let padded = padded_k::(k); let row_len = padded.value().get(); - // SAFETY: prepare constructs the view with this padded dimension. - let codes = unsafe { tile.codes.as_std_slice(padded) }; + // SAFETY: Preparation constructs the view with the recorded padded dimension. + let codes = unsafe { tile.codes.as_std_slice(tile.physical_k) }; assert_eq!(codes.len(), count * row_len); assert_eq!(tile.meta.len(), count); for i in 0..count { @@ -1775,6 +1866,7 @@ mod tests { ) where W: ExtraWide + LoadStore, Unsigned: Representation, + for<'a> PanelKernel<'a, W, MR, NR, PACK>: driver::PanelKernel, { let mut rng = StdRng::seed_from_u64(0x8f61d9c34a257be0); let mut shapes: Vec<_> = (1..NR).map(|n| (MR + 1, n, PACK + 1)).collect(); @@ -1801,41 +1893,6 @@ mod tests { } } - #[test] - fn scalar_driver_matches_reference_at_target_shapes() { - assert_driver_target_shapes::<_, 4, 16, 2, 4>(Scalar::new()); - } - - #[cfg(target_arch = "aarch64")] - #[test] - fn neon_driver_matches_reference_at_target_shapes() { - if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { - assert_driver_target_shapes::<_, 4, 16, 4, 4>(arch); - } else { - eprintln!("Neon/dotprod unavailable; target-shape regression not executed"); - } - } - - #[cfg(target_arch = "x86_64")] - #[test] - fn avx2_driver_matches_reference_at_target_shapes() { - if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { - assert_driver_target_shapes::<_, 4, 16, 4, 4>(arch); - } else { - eprintln!("AVX2 unavailable; target-shape regression not executed"); - } - } - - #[cfg(target_arch = "x86_64")] - #[test] - fn avx512_driver_matches_reference_at_target_shapes() { - if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked() { - assert_driver_target_shapes::<_, 4, 16, 8, 4>(arch); - } else { - eprintln!("V4 unavailable; target-shape regression not executed"); - } - } - fn assert_microkernel_nan_reduction( arch: W, ) where @@ -1848,9 +1905,9 @@ mod tests { let mut n = [0.0; MR]; n[MR - 1] = f32::NAN; let a_meta = Coefficients { - a: &[1.0; MR], - b: &[1.0; MR], - n: &n, + a: [1.0; MR], + b: [1.0; MR], + n, }; for (initial, score) in [ (f32::NEG_INFINITY, f32::NAN), @@ -1869,16 +1926,21 @@ mod tests { let mut c = [initial; MR]; let mut kernel = MicroKernel:: { arch, - // SAFETY: A is one full MR panel, padded to PACK for logical K=1. - a: unsafe { packed::Panel::new(Slice::new(&a_codes), k) }, - a_meta, - // SAFETY: B contains NR rows of PACK initialized elements. - b: unsafe { - unpacked::Panel::new(Slice::new(&b_codes), padded_k::(k)) + a: APanel { + // SAFETY: A is one full MR panel, padded to PACK for logical K=1. + codes: unsafe { packed::Panel::new(Slice::new(&a_codes), k) }, + meta: &a_meta, + }, + b: BPanel { + // SAFETY: B contains NR rows of PACK initialized codes. + codes: unsafe { + unpacked::Panel::new(Slice::new(&b_codes), padded_k::(k)) + }, + meta: &b_meta, + k, + physical_k: padded_k::(k), }, - b_meta: &b_meta, c: &mut c, - k, }; driver::MicroKernel::micro_kernel(&mut kernel); let mut expected = [initial.max(score); MR]; @@ -2050,85 +2112,51 @@ mod tests { }); } - fn assert_b_width< - W, - const B_BITS: usize, - const MR: usize, - const NR: usize, - const PACK: usize, - >( - arch: W, - rng: &mut StdRng, - ) where - W: ExtraWide + LoadStore, - Unsigned: Representation, - { - let max_a = u64::from(u8::MAX); + fn assert_dimension_limit() { + let max_a = (1u64 << A_BITS) - 1; let max_b = (1u64 << B_BITS) - 1; - let limit = max_k::<8, B_BITS>() as u64; + let limit = max_k::() as u64; assert!(limit * max_a * max_b <= i32::MAX as u64); assert!((limit + 1) * max_a * max_b > i32::MAX as u64); - - let (m, n, k) = (MR + 1, NR + 1, PACK + 1); - let a = random_minmax_matrix::<8>(m, k, rng); - let b = random_minmax_matrix::(n, k, rng); - let expected = reference_max_inner_products::<8, B_BITS>(a.as_view(), b.as_view()); - let packed = PackedA::::new(a.as_view()).unwrap(); - let mut guarded = vec![654321.0; m + 2]; - let mut kernel = checked_driver::( - arch, - &packed, - b.as_view(), - &mut guarded[1..=m], - ); - driver::Drive::drive(&mut kernel); - assert_scores_close( - kernel.c, - &expected, - format_args!("8x{B_BITS}: M={m}, N={n}, K={k}, PACK={PACK}"), - ); - assert_eq!(guarded[0], 654321.0); - assert_eq!(guarded[m + 1], 654321.0); - - let a = constant_minmax_matrix::<8>(1, k, max_a as i64, 1.0, 0.0); - let b = constant_minmax_matrix::(1, k, max_b as i64, 1.0, 0.0); - let packed = PackedA::::new(a.as_view()).unwrap(); - let mut c = [0.0]; - driver::Drive::drive(&mut checked_driver::( - arch, - &packed, - b.as_view(), - &mut c, - )); - assert_eq!(c, [(k as u64 * max_a * max_b) as f32]); } - macro_rules! check_widths { - ($arch:expr, $mr:literal, $nr:literal, $pack:literal; $($b:literal),+ $(,)?) => {{ - let arch = $arch; - let mut rng = StdRng::seed_from_u64(0x948cd3e02a671fb5); - assert_backend_code_range::<_, $mr, $pack>(arch); - $( - assert_b_width::<_, $b, $mr, $nr, $pack>(arch, &mut rng); - )+ - }}; + #[test] + fn dimension_limits_cover_all_input_widths() { + fn check_a() { + assert_dimension_limit::(); + assert_dimension_limit::(); + assert_dimension_limit::(); + assert_dimension_limit::(); + assert_dimension_limit::(); + assert_dimension_limit::(); + assert_dimension_limit::(); + assert_dimension_limit::(); + } + check_a::<1>(); + check_a::<2>(); + check_a::<3>(); + check_a::<4>(); + check_a::<5>(); + check_a::<6>(); + check_a::<7>(); + check_a::<8>(); } #[test] - fn scalar_supports_all_b_widths_and_layouts() { - check_widths!(Scalar::new(), 16, 2, 2; 1, 2, 3, 4, 5, 6, 7, 8); - check_widths!(Scalar::new(), 16, 2, 4; 1, 2, 3, 4, 5, 6, 7); - check_widths!(Scalar::new(), 8, 2, 2; 1, 2, 3, 4, 5, 6, 7, 8); + fn scalar_backends_cover_supported_code_ranges() { + assert_backend_code_range::<_, 16, 2>(Scalar::new()); + assert_backend_code_range::<_, 16, 4>(Scalar::new()); + assert_backend_code_range::<_, 8, 2>(Scalar::new()); } #[cfg(target_arch = "x86_64")] #[test] - fn avx2_supports_all_b_widths_and_layouts() { + fn avx2_backends_cover_supported_code_ranges() { if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { - check_widths!(arch, 16, 4, 2; 1, 2, 3, 4, 5, 6, 7, 8); - check_widths!(arch, 16, 4, 4; 1, 2, 3, 4, 5, 6); - check_widths!(arch, 8, 4, 2; 1, 2, 3, 4, 5, 6, 7, 8); - check_widths!(arch, 8, 4, 4; 1, 2, 3, 4, 5, 6); + assert_backend_code_range::<_, 16, 2>(arch); + assert_backend_code_range::<_, 16, 4>(arch); + assert_backend_code_range::<_, 8, 2>(arch); + assert_backend_code_range::<_, 8, 4>(arch); } else { eprintln!("AVX2 unavailable; all-width regression not executed"); } @@ -2136,10 +2164,10 @@ mod tests { #[cfg(target_arch = "x86_64")] #[test] - fn avx512_supports_all_b_widths_and_layouts() { + fn avx512_backends_cover_supported_code_ranges() { if let Some(arch) = diskann_wide::arch::x86_64::V4::new_checked() { - check_widths!(arch, 16, 8, 2; 1, 2, 3, 4, 5, 6, 7, 8); - check_widths!(arch, 16, 8, 4; 1, 2, 3, 4, 5, 6, 7); + assert_backend_code_range::<_, 16, 2>(arch); + assert_backend_code_range::<_, 16, 4>(arch); } else { eprintln!("V4 unavailable; all-width regression not executed"); } @@ -2147,10 +2175,10 @@ mod tests { #[cfg(target_arch = "aarch64")] #[test] - fn neon_supports_all_b_widths_and_layouts() { + fn neon_backends_cover_supported_code_ranges() { if let Some(arch) = diskann_wide::arch::aarch64::Neon::new_checked() { - check_widths!(arch, 16, 4, 4; 1, 2, 3, 4, 5, 6, 7, 8); - check_widths!(arch, 8, 4, 4; 1, 2, 3, 4, 5, 6, 7, 8); + assert_backend_code_range::<_, 16, 4>(arch); + assert_backend_code_range::<_, 8, 4>(arch); } else { eprintln!("Neon/dotprod unavailable; all-width regression not executed"); } @@ -2241,6 +2269,7 @@ mod tests { fn check_backend(arch: W) where W: ExtraWide<16, PACK> + LoadStore, + for<'a> PanelKernel<'a, W, 16, NR, PACK>: driver::PanelKernel, { assert_driver_contract::(arch); assert_driver_target_shapes::(arch); @@ -2250,6 +2279,7 @@ mod tests { fn assert_unsigned_byte_products(arch: W) where W: ExtraWide<16, PACK> + LoadStore, + for<'a> PanelKernel<'a, W, 16, NR, PACK>: driver::PanelKernel, { if cfg!(miri) { return;