Add staged (tiled_reduce_staged) multi-vector MaxSim kernel POC

[suri-kumkaran]

Introduce an experimental, generic cache-tiled reduction driver, tiled_reduce_staged, that computes multi-vector MaxSim/Chamfer for any element type and quantization by swapping four pluggable stages (StagedKernel, Postprocess, Reducer, StagedConvert) instead of forking the tiled loop nest per datatype.

Validated with two instantiations:

• f32: bit-identical to and on par with the hand-fused V3 kernel (selectable for A/B as MaxSimIsa::X86_64_V3_Staged).
• 4-bit MinMax-quantized i8: a new datatype added by swapping Stage A + the Acc type + Stage B only; 1.5-4.1x over the per-pair SIMD reference.

The driver owns all scratch, allocated from a caller-supplied ScopedAllocator; zero-allocation steady state is provided by a single-owner resettable bump arena (ResettableArena). A crisp design overview lives in the staged module README (diskann-quantization/.../kernels/staged/README.md).

Also adds benchmark examples (multi-vector-{staged,quant,3way}.json) and the quantized multi-vector benchmark backend.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 92.27226% with 151 lines in your changes missing coverage. Please review.
✅ Project coverage is 90.20%. Comparing base (745718c) to head (817c108).
⚠️ Report is 31 commits behind head on main.

Files with missing lines	Patch %	Lines
...ion/src/multi_vector/distance/kernels/staged/i8.rs	92.35%	35 Missing ⚠️
.../src/multi_vector/distance/kernels/tiler/tilers.rs	88.03%	25 Missing ⚠️
.../src/multi_vector/distance/kernels/tiler/minmax.rs	91.30%	22 Missing ⚠️
...ion/src/multi_vector/distance/kernels/tiler/f16.rs	89.41%	18 Missing ⚠️
...ion/src/multi_vector/distance/kernels/staged/v3.rs	91.97%	15 Missing ⚠️
...-quantization/src/multi_vector/distance/factory.rs	84.72%	11 Missing ⚠️
.../src/multi_vector/distance/kernels/staged/arena.rs	85.10%	7 Missing ⚠️
...n/src/multi_vector/distance/kernels/tiler/arena.rs	85.10%	7 Missing ⚠️
diskann-benchmark/src/inputs/multi_vector.rs	0.00%	2 Missing ⚠️
...kann-quantization/src/multi_vector/distance/isa.rs	0.00%	2 Missing ⚠️
... and 4 more

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1200      +/-   ##
==========================================
- Coverage   90.65%   90.20%   -0.46%     
==========================================
  Files         503      522      +19     
  Lines       95707    99208    +3501     
==========================================
+ Hits        86765    89492    +2727     
- Misses       8942     9716     +774     

Flag	Coverage Δ
miri	90.20% <92.27%> (-0.46%)	⬇️
unittests	89.88% <92.27%> (-0.74%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
diskann-benchmark/src/multi_vector/mod.rs	83.33% <ø> (-16.67%)	⬇️
...n-quantization/src/multi_vector/distance/kernel.rs	100.00% <ø> (ø)
...ntization/src/multi_vector/distance/kernels/mod.rs	100.00% <ø> (ø)
.../src/multi_vector/distance/kernels/tiler/leaves.rs	100.00% <100.00%> (ø)
...src/multi_vector/distance/kernels/staged/maxsim.rs	97.87% <97.87%> (ø)
diskann-benchmark/src/inputs/multi_vector.rs	19.04% <0.00%> (-0.63%)	⬇️
...kann-quantization/src/multi_vector/distance/isa.rs	0.00% <0.00%> (ø)
...src/multi_vector/distance/kernels/staged/driver.rs	98.41% <98.41%> (ø)
...on/src/multi_vector/distance/kernels/staged/mod.rs	96.92% <96.92%> (ø)
...ion/src/multi_vector/distance/kernels/tiler/mod.rs	98.58% <98.58%> (ø)
... and 8 more

... and 145 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Copilot]

Pull request overview

This PR introduces an experimental “staged” cache-tiled reduction driver for multi-vector MaxSim/Chamfer that composes computation from pluggable stages (kernel → postprocess → reducer), aiming to avoid forking the tiled loop nest per datatype/quantization. It also adds a separate coarse “tiler” pipeline experiment and wires up new x86_64 (V3) benchmark ops + example configs to A/B these kernels.

