Add Neon inner-product kernel for USlice<2> with spherical wiring - #1363
Conversation
There was a problem hiding this comment.
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
Target2implementation forInnerProductoverUSlice<2> × USlice<2>and updated the kernel support table and test bounds. - Updated spherical
dispatch_map!wiring to use Neon directly forAsData<2>andAsQuery<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.
979f6d7 to
867ff0d
Compare
Mark Hildebrand (hildebrandmw)
left a comment
There was a problem hiding this comment.
Thanks!
867ff0d to
a836b51
Compare
a836b51 to
a600f9c
Compare
|
Should of fixed cargo-fmt errors now. |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ 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
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
600c2b9
into
microsoft:main
# 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>
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