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 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;
@pfoxARM

pfoxARM commented Aug 25, 2026

Copy link
Copy Markdown
Author

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

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.

3 participants