Changes:

• Add staged kernel framework (tiled_reduce_staged) with V3 f32 parity path and a quantized 4-bit MinMax i8 POC path, including resettable arena scratch ownership for zero-allocation steady-state.
• Add a coarse tiler-based pipeline (including an f16 per-tile widen path) as a parallel experiment and expose standalone POC entry points on x86_64.
• Extend benchmarks/inputs/examples to exercise the staged ISA variant and new quantized/f16 A/B benchmark jobs.

Reviewed changes

Copilot reviewed 26 out of 26 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
diskann-quantization/src/multi_vector/distance/mod.rs	Re-export standalone x86_64 POC entry points for staged/tiler kernels.
diskann-quantization/src/multi_vector/distance/kernel.rs	Relax MaxSimKernel bounds (drops Sync) to allow interior-mutable per-call scratch.
diskann-quantization/src/multi_vector/distance/isa.rs	Add MaxSimIsa::X86_64_V3_Staged variant + display/support checks.
diskann-quantization/src/multi_vector/distance/factory.rs	Wire staged f32 kernel into build_max_sim and add reuse-regression test.
diskann-quantization/src/multi_vector/distance/kernels/mod.rs	Gate and expose staged and tiler modules + re-export POC types on x86_64.
diskann-quantization/src/multi_vector/distance/kernels/staged/README.md	Design write-up for staged kernel and its intended extensibility.
diskann-quantization/src/multi_vector/distance/kernels/staged/mod.rs	Define staged traits (kernel/postprocess/reducer/convert) and staging plan.
diskann-quantization/src/multi_vector/distance/kernels/staged/driver.rs	Implement staged tiled reduction driver allocating all scratch from ScopedAllocator.
diskann-quantization/src/multi_vector/distance/kernels/staged/arena.rs	Add resettable bump arena for reuse without heap allocation.
diskann-quantization/src/multi_vector/distance/kernels/staged/maxsim.rs	Identity postprocess + staged scratch container for f32 path.
diskann-quantization/src/multi_vector/distance/kernels/staged/v3.rs	V3 staged f32 kernel + reducer + bit-identical parity test vs fused kernel.
diskann-quantization/src/multi_vector/distance/kernels/staged/i8.rs	V3 quantized staged MinMax kernel (integer dot + metadata postprocess) + tests.
diskann-quantization/src/multi_vector/distance/kernels/tiler/mod.rs	Coarse tiler pipeline driver + generic traits for typed panel selection.
diskann-quantization/src/multi_vector/distance/kernels/tiler/arena.rs	Resettable arena copy for tiler experiment.
diskann-quantization/src/multi_vector/distance/kernels/tiler/tilers.rs	Implement identity and f16→f32 converting tile walks with lending semantics.
diskann-quantization/src/multi_vector/distance/kernels/tiler/leaves.rs	SIMD leaf kernels copied from staged to keep inner math identical.
diskann-quantization/src/multi_vector/distance/kernels/tiler/minmax.rs	4-bit MinMax instantiation of the coarse tiler driver + tests.
diskann-quantization/src/multi_vector/distance/kernels/tiler/f16.rs	f16 tiler instantiation (per-tile widen + f32 store kernel) + tests.
diskann-benchmark/src/multi_vector/mod.rs	Register new x86_64-only multi-vector benchmark ops.
diskann-benchmark/src/multi_vector/quant.rs	Add staged-vs-tiled-vs-reference quantized benchmark op.
diskann-benchmark/src/multi_vector/tiled_f16.rs	Add tiler-vs-preprocess benchmark op for f16.
diskann-benchmark/src/inputs/multi_vector.rs	Add new benchmark input types and ISA option for staged kernel.
diskann-benchmark/example/multi-vector-staged.json	Example job config including staged ISA variant.
diskann-benchmark/example/multi-vector-3way.json	Example job config comparing reference vs fused vs staged.
diskann-benchmark/example/multi-vector-quant.json	Example job config for quantized staged benchmark op.
diskann-benchmark/example/multi-vector-tiled-f16.json	Example job config for f16 tiler benchmark op.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.