Add Neon inner product U8xU1, U8xU2 and U8xU4 kernels

[pfoxARM]

What does this implement

Adds inner product 8-bit x 4-bit unsigned vectors using neon
Adds inner product 8-bit x 2-bit unsigned vectors using neon
Adds inner product 8-bit x 1-bit unsigned vectors using neon
Adds Zip and Unzip intrinsics into diskann-wide

[Copilot]

🟡 Changes recommended

The new Neon (USlice<8>, USlice<4>)Target2 impl is missing #[cfg(target_arch = "aarch64")], which will break non-aarch64 builds.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Adds AArch64/Neon-optimized inner product kernels for mixed-width unsigned bit-slices (U8×U4/U2/U1) and extends diskann-wide with Zip/Unzip support needed to efficiently unpack/interleave packed bit representations on Neon.

Changes:

• Added a ZipUnzip implementation for u8x16 on AArch64 to support efficient element interleaving/deinterleaving with 64-bit Neon zip/unzip intrinsics.
• Implemented Neon inner product kernels for USlice<8> × USlice<4|2|1> in the quantization bit-distance backend.
• Updated the bit-slice test bounds for (8, Neon) to exercise the new Neon kernel paths more thoroughly.

File summaries

File	Description
diskann-wide/src/arch/aarch64/double.rs	Adds ZipUnzip for u8x16 using vzip*/vuzp* to enable efficient interleaving used by packed-bit kernels.
diskann-quantization/src/bits/distances.rs	Adds Neon implementations for InnerProduct on (8,4), (8,2), (8,1) bit-slices and adjusts related documentation/test bounds.

Review details

• Files reviewed: 2/2 changed files
• Comments generated: 1
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Mark Hildebrand (hildebrandmw)]

Thanks - I have a bit of feedback on this on. Mainly questions on different approaches.

I also realized that the tests for heterogeneous kernels are woefully under powered when it comes to testing the neon kernels (as in, they aren't tested at all). Some work will be needed to refactor our tests. I'll try to get around to that today or so.

[Mark Hildebrand (hildebrandmw)]

Curious: did you try something like

let y = unsafe{ u8x16::load_simd(arch, py_u8.add(i))};letLoHi{lo: xodd,hi: xeven,} = unsafe{ u8x32::load_simd(arch, px_u8.add(2* i))}.unzip();
s0 = s0.dot_simd(y & mask, xodd).dot_simd((y >> 4)& mask, xeven);

(disclaimer: didn't quite test it due to our apparent testing gap 😢).

This would be about 3 loads, 3 integer ops, 2 shuffles, and 2 dots for 32 elements.

The current approach does 4 loads, 4 integer ops, 4 shuffles, and 2 dots for 32 elements.

Since vectors usually have dimensions that are a multiple of 2, I wouldn't worry too much about the extra remainder.

[Mark Hildebrand (hildebrandmw)]

Note that arch can come from input.arch(), so doesn't necessarily need to be passed separately.

[Mark Hildebrand (hildebrandmw)]

Question: did you try something like

pubfntest_function(arch:Neon,x:&u8,y:&u8) -> u8x16{let combined:u8x16 = LoHi::new(u8x8::splat(arch,*x), u8x8::splat(arch,*y)).join();let shifts = u8x16::from_array(arch,[0,1,2,3,4,5,6,7,0,1,2,3,4,5,6,7]);(combined >> shifts)& u8x16::splat(arch,1)}

for the bit-unpacking? With shifts and the mask amortized, this costs 5 instructions for 16 elements. I counted around 21 instructions per 64-elements in the approach taken in the PR.

The approach above does hit the load units more, but we might be able to introduce indexing operations for SIMDVector to amortize the loading of the 1-bit values.

[MustafaIdrisArm]

Thanks - I have a bit of feedback on this on. Mainly questions on different approaches.

I also realized that the tests for heterogeneous kernels are woefully under powered when it comes to testing the neon kernels (as in, they aren't tested at all). Some work will be needed to refactor our tests. I'll try to get around to that today or so.

Thanks for the feedback. With regards to the testing, I added the following code to all_ip_dispatches() in the heterogeneous_ip_tests_8xM macro for internal testing

#[cfg(target_arch = "aarch64")]ifletSome(arch) = diskann_wide::arch::aarch64::Neon::new_checked(){fuzz_heterogeneous_ip::<$M>(MAX_DIM,TRIALS_PER_DIM,
$max_val,&|x, y| arch.run2(InnerProduct, x, y),"neon",&mut rng,);}

With MAX_DIM set to 256 and TRIALS_PER_DIM set to 20, the coverage should be sufficient for the general use cases of the kernel. heterogeneous_ip_tests_8xM itself is expanded for all of our cases already (M = 4,2,1).

Note: There are also some corner case tests such as max_values() and edge_cases() where we could add neon dispatch so we test for those cases as well. The implementations are just as simple if not simpler than the snippet above, just let me know what fits best and I'll go for that.