[diskann-quantization] Add Neon F32 Minmax Kernels

[pfoxARM]

What does this implement

Adds AArch64 NEON full-precision query-to-data kernels for 1/2/4-bit quantized inner product, wires spherical dispatch to use them directly.

Kernels written by MustafaIdrisArm

Any other comments?

I added a commit to avoid expanding diskann-wide with an interleaved trait, instead taking a inline helper functions approach in diskann-quantization. However if there is an argument for interleaved traits in diskann-wide then we can revert my last commit.

[Copilot]

Copilot review overview

🔵 Needs a closer look

Add the missing conversion coverage and correct the Neon documentation classifications.

Review effort: Lite
Findings: None

What changed in this PR

Adds AArch64 NEON full-precision kernels for 1-, 2-, and 4-bit quantized inner products, with updated spherical dispatch.

Changes:

• Added NEON widening and integer-to-float conversions.
• Implemented quantized NEON kernels.
• Updated spherical dispatch and AArch64 code generation.

File	Summary
diskann-wide/​src/​arch/​aarch64/​u32x4_.rs	Added widening conversion support.
diskann-wide/​src/​arch/​aarch64/​f32x4_.rs	Added integer-to-float conversions. Moderate (1 vote): add tests for u32x4 → f32x4 and u16x4 conversion paths.
diskann-wide/​src/​arch/​aarch64/​double.rs	Added wider NEON conversion support and tests.
diskann-quantization/​src/​spherical/​iface.rs	Updated direct NEON dispatch.
diskann-quantization/​src/​spherical/​__codegen/​aarch64.rs	Added NEON kernel instantiations.
diskann-quantization/​src/​bits/​distances.rs	Added 1-, 2-, and 4-bit NEON kernels. Nit (1 vote): correct the Neon table classifications.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[pfoxARM]

Appreciate the comments Mark, I've done the following:

• Made use of SIMDCast , test_cast! and test_lossless_convert!
• Copied a similar approach to [diskann-wide] Add Neon vabdq_u8 and vabdq_s8. #1437 and moved the deinterleaved helpers into mod.rs with appropirate handlers for miri
• Added test_vld2q_f32() and test_vld4q_f32()
  I tested fmt, clippy and diskann-benchmark with minmax-quantization for these kernels too, I can post results shortly (running on wiki-1M with C100)

[Codecov Comments Bot (codecov-commenter)]

Codecov Report

❌ Patch coverage is 0% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 90.58%. Comparing base (49cdb88) to head (066e988).
⚠️ Report is 8 commits behind head on main.

Files with missing lines	Patch %	Lines
diskann-wide/src/reference.rs	0.00%	3 Missing ⚠️

❌ Your patch status has failed because the patch coverage (0.00%) is below the target coverage (90.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1431      +/-   ##
==========================================
- Coverage   91.51%   90.58%   -0.94%     
==========================================
  Files         527      567      +40     
  Lines      103055   112347    +9292     
==========================================
+ Hits        94309   101767    +7458     
- Misses       8746    10580    +1834     

Flag	Coverage Δ
miri	90.58% <0.00%> (-0.94%)	⬇️
unittests	90.34% <0.00%> (-0.92%)	⬇️

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

Files with missing lines	Coverage Δ
diskann-quantization/src/bits/distances.rs	92.20% <ø> (ø)
diskann-quantization/src/spherical/iface.rs	91.50% <ø> (+0.16%)	⬆️
diskann-wide/src/reference.rs	91.81% <0.00%> (-2.58%)	⬇️

... and 106 files with indirect coverage changes

🚀 New features to boost your workflow:

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

[pfoxARM]

These changes tested on wiki-1M with diskann-benchmark (exhaustive-minmax-quantization) using 8 threads

C100 VM

• f32xu1: 1.8021x (16.9801 -> 30.5998 qps)
• f32xu2: 2.4131x (14.5934 -> 35.2151 qps)
• f32xu4: 4.1833x (10.6926 -> 44.7303 qps)
  C200 VM
• f32xu1: 1.9875x (16.5900 -> 32.9731 qps)
• f32xu2: 2.4224x (14.5810 -> 35.3207 qps)
• f32xu4: 4.1394x (10.2360 -> 42.3707 qps)

Later I will test out the prefetching changes #1408 on top as well, measure cache miss rates.

[Mark Hildebrand (hildebrandmw)]

Thanks! Some housekeeping requests for diskann-wide and a quick question on the 1-bit kernels.

[pfoxARM]

PR is now ready for more review Mark Hildebrand (@hildebrandmw)

• Rewritten f32x1bit kernel using vtstq_u32 for bit-tests in each lane. This gives 300% QPS uplift compared to main and 100% speed up over the previous implementation in this PR (tested on 16-threads, QPS hit 113)
• Shuffled things around diskann-wide following advice, with inherent functions and avoiding expanding double.rs

Thanks!

[pfoxARM]

Ready for review again, no new changes since my last comment. Just fixed merge conflicts

[Mark Hildebrand (hildebrandmw)]

Thanks - a few small comments on my end. The biggest request is on the 1-bit kernel. LLVM appears to optimize sequences based on SIMDSelect, and & the same as when vtstq is used explicitly, so I'd prefer to keep the kernel in terms of more general machinery.

[Mark Hildebrand (hildebrandmw)]

For this PR - I think we can get away without introducing an intrinsic because LLVM kinda already knows how to handle these cases.

By this, I mean that the code generation between foo and bar in the linked compiler explorer pane is exactly the same, despite one using vtstq_u32 and the other using a more circuitous route using vandq, vceqq and vbslq. (Note that the LLVM version used by currently doesn't want to emit cmtst anyways, probably due to this issue)

The point is, we can get the same code generation using a simpler

let zero = f32x4::splat(arch, 0.0);

// This sequence below is effectively calls
// `bar` from the linked compiler explorer.
let mask0: u32x4 = y_word & bits0;
s0 = s0 + mask0.select(x0, zero); 

and save a lot of bespoke machinery.

[Mark Hildebrand (hildebrandmw)]

Can we add u8x8 => f32x8 and u8x16 => f32x16 as well for completeness?

[Mark Hildebrand (hildebrandmw)]

Can we add u8x8 => f32x8 and u8x16 => f32x16 as well for completeness?

[Mark Hildebrand (hildebrandmw)]

I know this seems kind of arbitrary, but since SIMDCast<f32> for u32x4 is implemented for the u32x4 type (i.e., it's on the right-hand side of the impl block), can test_cast!(u32x4 => f32x4, ...) go in u32x4_.rs?