diff --git a/diskann-quantization/src/multi_vector/distance/factory.rs b/diskann-quantization/src/multi_vector/distance/factory.rs index 5dcd4b8cd1..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`]. @@ -16,8 +18,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 +36,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 +60,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 +75,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 +99,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/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 a535c68dcf..ec89a834ca 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/f16.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/f16.rs @@ -1,52 +1,214 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. +/* + * Copyright (c) Microsoft Corporation. + * 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 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")] +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::{BlockTransposedTile, Cursor, RowMajorTile, contraction, tile_stride}; +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: NonZeroUsize, +} + +impl<'a, Arch: Architecture> Widen<'a, Arch> +where + SliceCast: for<'x> Target2, +{ + 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 { + 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 an f16 [`BlockTransposedRef`]. +/// +/// Widening is element-wise, so it preserves the block-transposed permutation. +struct BlockTransposedWiden<'a, Arch, const AR: usize>(Widen<'a, Arch>); + +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: 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 BlockTransposedWiden<'_, Arch, AR> { + type Tile = BlockTransposedTile<'t, f32, AR>; +} + +impl TileWalk for BlockTransposedWiden<'_, Arch, AR> +where + SliceCast: for<'x> Target2, +{ + fn next(&mut self) -> Option> { + let k = self.0.k; + self.0.next().map(|data| BlockTransposedTile::new(data, k)) + } + + fn reset(&mut self) { + self.0.cursor.reset(); + } +} -impl +/// Widens an f16 [`Standard`] matrix. +struct RowMajorWiden<'a, Arch, const BR: usize>(Widen<'a, Arch>); + +impl<'a, Arch: Architecture, const BR: usize> RowMajorWiden<'a, Arch, BR> +where + SliceCast: for<'x> Target2, +{ + 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 RowMajorWiden<'_, Arch, BR> { + type Tile = RowMajorTile<'t, f32, BR>; +} + +impl TileWalk for RowMajorWiden<'_, Arch, BR> +where + SliceCast: for<'x> Target2, +{ + fn next(&mut self) -> Option> { + let k = self.0.k; + self.0.next().map(|data| RowMajorTile::new(data, k)) + } + + fn reset(&mut self) { + self.0.cursor.reset(); + } +} + +/////////// +// Entry // +/////////// + +/// 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; + +#[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, + docs.num_vectors(), + query.padded_ncols(), + TileBudget::default(), + state, + |plan: Plan| { + ( + BlockTransposedWiden::new(arch, query, plan.a_panels), + RowMajorWiden::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, + docs.num_vectors(), + query.padded_ncols(), + TileBudget::default(), + state, + |plan: Plan| { + ( + BlockTransposedWiden::new(arch, query, plan.a_panels), + RowMajorWiden::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..d552ffe461 --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/float.rs @@ -0,0 +1,494 @@ +/* + * 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; +#[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::{ + 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}; + +/// Selects the f32 leaf for whichever architecture is in play. +pub(super) struct Kernel; + +/// Reduces a strip into the running per-A-row maxima, one entry per A row. +/// +/// 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 // +//////// + +#[cfg(target_arch = "x86_64")] +impl<'a, 'b, 's> + Accumulate< + V3, + BlockTransposedPanel<'a, f32, V3_A>, + RowMajorPanel<'b, f32, V3_B, Static>, + Slot<'s, f32, V3_A, V3_B>, + > for Kernel +{ + #[inline(always)] + fn accumulate( + &self, + arch: V3, + 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); + } +} + +#[cfg(target_arch = "x86_64")] +impl<'a, 'b, 's> + Accumulate< + V3, + BlockTransposedPanel<'a, f32, V3_A>, + RowMajorPanel<'b, f32, V3_B, Dynamic>, + Slot<'s, f32, V3_A, V3_B>, + > for Kernel +{ + #[inline(always)] + fn accumulate( + &self, + arch: V3, + 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. + 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 and its output cut never clamps. B is not, so the + // live width does. + let live = (b_panels.end * V3_B).min(self.nd) - b_panels.start * V3_B; + let rows = &mut self.out.as_chunks_mut::().0[a_panel]; + v3_leaf::max_into_rows(arch, scratch.columns(live), rows); + } +} + +////////////// +// Emulated // +////////////// + +impl<'a, 'b, 's> + Accumulate< + Scalar, + BlockTransposedPanel<'a, f32, SC_A>, + RowMajorPanel<'b, f32, SC_B, Static>, + Slot<'s, f32, SC_A, SC_B>, + > for Kernel +{ + #[inline(always)] + fn accumulate( + &self, + arch: Scalar, + 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); + } +} + +impl<'a, 'b, 's> + Accumulate< + Scalar, + BlockTransposedPanel<'a, f32, SC_A>, + RowMajorPanel<'b, f32, SC_B, Dynamic>, + Slot<'s, f32, SC_A, SC_B>, + > for Kernel +{ + #[inline(always)] + fn accumulate( + &self, + arch: Scalar, + a: BlockTransposedPanel<'a, f32, SC_A>, + b: RowMajorPanel<'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 live = (b_panels.end * SC_B).min(self.nd) - b_panels.start * SC_B; + let rows = &mut self.out.as_chunks_mut::().0[a_panel]; + scalar_leaf::max_into_rows(arch, scratch.columns(live), rows); + } +} + +/////////// +// Entry // +/////////// + +/// Plan, allocate the strip, and drive. +/// +/// `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 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-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 +/// instead of failing. +/// +/// # Panics +/// +/// 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, + k: usize, + budget: TileBudget, + state: &mut [f32], + walks: impl FnOnce(Plan) -> (AW, BW), +) where + Arch: Copy, + 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, + BlockTransposedPanel<'a, f32, AR>, + RowMajorPanel<'b, f32, BR, Static
>, + Slot<'s, f32, AR, BR>, + > + for<'a, 'b, 's> Accumulate< + Arch, + BlockTransposedPanel<'a, f32, AR>, + RowMajorPanel<'b, f32, BR, Dynamic>, + Slot<'s, f32, AR, BR>, + >, + for<'o, 'x> RawMax<'o>: Drain>, +{ + // The identity for max. + state.fill(f32::MIN); + + 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. 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()]; + + drive( + arch, + a_walk, + b_walk, + &mut Strip::new(&mut buf), + &Kernel, + &mut RawMax { out: state, nd }, + ); +} + +/// 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; + +#[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, + docs.num_vectors(), + query.padded_ncols(), + TileBudget::default(), + state, + |plan: Plan| { + ( + BlockTransposedWalk::new(query, plan.a_panels), + RowMajorWalk::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, + docs.num_vectors(), + query.padded_ncols(), + TileBudget::default(), + state, + |plan: Plan| { + ( + BlockTransposedWalk::new(query, plan.a_panels), + RowMajorWalk::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, + BlockTransposedPanel<'a, f32, AR>, + RowMajorPanel<'b, f32, BR, Static
>, + Slot<'s, f32, AR, BR>, + > + for<'a, 'b, 's> Accumulate< + Arch, + BlockTransposedPanel<'a, f32, AR>, + RowMajorPanel<'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, nd, k, budget, &mut state, |plan: Plan| { + ( + BlockTransposedWalk::new(bt.as_view(), plan.a_panels), + RowMajorWalk::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(), + 3, + 0, + TileBudget::default(), + &mut state, + |_: 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 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(), + 0, + 0, + TileBudget::default(), + &mut state, + |_: 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/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..6c73b8ef70 --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/leaves/mod.rs @@ -0,0 +1,23 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! Per-ISA micro-kernels. +//! +//! 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 +//! `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 `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 new file mode 100644 index 0000000000..e17d9441d6 --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/leaves/scalar.rs @@ -0,0 +1,123 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! 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 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; +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::{BlockTransposedPanel, RowMajorPanel}; + +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 `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: 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) } + + let k = a.k(); + let a_data = a.as_slice(); + let b_data = b.as_slice(); + assert_eq!(b_data.len(), UNROLL * k, "B panel extent"); + + 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.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()) }; + } + } +} + +/// 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 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); + for group in groups.by_ref() { + max_into_chains(arch, &mut chains, group); + } + // Fewer than WAYS columns remain, so each still lands on its own chain. + max_into_chains(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); + } + } + + 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 columns, one per chain. +#[inline(always)] +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 new file mode 100644 index 0000000000..3897deb125 --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/leaves/v3.rs @@ -0,0 +1,118 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! V3 (AVX2+FMA) f32 leaf, panels of 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::{BlockTransposedPanel, RowMajorPanel}; + +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 `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: 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) } + + let k = a.k(); + let a_data = a.as_slice(); + let b_data = b.as_slice(); + assert_eq!(b_data.len(), UNROLL * k, "B panel extent"); + + 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.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()) }; + } + } +} + +/// 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 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); + for group in groups.by_ref() { + max_into_chains(arch, &mut chains, group); + } + // Fewer than WAYS columns remain, so each still lands on its own chain. + max_into_chains(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); + } + } + + 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 columns, one per chain. +#[inline(always)] +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 55108698da..703bae017b 100644 --- a/diskann-quantization/src/multi_vector/distance/kernels/mod.rs +++ b/diskann-quantization/src/multi_vector/distance/kernels/mod.rs @@ -1,34 +1,45 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT license. +/* + * Copyright (c) Microsoft Corporation. + * 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. 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, not 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). +//! 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. -pub(super) mod f16; -pub(super) mod f32; -mod layouts; -mod reduce; -mod tiled_reduce; +use core::num::NonZeroUsize; +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 +59,315 @@ 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, never more +/// than the B-rows on hand. +#[derive(Debug, Clone, Copy)] +struct Plan { + 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 + /// 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 [`NonZeroUsize`] because they reach a divisor. The entries reject an + /// empty contraction before planning. + fn new( + a_row_bytes: NonZeroUsize, + b_row_bytes: NonZeroUsize, + b_rows: usize, + acc_bytes: usize, + budget: TileBudget, + ) -> Self { + // 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.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. + let b_panels = at_least_one_panel(cache_fit.min(b_rows.div_ceil(BR))); + + 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.get() * BR + } +} + +/////////////// +// Read side // +/////////////// -/// SIMD micro-kernel for the [`tiled_reduce`](tiled_reduce::tiled_reduce) loop. +/// 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. 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>; + fn reset(&mut self); +} + +/// An iterator whose short trailing element has its own type. /// -/// # Safety +/// `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; +} + +/// 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`], using the same implied-bound trick as [`TileAt`]. +trait SlotsAt<'s, B = &'s mut Self> { + /// Named here instead of being reached through the iterator, so [`drive`]'s bounds can + /// project it off the scratch. + type Slot; + /// Plain, not 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, 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. +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: narrower than the walk's B-tile") +} + +/// One A-panel against a whole B-tile. +/// +/// 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 + 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::*; + + 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. + // 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( + nz(64), + nz(64), + usize::MAX, + 4, + TileBudget { + l2_a: 40960, + l1_b: 36000, + }, + ); + assert_eq!((plan.a_panels.get(), plan.b_panels.get()), (40, 68)); + 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(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( + 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/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..93d89ad359 --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/strip.rs @@ -0,0 +1,131 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! The accumulator a fill writes and a drain reads. +//! +//! 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::mem; + +use super::{Scratch, SlotsAt}; + +/// 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; AR]], +} + +impl<'a, T, const AR: usize, const BR: usize> Strip<'a, T, AR, BR> { + /// Trailing elements beyond the last whole column are never touched. + pub(super) fn new(buf: &'a mut [T]) -> Self { + Self { + buf: buf.as_chunks_mut::().0, + } + } + + /// 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; AR]] { + &self.buf[..live] + } +} + +/// 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; AR]; BR], +} + +impl Slot<'_, T, AR, BR> { + pub(super) fn columns(&mut self) -> &mut [[T; AR]; BR] { + &mut *self.buf + } +} + +/// Cuts a [`Strip`] into disjoint [`Slot`]s, stopping short of a trailing partial tile. +/// +/// 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 { + // 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 }) + } +} + +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 { + rest: &mut *self.buf, + } + } +} + +#[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.columns().as_flattened_mut().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; 4]]); + } +} 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..d148e48146 --- /dev/null +++ b/diskann-quantization/src/multi_vector/distance/kernels/tiles.rs @@ -0,0 +1,392 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +//! The read ladder: walks lend tiles, tiles cut into panels. +//! +//! 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::num::NonZeroUsize; + +use super::{NoTail, Paneled, TailIterator, TileAt, TileWalk}; +use crate::bits::{Dynamic, Length, Static}; +use crate::multi_vector::{BlockTransposedRef, MatRef, Standard}; + +/////////////////////////// +// Block-transposed side // +/////////////////////////// + +/// A single block of a [`BlockTransposedRef`]: `AR` rows × `k` columns, column-major +/// within the block. +#[derive(Clone, Copy)] +pub(super) struct BlockTransposedPanel<'a, T, const AR: usize>(&'a [T]); + +impl<'a, T, const AR: usize> BlockTransposedPanel<'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 + } +} + +/// 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 BlockTransposedPanels<'a, T, AR> { + type Item = BlockTransposedPanel<'a, T, AR>; + + fn next(&mut self) -> Option { + let (panel, rest) = self.rest.split_at_checked(self.stride)?; + self.rest = rest; + Some(BlockTransposedPanel(panel)) + } +} + +impl TailIterator for BlockTransposedPanels<'_, T, AR> { + type Tail = NoTail; + + fn tail(self) -> Option { + None + } +} + +/// 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: NonZeroUsize, +} + +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: NonZeroUsize) -> Self { + debug_assert!(data.len().is_multiple_of(AR * k.get())); + Self { data, k } + } +} + +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 = BlockTransposedPanels<'a, T, AR>; + + fn panels(&self) -> Self::Panels { + BlockTransposedPanels { + rest: self.data, + stride: AR * self.k.get(), + } + } +} + +//////////////////// +// Row-major side // +//////////////////// + +/// 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 RowMajorPanel<'a, T, const BR: usize, L> { + data: &'a [T], + rows: L, +} + +impl<'a, T, const BR: usize, L: Length> RowMajorPanel<'a, T, BR, L> { + pub(super) fn as_slice(&self) -> &'a [T] { + self.data + } + + pub(super) fn rows(&self) -> usize { + self.rows.value() + } +} + +/// 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 RowMajorPanels<'a, T, BR> { + type Item = RowMajorPanel<'a, T, BR, Static
>; + + fn next(&mut self) -> Option { + let (data, rest) = self.rest.split_at_checked(BR * self.k.get())?; + self.rest = rest; + Some(RowMajorPanel { + data, + rows: Static::
, + }) + } +} + +impl<'a, T, const BR: usize> TailIterator for RowMajorPanels<'a, T, BR> { + type Tail = RowMajorPanel<'a, T, BR, Dynamic>; + + fn tail(self) -> Option { + 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()), + }) + } +} + +/// 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: NonZeroUsize, +} + +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: NonZeroUsize) -> Self { + debug_assert!(data.len().is_multiple_of(k.get())); + Self { data, k } + } +} + +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 { + RowMajorPanels { + rest: self.data, + k: self.k, + } + } +} + +/////////// +// 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. +pub(super) struct Cursor<'a, T> { + data: &'a [T], + stride: NonZeroUsize, + cur: usize, +} + +impl<'a, T> Cursor<'a, T> { + pub(super) fn new(data: &'a [T], stride: NonZeroUsize) -> Self { + Self { + data, + stride, + cur: 0, + } + } + + /// 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.get().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, which is the size a converting walk must stage. + pub(super) fn widest(&self) -> usize { + self.stride.get().min(self.data.len()) + } +} + +/// 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: NonZeroUsize, +} + +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: NonZeroUsize) -> Self { + let k = contraction(view.padded_ncols()); + Self { + cursor: Cursor::new(view.as_slice(), tile_stride(a_panels, AR, k)), + k, + } + } +} + +impl<'t, T: Copy, const AR: usize> TileAt<'t> for BlockTransposedWalk<'_, T, AR> { + type Tile = BlockTransposedTile<'t, T, AR>; +} + +impl TileWalk for BlockTransposedWalk<'_, T, AR> { + fn next(&mut self) -> Option> { + let k = self.k; + self.cursor + .next() + .map(|data| BlockTransposedTile::new(data, k)) + } + + fn reset(&mut self) { + self.cursor.reset(); + } +} + +/// 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: NonZeroUsize, +} + +impl<'a, T, const BR: usize> RowMajorWalk<'a, T, BR> { + /// # Panics + /// + /// 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(mat.as_slice(), tile_stride(b_panels, BR, k)), + k, + } + } +} + +impl<'t, T: Copy, const BR: usize> TileAt<'t> for RowMajorWalk<'_, T, BR> { + type Tile = RowMajorTile<'t, T, BR>; +} + +impl TileWalk for RowMajorWalk<'_, T, BR> { + fn next(&mut self) -> Option> { + let k = self.k; + self.cursor.next().map(|data| RowMajorTile::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 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::>>(), + 2 * size_of::() + ); + assert_eq!( + size_of::>(), + 3 * size_of::() + ); + } + + #[test] + 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 = RowMajorTile::::new(&data, NonZeroUsize::new(3).unwrap()); + + let mut panels = tile.panels(); + 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 = RowMajorTile::::new(&data, NonZeroUsize::new(3).unwrap()).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 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); + assert_eq!(walk.next().unwrap().data.len(), 3); + assert!(walk.next().is_none()); + + walk.reset(); + 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.