From c713e51febd15c78fe1ca170d8ff6d54759aff27 Mon Sep 17 00:00:00 2001 From: Suryansh Gupta Date: Fri, 14 Aug 2026 02:51:58 +0530 Subject: [PATCH 1/3] New Multi Vector Kernel desing --- .../src/multi_vector/distance/factory.rs | 33 +- .../src/multi_vector/distance/kernels/f16.rs | 207 ++++- .../multi_vector/distance/kernels/f32/mod.rs | 135 --- .../distance/kernels/f32/scalar.rs | 113 --- .../multi_vector/distance/kernels/f32/v3.rs | 112 --- .../multi_vector/distance/kernels/float.rs | 516 +++++++++++ .../multi_vector/distance/kernels/layouts.rs | 219 ----- .../distance/kernels/leaves/mod.rs | 21 + .../distance/kernels/leaves/scalar.rs | 129 +++ .../distance/kernels/leaves/v3.rs | 124 +++ .../src/multi_vector/distance/kernels/mod.rs | 351 ++++++-- .../multi_vector/distance/kernels/reduce.rs | 106 --- .../multi_vector/distance/kernels/strip.rs | 104 +++ .../distance/kernels/tiled_reduce.rs | 806 ------------------ .../multi_vector/distance/kernels/tiles.rs | 341 ++++++++ 15 files changed, 1700 insertions(+), 1617 deletions(-) delete mode 100644 diskann-quantization/src/multi_vector/distance/kernels/f32/mod.rs delete mode 100644 diskann-quantization/src/multi_vector/distance/kernels/f32/scalar.rs delete mode 100644 diskann-quantization/src/multi_vector/distance/kernels/f32/v3.rs create mode 100644 diskann-quantization/src/multi_vector/distance/kernels/float.rs delete mode 100644 diskann-quantization/src/multi_vector/distance/kernels/layouts.rs create mode 100644 diskann-quantization/src/multi_vector/distance/kernels/leaves/mod.rs create mode 100644 diskann-quantization/src/multi_vector/distance/kernels/leaves/scalar.rs create mode 100644 diskann-quantization/src/multi_vector/distance/kernels/leaves/v3.rs delete mode 100644 diskann-quantization/src/multi_vector/distance/kernels/reduce.rs create mode 100644 diskann-quantization/src/multi_vector/distance/kernels/strip.rs delete mode 100644 diskann-quantization/src/multi_vector/distance/kernels/tiled_reduce.rs create mode 100644 diskann-quantization/src/multi_vector/distance/kernels/tiles.rs diff --git a/diskann-quantization/src/multi_vector/distance/factory.rs b/diskann-quantization/src/multi_vector/distance/factory.rs index 5dcd4b8cd1..0539793acd 100644 --- a/diskann-quantization/src/multi_vector/distance/factory.rs +++ b/diskann-quantization/src/multi_vector/distance/factory.rs @@ -16,8 +16,7 @@ use diskann_wide::arch::x86_64::{V3, V4}; use super::isa::{MaxSimIsa, NotSupported}; use super::kernel::{Erase, MaxSimKernel}; -use super::kernels::f16::F16Entry; -use super::kernels::f32::F32Kernel; +use super::kernels::{MaxIp, MaxIpF16}; use super::max_sim::{MaxSim, MaxSimError}; use crate::multi_vector::distance::QueryMatRef; use crate::multi_vector::{BlockTransposed, BlockTransposedRef, Mat, MatRef, Standard}; @@ -35,7 +34,7 @@ struct Prepared { impl MaxSimKernel for Prepared> where A: Architecture, - F32Kernel: for<'a> diskann_wide::arch::Target3< + MaxIp: for<'a> diskann_wide::arch::Target3< A, (), BlockTransposedRef<'a, f32, GROUP>, @@ -59,14 +58,11 @@ where scores.fill(f32::MAX); return Ok(()); } - let mut scratch = vec![f32::MIN; self.prepared.padded_nrows()]; - self.arch.run3( - F32Kernel::, - self.prepared.reborrow(), - doc, - &mut scratch, - ); - for (dst, &src) in scores.iter_mut().zip(&scratch[..self.prepared.nrows()]) { + // `run` seeds the max itself, so the fill value here is arbitrary. + let mut state = vec![0.0; self.prepared.padded_nrows()]; + self.arch + .run3(MaxIp, self.prepared.reborrow(), doc, &mut state); + for (dst, &src) in scores.iter_mut().zip(&state[..self.prepared.nrows()]) { *dst = -src; } Ok(()) @@ -77,7 +73,7 @@ impl MaxSimKernel for Prepared> where A: Architecture, - F16Entry: for<'a> diskann_wide::arch::Target3< + MaxIpF16: for<'a> diskann_wide::arch::Target3< A, (), BlockTransposedRef<'a, half::f16, GROUP>, @@ -101,14 +97,11 @@ where scores.fill(f32::MAX); return Ok(()); } - let mut scratch = vec![f32::MIN; self.prepared.padded_nrows()]; - self.arch.run3( - F16Entry::, - self.prepared.reborrow(), - doc, - &mut scratch, - ); - for (dst, &src) in scores.iter_mut().zip(&scratch[..self.prepared.nrows()]) { + // `run` seeds the max itself, so the fill value here is arbitrary. + let mut state = vec![0.0; self.prepared.padded_nrows()]; + self.arch + .run3(MaxIpF16, self.prepared.reborrow(), doc, &mut state); + for (dst, &src) in scores.iter_mut().zip(&state[..self.prepared.nrows()]) { *dst = -src; } Ok(()) diff --git a/diskann-quantization/src/multi_vector/distance/kernels/f16.rs b/diskann-quantization/src/multi_vector/distance/kernels/f16.rs index a535c68dcf..febf4d962a 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/f16.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/f16.rs @@ -1,52 +1,197 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT license. -//! f16 dispatch adapter for block-transposed multi-vector distance. +//! f16 MaxSim. //! -//! Reuses the f32 micro-kernel family with tile-level f16→f32 conversion -//! via [`ConvertTo`](super::layouts::ConvertTo). No f16-specific micro-kernel -//! code is needed — the [`F32Kernel`](super::f32::F32Kernel) does all the -//! SIMD work after conversion. -//! -//! Conversion from f16 to f32 is performed at tile granularity via -//! [`SliceCast`](diskann_vector::conversion::SliceCast), dispatched through -//! the runtime architecture token — the same SIMD level used by the -//! micro-kernel. +//! There is no f16 leaf: f16 widens to f32 and reuses the f32 pipeline. Both sides widen a +//! tile at a time into a buffer the walk reuses, which is what the lending [`TileWalk`] +//! exists for — the whole query never has to be staged at once, and the staged copy stays +//! inside the cache level its tile was sized for. +use diskann_vector::conversion::SliceCast; use diskann_wide::Architecture; +#[cfg(target_arch = "x86_64")] +use diskann_wide::arch::x86_64::V3; +use diskann_wide::arch::{Scalar, Target2}; -use super::Kernel; -use super::TileBudget; -use super::f32::{F32Kernel, max_ip_kernel}; -use super::layouts; +use super::leaves::scalar::{A_PANEL as SC_A, B_PANEL as SC_B}; +#[cfg(target_arch = "x86_64")] +use super::leaves::v3::{A_PANEL as V3_A, B_PANEL as V3_B}; +use super::tiles::{Cursor, DocTile, QueryTile}; +use super::{Plan, TileAt, TileBudget, TileWalk, float}; use crate::multi_vector::{BlockTransposedRef, MatRef, Standard}; -pub(crate) struct F16Entry; +/// Stages one source tile at a time as f32. +struct Widen<'a, Arch> { + arch: Arch, + cursor: Cursor<'a, half::f16>, + buf: Vec, + k: usize, +} + +impl<'a, Arch: Architecture> Widen<'a, Arch> +where + SliceCast: for<'x> Target2, +{ + /// # Panics + /// + /// Panics if `k` is zero — the entry guards that case before any walk is built. + fn new(arch: Arch, src: &'a [half::f16], k: usize, stride: usize) -> Self { + assert!(k > 0, "widening walk requires a non-empty contraction"); + let cursor = Cursor::new(src, stride); + let buf = vec![0.0f32; cursor.widest()]; + Self { + arch, + cursor, + buf, + k, + } + } + + fn next(&mut self) -> Option<&[f32]> { + let arch = self.arch; + let src = self.cursor.next()?; + let len = src.len(); + arch.run2(SliceCast::new(), &mut self.buf[..len], src); + Some(&self.buf[..len]) + } +} + +/// Widens the padded storage of a block-transposed f16 query. +/// +/// Widening is element-wise, so it preserves the block-transposed permutation. +struct QueryWiden<'a, Arch, const AR: usize>(Widen<'a, Arch>); + +impl<'a, Arch: Architecture, const AR: usize> QueryWiden<'a, Arch, AR> +where + SliceCast: for<'x> Target2, +{ + fn new(arch: Arch, view: BlockTransposedRef<'a, half::f16, AR>, a_panels: usize) -> Self { + let k = view.padded_ncols(); + Self(Widen::new(arch, view.as_slice(), k, a_panels * AR * k)) + } +} + +impl<'t, Arch, const AR: usize> TileAt<'t> for QueryWiden<'_, Arch, AR> { + type Tile = QueryTile<'t, f32, AR>; +} + +impl TileWalk for QueryWiden<'_, Arch, AR> +where + SliceCast: for<'x> Target2, +{ + fn next(&mut self) -> Option> { + let k = self.0.k; + self.0.next().map(|data| QueryTile::new(data, k)) + } + + fn reset(&mut self) { + self.0.cursor.reset(); + } +} + +/// Widens a row-major f16 doc matrix. +struct DocWiden<'a, Arch, const BR: usize>(Widen<'a, Arch>); + +impl<'a, Arch: Architecture, const BR: usize> DocWiden<'a, Arch, BR> +where + SliceCast: for<'x> Target2, +{ + fn new(arch: Arch, docs: MatRef<'a, Standard>, b_panels: usize) -> Self { + let k = docs.vector_dim(); + Self(Widen::new(arch, docs.as_slice(), k, b_panels * BR * k)) + } +} + +impl<'t, Arch, const BR: usize> TileAt<'t> for DocWiden<'_, Arch, BR> { + type Tile = DocTile<'t, f32, BR>; +} + +impl TileWalk for DocWiden<'_, Arch, BR> +where + SliceCast: for<'x> Target2, +{ + fn next(&mut self) -> Option> { + let k = self.0.k; + self.0.next().map(|data| DocTile::new(data, k)) + } + + fn reset(&mut self) { + self.0.cursor.reset(); + } +} + +// ── Entry ──────────────────────────────────────────────────────── + +/// The f16 MaxSim entry — the f32 pipeline behind widening walks. +#[derive(Debug, Clone, Copy)] +pub(crate) struct MaxIpF16; -impl +#[cfg(target_arch = "x86_64")] +impl diskann_wide::arch::Target3< - A, + V3, (), - BlockTransposedRef<'_, half::f16, GROUP>, + BlockTransposedRef<'_, half::f16, V3_A>, MatRef<'_, Standard>, &mut [f32], - > for F16Entry -where - A: Architecture, - F32Kernel: Kernel, - layouts::BlockTransposed: layouts::ConvertTo as Kernel>::Left> - + layouts::Layout, - layouts::RowMajor: layouts::ConvertTo as Kernel>::Right> - + layouts::Layout, + > for MaxIpF16 +{ + #[inline(always)] + fn run( + self, + arch: V3, + query: BlockTransposedRef<'_, half::f16, V3_A>, + docs: MatRef<'_, Standard>, + state: &mut [f32], + ) { + float::run( + arch, + query.padded_nrows(), + docs.num_vectors(), + query.padded_ncols(), + TileBudget::default(), + state, + |plan: Plan| { + ( + QueryWiden::new(arch, query, plan.a_panels), + DocWiden::new(arch, docs, plan.b_panels), + ) + }, + ); + } +} + +impl + diskann_wide::arch::Target3< + Scalar, + (), + BlockTransposedRef<'_, half::f16, SC_A>, + MatRef<'_, Standard>, + &mut [f32], + > for MaxIpF16 { #[inline(always)] fn run( self, - arch: A, - lhs: BlockTransposedRef<'_, half::f16, GROUP>, - rhs: MatRef<'_, Standard>, - scratch: &mut [f32], + arch: Scalar, + query: BlockTransposedRef<'_, half::f16, SC_A>, + docs: MatRef<'_, Standard>, + state: &mut [f32], ) { - max_ip_kernel(arch, lhs, rhs, scratch, TileBudget::default()); + float::run( + arch, + query.padded_nrows(), + docs.num_vectors(), + query.padded_ncols(), + TileBudget::default(), + state, + |plan: Plan| { + ( + QueryWiden::new(arch, query, plan.a_panels), + DocWiden::new(arch, docs, plan.b_panels), + ) + }, + ); } } diff --git a/diskann-quantization/src/multi_vector/distance/kernels/f32/mod.rs b/diskann-quantization/src/multi_vector/distance/kernels/f32/mod.rs deleted file mode 100644 index a900ea3568..0000000000 --- a/diskann-quantization/src/multi_vector/distance/kernels/f32/mod.rs +++ /dev/null @@ -1,135 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -//! f32 micro-kernel family for block-transposed multi-vector distance. -//! -//! Provides: -//! -//! - `F32Kernel` — zero-sized marker type selecting the f32 micro-kernel -//! for `BlockTransposed` data. -//! - [`max_ip_kernel`] — architecture-, element-type-, and GROUP-generic entry point -//! for the reducing max-IP GEMM. Accepts any element type `T` for which -//! [`ConvertTo`](super::layouts::ConvertTo) impls exist (identity for f32, -//! SIMD-accelerated f16→f32, etc.). -//! -//! # Architecture-specific micro-kernels -//! -//! - `v3` (x86_64) — V3 (AVX2+FMA) 16×4 micro-kernel (GROUP=16). V4 delegates to V3 at dispatch. -//! - `scalar` — Emulated 8×2 micro-kernel (GROUP=8). Neon delegates to Scalar at dispatch. - -use diskann_wide::Architecture; - -use super::Kernel; -use super::TileBudget; -use super::layouts::{self, DescribeLayout}; -use super::tiled_reduce::tiled_reduce; -use crate::multi_vector::{BlockTransposedRef, MatRef, Standard}; - -mod scalar; -#[cfg(target_arch = "x86_64")] -mod v3; - -/// Zero-sized kernel type for f32 micro-kernels with block size `GROUP`. -pub(crate) struct F32Kernel; - -#[inline(never)] -#[cold] -#[allow(clippy::panic)] -fn max_ip_kernel_panic(scratch_len: usize, padded_nrows: usize, a_ncols: usize, b_dim: usize) { - panic!( - "max_ip_kernel: precondition failed: \ - scratch.len()={scratch_len} (expected {padded_nrows}), \ - a.ncols()={a_ncols}, b.vector_dim()={b_dim}" - ); -} - -/// Compute the reducing max-IP GEMM between a block-transposed A matrix and -/// a row-major B matrix, writing per-A-row max similarities into `scratch`. -/// -/// Thin wrapper over [`tiled_reduce`] using `F32Kernel` for the -/// requested architecture. The element type `T` can be any `Copy` type with -/// matching [`ConvertTo`](super::layouts::ConvertTo) impls (zero-cost for -/// `T = f32`; SIMD f16→f32 conversion once per tile for `T = half::f16`). -/// -/// `scratch` must have length [`BlockTransposedRef::padded_nrows()`] and be -/// initialized to `f32::MIN` before the first call. On return, `scratch[i]` -/// holds the maximum inner product between A row `i` and any B row. -/// -/// # Panics -/// -/// Panics if `scratch.len() != a.padded_nrows()` or `a.ncols() != b.vector_dim()`. -pub(super) fn max_ip_kernel( - arch: A, - a: BlockTransposedRef<'_, T, GROUP>, - b: MatRef<'_, Standard>, - scratch: &mut [f32], - budget: TileBudget, -) where - F32Kernel: Kernel, - layouts::BlockTransposed: - layouts::ConvertTo as Kernel>::Left> + layouts::Layout, - layouts::RowMajor: layouts::ConvertTo as Kernel>::Right> - + layouts::Layout, -{ - if scratch.len() != a.padded_nrows() || a.ncols() != b.vector_dim() { - max_ip_kernel_panic(scratch.len(), a.padded_nrows(), a.ncols(), b.vector_dim()); - } - - let k = a.ncols(); - let b_nrows = b.num_vectors(); - - // Compile-time: A_PANEL must equal GROUP for block-transposed layout correctness. - const { assert!( as Kernel>::A_PANEL == GROUP) } - - let ca = a.layout(); - let cb = b.layout(); - - // SAFETY: - // - a.as_ptr() is valid for a.padded_nrows() * k elements of T. - // - MatRef> stores nrows * ncols contiguous T elements. - // - scratch.len() == a.padded_nrows() (checked above). - // - a.padded_nrows() is always a multiple of GROUP, and the const assert above - // verifies A_PANEL == GROUP at compile time. - unsafe { - tiled_reduce::, _, _>( - arch, - &ca, - &cb, - a.as_ptr(), - a.padded_nrows(), - b.as_slice().as_ptr(), - b_nrows, - k, - scratch, - budget, - ); - } -} - -impl - diskann_wide::arch::Target3< - A, - (), - BlockTransposedRef<'_, f32, GROUP>, - MatRef<'_, Standard>, - &mut [f32], - > for F32Kernel -where - A: Architecture, - Self: Kernel, - layouts::BlockTransposed: - layouts::ConvertTo>::Left> + layouts::Layout, - layouts::RowMajor: - layouts::ConvertTo>::Right> + layouts::Layout, -{ - #[inline(always)] - fn run( - self, - arch: A, - lhs: BlockTransposedRef<'_, f32, GROUP>, - rhs: MatRef<'_, Standard>, - scratch: &mut [f32], - ) { - max_ip_kernel(arch, lhs, rhs, scratch, TileBudget::default()); - } -} diff --git a/diskann-quantization/src/multi_vector/distance/kernels/f32/scalar.rs b/diskann-quantization/src/multi_vector/distance/kernels/f32/scalar.rs deleted file mode 100644 index bd8fb1c4ab..0000000000 --- a/diskann-quantization/src/multi_vector/distance/kernels/f32/scalar.rs +++ /dev/null @@ -1,113 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -//! Scalar (emulated) f32 micro-kernel (8×2). -//! -//! Uses [`Emulated`](diskann_wide::Emulated) (aliased as `f32x8`) for -//! arithmetic — 8 multiply-accumulate operations per inner iteration, 8 scalar -//! comparisons per `max_simd`. Geometry is A_PANEL=8 (1 × f32x8), B_PANEL=2 -//! (matching the `Strategy2x1` pattern used by scalar distance functions -//! elsewhere in the codebase). -//! -//! The inner loop uses separate multiply and add (`a * b + acc`) instead of -//! `mul_add_simd` to avoid calling into libm's software `fma()` routine on -//! x86-64 targets without hardware FMA support. - -use diskann_wide::arch::Scalar; -use diskann_wide::{SIMDMinMax, SIMDVector}; - -use super::super::Kernel; -use super::super::layouts; -use super::super::reduce::Reduce; -use super::F32Kernel; - -diskann_wide::alias!(f32s = ::f32x8); - -// SAFETY: F32Kernel's `full_panel` and `partial_panel` only access -// A_PANEL(8) * k A elements, UNROLL * k B elements, and A_PANEL(8) -// scratch elements — all within the bounds guaranteed by `tiled_reduce`. -unsafe impl Kernel for F32Kernel<8> { - type Left = layouts::BlockTransposed; - type Right = layouts::RowMajor; - const A_PANEL: usize = 8; - const B_PANEL: usize = 2; - - #[inline(always)] - unsafe fn full_panel(arch: Scalar, a: *const f32, b: *const f32, k: usize, r: *mut f32) { - // SAFETY: pointer validity per Kernel contract. - unsafe { scalar_f32_microkernel::<{ Self::B_PANEL }>(arch, a, b, k, r) } - } - - #[inline(always)] - unsafe fn partial_panel( - arch: Scalar, - remainder: usize, - a: *const f32, - b: *const f32, - k: usize, - r: *mut f32, - ) { - // SAFETY: pointer validity per Kernel contract. - unsafe { - match remainder { - 1 => scalar_f32_microkernel::<1>(arch, a, b, k, r), - _ => unreachable!( - "unexpected remainder {remainder} for B_PANEL={}", - Self::B_PANEL - ), - } - } - } -} - -// ── Scalar f32 micro-kernel ────────────────────────────────────── - -/// Emulated micro-kernel: processes 8 A rows × `UNROLL` B rows. -/// -/// Uses separate multiply and add (`a * b + acc`) rather than `mul_add_simd` -/// to avoid calling libm's software `fma()` on x86-64 without hardware FMA. -/// A single register tile covers A_PANEL = 8 = f32s::LANES. B_PANEL=2 -/// follows the `Strategy2x1` pattern from scalar distance functions. -/// -/// # Safety -/// -/// 1. `a_packed` must point to `A_PANEL(8) × k` contiguous `f32` values. -/// 2. `b` must point to `UNROLL` rows of `k` contiguous `f32` values. -/// 3. `r` must point to at least `A_PANEL(8)` writable `f32` values. -#[inline(always)] -unsafe fn scalar_f32_microkernel( - arch: Scalar, - a_packed: *const f32, - b: *const f32, - k: usize, - r: *mut f32, -) where - [f32s; UNROLL]: Reduce, -{ - let op = |x: f32s, y: f32s| x.max_simd(y); - - let mut p0 = [f32s::default(arch); UNROLL]; - let offsets: [usize; UNROLL] = core::array::from_fn(|i| k * i); - - let a_stride = f32s::LANES; - - for i in 0..k { - // SAFETY: By preconditions 1 and 2; i < k and j < UNROLL. - unsafe { - let a0 = f32s::load_simd(arch, a_packed.add(a_stride * i)); - - for j in 0..UNROLL { - let bj = f32s::splat(arch, b.add(i + offsets[j]).read_unaligned()); - p0[j] = a0 * bj + p0[j]; - } - } - } - - // SAFETY: By precondition 3. - let mut r0 = unsafe { f32s::load_simd(arch, r) }; - - r0 = op(r0, p0.reduce(&op)); - - // SAFETY: By precondition 3. - unsafe { r0.store_simd(r) }; -} diff --git a/diskann-quantization/src/multi_vector/distance/kernels/f32/v3.rs b/diskann-quantization/src/multi_vector/distance/kernels/f32/v3.rs deleted file mode 100644 index b05195b1e6..0000000000 --- a/diskann-quantization/src/multi_vector/distance/kernels/f32/v3.rs +++ /dev/null @@ -1,112 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -//! V3 (AVX2+FMA) f32 micro-kernel (16×4). - -use diskann_wide::arch::x86_64::V3; -use diskann_wide::{SIMDMinMax, SIMDMulAdd, SIMDVector}; - -use super::super::Kernel; -use super::super::layouts; -use super::super::reduce::Reduce; -use super::F32Kernel; - -diskann_wide::alias!(f32s = ::f32x8); - -// SAFETY: F32Kernel's `full_panel` and `partial_panel` only access -// A_PANEL(16) * k A elements, UNROLL * k B elements, and A_PANEL(16) -// scratch elements — all within the bounds guaranteed by `tiled_reduce`. -unsafe impl Kernel for F32Kernel<16> { - type Left = layouts::BlockTransposed; - type Right = layouts::RowMajor; - const A_PANEL: usize = 16; - const B_PANEL: usize = 4; - - #[inline(always)] - unsafe fn full_panel(arch: V3, a: *const f32, b: *const f32, k: usize, r: *mut f32) { - // SAFETY: pointer validity per Kernel contract. - unsafe { f32_microkernel::<{ Self::B_PANEL }>(arch, a, b, k, r) } - } - - #[inline(always)] - unsafe fn partial_panel( - arch: V3, - remainder: usize, - a: *const f32, - b: *const f32, - k: usize, - r: *mut f32, - ) { - // SAFETY: pointer validity per Kernel contract. - unsafe { - match remainder { - 1 => f32_microkernel::<1>(arch, a, b, k, r), - 2 => f32_microkernel::<2>(arch, a, b, k, r), - 3 => f32_microkernel::<3>(arch, a, b, k, r), - _ => unreachable!( - "unexpected remainder {remainder} for B_PANEL={}", - Self::B_PANEL - ), - } - } - } -} - -// ── V3 f32 micro-kernel ───────────────────────────────────────── - -/// SIMD micro-kernel: processes 16 A rows × `UNROLL` B rows. -/// -/// Accumulates via FMA into two `f32x8` register tiles, reduces across the -/// `UNROLL` B lanes with `max_simd`, then merges into the scratch buffer `r`. -/// -/// # Safety -/// -/// 1. `a_packed` must point to `A_PANEL(16) × k` contiguous `f32` values. -/// 2. `b` must point to `UNROLL` rows of `k` contiguous `f32` values. -/// 3. `r` must point to at least `A_PANEL(16)` writable `f32` values. -#[inline(always)] -unsafe fn f32_microkernel( - arch: V3, - a_packed: *const f32, - b: *const f32, - k: usize, - r: *mut f32, -) where - [f32s; UNROLL]: Reduce, -{ - let op = |x: f32s, y: f32s| x.max_simd(y); - - let mut p0 = [f32s::default(arch); UNROLL]; - let mut p1 = [f32s::default(arch); UNROLL]; - let offsets: [usize; UNROLL] = core::array::from_fn(|i| k * i); - - let a_stride = 2 * f32s::LANES; - let a_stride_half = f32s::LANES; - - for i in 0..k { - // SAFETY: By preconditions 1 and 2; i < k and j < UNROLL. - unsafe { - let a0 = f32s::load_simd(arch, a_packed.add(a_stride * i)); - let a1 = f32s::load_simd(arch, a_packed.add(a_stride * i + a_stride_half)); - - for j in 0..UNROLL { - let bj = f32s::splat(arch, b.add(i + offsets[j]).read_unaligned()); - p0[j] = a0.mul_add_simd(bj, p0[j]); - p1[j] = a1.mul_add_simd(bj, p1[j]); - } - } - } - - // SAFETY: By precondition 3; LANES < A_PANEL so both halves are in-bounds. - let mut r0 = unsafe { f32s::load_simd(arch, r) }; - // SAFETY: By precondition 3; r.add(LANES) is still within the A_PANEL-sized scratch. - let mut r1 = unsafe { f32s::load_simd(arch, r.add(f32s::LANES)) }; - - r0 = op(r0, p0.reduce(&op)); - r1 = op(r1, p1.reduce(&op)); - - // SAFETY: By precondition 3. - unsafe { r0.store_simd(r) }; - // SAFETY: By precondition 3; r.add(LANES) is still within the A_PANEL-sized scratch. - unsafe { r1.store_simd(r.add(f32s::LANES)) }; -} diff --git a/diskann-quantization/src/multi_vector/distance/kernels/float.rs b/diskann-quantization/src/multi_vector/distance/kernels/float.rs new file mode 100644 index 0000000000..02d27cd639 --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/float.rs @@ -0,0 +1,516 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +//! f32 MaxSim: the accumulator is already the score, so the drain is a bare reduction. + +use core::ops::Range; + +use diskann_wide::arch::Scalar; +#[cfg(target_arch = "x86_64")] +use diskann_wide::arch::x86_64::V3; + +use super::leaves::scalar as scalar_leaf; +use super::leaves::scalar::{A_PANEL as SC_A, B_PANEL as SC_B}; +#[cfg(target_arch = "x86_64")] +use super::leaves::v3 as v3_leaf; +#[cfg(target_arch = "x86_64")] +use super::leaves::v3::{A_PANEL as V3_A, B_PANEL as V3_B}; +use super::strip::{Slot, Strip}; +use super::tiles::{DocPanel, DocTile, DocWalk, QueryPanel, QueryTile, QueryWalk}; +use super::{Accumulate, Drain, Plan, TileAt, TileBudget, TileWalk, drive}; +use crate::bits::{Dynamic, Static}; +use crate::multi_vector::{BlockTransposedRef, MatRef, Standard}; + +/// Selects the f32 leaf for whichever architecture is in play. +pub(super) struct Kernel; + +/// Reduces the strip straight into the running per-query-row maxima. +/// +/// Carries `nd` because the strip's trailing columns belong to doc rows past the end of +/// the matrix, and nothing else here knows where that end is. +pub(super) struct RawMax<'o> { + out: &'o mut [f32], + nd: usize, +} + +// ── V3 ─────────────────────────────────────────────────────────── + +#[cfg(target_arch = "x86_64")] +impl<'a, 'b, 's> + Accumulate< + V3, + QueryPanel<'a, f32, V3_A>, + DocPanel<'b, f32, V3_B, Static>, + Slot<'s, f32, V3_A, V3_B>, + > for Kernel +{ + #[inline(always)] + fn accumulate( + &self, + arch: V3, + a: QueryPanel<'a, f32, V3_A>, + b: DocPanel<'b, f32, V3_B, Static>, + out: Slot<'s, f32, V3_A, V3_B>, + ) { + v3_leaf::f32_store_microkernel::(arch, a, b, out); + } +} + +#[cfg(target_arch = "x86_64")] +impl<'a, 'b, 's> + Accumulate< + V3, + QueryPanel<'a, f32, V3_A>, + DocPanel<'b, f32, V3_B, Dynamic>, + Slot<'s, f32, V3_A, V3_B>, + > for Kernel +{ + #[inline(always)] + fn accumulate( + &self, + arch: V3, + a: QueryPanel<'a, f32, V3_A>, + b: DocPanel<'b, f32, V3_B, Dynamic>, + out: Slot<'s, f32, V3_A, V3_B>, + ) { + // Dispatch the runtime width onto a const the leaf can unroll for. + match b.rows() { + 3 => v3_leaf::f32_store_microkernel::<3, _>(arch, a, b, out), + 2 => v3_leaf::f32_store_microkernel::<2, _>(arch, a, b, out), + 1 => v3_leaf::f32_store_microkernel::<1, _>(arch, a, b, out), + other => unreachable!("tail width {other} outside 1..{V3_B}"), + } + } +} + +#[cfg(target_arch = "x86_64")] +impl Drain> for RawMax<'_> { + #[inline(always)] + fn drain( + &mut self, + arch: V3, + scratch: &Strip<'_, f32, V3_A, V3_B>, + a_panel: usize, + b_panels: Range, + ) { + // A is padded to whole panels, so the output cut never clamps; B is not, so the + // live width does. + let a0 = a_panel * V3_A; + let live = (b_panels.end * V3_B).min(self.nd) - b_panels.start * V3_B; + v3_leaf::fold_columns(arch, scratch.columns(live), &mut self.out[a0..][..V3_A]); + } +} + +// ── Emulated ───────────────────────────────────────────────────── + +impl<'a, 'b, 's> + Accumulate< + Scalar, + QueryPanel<'a, f32, SC_A>, + DocPanel<'b, f32, SC_B, Static>, + Slot<'s, f32, SC_A, SC_B>, + > for Kernel +{ + #[inline(always)] + fn accumulate( + &self, + arch: Scalar, + a: QueryPanel<'a, f32, SC_A>, + b: DocPanel<'b, f32, SC_B, Static>, + out: Slot<'s, f32, SC_A, SC_B>, + ) { + scalar_leaf::f32_store_microkernel::(arch, a, b, out); + } +} + +impl<'a, 'b, 's> + Accumulate< + Scalar, + QueryPanel<'a, f32, SC_A>, + DocPanel<'b, f32, SC_B, Dynamic>, + Slot<'s, f32, SC_A, SC_B>, + > for Kernel +{ + #[inline(always)] + fn accumulate( + &self, + arch: Scalar, + a: QueryPanel<'a, f32, SC_A>, + b: DocPanel<'b, f32, SC_B, Dynamic>, + out: Slot<'s, f32, SC_A, SC_B>, + ) { + match b.rows() { + 1 => scalar_leaf::f32_store_microkernel::<1, _>(arch, a, b, out), + other => unreachable!("tail width {other} outside 1..{SC_B}"), + } + } +} + +impl Drain> for RawMax<'_> { + #[inline(always)] + fn drain( + &mut self, + arch: Scalar, + scratch: &Strip<'_, f32, SC_A, SC_B>, + a_panel: usize, + b_panels: Range, + ) { + let a0 = a_panel * SC_A; + let live = (b_panels.end * SC_B).min(self.nd) - b_panels.start * SC_B; + scalar_leaf::fold_columns(arch, scratch.columns(live), &mut self.out[a0..][..SC_A]); + } +} + +// ── Entry ──────────────────────────────────────────────────────── + +/// Plan, allocate the strip, and drive. +/// +/// `k` is the *physical* row length both walks stride by, so it must be the query's padded +/// column count rather than its logical one. +/// +/// `walks` is built from the plan rather than passed in so the empty-contraction guard can +/// run before any walk exists — a zero-length row would give a walk a zero stride. +/// +/// On return `state` holds the per-query-row maximum inner product, one entry per padded +/// query row. Its incoming contents are ignored: seeding the max is this function's job, +/// not the caller's, because a caller that got it wrong would silently clamp the result +/// rather than fail. +/// +/// # Panics +/// +/// Panics unless `state` is exactly `nq` long. +pub(super) fn run( + arch: Arch, + nq: usize, + nd: usize, + k: usize, + budget: TileBudget, + state: &mut [f32], + walks: impl FnOnce(Plan) -> (AW, BW), +) where + Arch: Copy, + AW: TileWalk + for<'a> TileAt<'a, Tile = QueryTile<'a, f32, AR>>, + BW: TileWalk + for<'b> TileAt<'b, Tile = DocTile<'b, f32, BR>>, + Kernel: for<'a, 'b, 's> Accumulate< + Arch, + QueryPanel<'a, f32, AR>, + DocPanel<'b, f32, BR, Static
>, + Slot<'s, f32, AR, BR>, + > + for<'a, 'b, 's> Accumulate< + Arch, + QueryPanel<'a, f32, AR>, + DocPanel<'b, f32, BR, Dynamic>, + Slot<'s, f32, AR, BR>, + >, + for<'o, 'x> RawMax<'o>: Drain>, +{ + assert_eq!(state.len(), nq, "state is one entry per padded query row"); + + // The identity for max. + state.fill(f32::MIN); + + if k == 0 { + // Every inner product is the empty sum. + if nd > 0 { + state.fill(0.0); + } + return; + } + + // Both sides are f32 rows of `k`: whatever the source element type, the leaves + // consume f32. + let row_bytes = k * size_of::(); + let plan = Plan::::new(row_bytes, row_bytes, size_of::(), budget); + let (a_walk, b_walk) = walks(plan); + let mut buf = vec![0.0f32; plan.strip_len()]; + + drive( + arch, + a_walk, + b_walk, + &mut Strip::new(&mut buf), + &Kernel, + &mut RawMax { out: state, nd }, + ); +} + +/// The f32 MaxSim entry. Which leaf geometry applies follows from the architecture, so the +/// block size of `query` must match the [`Target3`](diskann_wide::arch::Target3) impl +/// selected. +#[derive(Debug, Clone, Copy)] +pub(crate) struct MaxIp; + +#[cfg(target_arch = "x86_64")] +impl + diskann_wide::arch::Target3< + V3, + (), + BlockTransposedRef<'_, f32, V3_A>, + MatRef<'_, Standard>, + &mut [f32], + > for MaxIp +{ + #[inline(always)] + fn run( + self, + arch: V3, + query: BlockTransposedRef<'_, f32, V3_A>, + docs: MatRef<'_, Standard>, + state: &mut [f32], + ) { + run( + arch, + query.padded_nrows(), + docs.num_vectors(), + query.padded_ncols(), + TileBudget::default(), + state, + |plan: Plan| { + ( + QueryWalk::new(query, plan.a_panels), + DocWalk::new(docs, plan.b_panels), + ) + }, + ); + } +} + +impl + diskann_wide::arch::Target3< + Scalar, + (), + BlockTransposedRef<'_, f32, SC_A>, + MatRef<'_, Standard>, + &mut [f32], + > for MaxIp +{ + #[inline(always)] + fn run( + self, + arch: Scalar, + query: BlockTransposedRef<'_, f32, SC_A>, + docs: MatRef<'_, Standard>, + state: &mut [f32], + ) { + run( + arch, + query.padded_nrows(), + docs.num_vectors(), + query.padded_ncols(), + TileBudget::default(), + state, + |plan: Plan| { + ( + QueryWalk::new(query, plan.a_panels), + DocWalk::new(docs, plan.b_panels), + ) + }, + ); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::multi_vector::BlockTransposed; + + fn sample(len: usize, phase: usize) -> Vec { + (0..len) + .map(|i| (((i * 7 + phase) % 23) as f32 - 11.0) / 4.0) + .collect() + } + + /// Reference in f64 so the comparison measures the kernel, not the reference. + fn naive(a: &[f32], nq: usize, b: &[f32], nd: usize, k: usize) -> Vec { + (0..nq) + .map(|i| { + (0..nd) + .map(|j| { + (0..k) + .map(|t| a[i * k + t] as f64 * b[j * k + t] as f64) + .sum::() + }) + .fold(f64::NEG_INFINITY, f64::max) + }) + .collect() + } + + fn check( + arch: Arch, + label: &str, + (nq, nd, k): (usize, usize, usize), + budget: TileBudget, + ) where + Arch: Copy, + Kernel: for<'a, 'b, 's> Accumulate< + Arch, + QueryPanel<'a, f32, AR>, + DocPanel<'b, f32, BR, Static
>, + Slot<'s, f32, AR, BR>, + > + for<'a, 'b, 's> Accumulate< + Arch, + QueryPanel<'a, f32, AR>, + DocPanel<'b, f32, BR, Dynamic>, + Slot<'s, f32, AR, BR>, + >, + for<'o, 'x> RawMax<'o>: Drain>, + { + let a = sample(nq * k, 0); + let b = sample(nd * k, 5); + let query = MatRef::new(Standard::new(nq, k).unwrap(), &a).unwrap(); + let docs = MatRef::new(Standard::new(nd, k).unwrap(), &b).unwrap(); + let bt = BlockTransposed::::from_matrix_view(query.as_matrix_view()); + + // Poisoned, not seeded: `run` owes the caller a full initialization, and every case + // below would fail loudly if it skipped one. + let mut state = vec![f32::MAX; bt.padded_nrows()]; + run( + arch, + bt.padded_nrows(), + nd, + k, + budget, + &mut state, + |plan: Plan| { + ( + QueryWalk::new(bt.as_view(), plan.a_panels), + DocWalk::new(docs, plan.b_panels), + ) + }, + ); + + for (i, &expected) in naive(&a, nq, &b, nd, k).iter().enumerate() { + let actual = state[i] as f64; + let tol = 1e-5 * expected.abs().max(1.0); + assert!( + (actual - expected).abs() < tol, + "[{label}] row {i} of ({nq},{nd},{k}): actual={actual}, expected={expected}", + ); + } + } + + /// Shapes chosen to cross every A-panel boundary (8 and 16), every B-panel remainder + /// class (mod 2 and mod 4), and both degenerate and prime contraction lengths. + const CASES: &[(usize, usize, usize)] = &[ + (1, 1, 1), + (1, 1, 64), + (5, 3, 5), + (8, 2, 8), + (8, 33, 127), + (9, 4, 3), + (15, 5, 16), + (16, 4, 64), + (16, 5, 128), + (16, 6, 64), + (16, 7, 256), + (17, 9, 65), + (32, 16, 256), + (33, 1, 2), + (64, 250, 64), + ]; + + #[test] + fn scalar_matches_naive() { + for &case in CASES { + check::<_, SC_A, SC_B>(Scalar::new(), "scalar", case, TileBudget::default()); + } + } + + /// Forces `a_panels == b_panels == 1`, so every panel is its own tile and the driver's + /// cross-tile ordinal carry is exercised on every shape. + #[test] + fn scalar_matches_naive_one_panel_per_tile() { + let budget = TileBudget { l2_a: 1, l1_b: 1 }; + for &case in CASES { + check::<_, SC_A, SC_B>(Scalar::new(), "scalar/tiny", case, budget); + } + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn v3_matches_naive() { + let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() else { + return; + }; + for &case in CASES { + check::<_, V3_A, V3_B>(arch, "x86-64-v3", case, TileBudget::default()); + } + } + + #[cfg(target_arch = "x86_64")] + #[test] + fn v3_matches_naive_one_panel_per_tile() { + let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() else { + return; + }; + let budget = TileBudget { l2_a: 1, l1_b: 1 }; + for &case in CASES { + check::<_, V3_A, V3_B>(arch, "x86-64-v3/tiny", case, budget); + } + } + + /// A zero-length contraction makes every inner product the empty sum, which the entry + /// has to answer without building a zero-stride walk. + #[test] + fn empty_contraction() { + let mut state = vec![f32::MAX; SC_A]; + run( + Scalar::new(), + SC_A, + 3, + 0, + TileBudget::default(), + &mut state, + |_: Plan| -> (QueryWalk<'_, f32, SC_A>, DocWalk<'_, f32, SC_B>) { + unreachable!("no walk is built for an empty contraction") + }, + ); + assert_eq!(state, vec![0.0; SC_A]); + + // With no docs there is no maximum, so the identity stands and the caller negates + // it into its empty-doc sentinel. + let mut state = vec![f32::MAX; SC_A]; + run( + Scalar::new(), + SC_A, + 0, + 0, + TileBudget::default(), + &mut state, + |_: Plan| -> (QueryWalk<'_, f32, SC_A>, DocWalk<'_, f32, SC_B>) { + unreachable!("no walk is built for an empty contraction") + }, + ); + assert_eq!(state, vec![f32::MIN; SC_A]); + } + + /// `state` sized from `nrows()` instead of `padded_nrows()` is the mistake the entry + /// check exists for: the two are derived independently, so they can disagree, and it + /// used to surface as a slice index panic deep inside the drain. + #[test] + #[should_panic(expected = "one entry per padded query row")] + fn state_sized_for_logical_rows_is_rejected() { + let a = sample(9 * 4, 0); + let b = sample(2 * 4, 5); + let query = MatRef::new(Standard::new(9, 4).unwrap(), &a).unwrap(); + let docs = MatRef::new(Standard::new(2, 4).unwrap(), &b).unwrap(); + let bt = BlockTransposed::::from_matrix_view(query.as_matrix_view()); + assert_eq!((bt.nrows(), bt.padded_nrows()), (9, 16)); + + let mut state = vec![0.0; bt.nrows()]; + run( + Scalar::new(), + bt.padded_nrows(), + 2, + 4, + TileBudget::default(), + &mut state, + |plan: Plan| { + ( + QueryWalk::new(bt.as_view(), plan.a_panels), + DocWalk::new(docs, plan.b_panels), + ) + }, + ); + } +} diff --git a/diskann-quantization/src/multi_vector/distance/kernels/layouts.rs b/diskann-quantization/src/multi_vector/distance/kernels/layouts.rs deleted file mode 100644 index e1ec8dd36e..0000000000 --- a/diskann-quantization/src/multi_vector/distance/kernels/layouts.rs +++ /dev/null @@ -1,219 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -//! Layout markers and tile-level conversion traits. -//! -//! - [`Layout`] — marker trait: memory layout + element type. -//! - [`BlockTransposed`] / [`RowMajor`] — zero-sized layout markers. -//! - [`DescribeLayout`] — bridges matrix types to layout markers. -//! - [`ConvertTo`] — tile-level conversion (blanket identity + f16→f32). - -use core::marker::PhantomData; - -use diskann_vector::conversion::SliceCast; -use diskann_wide::Architecture; -use diskann_wide::arch::Target2; - -// ── Layout trait ───────────────────────────────────── - -/// Memory layout and element type marker for tile data. -pub(super) trait Layout { - type Element: Copy; -} - -// ── Layout markers ─────────────────────────────────── - -/// Block-transposed tile layout: `GROUP` rows per block, `PACK` columns -/// interleaved. Matches [`BlockTransposedRef`](crate::multi_vector::BlockTransposedRef). -pub(super) struct BlockTransposed(PhantomData); - -impl BlockTransposed { - pub(super) fn new() -> Self { - Self(PhantomData) - } -} - -impl Copy for BlockTransposed {} - -impl Clone for BlockTransposed { - fn clone(&self) -> Self { - *self - } -} - -impl Layout for BlockTransposed { - type Element = T; -} - -/// Dense row-major tile layout. Matches [`MatRef>`](crate::multi_vector::MatRef). -pub(super) struct RowMajor(PhantomData); - -impl RowMajor { - pub(super) fn new() -> Self { - Self(PhantomData) - } -} - -impl Copy for RowMajor {} - -impl Clone for RowMajor { - fn clone(&self) -> Self { - *self - } -} - -impl Layout for RowMajor { - type Element = T; -} - -// ── DescribeLayout ─────────────────────────────────── - -/// Bridges a concrete matrix type to its [`Layout`] marker, enabling -/// type inference of [`ConvertTo`] parameters at call sites. -pub(super) trait DescribeLayout { - type Layout: Layout; - - fn layout(&self) -> Self::Layout; -} - -impl DescribeLayout - for crate::multi_vector::BlockTransposedRef<'_, T, GROUP, PACK> -{ - type Layout = BlockTransposed; - - fn layout(&self) -> Self::Layout { - BlockTransposed::new() - } -} - -impl DescribeLayout for crate::multi_vector::MatRef<'_, crate::multi_vector::Standard> { - type Layout = RowMajor; - - fn layout(&self) -> Self::Layout { - RowMajor::new() - } -} - -// ── ConvertTo trait ────────────────────────────────── - -/// Tile-level conversion from layout `Self` to layout `To`. -/// -/// The blanket identity impl covers every layout converting to itself -/// with `Buffer = ()` and zero cost. Explicit impls handle f16→f32 via -/// [`SliceCast`]. -/// -/// # Safety -/// -/// Implementors must ensure: -/// - `convert` reads at most `rows * k` source elements. -/// - `convert` writes only within `buf`. -/// - The returned pointer is valid until the next `&mut` access to `buf`. -pub(super) unsafe trait ConvertTo: Layout { - /// Staging buffer for converted tile data (`()` for identity conversions). - type Buffer; - - /// Allocate a buffer for up to `max_tile_rows` rows of dimension `k`. - fn new_buffer(&self, max_tile_rows: usize, k: usize) -> Self::Buffer; - - /// Convert `rows` rows of source data into `buf`, returning a read pointer. - /// - /// # Safety - /// - /// * `src` must point to `rows * k` valid elements in `Self`'s layout. - /// * `buf` must come from [`new_buffer`](Self::new_buffer) with the - /// same `k` and a `max_tile_rows >= rows`. - unsafe fn convert( - &self, - buf: &mut Self::Buffer, - arch: A, - src: *const Self::Element, - rows: usize, - k: usize, - ) -> *const To::Element; -} - -// ── Blanket identity ───────────────────────────────── - -/// Identity conversion: every layout converts to itself at zero cost. -// SAFETY: Identity conversion reads nothing beyond `src` and writes -// nothing into `buf`. The returned pointer is exactly `src`, which is -// valid for the lifetime guaranteed by the caller. -unsafe impl ConvertTo for L { - type Buffer = (); - - fn new_buffer(&self, _max_tile_rows: usize, _k: usize) {} - - unsafe fn convert( - &self, - _buf: &mut (), - _arch: A, - src: *const L::Element, - _rows: usize, - _k: usize, - ) -> *const L::Element { - src - } -} - -// ── f16 → f32 conversions ──────────────────────────── - -/// Block-transposed f16 → block-transposed f32 (element-wise, layout-preserving). -// SAFETY: `SliceCast` converts exactly `rows * k` f16 values from `src` -// into `rows * k` f32 values in `buf`. The returned pointer is -// `buf.as_ptr()`, valid until the next `&mut` access to `buf`. -unsafe impl - ConvertTo> for BlockTransposed -where - A: Architecture, - SliceCast: for<'a> Target2, -{ - type Buffer = Vec; - - fn new_buffer(&self, max_tile_rows: usize, k: usize) -> Vec { - vec![0.0f32; max_tile_rows * k] - } - - unsafe fn convert( - &self, - buf: &mut Vec, - arch: A, - src: *const half::f16, - rows: usize, - k: usize, - ) -> *const f32 { - let count = rows * k; - // SAFETY: Caller guarantees `src` points to `count` contiguous f16 values. - let src_slice = unsafe { std::slice::from_raw_parts(src, count) }; - arch.run2(SliceCast::new(), &mut buf[..count], src_slice); - buf.as_ptr() - } -} - -/// Row-major f16 → row-major f32 (element-wise, layout-preserving). -// SAFETY: Same as block-transposed variant — element-wise, layout-preserving. -unsafe impl ConvertTo> for RowMajor -where - A: Architecture, - SliceCast: for<'a> Target2, -{ - type Buffer = Vec; - - fn new_buffer(&self, max_tile_rows: usize, k: usize) -> Vec { - vec![0.0f32; max_tile_rows * k] - } - - unsafe fn convert( - &self, - buf: &mut Vec, - arch: A, - src: *const half::f16, - rows: usize, - k: usize, - ) -> *const f32 { - let count = rows * k; - // SAFETY: Caller guarantees `src` points to `count` contiguous f16 values. - let src_slice = unsafe { std::slice::from_raw_parts(src, count) }; - arch.run2(SliceCast::new(), &mut buf[..count], src_slice); - buf.as_ptr() - } -} diff --git a/diskann-quantization/src/multi_vector/distance/kernels/leaves/mod.rs b/diskann-quantization/src/multi_vector/distance/kernels/leaves/mod.rs new file mode 100644 index 0000000000..acb9f1b3f8 --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/leaves/mod.rs @@ -0,0 +1,21 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +//! Per-ISA micro-kernels. +//! +//! Panel geometry is fixed per ISA rather than derived from a lane count: the scalar leaf +//! is deliberately narrower than `2 × LANES` would suggest, because its "lanes" are a +//! loop, not silicon. +//! +//! Every leaf here **stores** its products into a slot and leaves the reduction to +//! `fold_columns`, which folds along independent max chains so the reduction does not +//! serialize on its own latency — it is not negligible work when the strip is wide and +//! shallow. + +pub(super) mod scalar; +#[cfg(target_arch = "x86_64")] +pub(super) mod v3; + +/// Independent max chains in `fold_columns`: enough to keep a multi-cycle max off its own +/// critical path, few enough that the chains stay in registers. +const WAYS: usize = 4; diff --git a/diskann-quantization/src/multi_vector/distance/kernels/leaves/scalar.rs b/diskann-quantization/src/multi_vector/distance/kernels/leaves/scalar.rs new file mode 100644 index 0000000000..13970793f8 --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/leaves/scalar.rs @@ -0,0 +1,129 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +//! Emulated f32 leaf — 8 × 2. +//! +//! Geometry is narrower than the V3 leaf because [`Emulated`](diskann_wide::Emulated) +//! lanes are an unrolled loop over scalars, not registers: 8 × 2 matches the `Strategy2x1` +//! shape used by the scalar distance functions elsewhere in the crate. +//! +//! The inner loop multiplies and adds separately rather than calling `mul_add_simd`, which +//! on x86-64 without hardware FMA would drop into libm's software `fma()`. + +use diskann_wide::arch::Scalar; +use diskann_wide::{SIMDMinMax, SIMDVector}; + +use super::WAYS; +use crate::bits::Length; +use crate::multi_vector::distance::kernels::strip::Slot; +use crate::multi_vector::distance::kernels::tiles::{DocPanel, QueryPanel}; + +diskann_wide::alias!(f32s = ::f32x8); + +pub(crate) const A_PANEL: usize = f32s::LANES; +pub(crate) const B_PANEL: usize = 2; + +/// Vector chunks spanned by one A-panel. +const REGS: usize = A_PANEL / f32s::LANES; + +/// `A_PANEL × UNROLL` inner products, stored column-major into `out`. +/// +/// # Panics +/// +/// Panics if `b` does not hold exactly `UNROLL` rows of `a`'s contraction length — the +/// extent every unchecked access below relies on. +#[inline(always)] +pub(crate) fn f32_store_microkernel( + arch: Scalar, + a: QueryPanel<'_, f32, A_PANEL>, + b: DocPanel<'_, f32, B_PANEL, L>, + mut out: Slot<'_, f32, A_PANEL, B_PANEL>, +) { + const { assert!(UNROLL >= 1 && UNROLL <= B_PANEL) } + + let k = a.k(); + let a_data = a.as_slice(); + let b_data = b.as_slice(); + assert_eq!(b_data.len(), UNROLL * k, "doc panel extent"); + + let out = out.as_mut_slice(); + debug_assert_eq!(out.len(), A_PANEL * B_PANEL); + + let ap = a_data.as_ptr(); + let bp = b_data.as_ptr(); + + let mut acc = [[f32s::default(arch); REGS]; UNROLL]; + let b_row: [usize; UNROLL] = core::array::from_fn(|j| k * j); + + for i in 0..k { + let a_col: [f32s; REGS] = core::array::from_fn(|r| { + // SAFETY: `r < REGS` and `A_PANEL == REGS * LANES`, so the load ends at or + // before `A_PANEL * (i + 1)`, hence at or before `A_PANEL * k`, which is at + // most `a_data.len()` because `k` is that length divided by `A_PANEL`. + unsafe { f32s::load_simd(arch, ap.add(A_PANEL * i + r * f32s::LANES)) } + }); + + for (acc_j, &row) in acc.iter_mut().zip(&b_row) { + // SAFETY: `row == k * j` for some `j < UNROLL`, so `i + row < UNROLL * k`, + // the asserted length of `b_data`. + let b = unsafe { bp.add(i + row).read() }; + + let bj = f32s::splat(arch, b); + for (acc_jr, a_r) in acc_j.iter_mut().zip(&a_col) { + *acc_jr = *a_r * bj + *acc_jr; + } + } + } + + for (acc_j, col) in acc.iter().zip(out.chunks_exact_mut(A_PANEL)) { + for (acc_jr, dst) in acc_j.iter().zip(col.chunks_exact_mut(f32s::LANES)) { + // SAFETY: `chunks_exact_mut` yields exactly `LANES` elements. + unsafe { acc_jr.store_simd(dst.as_mut_ptr()) }; + } + } +} + +/// Fold `acc`'s columns into the running per-A-row maxima in `state`, one A-panel wide. +#[inline(always)] +pub(crate) fn fold_columns(arch: Scalar, acc: &[f32], state: &mut [f32]) { + debug_assert_eq!(state.len(), A_PANEL); + debug_assert!(acc.len().is_multiple_of(A_PANEL)); + + let mut chains = [[f32s::splat(arch, f32::MIN); REGS]; WAYS]; + + let mut groups = acc.chunks_exact(WAYS * A_PANEL); + for group in groups.by_ref() { + fold_group(arch, &mut chains, group); + } + // Fewer than WAYS columns remain, so each still lands on its own chain. + fold_group(arch, &mut chains, groups.remainder()); + + let mut merged = chains[0]; + for chain in &chains[1..] { + for (m, c) in merged.iter_mut().zip(chain) { + *m = m.max_simd(*c); + } + } + + for (m, tile) in merged.iter().zip(state.chunks_exact_mut(f32s::LANES)) { + // SAFETY: `chunks_exact_mut` yields exactly `LANES` elements. + unsafe { + let sp = tile.as_mut_ptr(); + f32s::load_simd(arch, sp).max_simd(*m).store_simd(sp); + } + } +} + +/// Merge up to `WAYS` consecutive A-panels, one per chain. +#[inline(always)] +fn fold_group(arch: Scalar, chains: &mut [[f32s; REGS]; WAYS], src: &[f32]) { + for (chain, column) in chains.iter_mut().zip(src.chunks_exact(A_PANEL)) { + for (r, c) in chain.iter_mut().enumerate() { + // SAFETY: `chunks_exact` yields `A_PANEL == REGS * LANES` elements. + unsafe { + let v = f32s::load_simd(arch, column.as_ptr().add(r * f32s::LANES)); + *c = c.max_simd(v); + } + } + } +} diff --git a/diskann-quantization/src/multi_vector/distance/kernels/leaves/v3.rs b/diskann-quantization/src/multi_vector/distance/kernels/leaves/v3.rs new file mode 100644 index 0000000000..34980ca823 --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/leaves/v3.rs @@ -0,0 +1,124 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +//! V3 (AVX2+FMA) f32 leaf — 16 × 4. + +use diskann_wide::arch::x86_64::V3; +use diskann_wide::{SIMDMinMax, SIMDMulAdd, SIMDVector}; + +use super::WAYS; +use crate::bits::Length; +use crate::multi_vector::distance::kernels::strip::Slot; +use crate::multi_vector::distance::kernels::tiles::{DocPanel, QueryPanel}; + +diskann_wide::alias!(f32s = ::f32x8); + +/// Two `f32x8` register tiles per B-row: enough independent accumulators to cover FMA +/// latency at `B_PANEL = 4` without spilling. +pub(crate) const A_PANEL: usize = 2 * f32s::LANES; +pub(crate) const B_PANEL: usize = 4; + +/// Register tiles spanned by one A-panel. +const REGS: usize = A_PANEL / f32s::LANES; + +/// `A_PANEL × UNROLL` inner products, stored column-major into `out`. +/// +/// # Panics +/// +/// Panics if `b` does not hold exactly `UNROLL` rows of `a`'s contraction length — the +/// extent every unchecked access below relies on. +#[inline(always)] +pub(crate) fn f32_store_microkernel( + arch: V3, + a: QueryPanel<'_, f32, A_PANEL>, + b: DocPanel<'_, f32, B_PANEL, L>, + mut out: Slot<'_, f32, A_PANEL, B_PANEL>, +) { + const { assert!(UNROLL >= 1 && UNROLL <= B_PANEL) } + + let k = a.k(); + let a_data = a.as_slice(); + let b_data = b.as_slice(); + assert_eq!(b_data.len(), UNROLL * k, "doc panel extent"); + + let out = out.as_mut_slice(); + debug_assert_eq!(out.len(), A_PANEL * B_PANEL); + + let ap = a_data.as_ptr(); + let bp = b_data.as_ptr(); + + let mut acc = [[f32s::default(arch); REGS]; UNROLL]; + let b_row: [usize; UNROLL] = core::array::from_fn(|j| k * j); + + for i in 0..k { + let a_col: [f32s; REGS] = core::array::from_fn(|r| { + // SAFETY: `r < REGS` and `A_PANEL == REGS * LANES`, so the load ends at or + // before `A_PANEL * (i + 1)`, hence at or before `A_PANEL * k`, which is at + // most `a_data.len()` because `k` is that length divided by `A_PANEL`. + unsafe { f32s::load_simd(arch, ap.add(A_PANEL * i + r * f32s::LANES)) } + }); + + for (acc_j, &row) in acc.iter_mut().zip(&b_row) { + // SAFETY: `row == k * j` for some `j < UNROLL`, so `i + row < UNROLL * k`, + // the asserted length of `b_data`. + let b = unsafe { bp.add(i + row).read() }; + + let bj = f32s::splat(arch, b); + for (acc_jr, a_r) in acc_j.iter_mut().zip(&a_col) { + *acc_jr = a_r.mul_add_simd(bj, *acc_jr); + } + } + } + + for (acc_j, col) in acc.iter().zip(out.chunks_exact_mut(A_PANEL)) { + for (acc_jr, dst) in acc_j.iter().zip(col.chunks_exact_mut(f32s::LANES)) { + // SAFETY: `chunks_exact_mut` yields exactly `LANES` elements. + unsafe { acc_jr.store_simd(dst.as_mut_ptr()) }; + } + } +} + +/// Fold `acc`'s columns into the running per-A-row maxima in `state`, one A-panel wide. +#[inline(always)] +pub(crate) fn fold_columns(arch: V3, acc: &[f32], state: &mut [f32]) { + debug_assert_eq!(state.len(), A_PANEL); + debug_assert!(acc.len().is_multiple_of(A_PANEL)); + + let mut chains = [[f32s::splat(arch, f32::MIN); REGS]; WAYS]; + + let mut groups = acc.chunks_exact(WAYS * A_PANEL); + for group in groups.by_ref() { + fold_group(arch, &mut chains, group); + } + // Fewer than WAYS columns remain, so each still lands on its own chain. + fold_group(arch, &mut chains, groups.remainder()); + + let mut merged = chains[0]; + for chain in &chains[1..] { + for (m, c) in merged.iter_mut().zip(chain) { + *m = m.max_simd(*c); + } + } + + for (m, tile) in merged.iter().zip(state.chunks_exact_mut(f32s::LANES)) { + // SAFETY: `chunks_exact_mut` yields exactly `LANES` elements. + unsafe { + let sp = tile.as_mut_ptr(); + f32s::load_simd(arch, sp).max_simd(*m).store_simd(sp); + } + } +} + +/// Merge up to `WAYS` consecutive A-panels, one per chain. +#[inline(always)] +fn fold_group(arch: V3, chains: &mut [[f32s; REGS]; WAYS], src: &[f32]) { + for (chain, column) in chains.iter_mut().zip(src.chunks_exact(A_PANEL)) { + for (r, c) in chain.iter_mut().enumerate() { + // SAFETY: `chunks_exact` yields `A_PANEL == REGS * LANES` elements. + unsafe { + let v = f32s::load_simd(arch, column.as_ptr().add(r * f32s::LANES)); + *c = c.max_simd(v); + } + } + } +} diff --git a/diskann-quantization/src/multi_vector/distance/kernels/mod.rs b/diskann-quantization/src/multi_vector/distance/kernels/mod.rs index 55108698da..4f7cdcc33c 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/mod.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/mod.rs @@ -1,34 +1,38 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT license. -//! Block-transposed SIMD kernels for multi-vector distance computation. +//! Cache-tiled MaxSim: a [`TileWalk`] lends cache-sized tiles, each [`Paneled`] into the +//! panels one leaf call consumes, plus a typed tail. [`Accumulate`] folds one (A-panel, +//! B-panel) pair into a slot of the [`Scratch`]; [`Drain`] consumes the finished +//! accumulator. //! -//! SIMD-accelerated implementation that uses block-transposed memory layout -//! for **query** vectors, with documents remaining in row-major format. +//! Position is ordinal. [`drive`] counts the panels it passes and hands a [`Drain`] an +//! A-panel index and a B-panel range — never a stride or an address, so a drain scales +//! those ordinals by its own panel width and clamps the end against its own extent. A side +//! whose tail type is uninhabited ([`NoTail`]) need not clamp at all, and that is a +//! type-level fact rather than a convention. //! -//! # Memory Layout -//! -//! - **Query**: Block-transposed (`GROUP` vectors per block, dimensions contiguous -//! within each block). The block size is determined by the kernel's `A_PANEL`. -//! - **Document**: Row-major (standard [`MatRef`](crate::multi_vector::MatRef) format). +//! The driver speaks A/B; instantiations speak Query/Doc. -pub(super) mod f16; -pub(super) mod f32; -mod layouts; -mod reduce; -mod tiled_reduce; +use core::ops::Range; -// ── Tile budget ────────────────────────────────────────────────── +mod f16; +mod float; +mod leaves; +mod strip; +mod tiles; -/// Cache budgets fed to the tile planner. -/// -/// `Default` returns the production budgets derived from hardcoded L1/L2 -/// cache-size estimates and fixed fractions. +pub(crate) use f16::MaxIpF16; +pub(crate) use float::MaxIp; + +// ── Tile budget and planning ───────────────────────────────────── + +/// Cache budgets fed to [`Plan::new`]. #[derive(Debug, Clone, Copy)] struct TileBudget { - /// L2 budget in bytes reserved for A tiles. + /// L2 bytes reserved for A tiles. l2_a: usize, - /// L1 budget in bytes reserved for B tiles (before A-panel subtraction). + /// L1 bytes reserved for B tiles, before the resident A-panel is subtracted. l1_b: usize, } @@ -48,63 +52,260 @@ impl Default for TileBudget { } } -// ── Kernel trait ───────────────────────────────────────────────── +/// Panel counts per tile: `a_panels` A-panels resident in L2, and as many B-panels as +/// co-fit L1 alongside one A-panel *and* the accumulator columns they feed. +#[derive(Debug, Clone, Copy)] +struct Plan { + a_panels: usize, + b_panels: usize, +} + +impl Plan { + /// Row sizes are clamped to one byte so a degenerate row cannot divide by zero; the + /// entries reject an empty contraction before planning, so such a plan is never used. + fn new(a_row_bytes: usize, b_row_bytes: usize, acc_bytes: usize, budget: TileBudget) -> Self { + let a_row_bytes = a_row_bytes.max(1); + let b_row_bytes = b_row_bytes.max(1); + + let a_panels = (budget.l2_a / (a_row_bytes * AR)).max(1); + + // A B-row costs its own bytes plus the accumulator column it fills, and one + // A-panel stays resident alongside. + let per_b_row = b_row_bytes + AR * acc_bytes; + let b_budget = budget.l1_b.saturating_sub(AR * a_row_bytes); + let b_panels = ((b_budget / per_b_row) / BR).max(1); -/// SIMD micro-kernel for the [`tiled_reduce`](tiled_reduce::tiled_reduce) loop. + Self { a_panels, b_panels } + } + + /// Accumulator elements for one A-panel against the widest B-tile this plan allows. + /// + /// Sizing the scratch below this trips [`slots_exhausted`]. + fn strip_len(&self) -> usize { + AR * self.b_panels * BR + } +} + +// ── Read side ──────────────────────────────────────────────────── + +/// Per-lifetime half of [`TileWalk`]. /// -/// The kernel only sees already-converted data: storage-layout to -/// kernel-layout conversion is handled at tile boundaries by -/// [`ConvertTo`](layouts::ConvertTo), so implementors can assume their input -/// pointers reference `::Element` / -/// `::Element` directly. +/// The defaulted `B = &'a Self` carries the `Self: 'a` implied bound through +/// well-formedness; a plain GAT `where Self: 'a` collapses to `'static` under [`drive`]'s +/// `for<'a>` bound on stable. +trait TileAt<'a, B = &'a Self> { + type Tile: Paneled; +} + +/// A **lending** walk: `next` reborrows `&mut self`, so a tile may borrow a buffer the walk +/// reuses on the following call — which is what lets a walk convert as it goes. `reset` +/// rewinds, because B is re-walked once per A-tile. +trait TileWalk: for<'a> TileAt<'a> { + fn next(&mut self) -> Option<>::Tile>; + fn reset(&mut self); +} + +/// An iterator whose short trailing element has its own type. /// -/// # Safety +/// `tail` consumes the exhausted iterator, so the trailer comes off the cursor the loop was +/// already advancing rather than being recomputed from the source. +trait TailIterator: ExactSizeIterator { + type Tail; + fn tail(self) -> Option; +} + +/// A tile that knows how it breaks into panels. /// -/// Implementors must respect the per-method `# Safety` contracts on -/// [`full_panel`](Self::full_panel) and [`partial_panel`](Self::partial_panel). -unsafe trait Kernel { - /// Layout consumed by the A (left / query) side of the micro-kernel. - type Left: layouts::Layout; - /// Layout consumed by the B (right / document) side of the micro-kernel. - type Right: layouts::Layout; - - /// Number of A rows processed per micro-kernel invocation. - const A_PANEL: usize; - /// Number of B rows processed per micro-kernel invocation. - const B_PANEL: usize; - - /// Process one full `A_PANEL × B_PANEL` micro-panel pair. - /// - /// # Safety - /// - /// * `a` must point to `A_PANEL * k` contiguous elements of - /// `::Element`. - /// * `b` must point to `B_PANEL * k` contiguous elements of - /// `::Element`. - /// * `r` must point to at least `A_PANEL` writable `f32` values. - unsafe fn full_panel( - arch: A, - a: *const ::Element, - b: *const ::Element, - k: usize, - r: *mut f32, - ); - - /// Dispatch for `1..(B_PANEL-1)` remainder B rows. - /// - /// # Safety - /// - /// * `a` must point to `A_PANEL * k` contiguous elements of - /// `::Element`. - /// * `b` must point to `remainder * k` contiguous elements of - /// `::Element`. - /// * `r` must point to at least `A_PANEL` writable `f32` values. - unsafe fn partial_panel( - arch: A, - remainder: usize, - a: *const ::Element, - b: *const ::Element, - k: usize, - r: *mut f32, - ); +/// Extent only, never position: a tile says how much it holds, and the driver decides where +/// that lands. +trait Paneled { + type Panel: Copy; + type Tail: Copy; + type Panels: TailIterator; + + fn panels(&self) -> Self::Panels; +} + +/// [`Paneled::Tail`] for a side padded to whole panels. Uninhabited, so `tail()` provably +/// returns `None` and no consumer on that side needs to clamp. +#[derive(Clone, Copy)] +enum NoTail {} + +// ── Write side ─────────────────────────────────────────────────── + +/// Per-lifetime half of [`Scratch`] — same implied-bound trick as [`TileAt`]. +trait SlotsAt<'s, B = &'s mut Self> { + /// Named here rather than reached through the iterator so [`drive`]'s bounds can + /// project it off the scratch. + type Slot; + /// Plain rather than lending: slots partition one buffer within a single call, so none + /// of them borrows the cursor. + type Slots: Iterator; +} + +/// The write-side mirror of [`Paneled`]: memory a fill carves into one slot per B-panel, +/// each disjoint from the last. +trait Scratch: for<'s> SlotsAt<'s> { + fn slots(&mut self) -> >::Slots; +} + +// ── Compute side ───────────────────────────────────────────────── + +/// One A-panel × one B-panel → one accumulator slot. +/// +/// Pinned on all three as type parameters, so the walks' panel types select the impl. +trait Accumulate { + fn accumulate(&self, arch: Arch, a: A, b: B, out: O); +} + +/// [`NoTail`] is uninhabited, so this discharges [`drive`]'s A-tail bounds for every kernel +/// — and by coherence forbids any kernel from writing its own. +impl Accumulate for K { + #[inline(always)] + fn accumulate(&self, _: Arch, a: NoTail, _: B, _: O) { + match a {} + } +} + +/// Consume a finished accumulator. +/// +/// `a_panel` and `b_panels` are ordinals in [`drive`]'s visit order, not addresses. +trait Drain { + fn drain(&mut self, arch: Arch, scratch: &S, a_panel: usize, b_panels: Range); +} + +// ── Driver ─────────────────────────────────────────────────────── + +type PanelOf<'a, W> = <>::Tile as Paneled>::Panel; +type TailOf<'a, W> = <>::Tile as Paneled>::Tail; +type SlotOf<'s, S> = >::Slot; + +#[cold] +#[inline(never)] +fn slots_exhausted() -> ! { + unreachable!("scratch ran out of slots — it is narrower than the walk's B-tile") +} + +/// One A-panel against a whole B-tile. +/// +/// Returns the slots it filled — ground truth for how far B advanced, rather than the +/// prediction an `ExactSizeIterator::len` would give, which excludes the tail anyway. +#[inline(always)] +fn fill(arch: Arch, kernel: &K, a: A, b_tile: &BT, scratch: &mut S) -> usize +where + Arch: Copy, + A: Copy, + BT: Paneled, + S: Scratch, + K: for<'s> Accumulate> + + for<'s> Accumulate>, +{ + let mut panels = b_tile.panels(); + let mut slots = scratch.slots(); + let mut filled = 0; + + for b in panels.by_ref() { + let Some(out) = slots.next() else { + slots_exhausted() + }; + kernel.accumulate(arch, a, b, out); + filled += 1; + } + // The tail draws from the same cursor as the full panels. + if let Some(b) = panels.tail() { + let Some(out) = slots.next() else { + slots_exhausted() + }; + kernel.accumulate(arch, a, b, out); + filled += 1; + } + + filled +} + +/// Drive one A source against one B source, re-walking B once per A-tile. +/// +/// `scratch` precedes `kernel` deliberately: `S` must be resolved before the kernel's +/// [`Accumulate`] bounds are proved, or their slot type is still `<_ as SlotsAt>::Slot` +/// and no impl matches. +fn drive( + arch: Arch, + mut a_walk: AW, + mut b_walk: BW, + scratch: &mut S, + kernel: &K, + drain: &mut D, +) where + Arch: Copy, + AW: TileWalk, + BW: TileWalk, + S: Scratch, + K: for<'a, 'b, 's> Accumulate, PanelOf<'b, BW>, SlotOf<'s, S>> + + for<'a, 'b, 's> Accumulate, TailOf<'b, BW>, SlotOf<'s, S>> + + for<'a, 'b, 's> Accumulate, PanelOf<'b, BW>, SlotOf<'s, S>> + + for<'a, 'b, 's> Accumulate, TailOf<'b, BW>, SlotOf<'s, S>>, + D: Drain, +{ + let mut a_base = 0; + while let Some(a_tile) = a_walk.next() { + b_walk.reset(); + let mut b_base = 0; + // Last pass wins: every B-tile re-sweeps the same A-panels, so both counters are + // rewritten identically each pass and read after the last. An A-tile with no + // B-tiles advances neither, which is unobservable — no drain fires, and a B source + // empty for one A-tile is empty for all. + let mut a_end = a_base; + let mut b_used = 0; + + while let Some(b_tile) = b_walk.next() { + let mut a_panel = a_base; + let mut a_panels = a_tile.panels(); + + for panel in a_panels.by_ref() { + b_used = fill(arch, kernel, panel, &b_tile, scratch); + drain.drain(arch, scratch, a_panel, b_base..b_base + b_used); + a_panel += 1; + } + if let Some(panel) = a_panels.tail() { + b_used = fill(arch, kernel, panel, &b_tile, scratch); + drain.drain(arch, scratch, a_panel, b_base..b_base + b_used); + a_panel += 1; + } + + a_end = a_panel; + b_base += b_used; + } + a_base = a_end; + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn plan_reserves_l1_for_the_resident_a_panel_and_accumulator() { + // 64-byte rows, AR = 16, BR = 4, 4-byte accumulator. + // l2_a 40960 / (64 * 16) = 40 A-panels. + // l1_b 36000 - 16 * 64 = 34976 for B; per B-row = 64 + 16 * 4 = 128; + // 34976 / 128 = 273 rows -> 273 / 4 = 68 B-panels. + let plan = Plan::<16, 4>::new( + 64, + 64, + 4, + TileBudget { + l2_a: 40960, + l1_b: 36000, + }, + ); + assert_eq!((plan.a_panels, plan.b_panels), (40, 68)); + assert_eq!(plan.strip_len(), 16 * 68 * 4); + } + + #[test] + fn plan_clamps_to_one_panel_per_tile() { + let plan = Plan::<16, 4>::new(1024, 1024, 4, TileBudget { l2_a: 1, l1_b: 1 }); + assert_eq!((plan.a_panels, plan.b_panels), (1, 1)); + assert_eq!(plan.strip_len(), 16 * 4); + } } diff --git a/diskann-quantization/src/multi_vector/distance/kernels/reduce.rs b/diskann-quantization/src/multi_vector/distance/kernels/reduce.rs deleted file mode 100644 index d3dfe85d15..0000000000 --- a/diskann-quantization/src/multi_vector/distance/kernels/reduce.rs +++ /dev/null @@ -1,106 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -//! Compile-time unroll reduction over fixed-size accumulator arrays. -//! -//! Shared by every micro-kernel family (f32, f16, future u8/i8, …): each -//! kernel keeps `UNROLL` independent SIMD accumulators in the inner loop and -//! folds them down to a single value at the end with a caller-supplied binary -//! operator (e.g. `max_simd`). -//! -//! Implementations are provided for `[T; 1..=4]`, matching the unroll factors -//! currently used by the kernels. The 4-element fold is balanced (`(a⊕b)⊕(c⊕d)`) -//! to shorten the dependency chain; 2- and 3-element folds are left-associative. - -/// Compile-time unroll reduction over fixed-size arrays. -/// -/// Used by the micro-kernels to reduce `UNROLL` accumulators into a single -/// value using a caller-supplied binary operator (e.g. `max_simd`). -pub(super) trait Reduce { - type Element; - fn reduce(&self, f: &F) -> Self::Element - where - F: Fn(Self::Element, Self::Element) -> Self::Element; -} - -impl Reduce for [T; 1] { - type Element = T; - - #[inline(always)] - fn reduce(&self, _f: &F) -> T - where - F: Fn(T, T) -> T, - { - self[0] - } -} - -impl Reduce for [T; 2] { - type Element = T; - - #[inline(always)] - fn reduce(&self, f: &F) -> T - where - F: Fn(T, T) -> T, - { - f(self[0], self[1]) - } -} - -impl Reduce for [T; 3] { - type Element = T; - - #[inline(always)] - fn reduce(&self, f: &F) -> T - where - F: Fn(T, T) -> T, - { - f(f(self[0], self[1]), self[2]) - } -} - -impl Reduce for [T; 4] { - type Element = T; - - #[inline(always)] - fn reduce(&self, f: &F) -> T - where - F: Fn(T, T) -> T, - { - f(f(self[0], self[1]), f(self[2], self[3])) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn reduce_folds_correctly() { - let max = |a: f32, b: f32| a.max(b); - assert_eq!([5.0f32].reduce(&max), 5.0); - assert_eq!([1.0f32, 3.0].reduce(&max), 3.0); - assert_eq!([2.0f32, 1.0, 4.0].reduce(&max), 4.0); - assert_eq!([3.0f32, 1.0, 4.0, 2.0].reduce(&max), 4.0); - } - - /// Verify the exact fold order of each `Reduce` impl using a - /// non-commutative operator (subtraction). - /// - /// - `[a; 1]` → `a` - /// - `[a, b; 2]` → `a - b` - /// - `[a, b, c; 3]` → `(a - b) - c` (left fold) - /// - `[a, b, c, d; 4]` → `(a - b) - (c - d)` (balanced tree) - #[test] - fn reduce_fold_order() { - let sub = |a: f32, b: f32| a - b; - // [10] → 10 - assert_eq!([10.0f32].reduce(&sub), 10.0); - // [10, 3] → 10 - 3 = 7 - assert_eq!([10.0f32, 3.0].reduce(&sub), 7.0); - // [10, 3, 1] → (10 - 3) - 1 = 6 - assert_eq!([10.0f32, 3.0, 1.0].reduce(&sub), 6.0); - // [10, 3, 1, 2] → (10 - 3) - (1 - 2) = 7 - (-1) = 8 - assert_eq!([10.0f32, 3.0, 1.0, 2.0].reduce(&sub), 8.0); - } -} diff --git a/diskann-quantization/src/multi_vector/distance/kernels/strip.rs b/diskann-quantization/src/multi_vector/distance/kernels/strip.rs new file mode 100644 index 0000000000..30e3f94b7e --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/strip.rs @@ -0,0 +1,104 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +//! The accumulator a fill writes and a drain reads. +//! +//! Rows are query rows and one column is one doc — note that a doc arrives as a *row* of +//! the input and lands as a *column* here. Column-major over the whole strip: column `c` +//! occupies `[c * AR, (c + 1) * AR)`, and slot `p` covers columns `p * BR ..`, so a drain +//! can address a run of columns without knowing which slot produced them. + +use core::slice::ChunksExactMut; + +use super::{Scratch, SlotsAt}; + +/// Accumulator for one A-panel against one whole B-tile. +/// +/// Holds no cursor: [`Scratch::slots`] restarts from the front on every fill, which is +/// what lets the same memory be re-lent across tiles without clearing. +pub(super) struct Strip<'a, T, const AR: usize, const BR: usize> { + buf: &'a mut [T], +} + +impl<'a, T, const AR: usize, const BR: usize> Strip<'a, T, AR, BR> { + /// Trailing elements beyond the last whole slot are never touched. + pub(super) fn new(buf: &'a mut [T]) -> Self { + Self { buf } + } + + /// The first `live` columns — the only region the fill just performed wrote. + /// + /// # Panics + /// + /// Panics if `live` exceeds the strip's column capacity. + pub(super) fn columns(&self, live: usize) -> &[T] { + &self.buf[..live * AR] + } +} + +/// One `AR × BR` accumulator tile, column-major. +pub(super) struct Slot<'a, T, const AR: usize, const BR: usize> { + buf: &'a mut [T], +} + +impl Slot<'_, T, AR, BR> { + /// Exactly `AR * BR` elements: [`Slots`] only ever cuts whole tiles. + pub(super) fn as_mut_slice(&mut self) -> &mut [T] { + &mut *self.buf + } +} + +/// Cuts a strip into disjoint slots, stopping short of a trailing partial tile. +/// +/// Disjointness and that stopping rule are both structural, from `chunks_exact_mut`. +pub(super) struct Slots<'a, T, const AR: usize, const BR: usize>(ChunksExactMut<'a, T>); + +impl<'a, T, const AR: usize, const BR: usize> Iterator for Slots<'a, T, AR, BR> { + type Item = Slot<'a, T, AR, BR>; + + fn next(&mut self) -> Option { + self.0.next().map(|buf| Slot { buf }) + } +} + +impl<'s, T, const AR: usize, const BR: usize> SlotsAt<'s> for Strip<'_, T, AR, BR> { + type Slot = Slot<'s, T, AR, BR>; + type Slots = Slots<'s, T, AR, BR>; +} + +impl Scratch for Strip<'_, T, AR, BR> { + fn slots(&mut self) -> Slots<'_, T, AR, BR> { + Slots(self.buf.chunks_exact_mut(AR * BR)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn slots_partition_the_strip_and_stop_short_of_a_partial_tile() { + let mut buf = [0u32; 4 * 2 * 3 + 5]; + let mut strip = Strip::::new(&mut buf); + + for (n, mut slot) in strip.slots().enumerate() { + slot.as_mut_slice().fill(n as u32 + 1); + } + assert_eq!(strip.slots().count(), 3); + + // The 5 trailing elements are shorter than a slot, so no slot covers them. + assert_eq!(&buf[..8], &[1; 8]); + assert_eq!(&buf[16..24], &[3; 8]); + assert_eq!(&buf[24..], &[0; 5]); + } + + #[test] + fn columns_cut_ignores_slot_boundaries() { + let mut buf: [u32; 24] = core::array::from_fn(|i| i as u32); + let strip = Strip::::new(&mut buf); + + // Three columns straddle the first slot (2 columns) into the second. + assert_eq!(strip.columns(3), &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]); + assert_eq!(strip.columns(0), &[] as &[u32]); + } +} diff --git a/diskann-quantization/src/multi_vector/distance/kernels/tiled_reduce.rs b/diskann-quantization/src/multi_vector/distance/kernels/tiled_reduce.rs deleted file mode 100644 index 6e4b76e9c1..0000000000 --- a/diskann-quantization/src/multi_vector/distance/kernels/tiled_reduce.rs +++ /dev/null @@ -1,806 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. - -//! Generic tiling loop for reducing-GEMM micro-kernels. -//! -//! # Tiling Strategy -//! -//! This approach uses a reducing-GEMM pattern modeled after high-performance BLAS -//! implementations: -//! -//! - **L2 cache**: Tiles of A (conventionally the query) are sized to fit in L2. -//! - **L1 cache**: Tiles of B (conventionally the document) plus one micro-panel -//! of A are sized to fit in L1. -//! - **Micro-kernel**: An `A_PANEL × B_PANEL` micro-kernel (e.g. 16×4 for f32 on V3) -//! processes a panel of A rows against a panel of B rows per invocation, -//! accumulating max-IP into a scratch buffer. The panel sizes are determined -//! by the `Kernel` implementation for each element type. -//! -//! The loop itself is layout-agnostic: A and B are described by the generic -//! `LA`/`LB` parameters and converted to the kernel's expected layouts via -//! [`ConvertTo`] at tile boundaries. The current micro-kernels happen to want a -//! block-transposed A and a row-major B, but `tiled_reduce` does not require -//! either — any layout pair satisfying the `ConvertTo` bounds works. - -use diskann_wide::Architecture; - -use super::layouts::{ConvertTo, Layout}; -use super::{Kernel, TileBudget}; - -// ── Tile planner ───────────────────────────────────────────────── - -/// Tile-panel counts derived from cache budgets. -#[derive(Debug, Clone, Copy)] -struct FullReduce { - a_panels_per_tile: usize, - - b_panels_per_tile: usize, -} - -impl FullReduce { - /// Compute A-tile and B-tile panel counts from cache budgets. - /// - /// The L1 budget is reduced by one A micro-panel before splitting it into - /// B panels, since both must coexist in L1 during the inner loop. - fn new( - a_row_bytes: usize, - b_row_bytes: usize, - a_panel: usize, - b_panel: usize, - budget: TileBudget, - ) -> Self { - let a_row_bytes = a_row_bytes.max(1); - let b_row_bytes = b_row_bytes.max(1); - - let a_panels_per_tile = (budget.l2_a / (a_row_bytes * a_panel)).max(1); - - let a_panel_bytes = a_panel * a_row_bytes; - let b_tile_budget = budget.l1_b.saturating_sub(a_panel_bytes); - let b_panels_per_tile = (b_tile_budget / (b_row_bytes * b_panel)).max(1); - - Self { - a_panels_per_tile, - b_panels_per_tile, - } - } -} - -// ── Generic tiled reduce ───────────────────────────────────────── - -/// Execute the 5-level tiling loop with a pluggable SIMD micro-kernel and -/// tile-level layout converters. -/// -/// The loop nest is: -/// ```text -/// Loop 1: A tiles (sized to L2) — convert via `ca` -/// Loop 2: B tiles (sized to L1) — convert via `cb` -/// Loop 3: A panels (micro-panels within converted A tile) -/// Loop 4: B panels (micro-panels within converted B tile) -/// Loop 5: k (contraction dim, inside K::full_panel / K::partial_panel) -/// ``` -/// -/// Conversion from storage layout to kernel layout happens once per tile -/// (not per panel), amortizing cost over the entire tile. -/// -/// # Safety -/// -/// * `a_ptr` must be valid for `a_padded_nrows * k` elements of `AElem`. -/// * `a_padded_nrows` must be a multiple of `K::A_PANEL`. -/// * `b_ptr` must be valid for `b_nrows * k` elements of `BElem`. -/// * `scratch` must have length ≥ `a_padded_nrows` and be initialized by caller. -#[allow(clippy::too_many_arguments)] -pub(super) unsafe fn tiled_reduce( - arch: A, - ca: &LA, - cb: &LB, - a_ptr: *const LA::Element, - a_padded_nrows: usize, - b_ptr: *const LB::Element, - b_nrows: usize, - k: usize, - scratch: &mut [f32], - budget: TileBudget, -) where - A: Architecture, - K: Kernel, - LA: ConvertTo, - LB: ConvertTo, -{ - let a_row_bytes = k * std::mem::size_of::<::Element>(); - let b_row_bytes = k * std::mem::size_of::<::Element>(); - let plan = FullReduce::new(a_row_bytes, b_row_bytes, K::A_PANEL, K::B_PANEL, budget); - - let b_src_panel_stride = K::B_PANEL * k; - let b_src_tile_stride = b_src_panel_stride * plan.b_panels_per_tile; - - let a_kern_panel_stride = K::A_PANEL * k; - let b_kern_panel_stride = K::B_PANEL * k; - - let b_remainder = b_nrows % K::B_PANEL; - - assert_eq!( - a_padded_nrows % K::A_PANEL, - 0, - "a_padded_nrows ({a_padded_nrows}) must be a multiple of A_PANEL ({})", - K::A_PANEL, - ); - - // Zero-dimensional vectors have IP = 0 for every pair. Fill scratch and - // return to avoid zero-stride infinite loops in the tiling nest. - if k == 0 { - if b_nrows > 0 { - scratch[..a_padded_nrows].fill(0.0); - } - return; - } - - // Cap by actual data — planned tile can vastly exceed it, and some - // ConvertTo impls (e.g. f16) allocate per call. - let a_tile_rows = K::A_PANEL * plan.a_panels_per_tile; - let b_tile_rows = K::B_PANEL * plan.b_panels_per_tile; - let mut a_buf = ca.new_buffer(a_tile_rows.min(a_padded_nrows), k); - let mut b_buf = cb.new_buffer(b_tile_rows.min(b_nrows), k); - - // SAFETY: Caller guarantees b_ptr is valid for b_nrows * k elements. - let pb_end = unsafe { b_ptr.add(b_nrows * k) }; - // SAFETY: b_remainder < B_PANEL, so pb_end - b_remainder * k is within allocation. - let pb_full_end = unsafe { pb_end.sub(b_remainder * k) }; - - // SAFETY: All pointer arithmetic stays within the respective allocations. - unsafe { - let mut rows_done: usize = 0; - - // Loop 1: Tiles of `A`. - while rows_done < a_padded_nrows { - let tile_rows = a_tile_rows.min(a_padded_nrows - rows_done); - let pa_tile_src = a_ptr.add(rows_done * k); - // SAFETY: rows_done < a_padded_nrows (loop condition), so the - // pointer is in-bounds. - let pr_tile = scratch.as_mut_ptr().add(rows_done); - - // Convert A tile from storage layout to kernel layout. - let pa_tile = ca.convert(&mut a_buf, arch, pa_tile_src, tile_rows, k); - let pa_tile_end = pa_tile.add(tile_rows * k); - - let mut pb_tile_src = b_ptr; - - // ── Section A: Full B-tiles ───────────────────────────── - // Every panel in every B-tile here is complete; no - // remainder check inside Loop 4. - // - // Loop 2: Full B-tiles (every panel in the tile is complete). - // SAFETY: `pb_tile_src` is always in `[b_ptr, pb_full_end]` — both within - // the same allocation — so `offset_from` is well-defined. - while pb_full_end.offset_from(pb_tile_src) >= b_src_tile_stride as isize { - // Convert B tile from storage layout to kernel layout. - let pb_tile = cb.convert(&mut b_buf, arch, pb_tile_src, b_tile_rows, k); - let pb_tile_end = pb_tile.add(b_tile_rows * k); - - let mut pa_panel = pa_tile; - let mut pr_panel = pr_tile; - - // Loop 3: Micro-panels of `A`. - while pa_panel < pa_tile_end { - let mut pb_panel = pb_tile; - - // Loop 4: Micro-panels of `B` (all full, no remainder check). - while pb_panel < pb_tile_end { - K::full_panel(arch, pa_panel, pb_panel, k, pr_panel); - pb_panel = pb_panel.add(b_kern_panel_stride); - } - - pa_panel = pa_panel.add(a_kern_panel_stride); - pr_panel = pr_panel.add(K::A_PANEL); - } - pb_tile_src = pb_tile_src.add(b_src_tile_stride); - } - - // ── Section B: Peeled trailing B-tile ─────────────────── - // Holds whatever B-rows Section A could not consume: zero - // or more full B-panels followed by an optional partial - // panel of `b_remainder` rows. Skipped entirely when - // Section A consumed all of B. - // - // Peeled last B-tile: contains remaining full panels + remainder rows. - if pb_tile_src < pb_end { - let remaining_b_rows = b_nrows - ((pb_tile_src.offset_from(b_ptr) as usize) / k); - // Convert remaining B rows. - let pb_tile = cb.convert(&mut b_buf, arch, pb_tile_src, remaining_b_rows, k); - - let full_panels_in_remainder = remaining_b_rows / K::B_PANEL; - let pb_full_end_local = pb_tile.add(full_panels_in_remainder * b_kern_panel_stride); - - let mut pa_panel = pa_tile; - let mut pr_panel = pr_tile; - - // Loop 3 (peeled): Micro-panels of `A`. - while pa_panel < pa_tile_end { - let mut pb_panel = pb_tile; - - // Loop 4 (peeled): Full B-panels in the last tile. - while pb_panel < pb_full_end_local { - K::full_panel(arch, pa_panel, pb_panel, k, pr_panel); - pb_panel = pb_panel.add(b_kern_panel_stride); - } - - // Remainder dispatch: 1..(B_PANEL-1) leftover B-rows. - if b_remainder > 0 { - K::partial_panel(arch, b_remainder, pa_panel, pb_panel, k, pr_panel); - } - - pa_panel = pa_panel.add(a_kern_panel_stride); - pr_panel = pr_panel.add(K::A_PANEL); - } - } - - rows_done += tile_rows; - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - use diskann_wide::arch::Scalar; - - use super::super::f32::{F32Kernel, max_ip_kernel}; - use super::super::layouts; - use crate::multi_vector::{BlockTransposed, MatRef, Standard}; - - #[test] - fn basic_panel_counts() { - // 16 A-rows × 256 bytes/row = 4096 bytes per A-panel. - // L2 budget 40960 → 40960 / 4096 = 10 A-panels. - // One A-panel = 4096 bytes, L1 budget 36000 → 36000 - 4096 = 31904. - // 4 B-rows × 256 bytes/row = 1024 bytes per B-panel. - // 31904 / 1024 = 31 B-panels. - let plan = FullReduce::new( - 256, - 256, - 16, - 4, - TileBudget { - l2_a: 40960, - l1_b: 36000, - }, - ); - assert_eq!(plan.a_panels_per_tile, 10); - assert_eq!(plan.b_panels_per_tile, 31); - } - - #[test] - fn tiny_budget_clamps_to_one() { - // Budget too small for even one panel — clamp to 1. - let plan = FullReduce::new(1024, 1024, 16, 4, TileBudget { l2_a: 1, l1_b: 1 }); - assert_eq!(plan.a_panels_per_tile, 1); - assert_eq!(plan.b_panels_per_tile, 1); - } - - #[test] - fn zero_byte_rows_clamped() { - // Zero-byte rows (e.g. k=0) should not divide by zero. - // FullReduce clamps row bytes to max(1), so a_row_bytes=1, b_row_bytes=1. - let plan = FullReduce::new( - 0, - 0, - 16, - 4, - TileBudget { - l2_a: 100_000, - l1_b: 50_000, - }, - ); - // a_panels = 100_000 / (1 * 16) = 6250 - assert_eq!(plan.a_panels_per_tile, 6250); - // a_panel_bytes = 16 * 1 = 16. b_tile_budget = 50_000 - 16 = 49_984. - // b_panels = 49_984 / (1 * 4) = 12_496 - assert_eq!(plan.b_panels_per_tile, 12_496); - } - - #[test] - fn exact_fit_one_panel() { - // Budget exactly fits one A-panel (16 × 64 = 1024 bytes). - // No room for a second → a_panels = 1. - let plan = FullReduce::new( - 64, - 64, - 16, - 4, - TileBudget { - l2_a: 1024, - l1_b: 2048, - }, - ); - assert_eq!(plan.a_panels_per_tile, 1); - // L1: 2048 - 16*64(=1024) = 1024 for B. 4*64=256 per B-panel → 4 panels. - assert_eq!(plan.b_panels_per_tile, 4); - } - - #[test] - fn l1_saturated_by_a_panel() { - // A-panel alone exceeds L1 budget → b_tile_budget saturates to 0, - // b_panels_per_tile clamps to 1. - let plan = FullReduce::new( - 1024, - 64, - 16, - 4, - TileBudget { - l2_a: 100_000, - l1_b: 100, - }, - ); - assert_eq!(plan.b_panels_per_tile, 1); - } - - #[test] - #[should_panic(expected = "must be a multiple of A_PANEL")] - fn panics_on_unaligned_a_rows() { - let k = 4; - // 9 is not a multiple of A_PANEL (8). - let a = vec![0.0f32; 9 * k]; - let b = vec![0.0f32; 2 * k]; - let mut scratch = vec![f32::MIN; 16]; - - let ca = layouts::BlockTransposed::::new(); - let cb = layouts::RowMajor::::new(); - - // SAFETY: pointers and scratch are correctly sized; we expect a panic. - unsafe { - super::tiled_reduce::, _, _>( - Scalar::new(), - &ca, - &cb, - a.as_ptr(), - 9, - b.as_ptr(), - 2, - k, - &mut scratch, - TileBudget::default(), - ); - } - } - - #[test] - fn zero_dim_fills_scratch_and_returns() { - let a_rows = 8; - let b_rows = 3; - let k = 0; - - let a = Vec::::new(); - let b = Vec::::new(); - let mut scratch = vec![f32::MIN; a_rows]; - - let ca = layouts::BlockTransposed::::new(); - let cb = layouts::RowMajor::::new(); - - // SAFETY: k == 0 so no elements are read; pointers are never dereferenced. - unsafe { - super::tiled_reduce::, _, _>( - Scalar::new(), - &ca, - &cb, - a.as_ptr(), - a_rows, - b.as_ptr(), - b_rows, - k, - &mut scratch, - TileBudget::default(), - ); - } - - for &v in &scratch { - assert_eq!(v, 0.0, "zero-dim IP should be 0.0"); - } - } - - #[test] - fn zero_dim_zero_docs_leaves_scratch_untouched() { - let a_rows = 8; - let mut scratch = vec![f32::MIN; a_rows]; - - let ca = layouts::BlockTransposed::::new(); - let cb = layouts::RowMajor::::new(); - - // SAFETY: k == 0, b_nrows == 0; no elements read. - unsafe { - super::tiled_reduce::, _, _>( - Scalar::new(), - &ca, - &cb, - [].as_ptr(), - a_rows, - [].as_ptr(), - 0, - 0, - &mut scratch, - TileBudget::default(), - ); - } - - for &v in &scratch { - assert_eq!(v, f32::MIN, "zero docs should leave scratch untouched"); - } - } - - // Shape matrix for the breadth sweeps `tiled_reduce_f32_matches_naive` - // and `tiled_reduce_f16_matches_naive`. Covers degenerate shapes, - // prime `k`, exact and off-by-one `A_PANEL` boundaries on both - // `GROUP=8` and `GROUP=16`, and every B-row remainder class. Run - // under `TileBudget::default()`, which keeps every shape in a - // single tile (Section A skipped, Section B end-to-end). - // - // Multi-tile / multi-panel structural coverage lives in - // `tiled_reduce_all_loop_paths_match_naive`. - // - // (a_nrows, b_nrows, dim) - const NAIVE_CASES: &[(usize, usize, usize)] = &[ - (1, 1, 1), // Degenerate single-element - (1, 1, 2), // Minimal non-trivial - (1, 1, 4), // Single query, single doc - (1, 5, 8), // Single query, multiple docs - (5, 1, 8), // Multiple queries, single doc - (3, 2, 0), // Zero dimensions, both have rows - (3, 0, 4), // Zero docs - (3, 2, 3), // Prime k - (3, 4, 16), // General case - (5, 3, 5), // Prime k, A-remainder on aarch64 - (7, 7, 32), // Square case - (2, 3, 7), // k not divisible by SIMD lanes - (2, 3, 128), // Larger dimension - (8, 3, 4), // Single A-panel (Scalar), B remainder - (16, 5, 8), // Two A-panels (Scalar), B remainder - (16, 4, 64), // Two A-panels (Scalar), no B remainder; one (V3) - (17, 4, 64), // A-panel remainder on both Scalar and V3 - (32, 5, 16), // Multiple full A-panels, B remainder - (48, 3, 16), // 6 A-panels (Scalar) / 3 (V3) - (16, 6, 32), // V3 B remainder=2 - (16, 7, 32), // V3 B remainder=3 - (16, 8, 32), // No B remainder on either - ]; - - fn naive_max_ip_f32( - a: &[f32], - a_nrows: usize, - b: &[f32], - b_nrows: usize, - k: usize, - ) -> Vec { - (0..a_nrows) - .map(|i| { - (0..b_nrows) - .map(|j| (0..k).map(|d| a[i * k + d] * b[j * k + d]).sum::()) - .fold(f32::MIN, f32::max) - }) - .collect() - } - - fn naive_max_ip_f16( - a: &[half::f16], - a_nrows: usize, - b: &[half::f16], - b_nrows: usize, - k: usize, - ) -> Vec { - (0..a_nrows) - .map(|i| { - (0..b_nrows) - .map(|j| { - (0..k) - .map(|d| a[i * k + d].to_f32() * b[j * k + d].to_f32()) - .sum::() - }) - .fold(f32::MIN, f32::max) - }) - .collect() - } - - /// Run `max_ip_kernel::` against the naive reference - /// for one shape under `TileBudget::default()`. `arch_label` - /// identifies the arch branch in failure messages. - #[allow(clippy::too_many_arguments)] - fn check_kernel( - arch: A, - arch_label: &str, - tol: f32, - a_data: &[T], - a_nrows: usize, - b_data: &[T], - b_nrows: usize, - dim: usize, - expected: &[f32], - ) where - A: Architecture, - T: Copy + Default, - F32Kernel: Kernel, - layouts::BlockTransposed: - ConvertTo as Kernel>::Left> + Layout, - layouts::RowMajor: - ConvertTo as Kernel>::Right> + Layout, - { - let a_mat = MatRef::new(Standard::new(a_nrows, dim).unwrap(), a_data).unwrap(); - let a_bt = BlockTransposed::::from_matrix_view(a_mat.as_matrix_view()); - let b_mat = MatRef::new(Standard::new(b_nrows, dim).unwrap(), b_data).unwrap(); - - let mut scratch = vec![f32::MIN; a_bt.padded_nrows()]; - max_ip_kernel::( - arch, - a_bt.as_view(), - b_mat, - &mut scratch, - TileBudget::default(), - ); - - for i in 0..a_nrows { - let actual = scratch[i]; - let exp = expected[i]; - assert!( - (actual - exp).abs() < tol, - "[{arch_label}] row {i} mismatch for ({a_nrows},{b_nrows},{dim}): actual={actual}, expected={exp}", - ); - } - } - - /// Breadth sweep: every `NAIVE_CASES` shape, both f32 micro-kernels - /// (`F32Kernel<8>` Scalar; `F32Kernel<16>` V3 on x86_64 hosts with - /// AVX2+FMA), under `TileBudget::default()`. - /// - /// The V3 branch compiles on all targets but only executes on - /// x86_64 hosts that expose AVX2+FMA at runtime; CI's native - /// x86_64 runners and `sde-avx512-tests` cover this path. - #[test] - fn tiled_reduce_f32_matches_naive() { - for &(a_nrows, b_nrows, dim) in NAIVE_CASES { - let a_data: Vec = (0..a_nrows * dim).map(|i| (i + 1) as f32).collect(); - let b_data: Vec = (0..b_nrows * dim).map(|i| ((i + 1) * 2) as f32).collect(); - let expected = naive_max_ip_f32(&a_data, a_nrows, &b_data, b_nrows, dim); - - check_kernel::<_, f32, 8>( - Scalar::new(), - "scalar", - 1e-10, - &a_data, - a_nrows, - &b_data, - b_nrows, - dim, - &expected, - ); - - #[cfg(target_arch = "x86_64")] - if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { - check_kernel::<_, f32, 16>( - arch, - "x86-64-v3", - 1e-10, - &a_data, - a_nrows, - &b_data, - b_nrows, - dim, - &expected, - ); - } - } - } - - /// Structural coverage of every loop and section in `tiled_reduce`, - /// for every registered (arch, element-type) pair. - /// - /// For each [`LoopCoveragePlan`] in [`LOOP_COVERAGE_PLANS`], the - /// budget and shape are derived from `K::A_PANEL` / `K::B_PANEL` - /// so the realised tile plan matches the declared one (asserted - /// via [`FullReduce::new`]) and drives the loop nest down the - /// declared path. Across the three plans the following paths each - /// fire at least once on every (arch, T): - /// - /// * Loop 1 multi-A-tile (full + short). - /// * Loop 2 multi-iter (Section A consumes ≥2 full B-tiles). - /// * Section A's Loops 3 and 4 with ≥2 iters each. - /// * Section B fully-skipped. - /// * Section B with ≥2 leftover full B-panels. - /// * Section B's `partial_panel` dispatch. - /// - /// Asymmetric `(a_panels_per_tile, b_panels_per_tile)` rows catch - /// stride-confusion bugs that a square plan would mask. Running - /// both `T = f32` and `T = half::f16` re-enters the f16 per-tile - /// `ConvertTo` buffer on every multi-tile iteration. - #[test] - fn tiled_reduce_all_loop_paths_match_naive() { - check_tile_plan_paths::<_, f32, 8>(Scalar::new(), "scalar", gen_f32_data, naive_max_ip_f32); - check_tile_plan_paths::<_, half::f16, 8>( - Scalar::new(), - "scalar", - gen_f16_data, - naive_max_ip_f16, - ); - #[cfg(target_arch = "x86_64")] - if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { - check_tile_plan_paths::<_, f32, 16>(arch, "x86-64-v3", gen_f32_data, naive_max_ip_f32); - check_tile_plan_paths::<_, half::f16, 16>( - arch, - "x86-64-v3", - gen_f16_data, - naive_max_ip_f16, - ); - } - } - - /// One row of [`LOOP_COVERAGE_PLANS`]: target tile plan + trailing - /// remainder shape. See `tiled_reduce_all_loop_paths_match_naive` - /// for the full coverage story. - #[derive(Debug, Clone, Copy)] - struct LoopCoveragePlan { - /// Target `FullReduce::a_panels_per_tile` (Section A Loop 3 iters). - a_panels_per_tile: usize, - /// Target `FullReduce::b_panels_per_tile` (Section A Loop 4 iters). - b_panels_per_tile: usize, - /// Full B-tiles Section A's Loop 2 should consume. - section_a_b_tile_iters: usize, - /// Full B-panels left in Section B's trailing tile. - section_b_full_b_panels: usize, - /// Adds a 1-row B-remainder (drives `partial_panel`). - section_b_has_b_remainder: bool, - } - - /// Tile-plan rows exercised by - /// `tiled_reduce_all_loop_paths_match_naive`. Asymmetric - /// `(a_panels_per_tile, b_panels_per_tile)` pairs surface - /// stride-confusion bugs. - const LOOP_COVERAGE_PLANS: &[LoopCoveragePlan] = &[ - // Section A multi-panel + Section B with `partial_panel` only. - LoopCoveragePlan { - a_panels_per_tile: 2, - b_panels_per_tile: 2, - section_a_b_tile_iters: 1, - section_b_full_b_panels: 0, - section_b_has_b_remainder: true, - }, - // Section A's Loop 2 iterates twice; no Section B. - LoopCoveragePlan { - a_panels_per_tile: 2, - b_panels_per_tile: 3, - section_a_b_tile_iters: 2, - section_b_full_b_panels: 0, - section_b_has_b_remainder: false, - }, - // Section B carries multiple full B-panels plus a remainder. - LoopCoveragePlan { - a_panels_per_tile: 3, - b_panels_per_tile: 2, - section_a_b_tile_iters: 1, - section_b_full_b_panels: 2, - section_b_has_b_remainder: true, - }, - ]; - - fn gen_f32_data(len: usize, ceil: usize) -> Vec { - (0..len).map(|i| (i % ceil) as f32).collect() - } - - fn gen_f16_data(len: usize, ceil: usize) -> Vec { - (0..len) - .map(|i| diskann_wide::cast_f32_to_f16((i % ceil) as f32)) - .collect() - } - - type NaiveMaxIp = fn(&[T], usize, &[T], usize, usize) -> Vec; - - fn check_tile_plan_paths( - arch: A, - arch_label: &str, - gen_data: fn(usize, usize) -> Vec, - naive: NaiveMaxIp, - ) where - A: Architecture, - T: Copy + Default, - F32Kernel: Kernel, - layouts::BlockTransposed: - ConvertTo as Kernel>::Left> + Layout, - layouts::RowMajor: - ConvertTo as Kernel>::Right> + Layout, - { - let a_panel = as Kernel>::A_PANEL; - let b_panel = as Kernel>::B_PANEL; - let dim = 8usize; - let row_bytes = dim * std::mem::size_of::(); - - for &p in LOOP_COVERAGE_PLANS { - // Exact-fit budget for the declared tile plan; one A-panel - // is reserved in L1 alongside the B tile, matching - // `FullReduce::new`'s convention. - let budget = TileBudget { - l2_a: p.a_panels_per_tile * a_panel * row_bytes, - l1_b: a_panel * row_bytes + p.b_panels_per_tile * b_panel * row_bytes, - }; - - // Pin the realised plan so planner drift fails loudly - // rather than silently turning the test into a no-op. - let plan = FullReduce::new(row_bytes, row_bytes, a_panel, b_panel, budget); - assert_eq!( - plan.a_panels_per_tile, p.a_panels_per_tile, - "[{arch_label}] a_panels_per_tile for plan {p:?}", - ); - assert_eq!( - plan.b_panels_per_tile, p.b_panels_per_tile, - "[{arch_label}] b_panels_per_tile for plan {p:?}", - ); - - // a_nrows = a_panels_per_tile*A_PANEL + 1 → Loop 1 iterates twice. - // b_nrows packs: Section A consumption | Section B full panels | optional 1-row remainder. - let a_nrows = p.a_panels_per_tile * a_panel + 1; - let b_nrows = p.section_a_b_tile_iters * p.b_panels_per_tile * b_panel - + p.section_b_full_b_panels * b_panel - + usize::from(p.section_b_has_b_remainder); - - let ceil = dim; - let a_data = gen_data(a_nrows * dim, ceil); - let b_data = gen_data(b_nrows * dim, ceil); - let expected = naive(&a_data, a_nrows, &b_data, b_nrows, dim); - - let a_mat = MatRef::new(Standard::new(a_nrows, dim).unwrap(), &a_data).unwrap(); - let a_bt = BlockTransposed::::from_matrix_view(a_mat.as_matrix_view()); - let b_mat = MatRef::new(Standard::new(b_nrows, dim).unwrap(), &b_data).unwrap(); - let mut scratch = vec![f32::MIN; a_bt.padded_nrows()]; - max_ip_kernel::(arch, a_bt.as_view(), b_mat, &mut scratch, budget); - - for i in 0..a_nrows { - assert!( - (scratch[i] - expected[i]).abs() < 1e-10, - "[{arch_label}] plan={p:?} row {i}: actual={} expected={}", - scratch[i], - expected[i], - ); - } - } - } - - /// Breadth sweep for the f16 path (`F32Kernel` + `ConvertTo`): - /// every `NAIVE_CASES` shape, both Scalar and V3 (x86_64 only), - /// under `TileBudget::default()`. - #[test] - fn tiled_reduce_f16_matches_naive() { - for &(a_nrows, b_nrows, dim) in NAIVE_CASES { - // Use a small ceil so values stay exactly representable in f16 - // (bit-exact agreement with the f32 naive reference). - let ceil = dim.max(1); - let a_data: Vec = (0..a_nrows * dim) - .map(|i| diskann_wide::cast_f32_to_f16(((i + 1) % ceil) as f32)) - .collect(); - let b_data: Vec = (0..b_nrows * dim) - .map(|i| diskann_wide::cast_f32_to_f16((((i + 1) * 2) % ceil) as f32)) - .collect(); - let expected = naive_max_ip_f16(&a_data, a_nrows, &b_data, b_nrows, dim); - - check_kernel::<_, half::f16, 8>( - Scalar::new(), - "scalar", - 1e-10, - &a_data, - a_nrows, - &b_data, - b_nrows, - dim, - &expected, - ); - - #[cfg(target_arch = "x86_64")] - if let Some(arch) = diskann_wide::arch::x86_64::V3::new_checked() { - check_kernel::<_, half::f16, 16>( - arch, - "x86-64-v3", - 1e-10, - &a_data, - a_nrows, - &b_data, - b_nrows, - dim, - &expected, - ); - } - } - } -} diff --git a/diskann-quantization/src/multi_vector/distance/kernels/tiles.rs b/diskann-quantization/src/multi_vector/distance/kernels/tiles.rs new file mode 100644 index 0000000000..a9ae86763f --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/tiles.rs @@ -0,0 +1,341 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT license. + +//! The read ladder: walks lend tiles, tiles cut into panels. +//! +//! A query panel is one block-transposed block — `AR` rows interleaved column-major, so +//! the leaf loads `AR` contiguous rows per contraction step. A doc panel is `BR` row-major +//! rows. Neither carries `k`: it is `len() / AR` on the query side, and the leaf checks the +//! doc panel against it, so a mismatched pair is caught where it is used rather than +//! assumed. + +use core::slice::ChunksExact; + +use super::{NoTail, Paneled, TailIterator, TileAt, TileWalk}; +use crate::bits::{Dynamic, Length, Static}; +use crate::multi_vector::{BlockTransposedRef, MatRef, Standard}; + +// ── Query side ─────────────────────────────────────────────────── + +/// One block-transposed block: `AR` rows × `k` columns, column-major within the block. +#[derive(Clone, Copy)] +pub(super) struct QueryPanel<'a, T, const AR: usize>(&'a [T]); + +impl<'a, T, const AR: usize> QueryPanel<'a, T, AR> { + pub(super) fn as_slice(&self) -> &'a [T] { + self.0 + } + + /// The contraction length. + pub(super) fn k(&self) -> usize { + self.0.len() / AR + } +} + +pub(super) struct QueryPanels<'a, T, const AR: usize>(ChunksExact<'a, T>); + +impl<'a, T, const AR: usize> Iterator for QueryPanels<'a, T, AR> { + type Item = QueryPanel<'a, T, AR>; + + fn next(&mut self) -> Option { + self.0.next().map(QueryPanel) + } + + fn size_hint(&self) -> (usize, Option) { + self.0.size_hint() + } +} + +impl ExactSizeIterator for QueryPanels<'_, T, AR> {} + +impl TailIterator for QueryPanels<'_, T, AR> { + type Tail = NoTail; + + fn tail(self) -> Option { + None + } +} + +/// A run of whole blocks — block-transposed storage is padded to `AR`, hence [`NoTail`]. +pub(super) struct QueryTile<'a, T, const AR: usize> { + data: &'a [T], + k: usize, +} + +impl<'a, T, const AR: usize> QueryTile<'a, T, AR> { + /// # Panics + /// + /// Debug-only: panics unless `data` is a whole number of `AR × k` blocks, which is + /// what makes [`NoTail`] honest. + pub(super) fn new(data: &'a [T], k: usize) -> Self { + debug_assert!(k > 0 && data.len().is_multiple_of(AR * k)); + Self { data, k } + } +} + +impl<'a, T: Copy, const AR: usize> Paneled for QueryTile<'a, T, AR> { + type Panel = QueryPanel<'a, T, AR>; + type Tail = NoTail; + type Panels = QueryPanels<'a, T, AR>; + + fn panels(&self) -> Self::Panels { + QueryPanels(self.data.chunks_exact(AR * self.k)) + } +} + +// ── Doc side ───────────────────────────────────────────────────── + +/// Up to `BR` row-major rows of `k` elements. +/// +/// `L` is [`Static
`] for a whole panel and [`Dynamic`] for the trailing one, which is +/// how the tail reaches a leaf that can unroll for its width. +#[derive(Clone, Copy)] +pub(super) struct DocPanel<'a, T, const BR: usize, L> { + data: &'a [T], + rows: L, +} + +impl<'a, T, const BR: usize, L: Length> DocPanel<'a, T, BR, L> { + pub(super) fn as_slice(&self) -> &'a [T] { + self.data + } + + pub(super) fn rows(&self) -> usize { + self.rows.value() + } +} + +pub(super) struct DocPanels<'a, T, const BR: usize> { + chunks: ChunksExact<'a, T>, + k: usize, +} + +impl<'a, T, const BR: usize> Iterator for DocPanels<'a, T, BR> { + type Item = DocPanel<'a, T, BR, Static
>; + + fn next(&mut self) -> Option { + self.chunks.next().map(|data| DocPanel { + data, + rows: Static::
, + }) + } + + fn size_hint(&self) -> (usize, Option) { + self.chunks.size_hint() + } +} + +impl ExactSizeIterator for DocPanels<'_, T, BR> {} + +impl<'a, T, const BR: usize> TailIterator for DocPanels<'a, T, BR> { + type Tail = DocPanel<'a, T, BR, Dynamic>; + + fn tail(self) -> Option { + let data = self.chunks.remainder(); + (!data.is_empty()).then(|| DocPanel { + data, + rows: Dynamic(data.len() / self.k), + }) + } +} + +/// Row-major rows; only the final tile of a walk may be short. +pub(super) struct DocTile<'a, T, const BR: usize> { + data: &'a [T], + k: usize, +} + +impl<'a, T, const BR: usize> DocTile<'a, T, BR> { + /// # Panics + /// + /// Debug-only: panics unless `data` is a whole number of `k`-element rows. + pub(super) fn new(data: &'a [T], k: usize) -> Self { + debug_assert!(k > 0 && data.len().is_multiple_of(k)); + Self { data, k } + } +} + +impl<'a, T: Copy, const BR: usize> Paneled for DocTile<'a, T, BR> { + type Panel = DocPanel<'a, T, BR, Static
>; + type Tail = DocPanel<'a, T, BR, Dynamic>; + type Panels = DocPanels<'a, T, BR>; + + fn panels(&self) -> Self::Panels { + DocPanels { + chunks: self.data.chunks_exact(BR * self.k), + k: self.k, + } + } +} + +// ── Walks ──────────────────────────────────────────────────────── + +/// A cursor over a contiguous source, cut into tiles of `stride` elements. +/// +/// Shared by both sides — and by the widening walks — because the only difference is how +/// a tile is interpreted. +pub(super) struct Cursor<'a, T> { + data: &'a [T], + stride: usize, + cur: usize, +} + +impl<'a, T> Cursor<'a, T> { + pub(super) fn new(data: &'a [T], stride: usize) -> Self { + debug_assert!(stride > 0); + Self { + data, + stride, + cur: 0, + } + } + + pub(super) fn next(&mut self) -> Option<&'a [T]> { + let rest = self.data.get(self.cur..)?; + if rest.is_empty() { + return None; + } + let take = self.stride.min(rest.len()); + self.cur += take; + Some(&rest[..take]) + } + + pub(super) fn reset(&mut self) { + self.cur = 0; + } + + /// The longest tile this cursor can yield — the size a converting walk must stage. + pub(super) fn widest(&self) -> usize { + self.stride.min(self.data.len()) + } +} + +/// Walks the padded storage of a block-transposed query, `a_panels` blocks at a time. +pub(super) struct QueryWalk<'a, T, const AR: usize> { + cursor: Cursor<'a, T>, + k: usize, +} + +impl<'a, T: Copy, const AR: usize> QueryWalk<'a, T, AR> { + /// # Panics + /// + /// Panics if `view` has no columns. + pub(super) fn new(view: BlockTransposedRef<'a, T, AR>, a_panels: usize) -> Self { + let k = view.padded_ncols(); + assert!(k > 0, "QueryWalk requires a non-empty contraction"); + Self { + cursor: Cursor::new(view.as_slice(), a_panels * AR * k), + k, + } + } +} + +impl<'t, T: Copy, const AR: usize> TileAt<'t> for QueryWalk<'_, T, AR> { + type Tile = QueryTile<'t, T, AR>; +} + +impl TileWalk for QueryWalk<'_, T, AR> { + fn next(&mut self) -> Option> { + let k = self.k; + self.cursor.next().map(|data| QueryTile::new(data, k)) + } + + fn reset(&mut self) { + self.cursor.reset(); + } +} + +/// Walks a row-major doc matrix, `b_panels * BR` rows at a time. +pub(super) struct DocWalk<'a, T, const BR: usize> { + cursor: Cursor<'a, T>, + k: usize, +} + +impl<'a, T, const BR: usize> DocWalk<'a, T, BR> { + /// # Panics + /// + /// Panics if `docs` has zero-length rows. + pub(super) fn new(docs: MatRef<'a, Standard>, b_panels: usize) -> Self { + let k = docs.vector_dim(); + assert!(k > 0, "DocWalk requires a non-empty contraction"); + Self { + cursor: Cursor::new(docs.as_slice(), b_panels * BR * k), + k, + } + } +} + +impl<'t, T: Copy, const BR: usize> TileAt<'t> for DocWalk<'_, T, BR> { + type Tile = DocTile<'t, T, BR>; +} + +impl TileWalk for DocWalk<'_, T, BR> { + fn next(&mut self) -> Option> { + let k = self.k; + self.cursor.next().map(|data| DocTile::new(data, k)) + } + + fn reset(&mut self) { + self.cursor.reset(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Handles are passed by value into every leaf call; keeping them register-sized is + /// the reason `k` is derived rather than stored. + #[test] + fn handles_stay_thin() { + use core::mem::size_of; + assert_eq!(size_of::>(), 2 * size_of::()); + assert_eq!( + size_of::>>(), + 2 * size_of::() + ); + assert_eq!( + size_of::>(), + 3 * size_of::() + ); + } + + #[test] + fn doc_tile_splits_into_whole_panels_plus_a_dynamic_tail() { + let data: Vec = (0..7 * 3).map(|i| i as u8).collect(); + let tile = DocTile::::new(&data, 3); + + let mut panels = tile.panels(); + assert_eq!(panels.len(), 3); + assert_eq!(panels.next().unwrap().as_slice(), &[0, 1, 2, 3, 4, 5]); + let rest: Vec<_> = panels.by_ref().map(|p| p.rows()).collect(); + assert_eq!(rest, [2, 2]); + + let tail = panels.tail().unwrap(); + assert_eq!(tail.rows(), 1); + assert_eq!(tail.as_slice(), &[18, 19, 20]); + } + + #[test] + fn a_tile_that_divides_evenly_has_no_tail() { + let data: Vec = (0..6 * 3).map(|i| i as u8).collect(); + let mut panels = DocTile::::new(&data, 3).panels(); + assert_eq!(panels.by_ref().count(), 3); + assert!(panels.tail().is_none()); + } + + #[test] + fn a_walk_yields_short_final_tile_then_stops_and_rewinds() { + let data: Vec = (0..5 * 3).map(|i| i as u8).collect(); + let docs = MatRef::new(Standard::new(5, 3).unwrap(), &data).unwrap(); + let mut walk = DocWalk::::new(docs, 1); + + assert_eq!(walk.next().unwrap().data.len(), 6); + assert_eq!(walk.next().unwrap().data.len(), 6); + assert_eq!(walk.next().unwrap().data.len(), 3); + assert!(walk.next().is_none()); + + walk.reset(); + assert_eq!(walk.next().unwrap().data.len(), 6); + } +} From 7a2abcc009aca425383ae27faf1991f3ee0d4468 Mon Sep 17 00:00:00 2001 From: Suryansh Gupta Date: Sat, 15 Aug 2026 04:22:29 +0530 Subject: [PATCH 2/3] Restrict strip size in accordance to actual no of docs --- .../src/multi_vector/distance/kernels/f16.rs | 2 - .../multi_vector/distance/kernels/float.rs | 62 +++---------------- .../src/multi_vector/distance/kernels/mod.rs | 39 ++++++++++-- 3 files changed, 44 insertions(+), 59 deletions(-) diff --git a/diskann-quantization/src/multi_vector/distance/kernels/f16.rs b/diskann-quantization/src/multi_vector/distance/kernels/f16.rs index febf4d962a..d1da7fd78b 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/f16.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/f16.rs @@ -147,7 +147,6 @@ impl ) { float::run( arch, - query.padded_nrows(), docs.num_vectors(), query.padded_ncols(), TileBudget::default(), @@ -181,7 +180,6 @@ impl ) { float::run( arch, - query.padded_nrows(), docs.num_vectors(), query.padded_ncols(), TileBudget::default(), diff --git a/diskann-quantization/src/multi_vector/distance/kernels/float.rs b/diskann-quantization/src/multi_vector/distance/kernels/float.rs index 02d27cd639..cf6823c0c9 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/float.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/float.rs @@ -178,10 +178,10 @@ impl Drain> for RawMax<'_> { /// /// # Panics /// -/// Panics unless `state` is exactly `nq` long. +/// Panics if `state` is shorter than the query's padded row count: the drain indexes it +/// one A-panel at a time. pub(super) fn run( arch: Arch, - nq: usize, nd: usize, k: usize, budget: TileBudget, @@ -204,8 +204,6 @@ pub(super) fn run( >, for<'o, 'x> RawMax<'o>: Drain>, { - assert_eq!(state.len(), nq, "state is one entry per padded query row"); - // The identity for max. state.fill(f32::MIN); @@ -220,7 +218,7 @@ pub(super) fn run( // Both sides are f32 rows of `k`: whatever the source element type, the leaves // consume f32. let row_bytes = k * size_of::(); - let plan = Plan::::new(row_bytes, row_bytes, size_of::(), budget); + let plan = Plan::::new(row_bytes, row_bytes, nd, size_of::(), budget); let (a_walk, b_walk) = walks(plan); let mut buf = vec![0.0f32; plan.strip_len()]; @@ -260,7 +258,6 @@ impl ) { run( arch, - query.padded_nrows(), docs.num_vectors(), query.padded_ncols(), TileBudget::default(), @@ -294,7 +291,6 @@ impl ) { run( arch, - query.padded_nrows(), docs.num_vectors(), query.padded_ncols(), TileBudget::default(), @@ -364,20 +360,12 @@ mod tests { // Poisoned, not seeded: `run` owes the caller a full initialization, and every case // below would fail loudly if it skipped one. let mut state = vec![f32::MAX; bt.padded_nrows()]; - run( - arch, - bt.padded_nrows(), - nd, - k, - budget, - &mut state, - |plan: Plan| { - ( - QueryWalk::new(bt.as_view(), plan.a_panels), - DocWalk::new(docs, plan.b_panels), - ) - }, - ); + run(arch, nd, k, budget, &mut state, |plan: Plan| { + ( + QueryWalk::new(bt.as_view(), plan.a_panels), + DocWalk::new(docs, plan.b_panels), + ) + }); for (i, &expected) in naive(&a, nq, &b, nd, k).iter().enumerate() { let actual = state[i] as f64; @@ -456,7 +444,6 @@ mod tests { let mut state = vec![f32::MAX; SC_A]; run( Scalar::new(), - SC_A, 3, 0, TileBudget::default(), @@ -472,7 +459,6 @@ mod tests { let mut state = vec![f32::MAX; SC_A]; run( Scalar::new(), - SC_A, 0, 0, TileBudget::default(), @@ -483,34 +469,4 @@ mod tests { ); assert_eq!(state, vec![f32::MIN; SC_A]); } - - /// `state` sized from `nrows()` instead of `padded_nrows()` is the mistake the entry - /// check exists for: the two are derived independently, so they can disagree, and it - /// used to surface as a slice index panic deep inside the drain. - #[test] - #[should_panic(expected = "one entry per padded query row")] - fn state_sized_for_logical_rows_is_rejected() { - let a = sample(9 * 4, 0); - let b = sample(2 * 4, 5); - let query = MatRef::new(Standard::new(9, 4).unwrap(), &a).unwrap(); - let docs = MatRef::new(Standard::new(2, 4).unwrap(), &b).unwrap(); - let bt = BlockTransposed::::from_matrix_view(query.as_matrix_view()); - assert_eq!((bt.nrows(), bt.padded_nrows()), (9, 16)); - - let mut state = vec![0.0; bt.nrows()]; - run( - Scalar::new(), - bt.padded_nrows(), - 2, - 4, - TileBudget::default(), - &mut state, - |plan: Plan| { - ( - QueryWalk::new(bt.as_view(), plan.a_panels), - DocWalk::new(docs, plan.b_panels), - ) - }, - ); - } } diff --git a/diskann-quantization/src/multi_vector/distance/kernels/mod.rs b/diskann-quantization/src/multi_vector/distance/kernels/mod.rs index 4f7cdcc33c..410f30683e 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/mod.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/mod.rs @@ -53,7 +53,8 @@ impl Default for TileBudget { } /// Panel counts per tile: `a_panels` A-panels resident in L2, and as many B-panels as -/// co-fit L1 alongside one A-panel *and* the accumulator columns they feed. +/// co-fit L1 alongside one A-panel *and* the accumulator columns they feed, never more +/// than the B-rows on hand. #[derive(Debug, Clone, Copy)] struct Plan { a_panels: usize, @@ -61,9 +62,21 @@ struct Plan { } impl Plan { + /// `b_panels` is reconciled against `b_rows`, so a plan belongs to the B side it was + /// built for; reused against a longer one it still computes the right answer, but tiles + /// far more narrowly than the cache allows. `a_panels` needs no such reconciliation: it only + /// feeds a walk stride, which the cursor already bounds by the data it holds, whereas + /// `b_panels` also sizes the accumulator strip. + /// /// Row sizes are clamped to one byte so a degenerate row cannot divide by zero; the /// entries reject an empty contraction before planning, so such a plan is never used. - fn new(a_row_bytes: usize, b_row_bytes: usize, acc_bytes: usize, budget: TileBudget) -> Self { + fn new( + a_row_bytes: usize, + b_row_bytes: usize, + b_rows: usize, + acc_bytes: usize, + budget: TileBudget, + ) -> Self { let a_row_bytes = a_row_bytes.max(1); let b_row_bytes = b_row_bytes.max(1); @@ -73,7 +86,12 @@ impl Plan { // A-panel stays resident alongside. let per_b_row = b_row_bytes + AR * acc_bytes; let b_budget = budget.l1_b.saturating_sub(AR * a_row_bytes); - let b_panels = ((b_budget / per_b_row) / BR).max(1); + let cache_fit = (b_budget / per_b_row) / BR; + + // Never plan wider than the B-rows on hand: a cache-sized tile over a short B side + // would size accumulator columns no fill can reach. One panel at minimum, so a + // budget too small for even that still makes progress. + let b_panels = cache_fit.min(b_rows.div_ceil(BR)).max(1); Self { a_panels, b_panels } } @@ -292,6 +310,7 @@ mod tests { let plan = Plan::<16, 4>::new( 64, 64, + usize::MAX, 4, TileBudget { l2_a: 40960, @@ -302,9 +321,21 @@ mod tests { assert_eq!(plan.strip_len(), 16 * 68 * 4); } + #[test] + fn plan_never_tiles_past_the_b_rows_on_hand() { + // The cache would take 68 B-panels, but 10 B-rows fill three. + let budget = TileBudget { + l2_a: 40960, + l1_b: 36000, + }; + let plan = Plan::<16, 4>::new(64, 64, 10, 4, budget); + assert_eq!(plan.b_panels, 3); + assert_eq!(plan.strip_len(), 16 * 3 * 4); + } + #[test] fn plan_clamps_to_one_panel_per_tile() { - let plan = Plan::<16, 4>::new(1024, 1024, 4, TileBudget { l2_a: 1, l1_b: 1 }); + let plan = Plan::<16, 4>::new(1024, 1024, usize::MAX, 4, TileBudget { l2_a: 1, l1_b: 1 }); assert_eq!((plan.a_panels, plan.b_panels), (1, 1)); assert_eq!(plan.strip_len(), 16 * 4); } From 4683fd8732988cbc9833186d55f1530b4c3b53b0 Mon Sep 17 00:00:00 2001 From: Suryansh Gupta Date: Fri, 21 Aug 2026 04:37:15 +0530 Subject: [PATCH 3/3] Address review comments --- .../src/multi_vector/distance/factory.rs | 6 +- .../src/multi_vector/distance/fallback.rs | 6 +- .../src/multi_vector/distance/isa.rs | 6 +- .../src/multi_vector/distance/kernel.rs | 6 +- .../src/multi_vector/distance/kernels/f16.rs | 95 +++--- .../multi_vector/distance/kernels/float.rs | 162 +++++----- .../distance/kernels/leaves/mod.rs | 18 +- .../distance/kernels/leaves/scalar.rs | 72 +++-- .../distance/kernels/leaves/v3.rs | 70 ++--- .../src/multi_vector/distance/kernels/mod.rs | 133 +++++---- .../multi_vector/distance/kernels/strip.rs | 79 +++-- .../multi_vector/distance/kernels/tiles.rs | 279 +++++++++++------- .../src/multi_vector/distance/max_sim.rs | 6 +- .../src/multi_vector/distance/mod.rs | 6 +- .../multi_vector/distance/projected_eigen.rs | 6 +- 15 files changed, 552 insertions(+), 398 deletions(-) diff --git a/diskann-quantization/src/multi_vector/distance/factory.rs b/diskann-quantization/src/multi_vector/distance/factory.rs index 0539793acd..d6d773e001 100644 --- a/diskann-quantization/src/multi_vector/distance/factory.rs +++ b/diskann-quantization/src/multi_vector/distance/factory.rs @@ -1,5 +1,7 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ //! Factory + concrete `MaxSimKernel` impls for the multi-vector distance //! API. BYOTE entry point — see [`build_max_sim`]. diff --git a/diskann-quantization/src/multi_vector/distance/fallback.rs b/diskann-quantization/src/multi_vector/distance/fallback.rs index ed8da7a3e1..23b6209cf8 100644 --- a/diskann-quantization/src/multi_vector/distance/fallback.rs +++ b/diskann-quantization/src/multi_vector/distance/fallback.rs @@ -1,5 +1,7 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ //! Fallback kernel implementation of multi-vector distance computation. diff --git a/diskann-quantization/src/multi_vector/distance/isa.rs b/diskann-quantization/src/multi_vector/distance/isa.rs index d295438bc9..f94424c519 100644 --- a/diskann-quantization/src/multi_vector/distance/isa.rs +++ b/diskann-quantization/src/multi_vector/distance/isa.rs @@ -1,5 +1,7 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ //! Instruction Set Architecture (ISA) selector for the multi-vector MaxSim //! factory. diff --git a/diskann-quantization/src/multi_vector/distance/kernel.rs b/diskann-quantization/src/multi_vector/distance/kernel.rs index b292def54c..fd68428c1e 100644 --- a/diskann-quantization/src/multi_vector/distance/kernel.rs +++ b/diskann-quantization/src/multi_vector/distance/kernel.rs @@ -1,5 +1,7 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ //! Object-safe kernel boundary trait plus BYOTE visitor trait. diff --git a/diskann-quantization/src/multi_vector/distance/kernels/f16.rs b/diskann-quantization/src/multi_vector/distance/kernels/f16.rs index d1da7fd78b..ec89a834ca 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/f16.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/f16.rs @@ -1,13 +1,17 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ //! f16 MaxSim. //! //! There is no f16 leaf: f16 widens to f32 and reuses the f32 pipeline. Both sides widen a //! tile at a time into a buffer the walk reuses, which is what the lending [`TileWalk`] -//! exists for — the whole query never has to be staged at once, and the staged copy stays +//! exists for. The whole A side never has to be staged at once, and the staged copy stays //! inside the cache level its tile was sized for. +use core::num::NonZeroUsize; + use diskann_vector::conversion::SliceCast; use diskann_wide::Architecture; #[cfg(target_arch = "x86_64")] @@ -17,7 +21,7 @@ use diskann_wide::arch::{Scalar, Target2}; use super::leaves::scalar::{A_PANEL as SC_A, B_PANEL as SC_B}; #[cfg(target_arch = "x86_64")] use super::leaves::v3::{A_PANEL as V3_A, B_PANEL as V3_B}; -use super::tiles::{Cursor, DocTile, QueryTile}; +use super::tiles::{BlockTransposedTile, Cursor, RowMajorTile, contraction, tile_stride}; use super::{Plan, TileAt, TileBudget, TileWalk, float}; use crate::multi_vector::{BlockTransposedRef, MatRef, Standard}; @@ -26,18 +30,14 @@ struct Widen<'a, Arch> { arch: Arch, cursor: Cursor<'a, half::f16>, buf: Vec, - k: usize, + k: NonZeroUsize, } impl<'a, Arch: Architecture> Widen<'a, Arch> where SliceCast: for<'x> Target2, { - /// # Panics - /// - /// Panics if `k` is zero — the entry guards that case before any walk is built. - fn new(arch: Arch, src: &'a [half::f16], k: usize, stride: usize) -> Self { - assert!(k > 0, "widening walk requires a non-empty contraction"); + fn new(arch: Arch, src: &'a [half::f16], k: NonZeroUsize, stride: NonZeroUsize) -> Self { let cursor = Cursor::new(src, stride); let buf = vec![0.0f32; cursor.widest()]; Self { @@ -57,32 +57,41 @@ where } } -/// Widens the padded storage of a block-transposed f16 query. +/// Widens the padded storage of an f16 [`BlockTransposedRef`]. /// /// Widening is element-wise, so it preserves the block-transposed permutation. -struct QueryWiden<'a, Arch, const AR: usize>(Widen<'a, Arch>); +struct BlockTransposedWiden<'a, Arch, const AR: usize>(Widen<'a, Arch>); -impl<'a, Arch: Architecture, const AR: usize> QueryWiden<'a, Arch, AR> +impl<'a, Arch: Architecture, const AR: usize> BlockTransposedWiden<'a, Arch, AR> where SliceCast: for<'x> Target2, { - fn new(arch: Arch, view: BlockTransposedRef<'a, half::f16, AR>, a_panels: usize) -> Self { - let k = view.padded_ncols(); - Self(Widen::new(arch, view.as_slice(), k, a_panels * AR * k)) + fn new( + arch: Arch, + view: BlockTransposedRef<'a, half::f16, AR>, + a_panels: NonZeroUsize, + ) -> Self { + let k = contraction(view.padded_ncols()); + Self(Widen::new( + arch, + view.as_slice(), + k, + tile_stride(a_panels, AR, k), + )) } } -impl<'t, Arch, const AR: usize> TileAt<'t> for QueryWiden<'_, Arch, AR> { - type Tile = QueryTile<'t, f32, AR>; +impl<'t, Arch, const AR: usize> TileAt<'t> for BlockTransposedWiden<'_, Arch, AR> { + type Tile = BlockTransposedTile<'t, f32, AR>; } -impl TileWalk for QueryWiden<'_, Arch, AR> +impl TileWalk for BlockTransposedWiden<'_, Arch, AR> where SliceCast: for<'x> Target2, { - fn next(&mut self) -> Option> { + fn next(&mut self) -> Option> { let k = self.0.k; - self.0.next().map(|data| QueryTile::new(data, k)) + self.0.next().map(|data| BlockTransposedTile::new(data, k)) } fn reset(&mut self) { @@ -90,30 +99,35 @@ where } } -/// Widens a row-major f16 doc matrix. -struct DocWiden<'a, Arch, const BR: usize>(Widen<'a, Arch>); +/// Widens an f16 [`Standard`] matrix. +struct RowMajorWiden<'a, Arch, const BR: usize>(Widen<'a, Arch>); -impl<'a, Arch: Architecture, const BR: usize> DocWiden<'a, Arch, BR> +impl<'a, Arch: Architecture, const BR: usize> RowMajorWiden<'a, Arch, BR> where SliceCast: for<'x> Target2, { - fn new(arch: Arch, docs: MatRef<'a, Standard>, b_panels: usize) -> Self { - let k = docs.vector_dim(); - Self(Widen::new(arch, docs.as_slice(), k, b_panels * BR * k)) + fn new(arch: Arch, mat: MatRef<'a, Standard>, b_panels: NonZeroUsize) -> Self { + let k = contraction(mat.vector_dim()); + Self(Widen::new( + arch, + mat.as_slice(), + k, + tile_stride(b_panels, BR, k), + )) } } -impl<'t, Arch, const BR: usize> TileAt<'t> for DocWiden<'_, Arch, BR> { - type Tile = DocTile<'t, f32, BR>; +impl<'t, Arch, const BR: usize> TileAt<'t> for RowMajorWiden<'_, Arch, BR> { + type Tile = RowMajorTile<'t, f32, BR>; } -impl TileWalk for DocWiden<'_, Arch, BR> +impl TileWalk for RowMajorWiden<'_, Arch, BR> where SliceCast: for<'x> Target2, { - fn next(&mut self) -> Option> { + fn next(&mut self) -> Option> { let k = self.0.k; - self.0.next().map(|data| DocTile::new(data, k)) + self.0.next().map(|data| RowMajorTile::new(data, k)) } fn reset(&mut self) { @@ -121,9 +135,14 @@ where } } -// ── Entry ──────────────────────────────────────────────────────── +/////////// +// Entry // +/////////// -/// The f16 MaxSim entry — the f32 pipeline behind widening walks. +/// The f16 MaxSim entry: the f32 pipeline behind widening walks. +/// +/// Operand naming matches [`MaxIp`](super::MaxIp). The block-transposed A side is the +/// query and the row-major B side the documents. #[derive(Debug, Clone, Copy)] pub(crate) struct MaxIpF16; @@ -153,8 +172,8 @@ impl state, |plan: Plan| { ( - QueryWiden::new(arch, query, plan.a_panels), - DocWiden::new(arch, docs, plan.b_panels), + BlockTransposedWiden::new(arch, query, plan.a_panels), + RowMajorWiden::new(arch, docs, plan.b_panels), ) }, ); @@ -186,8 +205,8 @@ impl state, |plan: Plan| { ( - QueryWiden::new(arch, query, plan.a_panels), - DocWiden::new(arch, docs, plan.b_panels), + BlockTransposedWiden::new(arch, query, plan.a_panels), + RowMajorWiden::new(arch, docs, plan.b_panels), ) }, ); diff --git a/diskann-quantization/src/multi_vector/distance/kernels/float.rs b/diskann-quantization/src/multi_vector/distance/kernels/float.rs index cf6823c0c9..d552ffe461 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/float.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/float.rs @@ -1,8 +1,11 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ //! f32 MaxSim: the accumulator is already the score, so the drain is a bare reduction. +use core::num::NonZeroUsize; use core::ops::Range; use diskann_wide::arch::Scalar; @@ -16,7 +19,10 @@ use super::leaves::v3 as v3_leaf; #[cfg(target_arch = "x86_64")] use super::leaves::v3::{A_PANEL as V3_A, B_PANEL as V3_B}; use super::strip::{Slot, Strip}; -use super::tiles::{DocPanel, DocTile, DocWalk, QueryPanel, QueryTile, QueryWalk}; +use super::tiles::{ + BlockTransposedPanel, BlockTransposedTile, BlockTransposedWalk, RowMajorPanel, RowMajorTile, + RowMajorWalk, +}; use super::{Accumulate, Drain, Plan, TileAt, TileBudget, TileWalk, drive}; use crate::bits::{Dynamic, Static}; use crate::multi_vector::{BlockTransposedRef, MatRef, Standard}; @@ -24,23 +30,30 @@ use crate::multi_vector::{BlockTransposedRef, MatRef, Standard}; /// Selects the f32 leaf for whichever architecture is in play. pub(super) struct Kernel; -/// Reduces the strip straight into the running per-query-row maxima. +/// Reduces a strip into the running per-A-row maxima, one entry per A row. /// -/// Carries `nd` because the strip's trailing columns belong to doc rows past the end of -/// the matrix, and nothing else here knows where that end is. +/// This is the module's single [`Drain`], which makes it the one place the accumulator's +/// axes acquire meaning. A strip column holds one B row's inner products against every A +/// row. A B row therefore arrives as a *row* of the input and lands as a *column* of the +/// accumulator. +/// +/// Carries `nd` because a strip's trailing columns belong to B rows past the end of the +/// matrix, and nothing else here knows where that end is. pub(super) struct RawMax<'o> { out: &'o mut [f32], nd: usize, } -// ── V3 ─────────────────────────────────────────────────────────── +//////// +// V3 // +//////// #[cfg(target_arch = "x86_64")] impl<'a, 'b, 's> Accumulate< V3, - QueryPanel<'a, f32, V3_A>, - DocPanel<'b, f32, V3_B, Static>, + BlockTransposedPanel<'a, f32, V3_A>, + RowMajorPanel<'b, f32, V3_B, Static>, Slot<'s, f32, V3_A, V3_B>, > for Kernel { @@ -48,8 +61,8 @@ impl<'a, 'b, 's> fn accumulate( &self, arch: V3, - a: QueryPanel<'a, f32, V3_A>, - b: DocPanel<'b, f32, V3_B, Static>, + a: BlockTransposedPanel<'a, f32, V3_A>, + b: RowMajorPanel<'b, f32, V3_B, Static>, out: Slot<'s, f32, V3_A, V3_B>, ) { v3_leaf::f32_store_microkernel::(arch, a, b, out); @@ -60,8 +73,8 @@ impl<'a, 'b, 's> impl<'a, 'b, 's> Accumulate< V3, - QueryPanel<'a, f32, V3_A>, - DocPanel<'b, f32, V3_B, Dynamic>, + BlockTransposedPanel<'a, f32, V3_A>, + RowMajorPanel<'b, f32, V3_B, Dynamic>, Slot<'s, f32, V3_A, V3_B>, > for Kernel { @@ -69,8 +82,8 @@ impl<'a, 'b, 's> fn accumulate( &self, arch: V3, - a: QueryPanel<'a, f32, V3_A>, - b: DocPanel<'b, f32, V3_B, Dynamic>, + a: BlockTransposedPanel<'a, f32, V3_A>, + b: RowMajorPanel<'b, f32, V3_B, Dynamic>, out: Slot<'s, f32, V3_A, V3_B>, ) { // Dispatch the runtime width onto a const the leaf can unroll for. @@ -93,21 +106,23 @@ impl Drain> for RawMax<'_> { a_panel: usize, b_panels: Range, ) { - // A is padded to whole panels, so the output cut never clamps; B is not, so the + // A is padded to whole panels and its output cut never clamps. B is not, so the // live width does. - let a0 = a_panel * V3_A; let live = (b_panels.end * V3_B).min(self.nd) - b_panels.start * V3_B; - v3_leaf::fold_columns(arch, scratch.columns(live), &mut self.out[a0..][..V3_A]); + let rows = &mut self.out.as_chunks_mut::().0[a_panel]; + v3_leaf::max_into_rows(arch, scratch.columns(live), rows); } } -// ── Emulated ───────────────────────────────────────────────────── +////////////// +// Emulated // +////////////// impl<'a, 'b, 's> Accumulate< Scalar, - QueryPanel<'a, f32, SC_A>, - DocPanel<'b, f32, SC_B, Static>, + BlockTransposedPanel<'a, f32, SC_A>, + RowMajorPanel<'b, f32, SC_B, Static>, Slot<'s, f32, SC_A, SC_B>, > for Kernel { @@ -115,8 +130,8 @@ impl<'a, 'b, 's> fn accumulate( &self, arch: Scalar, - a: QueryPanel<'a, f32, SC_A>, - b: DocPanel<'b, f32, SC_B, Static>, + a: BlockTransposedPanel<'a, f32, SC_A>, + b: RowMajorPanel<'b, f32, SC_B, Static>, out: Slot<'s, f32, SC_A, SC_B>, ) { scalar_leaf::f32_store_microkernel::(arch, a, b, out); @@ -126,8 +141,8 @@ impl<'a, 'b, 's> impl<'a, 'b, 's> Accumulate< Scalar, - QueryPanel<'a, f32, SC_A>, - DocPanel<'b, f32, SC_B, Dynamic>, + BlockTransposedPanel<'a, f32, SC_A>, + RowMajorPanel<'b, f32, SC_B, Dynamic>, Slot<'s, f32, SC_A, SC_B>, > for Kernel { @@ -135,8 +150,8 @@ impl<'a, 'b, 's> fn accumulate( &self, arch: Scalar, - a: QueryPanel<'a, f32, SC_A>, - b: DocPanel<'b, f32, SC_B, Dynamic>, + a: BlockTransposedPanel<'a, f32, SC_A>, + b: RowMajorPanel<'b, f32, SC_B, Dynamic>, out: Slot<'s, f32, SC_A, SC_B>, ) { match b.rows() { @@ -155,31 +170,33 @@ impl Drain> for RawMax<'_> { a_panel: usize, b_panels: Range, ) { - let a0 = a_panel * SC_A; let live = (b_panels.end * SC_B).min(self.nd) - b_panels.start * SC_B; - scalar_leaf::fold_columns(arch, scratch.columns(live), &mut self.out[a0..][..SC_A]); + let rows = &mut self.out.as_chunks_mut::().0[a_panel]; + scalar_leaf::max_into_rows(arch, scratch.columns(live), rows); } } -// ── Entry ──────────────────────────────────────────────────────── +/////////// +// Entry // +/////////// /// Plan, allocate the strip, and drive. /// -/// `k` is the *physical* row length both walks stride by, so it must be the query's padded -/// column count rather than its logical one. +/// `k` is the *physical* row length both walks stride by, so it must be A's padded column +/// count, not its logical one. /// -/// `walks` is built from the plan rather than passed in so the empty-contraction guard can -/// run before any walk exists — a zero-length row would give a walk a zero stride. +/// `walks` is built from the plan instead of being passed in. That way the empty-contraction +/// guard runs before any walk exists. A zero-length row would give a walk a zero stride. /// -/// On return `state` holds the per-query-row maximum inner product, one entry per padded -/// query row. Its incoming contents are ignored: seeding the max is this function's job, +/// On return `state` holds the per-A-row maximum inner product, one entry per padded A +/// row. Its incoming contents are ignored: seeding the max is this function's job, /// not the caller's, because a caller that got it wrong would silently clamp the result -/// rather than fail. +/// instead of failing. /// /// # Panics /// -/// Panics if `state` is shorter than the query's padded row count: the drain indexes it -/// one A-panel at a time. +/// Panics if `state` is shorter than A's padded row count: the drain indexes it one +/// A-panel at a time. pub(super) fn run( arch: Arch, nd: usize, @@ -189,17 +206,17 @@ pub(super) fn run( walks: impl FnOnce(Plan) -> (AW, BW), ) where Arch: Copy, - AW: TileWalk + for<'a> TileAt<'a, Tile = QueryTile<'a, f32, AR>>, - BW: TileWalk + for<'b> TileAt<'b, Tile = DocTile<'b, f32, BR>>, + AW: TileWalk + for<'a> TileAt<'a, Tile = BlockTransposedTile<'a, f32, AR>>, + BW: TileWalk + for<'b> TileAt<'b, Tile = RowMajorTile<'b, f32, BR>>, Kernel: for<'a, 'b, 's> Accumulate< Arch, - QueryPanel<'a, f32, AR>, - DocPanel<'b, f32, BR, Static
>, + BlockTransposedPanel<'a, f32, AR>, + RowMajorPanel<'b, f32, BR, Static
>, Slot<'s, f32, AR, BR>, > + for<'a, 'b, 's> Accumulate< Arch, - QueryPanel<'a, f32, AR>, - DocPanel<'b, f32, BR, Dynamic>, + BlockTransposedPanel<'a, f32, AR>, + RowMajorPanel<'b, f32, BR, Dynamic>, Slot<'s, f32, AR, BR>, >, for<'o, 'x> RawMax<'o>: Drain>, @@ -207,17 +224,17 @@ pub(super) fn run( // The identity for max. state.fill(f32::MIN); - if k == 0 { + let Some(k) = NonZeroUsize::new(k) else { // Every inner product is the empty sum. if nd > 0 { state.fill(0.0); } return; - } + }; // Both sides are f32 rows of `k`: whatever the source element type, the leaves - // consume f32. - let row_bytes = k * size_of::(); + // consume f32. The clamp is an unreachable backstop, as in `tile_stride`. + let row_bytes = NonZeroUsize::new(k.get() * size_of::()).unwrap_or(NonZeroUsize::MIN); let plan = Plan::::new(row_bytes, row_bytes, nd, size_of::(), budget); let (a_walk, b_walk) = walks(plan); let mut buf = vec![0.0f32; plan.strip_len()]; @@ -232,9 +249,12 @@ pub(super) fn run( ); } -/// The f32 MaxSim entry. Which leaf geometry applies follows from the architecture, so the -/// block size of `query` must match the [`Target3`](diskann_wide::arch::Target3) impl -/// selected. +/// The f32 MaxSim entry, and with [`MaxIpF16`](super::MaxIpF16) one of the two places that +/// name the operands: the block-transposed A side is the query, the row-major B side the +/// documents. +/// +/// Which leaf geometry applies follows from the architecture, and the block size of the +/// query must match the [`Target3`](diskann_wide::arch::Target3) impl selected. #[derive(Debug, Clone, Copy)] pub(crate) struct MaxIp; @@ -264,8 +284,8 @@ impl state, |plan: Plan| { ( - QueryWalk::new(query, plan.a_panels), - DocWalk::new(docs, plan.b_panels), + BlockTransposedWalk::new(query, plan.a_panels), + RowMajorWalk::new(docs, plan.b_panels), ) }, ); @@ -297,8 +317,8 @@ impl state, |plan: Plan| { ( - QueryWalk::new(query, plan.a_panels), - DocWalk::new(docs, plan.b_panels), + BlockTransposedWalk::new(query, plan.a_panels), + RowMajorWalk::new(docs, plan.b_panels), ) }, ); @@ -340,13 +360,13 @@ mod tests { Arch: Copy, Kernel: for<'a, 'b, 's> Accumulate< Arch, - QueryPanel<'a, f32, AR>, - DocPanel<'b, f32, BR, Static
>, + BlockTransposedPanel<'a, f32, AR>, + RowMajorPanel<'b, f32, BR, Static
>, Slot<'s, f32, AR, BR>, > + for<'a, 'b, 's> Accumulate< Arch, - QueryPanel<'a, f32, AR>, - DocPanel<'b, f32, BR, Dynamic>, + BlockTransposedPanel<'a, f32, AR>, + RowMajorPanel<'b, f32, BR, Dynamic>, Slot<'s, f32, AR, BR>, >, for<'o, 'x> RawMax<'o>: Drain>, @@ -362,8 +382,8 @@ mod tests { let mut state = vec![f32::MAX; bt.padded_nrows()]; run(arch, nd, k, budget, &mut state, |plan: Plan| { ( - QueryWalk::new(bt.as_view(), plan.a_panels), - DocWalk::new(docs, plan.b_panels), + BlockTransposedWalk::new(bt.as_view(), plan.a_panels), + RowMajorWalk::new(docs, plan.b_panels), ) }); @@ -448,14 +468,15 @@ mod tests { 0, TileBudget::default(), &mut state, - |_: Plan| -> (QueryWalk<'_, f32, SC_A>, DocWalk<'_, f32, SC_B>) { - unreachable!("no walk is built for an empty contraction") - }, + |_: Plan| -> ( + BlockTransposedWalk<'_, f32, SC_A>, + RowMajorWalk<'_, f32, SC_B>, + ) { unreachable!("no walk is built for an empty contraction") }, ); assert_eq!(state, vec![0.0; SC_A]); - // With no docs there is no maximum, so the identity stands and the caller negates - // it into its empty-doc sentinel. + // With no B rows there is no maximum, so the identity stands and the caller + // negates it into its empty-input sentinel. let mut state = vec![f32::MAX; SC_A]; run( Scalar::new(), @@ -463,9 +484,10 @@ mod tests { 0, TileBudget::default(), &mut state, - |_: Plan| -> (QueryWalk<'_, f32, SC_A>, DocWalk<'_, f32, SC_B>) { - unreachable!("no walk is built for an empty contraction") - }, + |_: Plan| -> ( + BlockTransposedWalk<'_, f32, SC_A>, + RowMajorWalk<'_, f32, SC_B>, + ) { unreachable!("no walk is built for an empty contraction") }, ); assert_eq!(state, vec![f32::MIN; SC_A]); } diff --git a/diskann-quantization/src/multi_vector/distance/kernels/leaves/mod.rs b/diskann-quantization/src/multi_vector/distance/kernels/leaves/mod.rs index acb9f1b3f8..6c73b8ef70 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/leaves/mod.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/leaves/mod.rs @@ -1,21 +1,23 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ //! Per-ISA micro-kernels. //! -//! Panel geometry is fixed per ISA rather than derived from a lane count: the scalar leaf -//! is deliberately narrower than `2 × LANES` would suggest, because its "lanes" are a +//! Panel geometry is fixed per ISA, not derived from a lane count. The scalar leaf is +//! deliberately narrower than `2 × LANES` would suggest, because its "lanes" are a //! loop, not silicon. //! //! Every leaf here **stores** its products into a slot and leaves the reduction to -//! `fold_columns`, which folds along independent max chains so the reduction does not -//! serialize on its own latency — it is not negligible work when the strip is wide and -//! shallow. +//! `max_into_rows`. That reduction is not negligible work when the strip is wide and +//! shallow, so it runs several independent max chains rather than serializing on the +//! latency of a single one. pub(super) mod scalar; #[cfg(target_arch = "x86_64")] pub(super) mod v3; -/// Independent max chains in `fold_columns`: enough to keep a multi-cycle max off its own +/// Independent max chains in `max_into_rows`: enough to keep a multi-cycle max off its own /// critical path, few enough that the chains stay in registers. const WAYS: usize = 4; diff --git a/diskann-quantization/src/multi_vector/distance/kernels/leaves/scalar.rs b/diskann-quantization/src/multi_vector/distance/kernels/leaves/scalar.rs index 13970793f8..e17d9441d6 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/leaves/scalar.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/leaves/scalar.rs @@ -1,13 +1,15 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ -//! Emulated f32 leaf — 8 × 2. +//! Emulated f32 leaf, panels of 8 × 2. //! //! Geometry is narrower than the V3 leaf because [`Emulated`](diskann_wide::Emulated) //! lanes are an unrolled loop over scalars, not registers: 8 × 2 matches the `Strategy2x1` //! shape used by the scalar distance functions elsewhere in the crate. //! -//! The inner loop multiplies and adds separately rather than calling `mul_add_simd`, which +//! The inner loop multiplies and adds separately instead of calling `mul_add_simd`, which //! on x86-64 without hardware FMA would drop into libm's software `fma()`. use diskann_wide::arch::Scalar; @@ -16,7 +18,7 @@ use diskann_wide::{SIMDMinMax, SIMDVector}; use super::WAYS; use crate::bits::Length; use crate::multi_vector::distance::kernels::strip::Slot; -use crate::multi_vector::distance::kernels::tiles::{DocPanel, QueryPanel}; +use crate::multi_vector::distance::kernels::tiles::{BlockTransposedPanel, RowMajorPanel}; diskann_wide::alias!(f32s = ::f32x8); @@ -30,13 +32,13 @@ const REGS: usize = A_PANEL / f32s::LANES; /// /// # Panics /// -/// Panics if `b` does not hold exactly `UNROLL` rows of `a`'s contraction length — the -/// extent every unchecked access below relies on. +/// Panics if `b` does not hold `UNROLL` rows of `a`'s contraction length. That extent is +/// what every unchecked access below relies on. #[inline(always)] pub(crate) fn f32_store_microkernel( arch: Scalar, - a: QueryPanel<'_, f32, A_PANEL>, - b: DocPanel<'_, f32, B_PANEL, L>, + a: BlockTransposedPanel<'_, f32, A_PANEL>, + b: RowMajorPanel<'_, f32, B_PANEL, L>, mut out: Slot<'_, f32, A_PANEL, B_PANEL>, ) { const { assert!(UNROLL >= 1 && UNROLL <= B_PANEL) } @@ -44,10 +46,7 @@ pub(crate) fn f32_store_microkernel( let k = a.k(); let a_data = a.as_slice(); let b_data = b.as_slice(); - assert_eq!(b_data.len(), UNROLL * k, "doc panel extent"); - - let out = out.as_mut_slice(); - debug_assert_eq!(out.len(), A_PANEL * B_PANEL); + assert_eq!(b_data.len(), UNROLL * k, "B panel extent"); let ap = a_data.as_ptr(); let bp = b_data.as_ptr(); @@ -75,28 +74,29 @@ pub(crate) fn f32_store_microkernel( } } - for (acc_j, col) in acc.iter().zip(out.chunks_exact_mut(A_PANEL)) { - for (acc_jr, dst) in acc_j.iter().zip(col.chunks_exact_mut(f32s::LANES)) { - // SAFETY: `chunks_exact_mut` yields exactly `LANES` elements. + for (acc_j, col) in acc.iter().zip(out.columns()) { + let (tiles, _) = col.as_chunks_mut::<{ f32s::LANES }>(); + for (acc_jr, dst) in acc_j.iter().zip(tiles) { + // SAFETY: `as_chunks_mut` yields exactly `LANES` elements. unsafe { acc_jr.store_simd(dst.as_mut_ptr()) }; } } } -/// Fold `acc`'s columns into the running per-A-row maxima in `state`, one A-panel wide. +/// Merge each column of `acc` into the running per-A-row maxima in `state`. +/// +/// The maxima are held in registers for the whole sweep, so `state` is read once on entry +/// and written once on exit, not once per column. #[inline(always)] -pub(crate) fn fold_columns(arch: Scalar, acc: &[f32], state: &mut [f32]) { - debug_assert_eq!(state.len(), A_PANEL); - debug_assert!(acc.len().is_multiple_of(A_PANEL)); - +pub(crate) fn max_into_rows(arch: Scalar, acc: &[[f32; A_PANEL]], state: &mut [f32; A_PANEL]) { let mut chains = [[f32s::splat(arch, f32::MIN); REGS]; WAYS]; - let mut groups = acc.chunks_exact(WAYS * A_PANEL); + let mut groups = acc.chunks_exact(WAYS); for group in groups.by_ref() { - fold_group(arch, &mut chains, group); + max_into_chains(arch, &mut chains, group); } // Fewer than WAYS columns remain, so each still lands on its own chain. - fold_group(arch, &mut chains, groups.remainder()); + max_into_chains(arch, &mut chains, groups.remainder()); let mut merged = chains[0]; for chain in &chains[1..] { @@ -105,25 +105,19 @@ pub(crate) fn fold_columns(arch: Scalar, acc: &[f32], state: &mut [f32]) { } } - for (m, tile) in merged.iter().zip(state.chunks_exact_mut(f32s::LANES)) { - // SAFETY: `chunks_exact_mut` yields exactly `LANES` elements. - unsafe { - let sp = tile.as_mut_ptr(); - f32s::load_simd(arch, sp).max_simd(*m).store_simd(sp); - } + let (rows, _) = state.as_chunks_mut::<{ f32s::LANES }>(); + for (m, row) in merged.iter().zip(rows) { + *row = m.max_simd(f32s::from_array(arch, *row)).to_array(); } } -/// Merge up to `WAYS` consecutive A-panels, one per chain. +/// Merge up to [`WAYS`] consecutive columns, one per chain. #[inline(always)] -fn fold_group(arch: Scalar, chains: &mut [[f32s; REGS]; WAYS], src: &[f32]) { - for (chain, column) in chains.iter_mut().zip(src.chunks_exact(A_PANEL)) { - for (r, c) in chain.iter_mut().enumerate() { - // SAFETY: `chunks_exact` yields `A_PANEL == REGS * LANES` elements. - unsafe { - let v = f32s::load_simd(arch, column.as_ptr().add(r * f32s::LANES)); - *c = c.max_simd(v); - } +fn max_into_chains(arch: Scalar, chains: &mut [[f32s; REGS]; WAYS], src: &[[f32; A_PANEL]]) { + for (chain, column) in chains.iter_mut().zip(src) { + let (tiles, _) = column.as_chunks::<{ f32s::LANES }>(); + for (c, tile) in chain.iter_mut().zip(tiles) { + *c = c.max_simd(f32s::from_array(arch, *tile)); } } } diff --git a/diskann-quantization/src/multi_vector/distance/kernels/leaves/v3.rs b/diskann-quantization/src/multi_vector/distance/kernels/leaves/v3.rs index 34980ca823..3897deb125 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/leaves/v3.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/leaves/v3.rs @@ -1,7 +1,9 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ -//! V3 (AVX2+FMA) f32 leaf — 16 × 4. +//! V3 (AVX2+FMA) f32 leaf, panels of 16 × 4. use diskann_wide::arch::x86_64::V3; use diskann_wide::{SIMDMinMax, SIMDMulAdd, SIMDVector}; @@ -9,7 +11,7 @@ use diskann_wide::{SIMDMinMax, SIMDMulAdd, SIMDVector}; use super::WAYS; use crate::bits::Length; use crate::multi_vector::distance::kernels::strip::Slot; -use crate::multi_vector::distance::kernels::tiles::{DocPanel, QueryPanel}; +use crate::multi_vector::distance::kernels::tiles::{BlockTransposedPanel, RowMajorPanel}; diskann_wide::alias!(f32s = ::f32x8); @@ -25,13 +27,13 @@ const REGS: usize = A_PANEL / f32s::LANES; /// /// # Panics /// -/// Panics if `b` does not hold exactly `UNROLL` rows of `a`'s contraction length — the -/// extent every unchecked access below relies on. +/// Panics if `b` does not hold `UNROLL` rows of `a`'s contraction length. That extent is +/// what every unchecked access below relies on. #[inline(always)] pub(crate) fn f32_store_microkernel( arch: V3, - a: QueryPanel<'_, f32, A_PANEL>, - b: DocPanel<'_, f32, B_PANEL, L>, + a: BlockTransposedPanel<'_, f32, A_PANEL>, + b: RowMajorPanel<'_, f32, B_PANEL, L>, mut out: Slot<'_, f32, A_PANEL, B_PANEL>, ) { const { assert!(UNROLL >= 1 && UNROLL <= B_PANEL) } @@ -39,10 +41,7 @@ pub(crate) fn f32_store_microkernel( let k = a.k(); let a_data = a.as_slice(); let b_data = b.as_slice(); - assert_eq!(b_data.len(), UNROLL * k, "doc panel extent"); - - let out = out.as_mut_slice(); - debug_assert_eq!(out.len(), A_PANEL * B_PANEL); + assert_eq!(b_data.len(), UNROLL * k, "B panel extent"); let ap = a_data.as_ptr(); let bp = b_data.as_ptr(); @@ -70,28 +69,29 @@ pub(crate) fn f32_store_microkernel( } } - for (acc_j, col) in acc.iter().zip(out.chunks_exact_mut(A_PANEL)) { - for (acc_jr, dst) in acc_j.iter().zip(col.chunks_exact_mut(f32s::LANES)) { - // SAFETY: `chunks_exact_mut` yields exactly `LANES` elements. + for (acc_j, col) in acc.iter().zip(out.columns()) { + let (tiles, _) = col.as_chunks_mut::<{ f32s::LANES }>(); + for (acc_jr, dst) in acc_j.iter().zip(tiles) { + // SAFETY: `as_chunks_mut` yields exactly `LANES` elements. unsafe { acc_jr.store_simd(dst.as_mut_ptr()) }; } } } -/// Fold `acc`'s columns into the running per-A-row maxima in `state`, one A-panel wide. +/// Merge each column of `acc` into the running per-A-row maxima in `state`. +/// +/// The maxima are held in registers for the whole sweep, so `state` is read once on entry +/// and written once on exit, not once per column. #[inline(always)] -pub(crate) fn fold_columns(arch: V3, acc: &[f32], state: &mut [f32]) { - debug_assert_eq!(state.len(), A_PANEL); - debug_assert!(acc.len().is_multiple_of(A_PANEL)); - +pub(crate) fn max_into_rows(arch: V3, acc: &[[f32; A_PANEL]], state: &mut [f32; A_PANEL]) { let mut chains = [[f32s::splat(arch, f32::MIN); REGS]; WAYS]; - let mut groups = acc.chunks_exact(WAYS * A_PANEL); + let mut groups = acc.chunks_exact(WAYS); for group in groups.by_ref() { - fold_group(arch, &mut chains, group); + max_into_chains(arch, &mut chains, group); } // Fewer than WAYS columns remain, so each still lands on its own chain. - fold_group(arch, &mut chains, groups.remainder()); + max_into_chains(arch, &mut chains, groups.remainder()); let mut merged = chains[0]; for chain in &chains[1..] { @@ -100,25 +100,19 @@ pub(crate) fn fold_columns(arch: V3, acc: &[f32], state: &mut [f32]) { } } - for (m, tile) in merged.iter().zip(state.chunks_exact_mut(f32s::LANES)) { - // SAFETY: `chunks_exact_mut` yields exactly `LANES` elements. - unsafe { - let sp = tile.as_mut_ptr(); - f32s::load_simd(arch, sp).max_simd(*m).store_simd(sp); - } + let (rows, _) = state.as_chunks_mut::<{ f32s::LANES }>(); + for (m, row) in merged.iter().zip(rows) { + *row = m.max_simd(f32s::from_array(arch, *row)).to_array(); } } -/// Merge up to `WAYS` consecutive A-panels, one per chain. +/// Merge up to [`WAYS`] consecutive columns, one per chain. #[inline(always)] -fn fold_group(arch: V3, chains: &mut [[f32s; REGS]; WAYS], src: &[f32]) { - for (chain, column) in chains.iter_mut().zip(src.chunks_exact(A_PANEL)) { - for (r, c) in chain.iter_mut().enumerate() { - // SAFETY: `chunks_exact` yields `A_PANEL == REGS * LANES` elements. - unsafe { - let v = f32s::load_simd(arch, column.as_ptr().add(r * f32s::LANES)); - *c = c.max_simd(v); - } +fn max_into_chains(arch: V3, chains: &mut [[f32s; REGS]; WAYS], src: &[[f32; A_PANEL]]) { + for (chain, column) in chains.iter_mut().zip(src) { + let (tiles, _) = column.as_chunks::<{ f32s::LANES }>(); + for (c, tile) in chain.iter_mut().zip(tiles) { + *c = c.max_simd(f32s::from_array(arch, *tile)); } } } diff --git a/diskann-quantization/src/multi_vector/distance/kernels/mod.rs b/diskann-quantization/src/multi_vector/distance/kernels/mod.rs index 410f30683e..703bae017b 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/mod.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/mod.rs @@ -1,5 +1,7 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ //! Cache-tiled MaxSim: a [`TileWalk`] lends cache-sized tiles, each [`Paneled`] into the //! panels one leaf call consumes, plus a typed tail. [`Accumulate`] folds one (A-panel, @@ -7,13 +9,16 @@ //! accumulator. //! //! Position is ordinal. [`drive`] counts the panels it passes and hands a [`Drain`] an -//! A-panel index and a B-panel range — never a stride or an address, so a drain scales -//! those ordinals by its own panel width and clamps the end against its own extent. A side +//! A-panel index and a B-panel range, never a stride or an address. A drain scales those +//! ordinals by its own panel width and clamps the end against its own extent. A side //! whose tail type is uninhabited ([`NoTail`]) need not clamp at all, and that is a -//! type-level fact rather than a convention. +//! type-level fact, not a convention. //! -//! The driver speaks A/B; instantiations speak Query/Doc. +//! Naming follows the layer: `A` and `B` are positions in the contraction, type names say +//! which layout a thing is, and only the entry points and the [`Drain`] speak of queries +//! and documents. +use core::num::NonZeroUsize; use core::ops::Range; mod f16; @@ -25,7 +30,9 @@ mod tiles; pub(crate) use f16::MaxIpF16; pub(crate) use float::MaxIp; -// ── Tile budget and planning ───────────────────────────────────── +////////////////////////////// +// Tile budget and planning // +////////////////////////////// /// Cache budgets fed to [`Plan::new`]. #[derive(Debug, Clone, Copy)] @@ -57,41 +64,46 @@ impl Default for TileBudget { /// than the B-rows on hand. #[derive(Debug, Clone, Copy)] struct Plan { - a_panels: usize, - b_panels: usize, + a_panels: NonZeroUsize, + b_panels: NonZeroUsize, +} + +/// A panel count of at least one: a budget too small for even a single panel still has to +/// make progress, and a zero-wide tile would stall the walk it strides. +fn at_least_one_panel(panels: usize) -> NonZeroUsize { + NonZeroUsize::new(panels).unwrap_or(NonZeroUsize::MIN) } impl Plan { /// `b_panels` is reconciled against `b_rows`, so a plan belongs to the B side it was - /// built for; reused against a longer one it still computes the right answer, but tiles + /// built for. Reused against a longer one it still computes the right answer, but tiles /// far more narrowly than the cache allows. `a_panels` needs no such reconciliation: it only /// feeds a walk stride, which the cursor already bounds by the data it holds, whereas /// `b_panels` also sizes the accumulator strip. /// - /// Row sizes are clamped to one byte so a degenerate row cannot divide by zero; the - /// entries reject an empty contraction before planning, so such a plan is never used. + /// Row sizes are [`NonZeroUsize`] because they reach a divisor. The entries reject an + /// empty contraction before planning. fn new( - a_row_bytes: usize, - b_row_bytes: usize, + a_row_bytes: NonZeroUsize, + b_row_bytes: NonZeroUsize, b_rows: usize, acc_bytes: usize, budget: TileBudget, ) -> Self { - let a_row_bytes = a_row_bytes.max(1); - let b_row_bytes = b_row_bytes.max(1); - - let a_panels = (budget.l2_a / (a_row_bytes * AR)).max(1); + // Dividing by the row width and then the panel width, instead of by their product, + // leaves both divisors provably non-zero without a check, and leaves no + // intermediate that could overflow. + let a_panels = at_least_one_panel(budget.l2_a / a_row_bytes / AR); // A B-row costs its own bytes plus the accumulator column it fills, and one // A-panel stays resident alongside. - let per_b_row = b_row_bytes + AR * acc_bytes; - let b_budget = budget.l1_b.saturating_sub(AR * a_row_bytes); + let per_b_row = b_row_bytes.saturating_add(AR * acc_bytes); + let b_budget = budget.l1_b.saturating_sub(AR * a_row_bytes.get()); let cache_fit = (b_budget / per_b_row) / BR; // Never plan wider than the B-rows on hand: a cache-sized tile over a short B side - // would size accumulator columns no fill can reach. One panel at minimum, so a - // budget too small for even that still makes progress. - let b_panels = cache_fit.min(b_rows.div_ceil(BR)).max(1); + // would size accumulator columns no fill can reach. + let b_panels = at_least_one_panel(cache_fit.min(b_rows.div_ceil(BR))); Self { a_panels, b_panels } } @@ -100,23 +112,25 @@ impl Plan { /// /// Sizing the scratch below this trips [`slots_exhausted`]. fn strip_len(&self) -> usize { - AR * self.b_panels * BR + AR * self.b_panels.get() * BR } } -// ── Read side ──────────────────────────────────────────────────── +/////////////// +// Read side // +/////////////// /// Per-lifetime half of [`TileWalk`]. /// /// The defaulted `B = &'a Self` carries the `Self: 'a` implied bound through -/// well-formedness; a plain GAT `where Self: 'a` collapses to `'static` under [`drive`]'s +/// well-formedness. A plain GAT `where Self: 'a` collapses to `'static` under [`drive`]'s /// `for<'a>` bound on stable. trait TileAt<'a, B = &'a Self> { type Tile: Paneled; } /// A **lending** walk: `next` reborrows `&mut self`, so a tile may borrow a buffer the walk -/// reuses on the following call — which is what lets a walk convert as it goes. `reset` +/// reuses on the following call. That is what lets a walk convert as it goes. `reset` /// rewinds, because B is re-walked once per A-tile. trait TileWalk: for<'a> TileAt<'a> { fn next(&mut self) -> Option<>::Tile>; @@ -125,9 +139,9 @@ trait TileWalk: for<'a> TileAt<'a> { /// An iterator whose short trailing element has its own type. /// -/// `tail` consumes the exhausted iterator, so the trailer comes off the cursor the loop was -/// already advancing rather than being recomputed from the source. -trait TailIterator: ExactSizeIterator { +/// `tail` consumes the exhausted iterator. The trailer comes off the cursor the loop was +/// already advancing and is never recomputed from the source. +trait TailIterator: Iterator { type Tail; fn tail(self) -> Option; } @@ -149,15 +163,17 @@ trait Paneled { #[derive(Clone, Copy)] enum NoTail {} -// ── Write side ─────────────────────────────────────────────────── +//////////////// +// Write side // +//////////////// -/// Per-lifetime half of [`Scratch`] — same implied-bound trick as [`TileAt`]. +/// Per-lifetime half of [`Scratch`], using the same implied-bound trick as [`TileAt`]. trait SlotsAt<'s, B = &'s mut Self> { - /// Named here rather than reached through the iterator so [`drive`]'s bounds can + /// Named here instead of being reached through the iterator, so [`drive`]'s bounds can /// project it off the scratch. type Slot; - /// Plain rather than lending: slots partition one buffer within a single call, so none - /// of them borrows the cursor. + /// Plain, not lending. Slots partition one buffer within a single call, so none of them + /// borrows the cursor. type Slots: Iterator; } @@ -167,17 +183,20 @@ trait Scratch: for<'s> SlotsAt<'s> { fn slots(&mut self) -> >::Slots; } -// ── Compute side ───────────────────────────────────────────────── +////////////////// +// Compute side // +////////////////// /// One A-panel × one B-panel → one accumulator slot. /// -/// Pinned on all three as type parameters, so the walks' panel types select the impl. +/// Pinned on all three as type parameters, which lets the walks' panel types select the +/// impl. trait Accumulate { fn accumulate(&self, arch: Arch, a: A, b: B, out: O); } -/// [`NoTail`] is uninhabited, so this discharges [`drive`]'s A-tail bounds for every kernel -/// — and by coherence forbids any kernel from writing its own. +/// [`NoTail`] is uninhabited, so this discharges [`drive`]'s A-tail bounds for every +/// kernel, and by coherence forbids any kernel from writing its own. impl Accumulate for K { #[inline(always)] fn accumulate(&self, _: Arch, a: NoTail, _: B, _: O) { @@ -192,7 +211,9 @@ trait Drain { fn drain(&mut self, arch: Arch, scratch: &S, a_panel: usize, b_panels: Range); } -// ── Driver ─────────────────────────────────────────────────────── +//////////// +// Driver // +//////////// type PanelOf<'a, W> = <>::Tile as Paneled>::Panel; type TailOf<'a, W> = <>::Tile as Paneled>::Tail; @@ -201,13 +222,13 @@ type SlotOf<'s, S> = >::Slot; #[cold] #[inline(never)] fn slots_exhausted() -> ! { - unreachable!("scratch ran out of slots — it is narrower than the walk's B-tile") + unreachable!("scratch ran out of slots: narrower than the walk's B-tile") } /// One A-panel against a whole B-tile. /// -/// Returns the slots it filled — ground truth for how far B advanced, rather than the -/// prediction an `ExactSizeIterator::len` would give, which excludes the tail anyway. +/// Returns the slots it filled, which is ground truth for how far B advanced: a count of +/// whole panels would exclude the tail. #[inline(always)] fn fill(arch: Arch, kernel: &K, a: A, b_tile: &BT, scratch: &mut S) -> usize where @@ -270,7 +291,7 @@ fn drive( let mut b_base = 0; // Last pass wins: every B-tile re-sweeps the same A-panels, so both counters are // rewritten identically each pass and read after the last. An A-tile with no - // B-tiles advances neither, which is unobservable — no drain fires, and a B source + // B-tiles advances neither, which is unobservable. No drain fires, and a B source // empty for one A-tile is empty for all. let mut a_end = a_base; let mut b_used = 0; @@ -301,6 +322,10 @@ fn drive( mod tests { use super::*; + fn nz(n: usize) -> NonZeroUsize { + NonZeroUsize::new(n).unwrap() + } + #[test] fn plan_reserves_l1_for_the_resident_a_panel_and_accumulator() { // 64-byte rows, AR = 16, BR = 4, 4-byte accumulator. @@ -308,8 +333,8 @@ mod tests { // l1_b 36000 - 16 * 64 = 34976 for B; per B-row = 64 + 16 * 4 = 128; // 34976 / 128 = 273 rows -> 273 / 4 = 68 B-panels. let plan = Plan::<16, 4>::new( - 64, - 64, + nz(64), + nz(64), usize::MAX, 4, TileBudget { @@ -317,7 +342,7 @@ mod tests { l1_b: 36000, }, ); - assert_eq!((plan.a_panels, plan.b_panels), (40, 68)); + assert_eq!((plan.a_panels.get(), plan.b_panels.get()), (40, 68)); assert_eq!(plan.strip_len(), 16 * 68 * 4); } @@ -328,15 +353,21 @@ mod tests { l2_a: 40960, l1_b: 36000, }; - let plan = Plan::<16, 4>::new(64, 64, 10, 4, budget); - assert_eq!(plan.b_panels, 3); + let plan = Plan::<16, 4>::new(nz(64), nz(64), 10, 4, budget); + assert_eq!(plan.b_panels.get(), 3); assert_eq!(plan.strip_len(), 16 * 3 * 4); } #[test] fn plan_clamps_to_one_panel_per_tile() { - let plan = Plan::<16, 4>::new(1024, 1024, usize::MAX, 4, TileBudget { l2_a: 1, l1_b: 1 }); - assert_eq!((plan.a_panels, plan.b_panels), (1, 1)); + let plan = Plan::<16, 4>::new( + nz(1024), + nz(1024), + usize::MAX, + 4, + TileBudget { l2_a: 1, l1_b: 1 }, + ); + assert_eq!((plan.a_panels.get(), plan.b_panels.get()), (1, 1)); assert_eq!(plan.strip_len(), 16 * 4); } } diff --git a/diskann-quantization/src/multi_vector/distance/kernels/strip.rs b/diskann-quantization/src/multi_vector/distance/kernels/strip.rs index 30e3f94b7e..93d89ad359 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/strip.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/strip.rs @@ -1,63 +1,85 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ //! The accumulator a fill writes and a drain reads. //! -//! Rows are query rows and one column is one doc — note that a doc arrives as a *row* of -//! the input and lands as a *column* here. Column-major over the whole strip: column `c` -//! occupies `[c * AR, (c + 1) * AR)`, and slot `p` covers columns `p * BR ..`, so a drain -//! can address a run of columns without knowing which slot produced them. +//! Storage is a flat run of `AR`-element chunks. One chunk is what the rest of the module +//! calls a *column*. Read as memory the strip is therefore row-major with the axes +//! flipped: an `n × AR` matrix whose rows are the accumulator's columns. +//! +//! ```text +//! AR = 4, BR = 2: +//! +//! memory -> [ a0 a1 a2 a3 ][ a0 a1 a2 a3 ][ a0 a1 a2 a3 ][ a0 a1 a2 a3 ] +//! column 0 column 1 column 2 column 3 +//! \___________ slot 0 _________/\___________ slot 1 _________/ +//! ``` +//! +//! Column `c` occupies `[c * AR, (c + 1) * AR)` and slot `p` covers columns `p * BR ..`, +//! so [`Strip::columns`] can hand a drain a run that straddles slot boundaries. What the +//! two axes *mean* is the drain's business, not the strip's. See +//! [`RawMax`](super::float::RawMax). -use core::slice::ChunksExactMut; +use core::mem; use super::{Scratch, SlotsAt}; -/// Accumulator for one A-panel against one whole B-tile. +/// Accumulator for one A-panel against one whole B-tile, carved by its [`Scratch`] +/// impl into one [`Slot`] per B-panel. /// /// Holds no cursor: [`Scratch::slots`] restarts from the front on every fill, which is /// what lets the same memory be re-lent across tiles without clearing. pub(super) struct Strip<'a, T, const AR: usize, const BR: usize> { - buf: &'a mut [T], + buf: &'a mut [[T; AR]], } impl<'a, T, const AR: usize, const BR: usize> Strip<'a, T, AR, BR> { - /// Trailing elements beyond the last whole slot are never touched. + /// Trailing elements beyond the last whole column are never touched. pub(super) fn new(buf: &'a mut [T]) -> Self { - Self { buf } + Self { + buf: buf.as_chunks_mut::().0, + } } - /// The first `live` columns — the only region the fill just performed wrote. + /// The first `live` columns, the region the fill just performed wrote. /// /// # Panics /// /// Panics if `live` exceeds the strip's column capacity. - pub(super) fn columns(&self, live: usize) -> &[T] { - &self.buf[..live * AR] + pub(super) fn columns(&self, live: usize) -> &[[T; AR]] { + &self.buf[..live] } } -/// One `AR × BR` accumulator tile, column-major. +/// The `BR` consecutive columns of a [`Strip`] that one leaf call accumulates into. pub(super) struct Slot<'a, T, const AR: usize, const BR: usize> { - buf: &'a mut [T], + buf: &'a mut [[T; AR]; BR], } impl Slot<'_, T, AR, BR> { - /// Exactly `AR * BR` elements: [`Slots`] only ever cuts whole tiles. - pub(super) fn as_mut_slice(&mut self) -> &mut [T] { + pub(super) fn columns(&mut self) -> &mut [[T; AR]; BR] { &mut *self.buf } } -/// Cuts a strip into disjoint slots, stopping short of a trailing partial tile. +/// Cuts a [`Strip`] into disjoint [`Slot`]s, stopping short of a trailing partial tile. /// -/// Disjointness and that stopping rule are both structural, from `chunks_exact_mut`. -pub(super) struct Slots<'a, T, const AR: usize, const BR: usize>(ChunksExactMut<'a, T>); +/// Disjointness and that stopping rule are both structural, from `split_first_chunk_mut`. +pub(super) struct Slots<'a, T, const AR: usize, const BR: usize> { + rest: &'a mut [[T; AR]], +} impl<'a, T, const AR: usize, const BR: usize> Iterator for Slots<'a, T, AR, BR> { type Item = Slot<'a, T, AR, BR>; fn next(&mut self) -> Option { - self.0.next().map(|buf| Slot { buf }) + // A reborrow through `&mut self` cannot reach `'a`, so the remainder is moved out + // and put back. Too short a remainder leaves it empty, which is where it ends. + let (buf, rest) = mem::take(&mut self.rest).split_first_chunk_mut::
()?; + self.rest = rest; + Some(Slot { buf }) } } @@ -68,7 +90,9 @@ impl<'s, T, const AR: usize, const BR: usize> SlotsAt<'s> for Strip<'_, T, AR, B impl Scratch for Strip<'_, T, AR, BR> { fn slots(&mut self) -> Slots<'_, T, AR, BR> { - Slots(self.buf.chunks_exact_mut(AR * BR)) + Slots { + rest: &mut *self.buf, + } } } @@ -82,7 +106,7 @@ mod tests { let mut strip = Strip::::new(&mut buf); for (n, mut slot) in strip.slots().enumerate() { - slot.as_mut_slice().fill(n as u32 + 1); + slot.columns().as_flattened_mut().fill(n as u32 + 1); } assert_eq!(strip.slots().count(), 3); @@ -98,7 +122,10 @@ mod tests { let strip = Strip::::new(&mut buf); // Three columns straddle the first slot (2 columns) into the second. - assert_eq!(strip.columns(3), &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]); - assert_eq!(strip.columns(0), &[] as &[u32]); + assert_eq!( + strip.columns(3), + &[[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]] + ); + assert_eq!(strip.columns(0), &[] as &[[u32; 4]]); } } diff --git a/diskann-quantization/src/multi_vector/distance/kernels/tiles.rs b/diskann-quantization/src/multi_vector/distance/kernels/tiles.rs index a9ae86763f..d148e48146 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/tiles.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/tiles.rs @@ -1,27 +1,35 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ //! The read ladder: walks lend tiles, tiles cut into panels. //! -//! A query panel is one block-transposed block — `AR` rows interleaved column-major, so -//! the leaf loads `AR` contiguous rows per contraction step. A doc panel is `BR` row-major -//! rows. Neither carries `k`: it is `len() / AR` on the query side, and the leaf checks the -//! doc panel against it, so a mismatched pair is caught where it is used rather than -//! assumed. +//! The two sides read different layouts. A carries a [`BlockTransposedRef`], whose blocks +//! already interleave `AR` rows, so a panel is one block and a leaf loads `AR` contiguous +//! rows per contraction step. B carries a [`Standard`] matrix, plain row-major, so a panel +//! is `BR` consecutive rows. +//! +//! Neither panel stores `k`: it is `len() / AR` on the block-transposed side, and the leaf +//! checks the row-major panel against it, so a mismatched pair is caught where it is +//! used, not assumed. -use core::slice::ChunksExact; +use core::num::NonZeroUsize; use super::{NoTail, Paneled, TailIterator, TileAt, TileWalk}; use crate::bits::{Dynamic, Length, Static}; use crate::multi_vector::{BlockTransposedRef, MatRef, Standard}; -// ── Query side ─────────────────────────────────────────────────── +/////////////////////////// +// Block-transposed side // +/////////////////////////// -/// One block-transposed block: `AR` rows × `k` columns, column-major within the block. +/// A single block of a [`BlockTransposedRef`]: `AR` rows × `k` columns, column-major +/// within the block. #[derive(Clone, Copy)] -pub(super) struct QueryPanel<'a, T, const AR: usize>(&'a [T]); +pub(super) struct BlockTransposedPanel<'a, T, const AR: usize>(&'a [T]); -impl<'a, T, const AR: usize> QueryPanel<'a, T, AR> { +impl<'a, T, const AR: usize> BlockTransposedPanel<'a, T, AR> { pub(super) fn as_slice(&self) -> &'a [T] { self.0 } @@ -32,23 +40,23 @@ impl<'a, T, const AR: usize> QueryPanel<'a, T, AR> { } } -pub(super) struct QueryPanels<'a, T, const AR: usize>(ChunksExact<'a, T>); +/// The panel iterator of a [`BlockTransposedTile`]. +pub(super) struct BlockTransposedPanels<'a, T, const AR: usize> { + rest: &'a [T], + stride: usize, +} -impl<'a, T, const AR: usize> Iterator for QueryPanels<'a, T, AR> { - type Item = QueryPanel<'a, T, AR>; +impl<'a, T, const AR: usize> Iterator for BlockTransposedPanels<'a, T, AR> { + type Item = BlockTransposedPanel<'a, T, AR>; fn next(&mut self) -> Option { - self.0.next().map(QueryPanel) - } - - fn size_hint(&self) -> (usize, Option) { - self.0.size_hint() + let (panel, rest) = self.rest.split_at_checked(self.stride)?; + self.rest = rest; + Some(BlockTransposedPanel(panel)) } } -impl ExactSizeIterator for QueryPanels<'_, T, AR> {} - -impl TailIterator for QueryPanels<'_, T, AR> { +impl TailIterator for BlockTransposedPanels<'_, T, AR> { type Tail = NoTail; fn tail(self) -> Option { @@ -56,46 +64,55 @@ impl TailIterator for QueryPanels<'_, T, AR> { } } -/// A run of whole blocks — block-transposed storage is padded to `AR`, hence [`NoTail`]. -pub(super) struct QueryTile<'a, T, const AR: usize> { +/// A view over consecutive blocks of a [`BlockTransposedRef`](BlockTransposedRef). +/// +/// Its [`Paneled`] impl yields one [`BlockTransposedPanel`] per block. +pub(super) struct BlockTransposedTile<'a, T, const AR: usize> { data: &'a [T], - k: usize, + k: NonZeroUsize, } -impl<'a, T, const AR: usize> QueryTile<'a, T, AR> { +impl<'a, T, const AR: usize> BlockTransposedTile<'a, T, AR> { /// # Panics /// /// Debug-only: panics unless `data` is a whole number of `AR × k` blocks, which is /// what makes [`NoTail`] honest. - pub(super) fn new(data: &'a [T], k: usize) -> Self { - debug_assert!(k > 0 && data.len().is_multiple_of(AR * k)); + pub(super) fn new(data: &'a [T], k: NonZeroUsize) -> Self { + debug_assert!(data.len().is_multiple_of(AR * k.get())); Self { data, k } } } -impl<'a, T: Copy, const AR: usize> Paneled for QueryTile<'a, T, AR> { - type Panel = QueryPanel<'a, T, AR>; +impl<'a, T: Copy, const AR: usize> Paneled for BlockTransposedTile<'a, T, AR> { + type Panel = BlockTransposedPanel<'a, T, AR>; + /// Block-transposed storage is padded out to whole `AR`-row blocks, so a run of blocks + /// cannot end in a partial one. type Tail = NoTail; - type Panels = QueryPanels<'a, T, AR>; + type Panels = BlockTransposedPanels<'a, T, AR>; fn panels(&self) -> Self::Panels { - QueryPanels(self.data.chunks_exact(AR * self.k)) + BlockTransposedPanels { + rest: self.data, + stride: AR * self.k.get(), + } } } -// ── Doc side ───────────────────────────────────────────────────── +//////////////////// +// Row-major side // +//////////////////// -/// Up to `BR` row-major rows of `k` elements. +/// Up to `BR` consecutive rows of a [`Standard`] matrix, `k` elements each. /// /// `L` is [`Static
`] for a whole panel and [`Dynamic`] for the trailing one, which is /// how the tail reaches a leaf that can unroll for its width. #[derive(Clone, Copy)] -pub(super) struct DocPanel<'a, T, const BR: usize, L> { +pub(super) struct RowMajorPanel<'a, T, const BR: usize, L> { data: &'a [T], rows: L, } -impl<'a, T, const BR: usize, L: Length> DocPanel<'a, T, BR, L> { +impl<'a, T, const BR: usize, L: Length> RowMajorPanel<'a, T, BR, L> { pub(super) fn as_slice(&self) -> &'a [T] { self.data } @@ -105,84 +122,115 @@ impl<'a, T, const BR: usize, L: Length> DocPanel<'a, T, BR, L> { } } -pub(super) struct DocPanels<'a, T, const BR: usize> { - chunks: ChunksExact<'a, T>, - k: usize, +/// The panel iterator of a [`RowMajorTile`]. +pub(super) struct RowMajorPanels<'a, T, const BR: usize> { + rest: &'a [T], + /// The contraction length, not a panel stride. Only the tail needs a divisor, and what + /// it divides by is this, not `BR * k`. + k: NonZeroUsize, } -impl<'a, T, const BR: usize> Iterator for DocPanels<'a, T, BR> { - type Item = DocPanel<'a, T, BR, Static
>; +impl<'a, T, const BR: usize> Iterator for RowMajorPanels<'a, T, BR> { + type Item = RowMajorPanel<'a, T, BR, Static
>; fn next(&mut self) -> Option { - self.chunks.next().map(|data| DocPanel { + let (data, rest) = self.rest.split_at_checked(BR * self.k.get())?; + self.rest = rest; + Some(RowMajorPanel { data, rows: Static::
, }) } - - fn size_hint(&self) -> (usize, Option) { - self.chunks.size_hint() - } } -impl ExactSizeIterator for DocPanels<'_, T, BR> {} - -impl<'a, T, const BR: usize> TailIterator for DocPanels<'a, T, BR> { - type Tail = DocPanel<'a, T, BR, Dynamic>; +impl<'a, T, const BR: usize> TailIterator for RowMajorPanels<'a, T, BR> { + type Tail = RowMajorPanel<'a, T, BR, Dynamic>; fn tail(self) -> Option { - let data = self.chunks.remainder(); - (!data.is_empty()).then(|| DocPanel { - data, - rows: Dynamic(data.len() / self.k), + debug_assert!( + self.rest.len() < BR * self.k.get(), + "tail taken before the iterator was exhausted" + ); + (!self.rest.is_empty()).then(|| RowMajorPanel { + data: self.rest, + rows: Dynamic(self.rest.len() / self.k.get()), }) } } -/// Row-major rows; only the final tile of a walk may be short. -pub(super) struct DocTile<'a, T, const BR: usize> { +/// A view over consecutive rows of a [`Standard`] matrix. +/// +/// Its [`Paneled`] impl yields one [`RowMajorPanel`] per `BR` rows, plus a short tail. +pub(super) struct RowMajorTile<'a, T, const BR: usize> { data: &'a [T], - k: usize, + k: NonZeroUsize, } -impl<'a, T, const BR: usize> DocTile<'a, T, BR> { +impl<'a, T, const BR: usize> RowMajorTile<'a, T, BR> { /// # Panics /// /// Debug-only: panics unless `data` is a whole number of `k`-element rows. - pub(super) fn new(data: &'a [T], k: usize) -> Self { - debug_assert!(k > 0 && data.len().is_multiple_of(k)); + pub(super) fn new(data: &'a [T], k: NonZeroUsize) -> Self { + debug_assert!(data.len().is_multiple_of(k.get())); Self { data, k } } } -impl<'a, T: Copy, const BR: usize> Paneled for DocTile<'a, T, BR> { - type Panel = DocPanel<'a, T, BR, Static
>; - type Tail = DocPanel<'a, T, BR, Dynamic>; - type Panels = DocPanels<'a, T, BR>; +impl<'a, T: Copy, const BR: usize> Paneled for RowMajorTile<'a, T, BR> { + type Panel = RowMajorPanel<'a, T, BR, Static
>; + /// A row count need not divide `BR`, so a tile can end in a partial panel. + type Tail = RowMajorPanel<'a, T, BR, Dynamic>; + type Panels = RowMajorPanels<'a, T, BR>; fn panels(&self) -> Self::Panels { - DocPanels { - chunks: self.data.chunks_exact(BR * self.k), + RowMajorPanels { + rest: self.data, k: self.k, } } } -// ── Walks ──────────────────────────────────────────────────────── +/////////// +// Walks // +/////////// + +/// The contraction length as a [`NonZeroUsize`], which walks stride by and panels divide +/// by. Carrying the invariant in the type keeps those sites free of a zero check. +/// +/// # Panics +/// +/// Panics if `k` is zero. The entry guards that case before any walk is built. This is +/// the backstop for constructors reachable without it. +#[track_caller] +#[expect( + clippy::expect_used, + reason = "a walk cannot represent an empty contraction" +)] +pub(super) fn contraction(k: usize) -> NonZeroUsize { + NonZeroUsize::new(k).expect("walk requires a non-empty contraction") +} + +/// Elements spanned by `panels` panels of `width` rows, each row `k` long. +/// +/// Overflowing the product needs a contraction near `usize::MAX` and no entry admits one, +/// so a wrapped result is kept non-zero instead of being detected. A stride of zero would +/// leave [`Cursor`] handing out empty tiles without ever advancing. +pub(super) fn tile_stride(panels: NonZeroUsize, width: usize, k: NonZeroUsize) -> NonZeroUsize { + NonZeroUsize::new(panels.get() * width * k.get()).unwrap_or(NonZeroUsize::MIN) +} /// A cursor over a contiguous source, cut into tiles of `stride` elements. /// -/// Shared by both sides — and by the widening walks — because the only difference is how -/// a tile is interpreted. +/// Shared by both sides, and by the widening walks, because the only difference is how a +/// tile is interpreted. pub(super) struct Cursor<'a, T> { data: &'a [T], - stride: usize, + stride: NonZeroUsize, cur: usize, } impl<'a, T> Cursor<'a, T> { - pub(super) fn new(data: &'a [T], stride: usize) -> Self { - debug_assert!(stride > 0); + pub(super) fn new(data: &'a [T], stride: NonZeroUsize) -> Self { Self { data, stride, @@ -190,12 +238,13 @@ impl<'a, T> Cursor<'a, T> { } } + /// The next tile, clamped to what remains. The last one may be short. pub(super) fn next(&mut self) -> Option<&'a [T]> { let rest = self.data.get(self.cur..)?; if rest.is_empty() { return None; } - let take = self.stride.min(rest.len()); + let take = self.stride.get().min(rest.len()); self.cur += take; Some(&rest[..take]) } @@ -204,40 +253,41 @@ impl<'a, T> Cursor<'a, T> { self.cur = 0; } - /// The longest tile this cursor can yield — the size a converting walk must stage. + /// The longest tile this cursor can yield, which is the size a converting walk must stage. pub(super) fn widest(&self) -> usize { - self.stride.min(self.data.len()) + self.stride.get().min(self.data.len()) } } -/// Walks the padded storage of a block-transposed query, `a_panels` blocks at a time. -pub(super) struct QueryWalk<'a, T, const AR: usize> { +/// Walks the padded storage of a [`BlockTransposedRef`], `a_panels` blocks at a time. +pub(super) struct BlockTransposedWalk<'a, T, const AR: usize> { cursor: Cursor<'a, T>, - k: usize, + k: NonZeroUsize, } -impl<'a, T: Copy, const AR: usize> QueryWalk<'a, T, AR> { +impl<'a, T: Copy, const AR: usize> BlockTransposedWalk<'a, T, AR> { /// # Panics /// /// Panics if `view` has no columns. - pub(super) fn new(view: BlockTransposedRef<'a, T, AR>, a_panels: usize) -> Self { - let k = view.padded_ncols(); - assert!(k > 0, "QueryWalk requires a non-empty contraction"); + pub(super) fn new(view: BlockTransposedRef<'a, T, AR>, a_panels: NonZeroUsize) -> Self { + let k = contraction(view.padded_ncols()); Self { - cursor: Cursor::new(view.as_slice(), a_panels * AR * k), + cursor: Cursor::new(view.as_slice(), tile_stride(a_panels, AR, k)), k, } } } -impl<'t, T: Copy, const AR: usize> TileAt<'t> for QueryWalk<'_, T, AR> { - type Tile = QueryTile<'t, T, AR>; +impl<'t, T: Copy, const AR: usize> TileAt<'t> for BlockTransposedWalk<'_, T, AR> { + type Tile = BlockTransposedTile<'t, T, AR>; } -impl TileWalk for QueryWalk<'_, T, AR> { - fn next(&mut self) -> Option> { +impl TileWalk for BlockTransposedWalk<'_, T, AR> { + fn next(&mut self) -> Option> { let k = self.k; - self.cursor.next().map(|data| QueryTile::new(data, k)) + self.cursor + .next() + .map(|data| BlockTransposedTile::new(data, k)) } fn reset(&mut self) { @@ -245,34 +295,33 @@ impl TileWalk for QueryWalk<'_, T, AR> { } } -/// Walks a row-major doc matrix, `b_panels * BR` rows at a time. -pub(super) struct DocWalk<'a, T, const BR: usize> { +/// Walks a [`Standard`] matrix, `b_panels * BR` rows at a time. +pub(super) struct RowMajorWalk<'a, T, const BR: usize> { cursor: Cursor<'a, T>, - k: usize, + k: NonZeroUsize, } -impl<'a, T, const BR: usize> DocWalk<'a, T, BR> { +impl<'a, T, const BR: usize> RowMajorWalk<'a, T, BR> { /// # Panics /// - /// Panics if `docs` has zero-length rows. - pub(super) fn new(docs: MatRef<'a, Standard>, b_panels: usize) -> Self { - let k = docs.vector_dim(); - assert!(k > 0, "DocWalk requires a non-empty contraction"); + /// Panics if `mat` has zero-length rows. + pub(super) fn new(mat: MatRef<'a, Standard>, b_panels: NonZeroUsize) -> Self { + let k = contraction(mat.vector_dim()); Self { - cursor: Cursor::new(docs.as_slice(), b_panels * BR * k), + cursor: Cursor::new(mat.as_slice(), tile_stride(b_panels, BR, k)), k, } } } -impl<'t, T: Copy, const BR: usize> TileAt<'t> for DocWalk<'_, T, BR> { - type Tile = DocTile<'t, T, BR>; +impl<'t, T: Copy, const BR: usize> TileAt<'t> for RowMajorWalk<'_, T, BR> { + type Tile = RowMajorTile<'t, T, BR>; } -impl TileWalk for DocWalk<'_, T, BR> { - fn next(&mut self) -> Option> { +impl TileWalk for RowMajorWalk<'_, T, BR> { + fn next(&mut self) -> Option> { let k = self.k; - self.cursor.next().map(|data| DocTile::new(data, k)) + self.cursor.next().map(|data| RowMajorTile::new(data, k)) } fn reset(&mut self) { @@ -284,29 +333,31 @@ impl TileWalk for DocWalk<'_, T, BR> { mod tests { use super::*; - /// Handles are passed by value into every leaf call; keeping them register-sized is - /// the reason `k` is derived rather than stored. + /// Handles are passed by value into every leaf call. Keeping them register-sized is why + /// `k` is derived instead of stored. #[test] fn handles_stay_thin() { use core::mem::size_of; - assert_eq!(size_of::>(), 2 * size_of::()); assert_eq!( - size_of::>>(), + size_of::>(), + 2 * size_of::() + ); + assert_eq!( + size_of::>>(), 2 * size_of::() ); assert_eq!( - size_of::>(), + size_of::>(), 3 * size_of::() ); } #[test] - fn doc_tile_splits_into_whole_panels_plus_a_dynamic_tail() { + fn row_major_tile_splits_into_whole_panels_plus_a_dynamic_tail() { let data: Vec = (0..7 * 3).map(|i| i as u8).collect(); - let tile = DocTile::::new(&data, 3); + let tile = RowMajorTile::::new(&data, NonZeroUsize::new(3).unwrap()); let mut panels = tile.panels(); - assert_eq!(panels.len(), 3); assert_eq!(panels.next().unwrap().as_slice(), &[0, 1, 2, 3, 4, 5]); let rest: Vec<_> = panels.by_ref().map(|p| p.rows()).collect(); assert_eq!(rest, [2, 2]); @@ -319,7 +370,7 @@ mod tests { #[test] fn a_tile_that_divides_evenly_has_no_tail() { let data: Vec = (0..6 * 3).map(|i| i as u8).collect(); - let mut panels = DocTile::::new(&data, 3).panels(); + let mut panels = RowMajorTile::::new(&data, NonZeroUsize::new(3).unwrap()).panels(); assert_eq!(panels.by_ref().count(), 3); assert!(panels.tail().is_none()); } @@ -327,8 +378,8 @@ mod tests { #[test] fn a_walk_yields_short_final_tile_then_stops_and_rewinds() { let data: Vec = (0..5 * 3).map(|i| i as u8).collect(); - let docs = MatRef::new(Standard::new(5, 3).unwrap(), &data).unwrap(); - let mut walk = DocWalk::::new(docs, 1); + let mat = MatRef::new(Standard::new(5, 3).unwrap(), &data).unwrap(); + let mut walk = RowMajorWalk::::new(mat, NonZeroUsize::new(1).unwrap()); assert_eq!(walk.next().unwrap().data.len(), 6); assert_eq!(walk.next().unwrap().data.len(), 6); diff --git a/diskann-quantization/src/multi_vector/distance/max_sim.rs b/diskann-quantization/src/multi_vector/distance/max_sim.rs index d9a4fb541c..ad5186cd4e 100644 --- a/diskann-quantization/src/multi_vector/distance/max_sim.rs +++ b/diskann-quantization/src/multi_vector/distance/max_sim.rs @@ -1,5 +1,7 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ //! MaxSim and Chamfer distance types for multi-vector representations. diff --git a/diskann-quantization/src/multi_vector/distance/mod.rs b/diskann-quantization/src/multi_vector/distance/mod.rs index 178e8cad08..3986234bdd 100644 --- a/diskann-quantization/src/multi_vector/distance/mod.rs +++ b/diskann-quantization/src/multi_vector/distance/mod.rs @@ -1,5 +1,7 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ //! Distance computation for multi-vector representations. //! diff --git a/diskann-quantization/src/multi_vector/distance/projected_eigen.rs b/diskann-quantization/src/multi_vector/distance/projected_eigen.rs index cd79575e8e..95889fe63f 100644 --- a/diskann-quantization/src/multi_vector/distance/projected_eigen.rs +++ b/diskann-quantization/src/multi_vector/distance/projected_eigen.rs @@ -1,5 +1,7 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ //! Projected-eigen distance type for multi-vector representations.