Add Neon Inner product u4*u4 kernel - #1353
Conversation
There was a problem hiding this comment.
Pull request overview
This PR adds an AArch64-specific Neon+dotprod inner-product kernel for 4-bit (USlice<4> × USlice<4>) paths in spherical quantization, and wires it into the spherical quantizer’s architecture dispatch so the Neon implementation can be selected where available.
Changes:
- Add
aarch64spherical__codegeninstantiations for the 4-bit Neon inner-product paths. - Enable Neon dispatch for spherical quantization
AsData<4>andAsQuery<4>without downcasting to Scalar. - Implement an AArch64 Neon
InnerProductkernel forUSlice<4> × USlice<4>and adjust retargeting to avoid overlap.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
diskann-quantization/src/spherical/iface.rs |
Updates dispatch mapping so 4-bit spherical paths can use Neon directly (no downcast). |
diskann-quantization/src/spherical/__codegen/mod.rs |
Adds an AArch64 codegen module behind cfg(target_arch = "aarch64"). |
diskann-quantization/src/spherical/__codegen/aarch64.rs |
New AArch64 instantiation helpers for the 4-bit Neon inner-product distance computer. |
diskann-quantization/src/bits/distances.rs |
Adds the Neon USlice<4> × USlice<4> inner-product implementation and updates retargeting accordingly. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
|
@microsoft-github-policy-service agree company="Arm" |
Mark Hildebrand (hildebrandmw)
left a comment
There was a problem hiding this comment.
Thanks! Looks good to me - great to start having Neon kernels!
Outside of the small tweak to the test bounds, please start a new aarch64.rs file in diskann-quantization/src/__codegen.
5a885d5 to
a3506ac
Compare
a3506ac to
a654681
Compare
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1353 +/- ##
==========================================
- Coverage 91.55% 91.55% -0.01%
==========================================
Files 521 521
Lines 100371 100371
==========================================
- Hits 91898 91895 -3
- Misses 8473 8476 +3
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
What does this implement?
Exclusive aarch64 USlice4 * USlice 4 Inner Product kernel using Neon and dotprod. We also add quantization instantiation for spherical quantization.
Any other comments?
We use dot_simd() heavily in this kernel, for performance we rely on dotprod feature.