[multi vector] Distance Kernels - #1368
Conversation
|
Hi Mark, I completed a broader performance investigation of PR #1368 at I tested the Scalar and V3 paths locally on an Intel Core Ultra 7 268V. Since that machine does not support AVX-512, I ran the V4 investigation separately on an Intel Xeon Platinum 8370C that passes DiskANN’s complete Both sets of benchmarks used the same nine shapes, three passes per configuration, and 50 measurements per pass. The primary metric was minimum ns/IP across all 150 measurements. Negative percentages mean faster execution. Scalar and V3 resultsThe V3 results look encouraging:
The PR V3 implementation was faster than the baseline on all nine shapes for both element types. However, the
Scalar investigationI traced the Scalar regression to core::array::from_fn(|i| a[i].mul_add_simd(b, acc[i]))The emulated Scalar implementation invokes core::array::from_fn(|i| a[i] * b + acc[i])I tested only this one-line change while leaving the rest of #1368 unchanged. Scalar performance returned to approximately baseline level:
All selected MaxSim tests still passed. This suggests that Scalar should probably specialize this operation to use separate multiplication and addition, while V3 and V4 continue using hardware Unmodified #1368: V4 versus V3Before making any experimental V4 changes, I benchmarked the original V4 and V3 implementations from the same #1368 commit and the same release binary on the Xeon Platinum 8370C. All V4 correctness tests executed successfully.
In other words, the unmodified V4 implementation had approximately 36× the V3 latency for f32 and 32× the V3 latency for f16. These results come directly from the unmodified PR code, before applying any experimental microkernel or changing the MR/NR geometry. The regression was highly stable rather than thermal noise: the maximum V4 pass-to-pass spread was below 0.5%. V4 code-generation investigationDisassembly showed that the generic Each broadcast, FMA, and max operation became an out-of-line call: These calls spilled ZMM values through memory and also executed I tried wrapping the outer Experimental V4 code-generation fixAs an experiment, I added a concrete V4 microkernel inside #[target_feature(enable = "avx512f")]The concrete microkernel invokes After this change, the generated hot loop contained direct ZMM loads, broadcast-FMAs, and max instructions without per-operation calls. Compared with the original unmodified V4 implementation, this experimental fix improved performance by:
These speedups are relative to the original broken V4 code-generation path, not relative to V3. All 12 MaxSim tests passed, including the V4 driver, panel, and microkernel tests. V4 MR/NR explorationAfter fixing the V4 code-generation problem, I benchmarked these tile configurations: Each configuration used the same nine shapes and 150 measurements per shape. MR=16Relative to the existing
For f32, it improved all nine shapes relative to For f16, it was also the fastest tested tile on six shapes, although one MR=32MR32 improved some small-D or low-dimensional shapes by approximately 40–50%, but it regressed sustained high-dimensional workloads by more than 2×:
MR32 therefore does not appear suitable as an unconditional default on Ice Lake. The sustained regressions may be related to register pressure and/or AVX-512 throughput and frequency behavior. It may only be useful through shape- or microarchitecture-dependent dispatch. Best fixed V4 versus V3Finally, I compared the best general V4 tile,
Therefore, the concrete microkernel removes the catastrophic code-generation regression, and Current conclusions
For Scalar, would it make sense to specialize For V4, does a dedicated target-feature leaf microkernel fit the direction you had in mind, or would you prefer to address the target-feature/inlining boundary inside I have the experimental patches, disassembly, complete per-shape tables, and raw benchmark outputs available if useful. |
|
Thanks for the write up Junkui. The issues around scalar and AVX-512 are known and will be addressed before this PR is moved out of draft state. For all architecture, a strategic call to Thanks for the exploration on MR vs NR on AVX-512. I'm a little surprised 32x6 performed poorly as this is the blocking used by many BLAS implementations. This PR isn't quite at the performance tuning phase though, and I suspect our cache sizing is incorrect. |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #1368 +/- ##
==========================================
+ Coverage 91.55% 92.60% +1.04%
==========================================
Files 521 524 +3
Lines 100371 102090 +1719
==========================================
+ Hits 91895 94540 +2645
+ Misses 8476 7550 -926
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
🟡 Changes recommended
The new factory-path kernels return 0.0 (instead of f32::MAX) for the edge case of 0-dimensional queries with empty documents, diverging from the existing/reference semantics.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
This PR is a substantial refactor of multi-vector max-sim (“maxsim”) distance computation in diskann-quantization, replacing the prior distance/kernels/* tiling + micro-kernel stack with a new matrix_kernels infrastructure that supports more architectures (AVX-512, Neon), multiple kernel shapes, and removes the need for output padding. It also adds some supporting utilities (e.g., matrix transpose) and improves error reporting for dimension mismatches.
Changes:
- Replaced the prior
multi_vector/distance/kernels/*implementation with the newmatrix_kernelsmodule and wired the maxsim factory paths to the new drivers. - Added explicit dimension-mismatch error reporting (
MaxSimError::UnequalDim) and enforced it in both SIMD and fallback implementations. - Improved AArch64 Neon min/max behavior under Miri by using an emulated path.
File summaries
| File | Description |
|---|---|
| diskann-wide/src/arch/aarch64/f32x4_.rs | Routes min/max to standard variants and uses emulation under Miri to improve testability. |
| diskann-utils/src/views.rs | Adds transpose() for MatrixBase plus unit tests. |
| diskann-quantization/src/multi_vector/distance/mod.rs | Removes the old kernels module from the distance module tree. |
| diskann-quantization/src/multi_vector/distance/max_sim.rs | Adds MaxSimError::UnequalDim for dimension mismatch reporting. |
| diskann-quantization/src/multi_vector/distance/kernels/tiled_reduce.rs | Deletes the old generic tiling loop implementation (replaced by matrix_kernels). |
| diskann-quantization/src/multi_vector/distance/kernels/reduce.rs | Deletes the old accumulator reduction helper trait (replaced by matrix_kernels::util::Fold). |
| diskann-quantization/src/multi_vector/distance/kernels/mod.rs | Deletes the old kernels module root. |
| diskann-quantization/src/multi_vector/distance/kernels/layouts.rs | Deletes the old layout + conversion traits (replaced by matrix_kernels::blocks/* + utilities). |
| diskann-quantization/src/multi_vector/distance/kernels/f32/v3.rs | Deletes the old AVX2+FMA f32 micro-kernel implementation. |
| diskann-quantization/src/multi_vector/distance/kernels/f32/scalar.rs | Deletes the old scalar/emulated f32 micro-kernel implementation. |
| diskann-quantization/src/multi_vector/distance/kernels/f32/mod.rs | Deletes the old f32 kernel family entry point. |
| diskann-quantization/src/multi_vector/distance/kernels/f16.rs | Deletes the old f16 adapter entry point. |
| diskann-quantization/src/multi_vector/distance/kernel.rs | Updates trait docs to include the new UnequalDim error condition. |
| diskann-quantization/src/multi_vector/distance/fallback.rs | Enforces dimension equality and returns UnequalDim for fallback computation. |
| diskann-quantization/src/multi_vector/distance/factory.rs | Rewires the factory to use matrix_kernels maxsim drivers and updates query preparation logic. |
| diskann-quantization/src/matrix_kernels/mod.rs | Introduces the new “GEMM-lite” kernel framework and cache model. |
| diskann-quantization/src/matrix_kernels/driver.rs | Adds semantic driver/panel/micro-kernel traits used by the new kernel stack. |
| diskann-quantization/src/matrix_kernels/bounds.rs | Adds debug-only bounds tracking to support safer low-level pointer math in kernels. |
| diskann-quantization/src/matrix_kernels/num.rs | Adds strongly typed dimensions/offset units (DimK, Elements, Bytes). |
| diskann-quantization/src/matrix_kernels/ptr.rs | Adds slice wrappers with debug-only length tracking for kernel pointer traversal. |
| diskann-quantization/src/matrix_kernels/util.rs | Adds conversion, load/store helpers, and compile-time fold utilities used by kernels. |
| diskann-quantization/src/matrix_kernels/test_util.rs | Adds panic-capture helpers and test distributions for matrix-kernel tests. |
| diskann-quantization/src/matrix_kernels/blocks/mod.rs | Adds packed/unpacked block-view module root for kernel argument views. |
| diskann-quantization/src/matrix_kernels/blocks/packed.rs | Adds packed (block-transposed) views/panels over query-like data. |
| diskann-quantization/src/matrix_kernels/blocks/unpacked.rs | Adds unpacked views/panels over row-major/column-major-like banded data. |
| diskann-quantization/src/matrix_kernels/maxsim/mod.rs | Adds maxsim kernel module root for the new framework. |
| diskann-quantization/src/matrix_kernels/maxsim/test.rs | Adds shared test-case generation for maxsim reference comparisons. |
| diskann-quantization/src/matrix_kernels/maxsim/packed_f32_x_unpacked_f32.rs | Adds the new packed-f32 × unpacked-f32 maxsim kernel stack (drivers + kernels + tests). |
| diskann-quantization/src/matrix_kernels/maxsim/packed_f32_x_unpacked_f16.rs | Adds the new packed-f32 × unpacked-f16 driver with per-tile f16→f32 conversion and tests. |
| diskann-quantization/src/lib.rs | Registers the new matrix_kernels module in the crate. |
Review details
Suppressed comments (1)
diskann-quantization/src/multi_vector/distance/factory.rs:134
- Same edge case as the f32 path: for 0-dimensional vectors, this fills
scoreswith 0.0 even whendoc.num_vectors() == 0, but the reference semantics for empty docs aref32::MAX.
let Some(k) = NonZeroUsize::new(self.prepared.ncols()).map(mk::DimK::new) else {
scores.fill(0.0);
return Ok(());
};
- Files reviewed: 30/30 changed files
- Comments generated: 3
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| let Some(k) = NonZeroUsize::new(self.prepared.ncols()).map(mk::DimK::new) else { | ||
| scores.fill(0.0); | ||
| return Ok(()); | ||
| }; |
There was a problem hiding this comment.
Saving this for a later cleanup pass.
| bounds::check_eq!(a.k(), k, "constraction dimensions to not agree"); | ||
| bounds::check_eq!(b.k(), k, "constraction dimensions to not agree"); | ||
| bounds::check_eq!( | ||
| bounds::Bound::new(a.blocks().get()), | ||
| c.len().div_ceil(MR), | ||
| "output length must occupiy exactly the packed A blocks", | ||
| ); |
| bounds::check_eq!(a.k(), k, "constraction dimensions to not agree"); | ||
| bounds::check_eq!(b.k(), k, "constraction dimensions to not agree"); | ||
| bounds::check_eq!( | ||
| bounds::Bound::new(a.blocks().get()), | ||
| c.len().div_ceil(MR), | ||
| "output length must occupiy exactly the packed A blocks", | ||
| ); |
| { | ||
| /// Load up to the first `src.len()` and return the results in an array. | ||
| /// | ||
| /// The remaining items items should be left in a default state. |
There was a problem hiding this comment.
typo: “items” is repeated twice.
| || { | ||
| // Pre-fill `c`. | ||
| self.c.fill(f32::NEG_INFINITY); | ||
|
|
||
| // We allow `c` to be slightly under-filled. | ||
| // | ||
| // These variables track if under-fill is happening. | ||
| let remainder = self.c.len() % MR; | ||
| let last_a_block = self.a.blocks().get() - 1; | ||
|
|
||
| let mut c = MutSlice::new(self.c); | ||
|
|
||
| let on_a_panels = |a_panels: packed::View<'_, f32, MR>, a_block_base| { | ||
| let on_b_panels = |b_panels: unpacked::View<'_, f16>, _| { | ||
| // Convert `f16` to `f32`. | ||
| // | ||
| // SAFETY: Class invariant - `self.b.k()` is equal to `self.k`. | ||
| let b_flat = unsafe { b_panels.as_std_slice(self.k) }; | ||
| let b_converted = &mut self.b_converted[..b_flat.len()]; | ||
| Converter::new(self.arch).convert(b_converted, b_flat); | ||
|
|
||
| // SAFETY: `b_converted` has length `b_panels.extent() * self.k`. | ||
| let b_panels_converted = unsafe { | ||
| unpacked::View::new(Slice::new(b_converted), b_panels.extent(), self.k) | ||
| }; | ||
|
|
||
| let panel_kernel = |a_panel: packed::Panel<'_, f32, MR>, a_block_offset| { | ||
| // If we are in the very last block and we need to sub-fill, do that. | ||
| // Otherwise, reference the output in place. | ||
| let a_block = a_block_base + a_block_offset; | ||
| let handling_tail = a_block == last_a_block && remainder != 0; | ||
|
|
||
| let bound = | ||
| bounds::Bound::from_fn( | ||
| || if handling_tail { remainder } else { MR }, | ||
| ); | ||
|
|
||
| // SAFETY: By class invariant, | ||
| // | ||
| // `MR * (self.a.blocks() - 1) < c.len() <= MR * self.a.blocks()`. | ||
| // | ||
| // From the visitor, `a_block <= self.a.blocks()`. | ||
| let mut region = unsafe { c.subslice(MR * a_block, bound) }; | ||
| let c = if handling_tail { | ||
| util::LoadStore::<f32, MR>::load( | ||
| self.arch, | ||
| // SAFETY: `region` as length exactly `remainder`. | ||
| unsafe { region.as_std_slice(remainder) }, | ||
| ) | ||
| } else { | ||
| // SAFETY: `region` has length exactly `MR`. | ||
| unsafe { *region.as_array::<MR>() } | ||
| }; | ||
|
|
||
| // Run the kernel | ||
| // | ||
| // SAFETY: By class invariant, `a_panel.k()` and | ||
| // `b_panels_converted.k()` are equal to `self.k`. | ||
| let mut kernel = unsafe { | ||
| PanelKernel::new(self.arch, a_panel, b_panels_converted, c, self.k) | ||
| }; | ||
|
|
||
| driver::PanelKernel::panel_kernel(&mut kernel); | ||
|
|
||
| let c_final = kernel.take(); | ||
|
|
||
| // Put back `C`. | ||
| if handling_tail { | ||
| util::LoadStore::<f32, MR>::store( | ||
| self.arch, | ||
| c_final, | ||
| // SAFETY: `region` has length exactly `remainder`. | ||
| unsafe { region.as_std_mut_slice(remainder) }, | ||
| ); | ||
| } else { | ||
| // SAFETY: `region` has length exactly `MR`. | ||
| unsafe { *region.as_array::<MR>() = c_final }; | ||
| } | ||
| }; | ||
|
|
||
| // SAFETY: By class invariant, `a_panels.k() == self.k`. | ||
| unsafe { | ||
| a_panels.visit_panels(self.k, panel_kernel); | ||
| } | ||
| }; | ||
|
|
||
| // SAFETY: By class invariant, `self.b.k() == self.k`. | ||
| unsafe { | ||
| self.b | ||
| .visit_sub_views(self.params.b_cols_in_l1, self.k, on_b_panels); | ||
| } | ||
| }; | ||
|
|
||
| // SAFETY: By class invariant, `self.a.k() == self.k`. | ||
| unsafe { | ||
| self.a | ||
| .visit_sub_views(self.params.a_panels_in_l2, self.k, on_a_panels) | ||
| }; | ||
| }, |
There was a problem hiding this comment.
The A tile loop is outermost, so each A tile walks all of B. With small queries there is only ever one A tile, so this makes no difference today.
But if the query grows past one A tile, B gets re-walked, and in the f16 case re-converted, once per A tile. Is B on the outside worth exploring there?
Not claiming it is a win. Flipping it means A gets re-read instead, you would want to flip the tile sizing too, and the converted buffer stops being small scratch.
Mainly I want to understand what factors should decide this loop order. I can see a few candidates, like the relative sizes of the two operands, which one is cheaper to re-read given f16 halves B, and how the cache budget gets split between them. But I cannot tell which of these actually dominates, so any pointers would help.
There was a problem hiding this comment.
This loop order was chosen to get us to parity with the current code. I suspect that a less-than-optimal tiling is just one of its many flaws. It is meant to work as a baseline for exploring other orderings. We are woefully lacking infrastructure to properly sweep and benchmark different orderings.
Suryansh Gupta (suri-kumkaran)
left a comment
There was a problem hiding this comment.
Small thing, a few typos that show up in panic messages users would see:
-
"constraction" should be "contraction", and "to not agree" should be "do not
agree". 8 places:f32lines:110, 111, 143, 144andf16lines:66, 67, 97, 98. -
"occupiy" should be "occupy". 4 places:
f32lines:115, 148andf16lines:71, 102. -
A few in doc comments too: "oclumns" in
matrix_kernels/mod.rs:26, "matrics" in
num.rs:117, and "the pointer offset if valid" should be "is valid" in
blocks/packed.rs:155and194andblocks/unpacked.rs:185.
| /// Shapes for the `chamfer_matches_fallback` / `max_sim_matches_fallback` | ||
| /// agreement checks: `(num_queries, num_docs, dim)`. | ||
| /// | ||
| /// Targets the factory wiring (query setup, score writeback) above the | ||
| /// kernel layer; exhaustive panel/remainder coverage is pinned in | ||
| /// `kernels::tiled_reduce::tests`. | ||
| const TEST_CASES: &[(usize, usize, usize)] = &[ | ||
| (1, 1, 4), // Degenerate | ||
| (5, 3, 5), // Prime k; nq > 1 and nd > 1 exercise per-row writeback | ||
| (17, 4, 64), // A-panel remainder crossing both Scalar and V3 panel widths | ||
| (16, 6, 32), // B-remainder ≠ 1 (V3 b_remainder = 2) | ||
| ]; |
There was a problem hiding this comment.
nit: The zero-doc and zero-query cases are covered well below, and I like that those loop over both Auto and Reference.
The one I could not find anywhere is dim == 0. This list starts at (1,1,4), and
the kernel-level dims in maxsim/test.rs start at [1,1,1,1,1], so nothing
exercises a zero contraction dimension. That is the case this commit changed
from f32::MAX to 0.0.
| } | ||
|
|
||
| /// Transpose the elements in `self`. | ||
| pub fn transpose(&self) -> Matrix<T::Elem> |
There was a problem hiding this comment.
Is this meant to be public API? Every caller I can find is inside a #[cfg(test)]
block, both the ones here and the ones in the maxsim kernel tests.
| @@ -199,9 +263,12 @@ | |||
| { | |||
| fn run(self, arch: V4, query: MatRef<'_, Standard<f32>>) -> E::Output { | |||
| // V4 dispatches to V3 (no V4-specific kernel). | |||
There was a problem hiding this comment.
stale
Infrastructure investment for multi-vector distances.
This PR replaces the implementation of
f32andf16max-sim kernels with one that takes approximately 2.8 times the amount of code! 😢Some things it does:
Some things it doesn't do:
Why
We eventually want to support quantized multi-vector operations. The hope with this PR is that it introduces enough design ideas and infrastructure to make the process easier.
The design in the PR follows several principles:
From a performance standpoint, data structures representing memory spans and 2D matrix views are kept as lean as possible. This means discarding information that is expected to be known by something higher in the stack. For example, the contraction dimension
Kin aMxK . KxNmatrix-matrix multiplication is shared by both the left-hand and right-hand arguments. Memory views thus shouldn't independently track this information because:Kexternally helps LLVM with code-generation.For better debugging, these dimensions above that get thrown away are retained and tracked closely. I found this helps tremendously when working on corner cases and helps document the assumptions about the intended memory use.
Strongly typed data structures for representing memory layouts. In release builds, these boil down to basically a pointer and one dimension. In addition, these representations allow for some sharing of infrastructure when it comes to traversing memory.
Suggested Reviewing Order
Most of this rather large diff is infrastructure. The kernels themselves are a relatively small fraction. From lower-level to higher level:
matrix_kernelsutil.rs/num.rs: Small utilities and strongly typed integers.bounds.rs: This is the main debugging aid. ABoundis an integer underdebug_assertionsand test builds, and a ZST otherwise. It is used in the contexts mentioned above where the true lengths and sizes of data structures should be known by another component and are hence not tracked in optimized build, but exist in debug builds to double-check indexing.ptr.rs: Basically a&[T]and&mut [T]but with its length aBoundrather than a trueusize. These effectively become pointers in release builds, but are full slices in debug ubilds.driver.rs: Simple semantic traits.Drive: The entry point for full GEMM operations. In this PR,Driveimplementations involve multiple invocations of a lower-levelPanelKernel.PanelKernel: A step working on a sub-portion of the full operation with arguments fixed in the L1 and L2 cache. In this PR,PanelKernelimplementations involve multiple invocations of a lower-levelMicroKernel.MicroKernel: A low level operation where two operands in memory are (hopefully) in the L1 cache. This contains the broadcasting FMAs.blocks/unpacked.rs: A view over a row-major or column-major matrix. The terminology "unpacked" means that the inner dimensionkfor each row/column is the contiguous one in memory.This has a few refinements:
Panel: A view where the number of rows/columns (not thekdimension) is a compile time constant. This is used to parameterize micro-kernels.Remainder: A view with a struct upper-bound on the number of rows/columns (not thekdimension). This is used to represent the tail portion of visitation across panels that make up a parent view.blocks/packed.rs: The equivalent of the unpacked view for transposed matrices. This representation is key the performance of a GEMM operation. Typically, packing into this representation is done incrementally, but both our current and this implementation pre-pack the query into this layout.maxsim/packed_f32_x_unpacked_f32.rs: The entry point for thef32 x f32maxsim. This follows a hierarchy ofDriver->PanelKernel->MicroKernelwhere each level tries to take advantage of memory locality.maxsim/packed_f32_x_unpacked_f16.rs: The entry point for thef32 x f16maxsim. This incrementally convers the RHS tof32before entering thePanelKernelfor thef32 x f32case.The top level drivers are currently parameterized by the blocking factors used in the micro-kernel. Blocking factors are represented as
MR x NRwhere the left-handAmatrix is blocked byMRand the right-handBmatrix is blocked byNR. Micro-kernels compute inner products forMRrows ofAwithNRcolumns ofB. There number of SIMD registers needed to hold the partial accumulators is given bySo a 16x6 blocked kernel for AVX2 requires
(16 / 8) * 6 = 12accumulators (close to the architectural limit of 16 registers). A32x6AVX-512 kernel also requires 12 accumulators since AVX-512 registers are twice as wide.multi_vectorThe factory methods are wired up to reasonable defaults for benchmarking.
Benchmarks
Some benchmark numbers are listed below. Values are speedup relative to
main.f32f16f32f16f32There is still plenty of work to be done with respect to performance tuning (better cache blocking, kernel selection, maybe prefetching etc). This PR intentionally keeps the same cache size heuristics as
main. But there is obviously still performance on the table. For example, theQ = 32, D = 32, DIM = 512case on my laptop gets 10.0 ns/IP where-as MKL gets 9.99. Since we have the benefit of pre-packingAand can avoid materializingC, we should be able to do better.