Repository navigation
[muilt-vetor] Add standalone MinMax8 x MinMax4 MaxSim kernel - #1494
juchen-ms (partychen) wants to merge 2 commits into
Conversation
There was a problem hiding this comment.
🔵 Needs a closer look
The new unsafe SIMD paths require cross-platform validation and final human review.
0 open findings
What changed in this PR
Adds a standalone quantized MaxSim kernel to DiskANN’s matrix kernels, leaving distance API integration for a follow-up.
Changes:
- Implements packed-A, B-tiled computation for Scalar, V3, V4, and Neon.
- Adds compensated inner products and NaN-aware maximum reduction.
- Adds supporting SIMD conversions and regression coverage.
| File | Description |
|---|---|
| diskann-wide/src/reference.rs | Adds reference u32-to-f32 conversion. |
| diskann-wide/src/emulated.rs | Adds and tests 16-lane i32-to-f32 conversion. |
| diskann-wide/src/arch/x86_64/v4/conversion.rs | Adds and tests AVX-512 conversion. |
| diskann-wide/src/arch/x86_64/v3/conversion.rs | Adds and tests split-vector conversion. |
| diskann-wide/src/arch/aarch64/u32x4_.rs | Adds and tests Neon conversion and reinterpretation. |
| diskann-quantization/src/matrix_kernels/maxsim/packed_minmax8_x_unpacked_minmax4.rs | Implements the kernel, backends, and regression tests. |
| diskann-quantization/src/matrix_kernels/maxsim/mod.rs | Registers the kernel pending integration. |
🧠 Review effort: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #1494 +/- ##
==========================================
+ Coverage 90.94% 92.13% +1.18%
==========================================
Files 586 587 +1
Lines 118391 121206 +2815
==========================================
+ Hits 107675 111677 +4002
+ Misses 10716 9529 -1187
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
Support all 1-8 bit A/B combinations through shared packing, tiled execution, compensation and maximum-inner-product reduction. Retain the tuned byte backends alongside full-range i16 paths, keep quantization widths out of the numeric microkernels, and add the required SIMD conversions and correctness coverage. Leave factory, fallback and public distance integration to follow-up PRs. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 74814a43-c576-493a-8783-585ef48c3f1c
f074880 to
e465898
Compare
juchen-ms (partychen)
changed the title
Mark Hildebrand (hildebrandmw)
left a comment
Mark Hildebrand (hildebrandmw)
left a comment
There was a problem hiding this comment.
Thanks Junkui. This is a simplification over the earlier PR! There still is some coupling between the lower-level kernels and some of the higher data preparation stuff.
A few other things that stand out to me:
- It makes heavy use of post-monomorphization errors via
const { assert!(...) }. These aren't necessarily wrong, but they are kind of hard to debug when they trigger, don't compose with higher level code well, and to me often indicate that there is another way of structuring the code so they aren't needed. I'm not saying get rid of them, but that over reliance could indicate a missing abstraction. - Why are 16-wide accumulators used for all architectures? This seems like it could be problematic for Neon or Scalar. Have you checked that the generated code isn't spilling too much and is reasonable?
I know a few other folks are planning on reviewing this as well - let's see what they have to say.
| for (i, out) in dst.iter_mut().enumerate() { | ||
| *out = W::convert_b(codes.get(i).expect("input and output lengths match") as u8); | ||
| } | ||
| } |
There was a problem hiding this comment.
I played around with the assembly - this does result in better code generation than pulling items directly out of the BitSlice, and that is unfortunate.
However, I do not like that this reinvents unpacking. When we find things like this, I would really prefer we create an issue to fix problems at the root (in this case, improving extraction from bitslice with iterators or something) rather than patching around it locally and forgetting. The former benefits the entire code base.
As much as it pains me, please don't special case certain bit-widths and let's see if we can improve BitSlice instead.
There was a problem hiding this comment.
Moved bulk decoding into BitSlice::decode_into and removed the local bit-width special cases. All widths now use the same decoding algorithm.
| Unsigned: Representation<A_BITS>, | ||
| { | ||
| codes: BlockTransposed<T, MR, PACK>, | ||
| meta: Coefficients<Vec<f32>>, |
There was a problem hiding this comment.
This is a bit of a digression from the f32, f16 and i8 cases where the query layout is owned externally. This does make it a bit harder to potentially reuse, especially in an FFI context. But if this remains pub(crate), it's probably fine for now.
There was a problem hiding this comment.
Keeping PackedA as pub(crate) for now. I agree that exposing this publicly or through FFI would warrant revisiting the ownership boundary.
| } | ||
| )+}; | ||
| } | ||
| remainder!(1, 2, 3, 4, 5, 6, 7); |
There was a problem hiding this comment.
This seems aggressive even when NR == 2. Does dead code elimination eliminate the extra remainders?
There was a problem hiding this comment.
Added explicit $n < NR guards. These are compile-time constants, so for NR == 2, only the one-column remainder is reachable.
| const { | ||
| assert!(W::Float::LANES == MR); | ||
| assert!(W::Wide::LANES == MR * PACK); | ||
| } |
There was a problem hiding this comment.
This const-asserts something that the other micro-kernels ensure by construction. This is in part because the heavy use of constraints on associated types on ExtraWide.
Not necessarily a blocker, but it does make me wonder if there is a better decomposition.
There was a problem hiding this comment.
Refactored ExtraWide around backend operations. Concrete array conversions now establish load/store sizes, replacing the generic lane assertions.
| } | ||
|
|
||
| // SAFETY: Float has exactly MR lanes, and C is an MR-element array. | ||
| let mut best = unsafe { W::Float::load_simd(self.arch, self.c.as_ptr()) }; |
There was a problem hiding this comment.
I did some cursory inspection of the V3 generated code - this seems to spill pretty heavily to the stack. Is that intentional?
There was a problem hiding this comment.
We tried smaller tiles, delayed coefficient loads, and different compensation schedules. Some reduced spilling, but the performance gains were modest or inconsistent and came with additional complexity.
Among the variants we tested, the current implementation offers the better balance of performance and simplicity, though compensation still spills. Do you have any suggestions for a different approach worth exploring?
| /// | ||
| /// `valid <= PACK` and the first `valid` elements of B must be readable. | ||
| #[inline(always)] | ||
| unsafe fn group<T: Copy + Default, const PACK: usize>(b: Slice<'_, T>, valid: usize) -> [T; PACK] { |
There was a problem hiding this comment.
Since we're preparing B anyways, couldn't we add a little bit of zero padding and remove the need for this?
In the main body of the microkernel, all of this is constant propagated and inlined, but we could turn it into a non-issue entirely.
There was a problem hiding this comment.
B rows are now zero-padded to PACK during preparation. The kernel uses one loop over all padded packs; compensation still uses the logical K.
| const MR: usize, | ||
| const NR: usize, | ||
| const PACK: usize, | ||
| >( |
There was a problem hiding this comment.
What is this testing? PreparedA doesn't have a lifetime and is treated as const, so why would reuse not work?
There was a problem hiding this comment.
Agreed—removed the redundant reuse test. Kept the tests that verify packed codes, metadata, and padding.
|
|
||
| #[cfg(target_arch = "x86_64")] | ||
| #[test] | ||
| fn avx2_supports_all_bit_pairs_and_byte_layouts() { |
There was a problem hiding this comment.
This test and corresponding scalar tests are now the longest running tests when running code coverage. Is this really necessary? Can't more directed tests be used?
There was a problem hiding this comment.
Replaced the Cartesian sweep with directed tests for packing, supported B widths, and boundaries. Local coverage-instrumented runtime dropped from 10.31s to 0.95s, with no loss of covered production lines at that cleanup stage.
| } | ||
| } | ||
|
|
||
| fn assert_nan_candidates_preserve_maxima< |
There was a problem hiding this comment.
This seems like quite an expensive way to check that SIMDMinMax is used properly in the final data preparation.
There was a problem hiding this comment.
Replaced the full-driver fixtures with small, direct microkernel tests covering NaN candidates, mixed lanes, and preservation of existing maxima.
75183b4 to
e360602
Compare
Move bulk decoding into BitSlice and zero-pad prepared B rows for a unified pack loop. Decouple backend representations behind fixed-size operations, retain smaller row-panel choices, and replace redundant test sweeps with directed coverage. Make B preparation independent of backend and row-panel types by accepting a code converter. Keep backend capability validation at the Driver construction boundary. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
e360602 to
2538cc7
Compare
Thanks for the feedback. I’ve moved SIMD-specific operations behind I’ve also added MR=8 alternatives for Neon, V3, and Scalar’s i16 path rather than requiring 16-wide accumulators everywhere. Earlier V3 assembly inspection showed register-resident accumulators in the full contraction loop, but stack traffic around compensation and tails; I’m not claiming all configurations are spill-free.
|
No new dependencies. Public additions are backward-compatible: dense
BitSlicebulk decoding and supporting SIMD trait implementations. Existing distance APIs and canonical MinMax storage remain unchanged.Reference Issues/PRs
Reworks the matrix-kernel portion of #1394 into a standalone implementation. Uses the SIMD reinterpretation and
maddubssupport from #1445 and #1479.What does this implement/fix? Briefly explain your changes.
Adds an internal MaxSim matrix kernel for MinMax A/B inputs with 1-8 bits per code, supporting all 64 bit-width combinations through the available numeric layouts.
Execution
PackedAowns the prepacked query codes and compensation coefficients and can be reused across B comparisons.Drive -> PanelKernel -> MicroKernelstructure decodes each B tile once and reuses it across A panels. Compensation precedes max reduction; no full pairwise score matrix is materialized.BitSlice::decode_into, rather than MinMax-local 4/8-bit decoding branches. Output length is validated before writing or invoking the conversion callback.validargument or separate partial-K contraction.ExtraWideexposes backend operations for loads, accumulation, score storage, and compensated reduction. Concrete array conversions replace the generic microkernel's lane-count const assertions. Construction-time checks for legal NR/backend bit widths remain.Available layouts
All layouts support A_BITS=1-8. MR counts A rows per microkernel; NR is a separate parameter in 1-8. These are supported layouts, not a universal optimal configuration or automatic selection policy.
u8/i8i16/i16u8/i8i16/i16u8/i8i16/i16u8/u8The byte limits prevent signed-byte misinterpretation and, on V3, saturation of the intermediate i16 sum. Full-range eight-bit B uses the i16 path on Scalar/x86 and unsigned byte dots on Neon.
Behavior
Each drive overwrites its output with unnegated maximum inner products. NaNs are ignored; empty B or all-NaN candidates leave negative infinity. The internal driver requires nonempty A, positive K, and matching dimensions. K is bounded by
i32::MAX / ((2^A_BITS - 1) * (2^B_BITS - 1)); conversion to f32 may still round.Tests separately cover decoding, packing, arithmetic limits, supported bit widths, partial panels, tile reuse, smaller MR layouts, compensation, and NaN reduction. Repeated whole-kernel combinations have been reduced in favor of focused cases.
Any other comments?
Scope: Factory/ISA selection policy, public prepared-query distance integration, FFI/externally owned query views, and benchmark integration remain follow-up work.
PackedAstayspub(crate).Performance: MR=16 is retained alongside narrower layouts, not claimed optimal for every architecture or query size. Earlier exploratory timings and assembly observations are not validation of this head; current-head native x86 performance and spill analysis still need confirmation.
Validation at
2538cc7-- Windows ARM64 host:cargo fmt --all --checkcargo test -p diskann-quantization --release --lib -- bits::slice::tests matrix_kernels::maxsim::packed_minmax_x_unpacked_minmax --test-threads=1cargo clippy -p diskann-quantization -p diskann-wide --all-targets -- -D warningscargo clippy -p diskann-quantization -p diskann-wide --all-targets --target x86_64-pc-windows-msvc -- -D warningsNative V3/V4 execution, SDE/Miri, and coverage were not rerun as part of these checks. Cross-compilation is not runtime validation.