perf: add SIMD-accelerated u8 dot product for SQ distance#6506
Merged
westonpace merged 11 commits intoApr 16, 2026
Conversation
Add AVX-512 VNNI and AVX2 backends for unsigned int8 dot product, with runtime CPU feature detection and automatic fallback to scalar. This accelerates Scalar Quantization (SQ) dot-product distance computation which operates on u8-encoded vectors. Backends (selected at runtime, best available wins): - AVX-512 VNNI: VPDPBUSD with XOR-0x80 bias trick, 64 elements/iter - AVX2: zero-extend to u16 + VPMADDWD, 32 elements/iter - Scalar: portable fallback The VNNI path uses the algebraic trick of biasing one operand into the signed domain (XOR 0x80) and correcting with 128*sum(a) at the end. The correction uses VPSADBW which runs on a different execution port than VPDPBUSD, making it effectively free. Replaces the unoptimized scalar Dot<u8> impl (which had an explicit TODO for SIMD optimization). All existing callers including SQ distance computation benefit automatically. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
Add two aarch64 backends to the u8 dot product dispatch: - NEON dotprod (UDOT): native u8×u8→u32 via inline asm, 16 elem/iter. Available on Apple M1+, Graviton2+, Cortex-A76+ (ARMv8.2-A+dotprod). Uses inline asm because vdotq_u32 is not yet stable in Rust std::arch. - Base NEON: widening multiply (vmull_u8→u16) + pairwise accumulate (vpadalq_u16→u32), 16 elem/iter. Always available on aarch64. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
Trim per-function doc comments and inline comments to match the concise style used in sibling distance modules, while preserving the detailed module-level documentation explaining the VNNI bias trick and backend selection. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
LLVM auto-vectorizes the scalar loop with NEON on aarch64, so explicit NEON intrinsics provide no benefit on Apple Silicon. The real win is AVX-512 VNNI on Intel where the compiler cannot auto-generate VPDPBUSD. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
Cast to *const __m512i instead of *const i32 for _mm512_loadu_si512, fixing compilation on CI (MSRV, linux, windows). Also fix cargo fmt. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
Compare scalar vs SIMD-dispatched u8 dot product across 128/256/512/1024 dimensions to measure gains from AVX2 and AVX-512 VNNI backends. Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
This was referenced Apr 26, 2026
feat(lance-linalg): runtime SIMD dispatch for pre-Haswell x86_64 from-source builds
tobocop2/lance#2
Draft
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Dot<u8>impl (which had an explicit TODO) with dispatched SIMD kernelDetails
New file:
rust/lance-linalg/src/distance/dot_u8.rsThree backends selected at runtime via
OnceLock+is_x86_feature_detected!:VPDPBUSD+ XOR-0x80 bias trickVPMADDWDon zero-extended u16The VNNI bias trick
VPDPBUSDexpects one unsigned and one signed operand, but SQ vectors are u8×u8. We XOR one operand with 0x80 to map it to the signed domain, then correct by adding128·Σaat the end. The correction usesVPSADBWwhich runs on execution port 5 whileVPDPBUSDruns on port 0 — they execute in parallel every cycle, making the correction effectively free.SQ integration (automatic)
SQDistCalculator::distance()already callsdot_distance()→u8::dot()for Dot distance type. Replacing theDot<u8>body is the only change needed.Benchmarks
Ryzen 4500 (AVX2, no VNNI)
1M total u8 elements, varying vector dimension. Scalar baseline vs AVX2-dispatched path:
AVX2 delivers up to 1.63x throughput at dim=1024. At dim=128 the
OnceLockdispatch and AVX2 loop setup overhead exceeds the SIMD gains on short vectors. AVX-512 VNNI (Ice Lake+ / Zen 4+) is expected to show larger gains with 64 elements/iter.Apple M4 (ARM64, scalar fallback)
On ARM64 the dispatch falls back to scalar, so both paths perform identically (~13 µs at dim=1024). A follow-up ARM NEON
UDOTpath would bring SIMD gains to Apple Silicon.Out of scope (follow-up)
Σ(a-b)²)VPDPBUUD(unsigned×unsigned, Sierra Forest+) — too new for stable RustUDOTpathTest plan
#[cfg]guards for missing CPU features)dottests continue to pass (9/9)cargo clippy -p lance-linalg --tests --benches -- -D warningscleancargo bench --bench dot -p lance-linalg -- "Dot\(u8"🤖 Generated with Claude Code