[diskann-wide] AVX512/Neon popcount - #1310
Conversation
There was a problem hiding this comment.
Pull request overview
Adds a new SIMDPopcount trait to diskann-wide and wires it through emulated/reference paths plus selected x86_64 AVX-512 and AArch64 Neon backends, enabling lane-wise popcount where hardware instructions exist (supporting upcoming 1-bit quantization work).
Changes:
- Introduce and export
SIMDPopcount, with implementations forEmulatedandDoubledvector wrappers. - Add reference scalar “integer ops” helpers including expected popcount, and new popcount correctness tests.
- Implement instruction-backed popcount for AVX-512 (
avx512bitalg/avx512vpopcntdq) and Neon (vcnt*) vector types.
Reviewed changes
Copilot reviewed 28 out of 28 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| diskann-wide/src/traits.rs | Adds the public SIMDPopcount trait. |
| diskann-wide/src/lib.rs | Re-exports SIMDPopcount from the crate root. |
| diskann-wide/src/reference.rs | Renames/expands reference integer ops to include popcount expected results. |
| diskann-wide/src/emulated.rs | Implements SIMDPopcount for the emulated backend via reference ops. |
| diskann-wide/src/doubled.rs | Implements SIMDPopcount for Doubled<T> by delegating to halves. |
| diskann-wide/src/test_utils/ops.rs | Adds test_popcount! and updates integer test bounds to ReferenceIntegerOps. |
| diskann-wide/src/arch/x86_64/v4/u8x64_.rs | Adds AVX-512 byte popcount and tests. |
| diskann-wide/src/arch/x86_64/v4/u8x32_.rs | Adds AVX-512VL byte popcount and tests. |
| diskann-wide/src/arch/x86_64/v4/u8x16_.rs | Adds AVX-512VL byte popcount and tests. |
| diskann-wide/src/arch/x86_64/v4/i8x64_.rs | Adds AVX-512 byte popcount and tests. |
| diskann-wide/src/arch/x86_64/v4/i8x32_.rs | Adds AVX-512VL byte popcount and tests. |
| diskann-wide/src/arch/x86_64/v4/i8x16_.rs | Adds AVX-512VL byte popcount and tests. |
| diskann-wide/src/arch/x86_64/v4/i16x32_.rs | Adds AVX-512 word popcount and tests. |
| diskann-wide/src/arch/x86_64/v4/i16x16_.rs | Adds AVX-512VL word popcount and tests. |
| diskann-wide/src/arch/x86_64/v4/i16x8_.rs | Adds AVX-512VL word popcount and tests. |
| diskann-wide/src/arch/x86_64/v4/u32x16_.rs | Adds AVX-512 dword popcount and tests. |
| diskann-wide/src/arch/x86_64/v4/u32x8_.rs | Adds AVX-512VL dword popcount and tests. |
| diskann-wide/src/arch/x86_64/v4/u32x4_.rs | Adds AVX-512VL dword popcount and tests. |
| diskann-wide/src/arch/x86_64/v4/i32x16_.rs | Adds AVX-512 dword popcount and tests. |
| diskann-wide/src/arch/x86_64/v4/i32x8_.rs | Adds AVX-512VL dword popcount and tests. |
| diskann-wide/src/arch/x86_64/v4/i32x4_.rs | Adds AVX-512VL dword popcount and tests. |
| diskann-wide/src/arch/x86_64/v4/u64x4_.rs | Adds AVX-512VL qword popcount and tests. |
| diskann-wide/src/arch/x86_64/v4/u64x2_.rs | Adds AVX-512VL qword popcount and tests. |
| diskann-wide/src/arch/aarch64/u8x8_.rs | Adds Neon byte popcount and tests. |
| diskann-wide/src/arch/aarch64/u8x16_.rs | Adds Neon byte popcount and tests. |
| diskann-wide/src/arch/aarch64/i8x8_.rs | Adds Neon byte popcount and tests. |
| diskann-wide/src/arch/aarch64/i8x16_.rs | Adds Neon byte popcount and tests. |
| diskann-wide/src/arch/aarch64/double.rs | Adds popcount tests for doubled Neon vectors. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| /// Count the set bits in each lane. | ||
| /// | ||
| /// This trait is only implemented for vector/architecture pairs that have native support | ||
| /// for popcounts. | ||
| pub trait SIMDPopcount { |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 28 out of 28 changed files in this pull request and generated no new comments.
Suppressed comments (1)
diskann-wide/src/traits.rs:531
- The
SIMDPopcountdocs say it is only implemented for vector/architecture pairs with native popcount support, but the PR also implements it forEmulated<..>(and thereforeDoubled<Emulated<..>>), which is not native SIMD popcount. Either adjust the docs to match the actual impl set, or remove the emulated impl if the intent is truly “native only”.
/// Count the set bits in each lane.
///
/// This trait is only implemented for vector/architecture pairs that have native support
/// for popcounts.
pub trait SIMDPopcount {
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1310 +/- ##
=======================================
Coverage 91.46% 91.47%
=======================================
Files 516 516
Lines 98276 98305 +29
=======================================
+ Hits 89891 89922 +31
+ Misses 8385 8383 -2
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
Adds limited support for SIMD popcounts. These instructions are weird: AVX512-BITALG has great support, Neon only supports
i8/u8, and AVX2 supports nothing. The decision made here is to only implement theSIMDPopcounttrait when it maps to an instruction. Trying to emulate it on an op-level basis is disastrous for performance.This is part of a small line of work adding better AVX-512 1-bit quantization.
AI Disclosure: the author is unable to efficiently type at the moment so had an agent do most of the work. The change set is largely mechanical.