Skip to content

Add Neon inner-product kernel for USlice<2> with spherical wiring - #1363

Merged
Mark Hildebrand (hildebrandmw) merged 1 commit into
microsoft:mainfrom
pfoxARM:pfoxARM/neon_ip_u2
Sep 2, 2026
Merged

Add Neon inner-product kernel for USlice<2> with spherical wiring#1363
Mark Hildebrand (hildebrandmw) merged 1 commit into
microsoft:mainfrom
pfoxARM:pfoxARM/neon_ip_u2

Conversation

@pfoxARM

@pfoxARM pfoxARM commented Aug 29, 2026

Copy link
Copy Markdown
Contributor

What does this implement?

Adds Neon Inner product kernel USlice<2> * USlice<2> into diskann-quantization/src/bits/distances.rs
Adds 2-bit spherical wiring into diskann-quantization/src/spherical
Updates unit test bounds

@pfoxARM
pfoxARM requested review from a team and a lite review from Copilot August 29, 2026 07:15

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 extends DiskANN’s quantized distance infrastructure by adding an AArch64/Neon-optimized inner-product kernel for 2-bit unsigned bit-slices and wiring it into the spherical quantizer dispatch so 2-bit spherical distances can run on Neon instead of retargeting to scalar.

Changes:

  • Added a Neon Target2 implementation for InnerProduct over USlice<2> × USlice<2> and updated the kernel support table and test bounds.
  • Updated spherical dispatch_map! wiring to use Neon directly for AsData<2> and AsQuery<2> (no downcast/retarget).
  • Added explicit AArch64 codegen instantiations for the 2-bit spherical inner-product paths.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File Description
diskann-quantization/src/spherical/iface.rs Enables Neon dispatch for 2-bit spherical data/query layouts by removing the downcast/retarget.
diskann-quantization/src/spherical/__codegen/aarch64.rs Adds explicit codegen instantiations for the 2-bit Neon spherical inner-product paths.
diskann-quantization/src/bits/distances.rs Implements the 2-bit Neon inner-product kernel, updates retargeting, and adjusts test bounds/documentation.

💡 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/spherical/__codegen/aarch64.rs Outdated
Comment thread diskann-quantization/src/bits/distances.rs Outdated
@pfoxARM
pfoxARM force-pushed the pfoxARM/neon_ip_u2 branch from 979f6d7 to 867ff0d Compare August 29, 2026 07:26

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!

Comment thread diskann-quantization/src/bits/distances.rs Outdated
@pfoxARM

pfoxARM commented Sep 2, 2026

Copy link
Copy Markdown
Contributor Author

Should of fixed cargo-fmt errors now.

@codecov-commenter

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 91.54%. Comparing base (fabcb9b) to head (a600f9c).
⚠️ Report is 3 commits behind head on main.

Additional details and impacted files

Impacted file tree graph

@@            Coverage Diff             @@
##             main    #1363      +/-   ##
==========================================
- Coverage   91.55%   91.54%   -0.01%     
==========================================
  Files         521      521              
  Lines      100371   100302      -69     
==========================================
- Hits        91895    91825      -70     
- Misses       8476     8477       +1     
Flag Coverage Δ
miri 91.54% <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 16 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 600c2b9 into microsoft:main Sep 2, 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