Add Neon Inner product u4*u4 kernel - #1353
Open
pfoxARM wants to merge 1 commit into
Open
Conversation
Contributor
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.
Comment on lines
+1571
to
+1575
| #[cfg(target_arch = "aarch64")] | ||
| impl | ||
| Target2< | ||
| diskann_wide::arch::aarch64::Neon, | ||
| MathematicalResult<u32>, |
Comment on lines
+12
to
+14
|
|
||
| #[cfg(target_arch = "aarch64")] | ||
| pub mod aarch64; |
Author
|
@microsoft-github-policy-service agree company="Arm" |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.