Skip to content

Add Neon Inner product u4*u4 kernel - #1353

Merged
Mark Hildebrand (hildebrandmw) merged 1 commit into
microsoft:mainfrom
pfoxARM:pfoxARM/neon_ip_u4
Aug 28, 2026
Merged

Add Neon Inner product u4*u4 kernel#1353
Mark Hildebrand (hildebrandmw) merged 1 commit into
microsoft:mainfrom
pfoxARM:pfoxARM/neon_ip_u4

Conversation

@pfoxARM

@pfoxARM pfoxARM commented Aug 25, 2026

Copy link
Copy Markdown
Contributor
  • 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
Contributor 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.

@hildebrandmw
Mark Hildebrand (hildebrandmw) merged commit fabcb9b into microsoft:main Aug 28, 2026
30 checks passed
MustafaIdrisArm pushed a commit to MustafaIdrisArm/DiskANN that referenced this pull request Sep 11, 2026
# Breaking Changes

Flat search visitors are now query-aware (microsoft#1359)
* `flat::FlatIndex`  has been removed. The search entry point is now the
free function `flat::knn_search`, and the `DistancesUnordered` visitor
is constructed per-query rather than reused. The `ElementRef`,
`QueryComputer`, and `QueryComputerError` associated types and the
visitor GAT have also been removed. `DistancesUnordered` now yields
`(id, distance)` pairs directly.

Migration: Callers of the old `FlatIndex`/visitor API should:

1. Drop `FlatIndex` and construct your `DistancesUnordered` visitor
directly for the query being searched.
2. Replace calls into the removed wrapper with `flat::knn_search(&mut
visitor, k, processor, query, &mut output)`.
3. Remove any `ElementRef`/`QueryComputer` implementation, fuse scanning
and distance computation directly in your
`DistancesUnordered::distances_unordered` implementation.

# All Changes

* Make `UnalignedSlice` Send and Sync. by @hildebrandmw in
microsoft#1348
* Bump actions/checkout from 4.4.0 to 7.0.1 in the github-actions group
by @dependabot[bot] in microsoft#1344
* Deduplicate virtual start point edges during disk serialization by
@partychen in microsoft#1350
* Fix alpha pruning documentation by @xinyuwen2 in
microsoft#1351
* Bump the github-actions group with 4 updates by @dependabot[bot] in
microsoft#1356
* Add Neon Inner product u4*u4 kernel by @pfoxARM in
microsoft#1353
* Strengthen arguments to `robust_prune`. by @hildebrandmw in
microsoft#1358
* Allow inspection of the paged search accessor. by @hildebrandmw in
microsoft#1364
* Make flat search visitors query-aware by @partychen in
microsoft#1359
* Add Neon inner-product kernel for USlice<2> with spherical wiring by
@pfoxARM in microsoft#1363

## New Contributors
* @pfoxARM made their first contribution in
microsoft#1353

**Full Changelog**:
microsoft/DiskANN@v0.56.0...v0.57.0

Co-authored-by: Mark Hildebrand <mhildebrand@microsoft.com>
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