Skip to content

Add Neon Inner product u4*u4 kernel - #1353

Open
pfoxARM wants to merge 1 commit into
microsoft:mainfrom
pfoxARM:pfoxARM/neon_ip_u4
Open

Add Neon Inner product u4*u4 kernel#1353
pfoxARM wants to merge 1 commit into
microsoft:mainfrom
pfoxARM:pfoxARM/neon_ip_u4

Conversation

@pfoxARM

@pfoxARM pfoxARM commented Aug 25, 2026

Copy link
Copy Markdown
  • Does this PR have a descriptive title that could go in our release notes? Yes
  • Does this PR add any new dependencies? No
  • Does this PR modify any existing APIs? No
  • Is the change to the API backwards compatible? No
  • Should this result in any changes to our documentation, either updating existing docs or adding new ones? No

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.

@pfoxARM
pfoxARM requested review from a team and a lite review from Copilot August 25, 2026 19:46

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 aarch64 spherical __codegen instantiations for the 4-bit Neon inner-product paths.
  • Enable Neon dispatch for spherical quantization AsData<4> and AsQuery<4> without downcasting to Scalar.
  • Implement an AArch64 Neon InnerProduct kernel for USlice<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 thread diskann-quantization/src/bits/distances.rs
Comment thread diskann-quantization/src/spherical/__codegen/mod.rs
@pfoxARM

pfoxARM commented Aug 25, 2026

Copy link
Copy Markdown
Author

@microsoft-github-policy-service agree company="Arm"

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread diskann-quantization/src/bits/distances.rs Outdated
@pfoxARM
pfoxARM force-pushed the pfoxARM/neon_ip_u4 branch from 5a885d5 to a3506ac Compare August 27, 2026 11:02
@pfoxARM
pfoxARM force-pushed the pfoxARM/neon_ip_u4 branch from a3506ac to a654681 Compare August 27, 2026 13:24
@codecov-commenter

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 91.55%. Comparing base (158126e) to head (a654681).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

Impacted file tree graph

@@            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     
Flag Coverage Δ
miri 91.55% <100.00%> (-0.01%) ⬇️
unittests 91.23% <100.00%> (-0.01%) ⬇️

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 99.59% <100.00%> (ø)
diskann-quantization/src/spherical/iface.rs 90.19% <ø> (ø)

... and 4 files with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants