perf: add SIMD-accelerated u8 L2 and cosine distance kernels

[justinrmiller]

Summary

• Add hand-written AVX2 and AVX-512 VNNI backends for u8 squared L2 distance (Σ(a-b)²) in new l2_u8.rs
• Add fused single-pass u8 cosine distance kernel in new cosine_u8.rs — computes dot(a,b), ‖a‖², ‖b‖² simultaneously, halving memory traffic vs the previous 2-3 pass approach
• Wire both into the L2 for u8 and Cosine for u8 trait impls
• Add benchmarks comparing scalar vs SIMD for both kernels

Algorithmic approach (adapted from NumKong)

L2 (AVX2): Saturating subtraction for |a-b|, zero-extend u8→i16, VPMADDWD(diff, diff) to square and accumulate into i32. 32 elements/iter.

L2 (AVX-512 VNNI): Same abs-diff approach with VPDPWSSD for fused square-accumulate. 64 elements/iter.

Cosine (AVX2): Zero-extend both vectors to i16, triple VPMADDWD per half (a·b, a·a, b·b). 32 elements/iter, single pass.

Cosine (AVX-512 VNNI): Same three-accumulator approach with VPDPWSSD. 64 elements/iter.

Both kernels use OnceLock-based runtime CPU dispatch, falling back to portable scalar on non-x86 platforms.

Benchmarks

1M × 1024-dim u8 vectors.

x86_64 — AMD Ryzen 5 4500 6-Core (AVX2, no AVX-512)

Kernel	Scalar	SIMD	Speedup
L2(u8)	73.5 ms	58.2 ms	1.26x
Cosine(u8)	122.2 ms	82.1 ms	1.49x

L2 auto-vectorization baseline was 91.5 ms, so SIMD is 1.57x faster than that path.

aarch64 — Apple Silicon M3 Max (no AVX2, scalar fallback)

Kernel	Scalar	SIMD (dispatch)
L2(u8)	26.8 ms	27.3 ms
Cosine(u8)	90.1 ms	90.4 ms

On aarch64 the SIMD path falls through to scalar (no AVX2), so times are identical — confirms no regression on non-x86 platforms. AVX-512 VNNI systems (Ice Lake+, Zen 4+) should see larger gains.

Test plan

• All 11 new tests pass: SIMD backends verified against scalar reference across 18 vector sizes (0–4097), boundary values (0/255), alternating patterns, random seeds
• All 63 existing lance-linalg tests pass (no regressions)
• Clippy clean, fmt clean
• Benchmarked on x86_64 AVX2 (AMD Ryzen 5 4500) — L2 1.26x, Cosine 1.49x faster
• Verify on AVX-512 VNNI system for additional speedup data

🤖 Generated with Claude Code

[BubbleCal]

Followup: switch to the new kernels for SQ

[codecov]

Codecov Report

❌ Patch coverage is 73.94541% with 105 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
rust/lance-linalg/src/distance/cosine_u8.rs	75.21%	56 Missing and 2 partials ⚠️
rust/lance-linalg/src/distance/l2_u8.rs	73.93%	41 Missing and 2 partials ⚠️
rust/lance-linalg/src/distance/cosine.rs	0.00%	3 Missing ⚠️
rust/lance-linalg/src/distance/l2.rs	0.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!