diff --git a/lib/quantization/Cargo.toml b/lib/quantization/Cargo.toml index 3e541632ee..77a8a24c86 100644 --- a/lib/quantization/Cargo.toml +++ b/lib/quantization/Cargo.toml @@ -64,3 +64,7 @@ harness = false [[bench]] name = "turboquant" harness = false + +[[bench]] +name = "turbo_simd" +harness = false diff --git a/lib/quantization/benches/turbo_simd.rs b/lib/quantization/benches/turbo_simd.rs new file mode 100644 index 0000000000..cb44213004 --- /dev/null +++ b/lib/quantization/benches/turbo_simd.rs @@ -0,0 +1,669 @@ +use std::hint::black_box; + +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use quantization::encoded_vectors_binary::BitsStoreType; +use quantization::turboquant::simd::{ + Query1bitSimd, Query2bitSimd, Query4bitSimd, score_1bit_internal, score_1bit_internal_scalar, + score_2bit_internal, score_2bit_internal_scalar, score_4bit_internal, + score_4bit_internal_scalar, +}; +#[cfg(target_arch = "x86_64")] +use quantization::turboquant::simd::{ + score_1bit_internal_avx2, score_1bit_internal_avx512_vpopcntdq, score_1bit_internal_sse, + score_2bit_internal_avx2, score_2bit_internal_avx512_vnni, score_2bit_internal_sse, + score_4bit_internal_avx2, score_4bit_internal_avx512_vnni, score_4bit_internal_sse, +}; +#[cfg(all(target_arch = "aarch64", target_feature = "neon"))] +use quantization::turboquant::simd::{ + score_1bit_internal_neon, score_2bit_internal_neon, score_2bit_internal_neon_sdot, + score_4bit_internal_neon, score_4bit_internal_neon_sdot, +}; +use rand::prelude::StdRng; +use rand::seq::SliceRandom; +use rand::{RngExt, SeedableRng}; + +/// Two "nicely aligned" dims — `128` (single SIMD block) and `1536` (large, +/// divisible by every chunk size we ship) — plus a per-bit-width "ugly" dim +/// that hits the worst-case tail for that pipeline: +/// • odd chunk count → SDOT / AVX2 / AVX-512 leftover branch fires. +/// • `dim % chunk_dim` at its maximum → scalar tail helper does the most work. +const DIMS_4BIT: &[usize] = &[128, 1534, 1536]; // 1534 = 95 chunks (odd) + 14-dim tail +const DIMS_2BIT: &[usize] = &[128, 1532, 1536]; // 1532 = 95 chunks (odd) + 12-dim tail +const DIMS_1BIT: &[usize] = &[128, 1528, 1536]; // 1528 = 11 blocks (odd) + 120-dim tail + +/// Pool size ≫ L2. Indices are shuffled so the hardware prefetcher can't stream +/// vectors into cache — each iteration pays a real DRAM fetch. +const POOL_BYTES: usize = 64 * 1024 * 1024; + +struct VectorPool { + buf: Vec, + indices: Vec, + /// Packed bytes per vector (= dim / 2, since two 4-bit codes share a byte). + packed_bytes: usize, +} + +impl VectorPool { + fn with_packed_bytes(packed_bytes: usize, seed: u64) -> Self { + let count = (POOL_BYTES / packed_bytes).max(1024); + let mut rng = StdRng::seed_from_u64(seed); + let buf: Vec = (0..count * packed_bytes) + .map(|_| rng.random_range(0..=u8::MAX)) + .collect(); + let mut indices: Vec = (0..count as u32).collect(); + indices.shuffle(&mut rng); + Self { + buf, + indices, + packed_bytes, + } + } + + /// 4-bit PQ pool: two codes per byte → `dim / 2` packed bytes per vector. + fn new_4bit(dim: usize, seed: u64) -> Self { + assert!(dim.is_multiple_of(2)); + Self::with_packed_bytes(dim / 2, seed) + } + + /// 2-bit PQ pool: four codes per byte → `dim / 4` packed bytes per vector. + fn new_2bit(dim: usize, seed: u64) -> Self { + assert!(dim.is_multiple_of(4)); + Self::with_packed_bytes(dim / 4, seed) + } + + /// 1-bit PQ pool: 8 codes per byte → `dim / 8` packed bytes per vector. + fn new_1bit(dim: usize, seed: u64) -> Self { + assert!(dim.is_multiple_of(8)); + Self::with_packed_bytes(dim / 8, seed) + } + + #[inline] + fn vector(&self, cursor: usize) -> &[u8] { + let idx = self.indices[cursor % self.indices.len()] as usize; + &self.buf[idx * self.packed_bytes..(idx + 1) * self.packed_bytes] + } +} + +fn make_query(dim: usize) -> Vec { + let mut rng = StdRng::seed_from_u64(42); + (0..dim).map(|_| rng.random_range(-1.0_f32..1.0)).collect() +} + +fn bench_dotprod_cold(c: &mut Criterion) { + let mut group = c.benchmark_group("query4bit_dotprod_cold"); + for &dim in DIMS_4BIT { + let q = make_query(dim); + let query = Query4bitSimd::new(&q); + let pool = VectorPool::new_4bit(dim, 7); + + group.throughput(Throughput::Elements(dim as u64)); + + group.bench_with_input(BenchmarkId::new("scalar", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let v = pool.vector(cursor); + cursor = cursor.wrapping_add(1); + black_box(&query).dotprod_raw(black_box(v)) + }); + }); + + // Full public path: `dotprod_raw_best` (best SIMD backend) + the + // `sum_codebook_over_vector` bias correction + `as f32` reconstruction. + // This is what a real caller pays. Comparing against the best raw + // backend for the current CPU shows the bias-correction overhead. + group.bench_with_input(BenchmarkId::new("dotprod", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let v = pool.vector(cursor); + cursor = cursor.wrapping_add(1); + black_box(&query).dotprod(black_box(v)) + }); + }); + + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + { + group.bench_with_input(BenchmarkId::new("neon", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let v = pool.vector(cursor); + cursor = cursor.wrapping_add(1); + unsafe { black_box(&query).dotprod_raw_neon(black_box(v)) } + }); + }); + + if std::arch::is_aarch64_feature_detected!("dotprod") { + group.bench_with_input(BenchmarkId::new("neon_sdot", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let v = pool.vector(cursor); + cursor = cursor.wrapping_add(1); + unsafe { black_box(&query).dotprod_raw_neon_sdot(black_box(v)) } + }); + }); + } + } + + #[cfg(target_arch = "x86_64")] + { + if std::is_x86_feature_detected!("sse4.1") && std::is_x86_feature_detected!("ssse3") { + group.bench_with_input(BenchmarkId::new("sse", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let v = pool.vector(cursor); + cursor = cursor.wrapping_add(1); + unsafe { black_box(&query).dotprod_raw_sse(black_box(v)) } + }); + }); + } + + if std::is_x86_feature_detected!("avx2") { + group.bench_with_input(BenchmarkId::new("avx2", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let v = pool.vector(cursor); + cursor = cursor.wrapping_add(1); + unsafe { black_box(&query).dotprod_raw_avx2(black_box(v)) } + }); + }); + } + + if std::is_x86_feature_detected!("avx512f") + && std::is_x86_feature_detected!("avx512bw") + && std::is_x86_feature_detected!("avx512vnni") + { + group.bench_with_input(BenchmarkId::new("avx512_vnni", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let v = pool.vector(cursor); + cursor = cursor.wrapping_add(1); + unsafe { black_box(&query).dotprod_raw_avx512_vnni(black_box(v)) } + }); + }); + } + } + } + group.finish(); +} + +/// Benchmarks [`score_4bit_internal`] (both vectors already PQ-encoded, both +/// cold from DRAM). Every iteration picks two *different* pool indices so +/// each call pays two independent random cache misses — this models the +/// HNSW-internal case where we score one PQ vector against another stored +/// PQ vector, not against a hot query. +fn bench_score_cold(c: &mut Criterion) { + let mut group = c.benchmark_group("query4bit_score_cold"); + for &dim in DIMS_4BIT { + let pool = VectorPool::new_4bit(dim, 7); + + group.throughput(Throughput::Elements(dim as u64)); + + group.bench_with_input(BenchmarkId::new("scalar", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let va = pool.vector(cursor); + let vb = pool.vector(cursor + 1); + cursor = cursor.wrapping_add(2); + score_4bit_internal_scalar(black_box(va), black_box(vb)) + }); + }); + + // Public dispatcher — picks best available SIMD at runtime. + group.bench_with_input(BenchmarkId::new("dispatch", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let va = pool.vector(cursor); + let vb = pool.vector(cursor + 1); + cursor = cursor.wrapping_add(2); + score_4bit_internal(black_box(va), black_box(vb)) + }); + }); + + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + { + group.bench_with_input(BenchmarkId::new("neon", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let va = pool.vector(cursor); + let vb = pool.vector(cursor + 1); + cursor = cursor.wrapping_add(2); + unsafe { score_4bit_internal_neon(black_box(va), black_box(vb)) } + }); + }); + + if std::arch::is_aarch64_feature_detected!("dotprod") { + group.bench_with_input(BenchmarkId::new("neon_sdot", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let va = pool.vector(cursor); + let vb = pool.vector(cursor + 1); + cursor = cursor.wrapping_add(2); + unsafe { score_4bit_internal_neon_sdot(black_box(va), black_box(vb)) } + }); + }); + } + } + + #[cfg(target_arch = "x86_64")] + { + if std::is_x86_feature_detected!("sse4.1") && std::is_x86_feature_detected!("ssse3") { + group.bench_with_input(BenchmarkId::new("sse", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let va = pool.vector(cursor); + let vb = pool.vector(cursor + 1); + cursor = cursor.wrapping_add(2); + unsafe { score_4bit_internal_sse(black_box(va), black_box(vb)) } + }); + }); + } + + if std::is_x86_feature_detected!("avx2") { + group.bench_with_input(BenchmarkId::new("avx2", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let va = pool.vector(cursor); + let vb = pool.vector(cursor + 1); + cursor = cursor.wrapping_add(2); + unsafe { score_4bit_internal_avx2(black_box(va), black_box(vb)) } + }); + }); + } + + if std::is_x86_feature_detected!("avx512f") + && std::is_x86_feature_detected!("avx512bw") + && std::is_x86_feature_detected!("avx512vnni") + { + group.bench_with_input(BenchmarkId::new("avx512_vnni", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let va = pool.vector(cursor); + let vb = pool.vector(cursor + 1); + cursor = cursor.wrapping_add(2); + unsafe { score_4bit_internal_avx512_vnni(black_box(va), black_box(vb)) } + }); + }); + } + } + } + group.finish(); +} + +/// Cold-cache benchmarks for [`score_1bit_internal`] — vector-vs-vector +/// XOR+popcount scoring. Both operands are drawn from different shuffled +/// pool indices so each call pays two independent DRAM fetches. +fn bench_score_1bit_cold(c: &mut Criterion) { + let mut group = c.benchmark_group("query1bit_score_cold"); + for &dim in DIMS_1BIT { + let pool = VectorPool::new_1bit(dim, 7); + + group.throughput(Throughput::Elements(dim as u64)); + + group.bench_with_input(BenchmarkId::new("scalar", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let va = pool.vector(cursor); + let vb = pool.vector(cursor + 1); + cursor = cursor.wrapping_add(2); + score_1bit_internal_scalar(black_box(va), black_box(vb)) + }); + }); + + group.bench_with_input(BenchmarkId::new("dispatch", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let va = pool.vector(cursor); + let vb = pool.vector(cursor + 1); + cursor = cursor.wrapping_add(2); + score_1bit_internal(black_box(va), black_box(vb)) + }); + }); + + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + { + group.bench_with_input(BenchmarkId::new("neon", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let va = pool.vector(cursor); + let vb = pool.vector(cursor + 1); + cursor = cursor.wrapping_add(2); + unsafe { score_1bit_internal_neon(black_box(va), black_box(vb)) } + }); + }); + } + + #[cfg(target_arch = "x86_64")] + { + if std::is_x86_feature_detected!("sse4.1") && std::is_x86_feature_detected!("ssse3") { + group.bench_with_input(BenchmarkId::new("sse", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let va = pool.vector(cursor); + let vb = pool.vector(cursor + 1); + cursor = cursor.wrapping_add(2); + unsafe { score_1bit_internal_sse(black_box(va), black_box(vb)) } + }); + }); + } + + if std::is_x86_feature_detected!("avx2") { + group.bench_with_input(BenchmarkId::new("avx2", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let va = pool.vector(cursor); + let vb = pool.vector(cursor + 1); + cursor = cursor.wrapping_add(2); + unsafe { score_1bit_internal_avx2(black_box(va), black_box(vb)) } + }); + }); + } + + if std::is_x86_feature_detected!("avx512f") + && std::is_x86_feature_detected!("avx512vpopcntdq") + { + group.bench_with_input(BenchmarkId::new("avx512_vpopcntdq", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let va = pool.vector(cursor); + let vb = pool.vector(cursor + 1); + cursor = cursor.wrapping_add(2); + unsafe { + score_1bit_internal_avx512_vpopcntdq(black_box(va), black_box(vb)) + } + }); + }); + } + } + } + group.finish(); +} + +/// Query-against-data benchmarks: a single hot query is scored against cold +/// 1-bit PQ data vectors. Mirrors the HNSW scoring pattern. +/// +/// Compares our `Query1bitSimd<{8,12,16}>` (signed bit-plane transpose + +/// AND-popcount) against the existing BQ `Scalar8bits` path +/// (`BitsStoreType::xor_popcnt_scalar` with `bits_count=8`) — BQ stays at +/// 8 bits (its only supported scalar width) and serves as the baseline. +/// 12/16 rows show the linear cost of widening the query. +fn bench_query1bit_vs_bq_hot(c: &mut Criterion) { + let mut group = c.benchmark_group("query1bit_vs_bq_scalar8bits"); + let mut rng_seed = StdRng::seed_from_u64(42); + for &dim in DIMS_1BIT { + let pool = VectorPool::new_1bit(dim, 7); + let query_floats: Vec = (0..dim) + .map(|_| rng_seed.random_range(-1.0_f32..1.0)) + .collect(); + + let q_our_8 = Query1bitSimd::<8>::new(&query_floats); + let q_our_12 = Query1bitSimd::<12>::new(&query_floats); + let q_our_16 = Query1bitSimd::<16>::new(&query_floats); + let q_bq = encode_bq_scalar8bits(&query_floats); + + group.throughput(Throughput::Elements(dim as u64)); + + group.bench_with_input(BenchmarkId::new("query1bit_8bit", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let v = pool.vector(cursor); + cursor = cursor.wrapping_add(1); + q_our_8.dotprod(black_box(v)) + }); + }); + + group.bench_with_input(BenchmarkId::new("query1bit_12bit", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let v = pool.vector(cursor); + cursor = cursor.wrapping_add(1); + q_our_12.dotprod(black_box(v)) + }); + }); + + group.bench_with_input(BenchmarkId::new("query1bit_16bit", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let v = pool.vector(cursor); + cursor = cursor.wrapping_add(1); + q_our_16.dotprod(black_box(v)) + }); + }); + + group.bench_with_input(BenchmarkId::new("bq_scalar8bits", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let v = pool.vector(cursor); + cursor = cursor.wrapping_add(1); + ::xor_popcnt_scalar(black_box(v), black_box(&q_bq), 8) + }); + }); + } + group.finish(); +} + +/// Minimal reproduction of the BQ `Scalar8bits` encoding (u8 store, +/// `bits_count = 8`): per group of 8 consecutive query dims, emit 8 +/// consecutive bytes — one per bit-plane — containing those 8 dims' bits. +/// See `_encode_scalar_query_vector` in `encoded_vectors_binary.rs` for the +/// canonical version. +fn encode_bq_scalar8bits(query: &[f32]) -> Vec { + assert!(query.len().is_multiple_of(8)); + let max_abs = query + .iter() + .map(|x| x.abs()) + .fold(0.0_f32, f32::max) + .max(f32::EPSILON); + let min = -max_abs; + let delta = 2.0 * max_abs / 255.0; + + let mut encoded = vec![0_u8; query.len()]; + for (chunk_idx, chunk) in query.chunks(8).enumerate() { + for (shift, &value) in chunk.iter().enumerate() { + let q = ((value - min) / delta).round().clamp(0.0, 255.0) as usize; + for b in 0..8 { + let bit = ((q >> b) & 1) as u8; + encoded[8 * chunk_idx + b] |= bit << shift; + } + } + } + encoded +} + +/// Cold-cache query-vs-vector dotprod for 2-bit PQ. Mirrors +/// [`bench_dotprod_cold`] for 4-bit — a hot `Query2bitSimd` against cold +/// data vectors drawn from a shuffled 64 MB pool. +fn bench_dotprod_2bit_cold(c: &mut Criterion) { + let mut group = c.benchmark_group("query2bit_dotprod_cold"); + for &dim in DIMS_2BIT { + let q = make_query(dim); + let query = Query2bitSimd::new(&q); + let pool = VectorPool::new_2bit(dim, 7); + + group.throughput(Throughput::Elements(dim as u64)); + + group.bench_with_input(BenchmarkId::new("scalar", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let v = pool.vector(cursor); + cursor = cursor.wrapping_add(1); + black_box(&query).dotprod_raw(black_box(v)) + }); + }); + + group.bench_with_input(BenchmarkId::new("dotprod", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let v = pool.vector(cursor); + cursor = cursor.wrapping_add(1); + black_box(&query).dotprod(black_box(v)) + }); + }); + + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + { + group.bench_with_input(BenchmarkId::new("neon", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let v = pool.vector(cursor); + cursor = cursor.wrapping_add(1); + unsafe { black_box(&query).dotprod_raw_neon(black_box(v)) } + }); + }); + + if std::arch::is_aarch64_feature_detected!("dotprod") && dim.is_multiple_of(32) { + group.bench_with_input(BenchmarkId::new("neon_sdot", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let v = pool.vector(cursor); + cursor = cursor.wrapping_add(1); + unsafe { black_box(&query).dotprod_raw_neon_sdot(black_box(v)) } + }); + }); + } + } + + #[cfg(target_arch = "x86_64")] + { + if std::is_x86_feature_detected!("sse4.1") && std::is_x86_feature_detected!("ssse3") { + group.bench_with_input(BenchmarkId::new("sse", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let v = pool.vector(cursor); + cursor = cursor.wrapping_add(1); + unsafe { black_box(&query).dotprod_raw_sse(black_box(v)) } + }); + }); + } + if std::is_x86_feature_detected!("avx2") { + group.bench_with_input(BenchmarkId::new("avx2", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let v = pool.vector(cursor); + cursor = cursor.wrapping_add(1); + unsafe { black_box(&query).dotprod_raw_avx2(black_box(v)) } + }); + }); + } + if std::is_x86_feature_detected!("avx512f") + && std::is_x86_feature_detected!("avx512bw") + && std::is_x86_feature_detected!("avx512vnni") + { + group.bench_with_input(BenchmarkId::new("avx512_vnni", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let v = pool.vector(cursor); + cursor = cursor.wrapping_add(1); + unsafe { black_box(&query).dotprod_raw_avx512_vnni(black_box(v)) } + }); + }); + } + } + } + group.finish(); +} + +/// Cold-cache vector-vs-vector score for 2-bit PQ. Mirrors +/// [`bench_score_cold`] for 4-bit. +fn bench_score_2bit_cold(c: &mut Criterion) { + let mut group = c.benchmark_group("query2bit_score_cold"); + for &dim in DIMS_2BIT { + let pool = VectorPool::new_2bit(dim, 7); + + group.throughput(Throughput::Elements(dim as u64)); + + group.bench_with_input(BenchmarkId::new("scalar", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let va = pool.vector(cursor); + let vb = pool.vector(cursor + 1); + cursor = cursor.wrapping_add(2); + score_2bit_internal_scalar(black_box(va), black_box(vb)) + }); + }); + + group.bench_with_input(BenchmarkId::new("dispatch", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let va = pool.vector(cursor); + let vb = pool.vector(cursor + 1); + cursor = cursor.wrapping_add(2); + score_2bit_internal(black_box(va), black_box(vb)) + }); + }); + + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + { + group.bench_with_input(BenchmarkId::new("neon", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let va = pool.vector(cursor); + let vb = pool.vector(cursor + 1); + cursor = cursor.wrapping_add(2); + unsafe { score_2bit_internal_neon(black_box(va), black_box(vb)) } + }); + }); + + if std::arch::is_aarch64_feature_detected!("dotprod") { + group.bench_with_input(BenchmarkId::new("neon_sdot", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let va = pool.vector(cursor); + let vb = pool.vector(cursor + 1); + cursor = cursor.wrapping_add(2); + unsafe { score_2bit_internal_neon_sdot(black_box(va), black_box(vb)) } + }); + }); + } + } + + #[cfg(target_arch = "x86_64")] + { + if std::is_x86_feature_detected!("sse4.1") && std::is_x86_feature_detected!("ssse3") { + group.bench_with_input(BenchmarkId::new("sse", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let va = pool.vector(cursor); + let vb = pool.vector(cursor + 1); + cursor = cursor.wrapping_add(2); + unsafe { score_2bit_internal_sse(black_box(va), black_box(vb)) } + }); + }); + } + if std::is_x86_feature_detected!("avx2") { + group.bench_with_input(BenchmarkId::new("avx2", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let va = pool.vector(cursor); + let vb = pool.vector(cursor + 1); + cursor = cursor.wrapping_add(2); + unsafe { score_2bit_internal_avx2(black_box(va), black_box(vb)) } + }); + }); + } + if std::is_x86_feature_detected!("avx512f") + && std::is_x86_feature_detected!("avx512bw") + && std::is_x86_feature_detected!("avx512vnni") + { + group.bench_with_input(BenchmarkId::new("avx512_vnni", dim), &dim, |b, _| { + let mut cursor = 0usize; + b.iter(|| { + let va = pool.vector(cursor); + let vb = pool.vector(cursor + 1); + cursor = cursor.wrapping_add(2); + unsafe { score_2bit_internal_avx512_vnni(black_box(va), black_box(vb)) } + }); + }); + } + } + } + group.finish(); +} + +criterion_group!( + benches, + bench_dotprod_cold, + bench_score_cold, + bench_dotprod_2bit_cold, + bench_score_2bit_cold, + bench_score_1bit_cold, + bench_query1bit_vs_bq_hot, +); +criterion_main!(benches); diff --git a/lib/quantization/src/turboquant/encoding.rs b/lib/quantization/src/turboquant/encoding.rs index 22a753dfc6..edc97a87bb 100644 --- a/lib/quantization/src/turboquant/encoding.rs +++ b/lib/quantization/src/turboquant/encoding.rs @@ -59,6 +59,15 @@ impl TurboQuantizer { out } + /// Split an encoded vector into its dimension bytes and the decoded extras. + /// Shared between [`Self::unpack_vector`] (scalar scoring) and the SIMD + /// scoring paths that consume the raw packed bytes directly. + pub(super) fn split_vector<'a>(&self, vec: &'a [u8]) -> (TqVectorExtras, &'a [u8]) { + let extra_len = TqVectorExtras::size_for(self.bits, self.distance, self.mode); + let (dim_part, extra_part) = vec.split_at(vec.len() - extra_len); + (self.unpack_extras_from(extra_part), dim_part) + } + /// Unpacks `vec` into an iterator of `dim` centroid values and returns any /// `Extras` stored alongside. /// @@ -69,11 +78,7 @@ impl TurboQuantizer { /// Does not apply the inverse rotation — the iterator yields values in the /// rotated space. pub fn unpack_vector(&self, vec: &[u8]) -> (TqVectorExtras, impl Iterator) { - let extra_len = TqVectorExtras::size_for(self.bits, self.distance, self.mode); - - let (dim_part, extra_part) = vec.split_at(vec.len() - extra_len); - - let extras = self.unpack_extras_from(extra_part); + let (extras, dim_part) = self.split_vector(vec); let centroids = self.bits.get_centroids(); let mut reader = BitReader::new(dim_part); diff --git a/lib/quantization/src/turboquant/mod.rs b/lib/quantization/src/turboquant/mod.rs index 1fa28da6a6..e353ed79eb 100644 --- a/lib/quantization/src/turboquant/mod.rs +++ b/lib/quantization/src/turboquant/mod.rs @@ -3,6 +3,7 @@ pub mod lloyd_max; mod permutation; pub mod quantization; pub mod rotation; +pub mod simd; use std::alloc::Layout; use std::borrow::Cow; @@ -20,7 +21,8 @@ use serde::{Deserialize, Serialize}; use crate::EncodingError; use crate::encoded_storage::{EncodedStorage, EncodedStorageBuilder}; use crate::encoded_vectors::{EncodedVectors, VectorParameters, validate_vector_parameters}; -use crate::turboquant::quantization::{Precomputed, TurboQuantizer}; +use crate::turboquant::quantization::TurboQuantizer; +use crate::turboquant::simd::{Query1bitSimd, Query2bitSimd, Query4bitSimd}; #[derive(Serialize, Deserialize, Clone, Copy, Debug, PartialEq, Eq, Hash)] #[serde(rename_all = "snake_case")] @@ -73,9 +75,15 @@ pub struct EncodedQueryTQ { // TODO(turbo): add precomputed extras here when needed } +/// SIMD-ready encoded query, one variant per supported bit-width. Each +/// variant wraps the bit-width's [`simd::Query{N}bitSimd`] precomputation +/// (rotation-applied query, quantized to the SIMD-friendly integer form); +/// on architectures without a matching SIMD instruction set the scalar +/// reference kernel inside each type takes over automatically. pub enum EncodedQueryTQData { - Native(Precomputed), - // TODO(turbo): add other variants for SIMD-optimized precomputations, etc. + Bits1(Query1bitSimd), + Bits2(Query2bitSimd), + Bits4(Query4bitSimd), } #[derive(Serialize, Deserialize)] diff --git a/lib/quantization/src/turboquant/quantization.rs b/lib/quantization/src/turboquant/quantization.rs index 277270ac91..79c103be53 100644 --- a/lib/quantization/src/turboquant/quantization.rs +++ b/lib/quantization/src/turboquant/quantization.rs @@ -1,5 +1,9 @@ use crate::DistanceType; use crate::turboquant::rotation::HadamardRotation; +use crate::turboquant::simd::{ + Query1bitSimd, Query2bitSimd, Query4bitSimd, score_1bit_internal, score_2bit_internal, + score_4bit_internal, +}; use crate::turboquant::{EncodedQueryTQ, EncodedQueryTQData, Metadata, TQBits, TQMode}; /// Quantize vectors using TurboQuant. @@ -14,20 +18,6 @@ pub struct TurboQuantizer { dim_sqrt: f32, } -/// A query with the Hadamard rotation already applied. Built via -/// [`TurboQuantizer::precompute_query`] and consumed by -/// [`TurboQuantizer::score_precomputed`] to amortize the rotation cost across -/// many scoring calls. -pub struct Precomputed(Vec); - -impl Precomputed { - /// Borrow the rotated query components. - #[inline] - pub fn as_slice(&self) -> &[f64] { - &self.0 - } -} - impl TurboQuantizer { /// Initialize a new TurboQuantizer. pub fn new(dim: usize, bits: TQBits, mode: TQMode, distance: DistanceType) -> Self { @@ -97,9 +87,19 @@ impl TurboQuantizer { /// For asymmetric scoring (original vector against quantized vector), refer to [`Self::score_precomputed`] /// and precompute the query first. pub fn score_symmetric(&self, v1: &[u8], v2: &[u8]) -> f32 { - let (extra_v1, iter1) = self.unpack_vector(v1); - let (extra_v2, iter2) = self.unpack_vector(v2); - let raw_dot = dot_impl(iter1, iter2); + let (extra_v1, data_v1) = self.split_vector(v1); + let (extra_v2, data_v2) = self.split_vector(v2); + + // `score_{N}bit_internal` computes `Σ centroid(a_k) · centroid(b_k)` + // directly on packed bytes — same quantity the old + // `dot_impl(unpack_vector, unpack_vector)` f64 path produced, now via + // the bit-width's SIMD kernel (with a scalar fallback built in). + let raw_dot = match self.bits { + TQBits::Bits1 => score_1bit_internal(data_v1, data_v2), + TQBits::Bits1_5 => score_1bit_internal(data_v1, data_v2), + TQBits::Bits2 => score_2bit_internal(data_v1, data_v2), + TQBits::Bits4 => score_4bit_internal(data_v1, data_v2), + }; let v1_l2 = extra_v1.l2_length.unwrap_or(1.0); let v2_l2 = extra_v2.l2_length.unwrap_or(1.0); @@ -129,8 +129,9 @@ impl TurboQuantizer { } } - /// Precompute the Hadamard rotation of `query` so subsequent - /// [`Self::score_precomputed`] calls skip the per-call rotation. + /// Precompute the Hadamard rotation of `query` and hand it to the + /// bit-width's SIMD encoder. Subsequent [`Self::score_precomputed`] + /// calls reuse this precomputation — rotation runs once, not per score. pub fn precompute_query(&self, query: &[f32]) -> EncodedQueryTQ { debug_assert!(query.len() <= self.padded_dim); @@ -140,7 +141,6 @@ impl TurboQuantizer { .chain(std::iter::repeat(0.0)) .take(self.padded_dim) .collect(); - self.rotation.apply(&mut rotated); let l2_norm = match self.distance { @@ -155,8 +155,19 @@ impl TurboQuantizer { DistanceType::Cosine | DistanceType::Dot | DistanceType::L2 => None, }; + // SIMD encoders consume f32 (the quantization step will re-bucket into + // integer codebooks anyway, so the extra f64 precision from rotation + // has no downstream benefit here). + let rotated_f32: Vec = rotated.iter().map(|&x| x as f32).collect(); + + let data = match self.bits { + TQBits::Bits1 => EncodedQueryTQData::Bits1(Query1bitSimd::new(&rotated_f32)), + TQBits::Bits1_5 => EncodedQueryTQData::Bits1(Query1bitSimd::new(&rotated_f32)), + TQBits::Bits2 => EncodedQueryTQData::Bits2(Query2bitSimd::new(&rotated_f32)), + TQBits::Bits4 => EncodedQueryTQData::Bits4(Query4bitSimd::new(&rotated_f32)), + }; EncodedQueryTQ { - data: EncodedQueryTQData::Native(Precomputed(rotated)), + data, l2_norm, query, } @@ -166,11 +177,11 @@ impl TurboQuantizer { /// [`Self::precompute_query`]. Returns an approximate `` for Dot /// and `cos(θ)` for Cosine. pub fn score_precomputed(&self, query: &EncodedQueryTQ, vec: &[u8]) -> f32 { - let (vector_extras, unpacked) = self.unpack_vector(vec); + let (vector_extras, data_bytes) = self.split_vector(vec); let dot = match &query.data { - EncodedQueryTQData::Native(precomputed) => { - dot_impl(precomputed.as_slice().iter().copied(), unpacked) - } // TODO(turbo): add other variants for SIMD-optimized precomputations, etc. + EncodedQueryTQData::Bits1(q) => q.dotprod(data_bytes), + EncodedQueryTQData::Bits2(q) => q.dotprod(data_bytes), + EncodedQueryTQData::Bits4(q) => q.dotprod(data_bytes), }; match self.distance { @@ -206,16 +217,6 @@ impl TurboQuantizer { } } -/// Raw dot implementation between two vectors, `left` and `right`. -#[inline] -fn dot_impl(left: I, right: J) -> f32 -where - I: Iterator, - J: Iterator, -{ - left.zip(right).map(|(q, u)| q * u).sum::() as f32 -} - #[cfg(test)] mod tests { use common::bitpacking::BitReader; @@ -841,4 +842,64 @@ mod tests { } } } + + /// Sanity-check that [`TurboQuantizer::precompute_query`] + + /// [`TurboQuantizer::score_precomputed`] dispatch works for every + /// supported bit width. The precision-oriented tests above lock + /// `TQBits::Bits4`, so the `EncodedQueryTQData::Bits1`/`Bits2` dispatch + /// arms (and the Query1bitSimd / Query2bitSimd wiring behind them) would + /// otherwise only be exercised by each kernel's own module-level parity + /// tests, never through the integrated `precompute_query` path. + /// + /// The checks here are deliberately loose — 1-bit scoring discards almost + /// all amplitude info so tight `|got − truth| < ε` asserts would be + /// meaningless. We require: finite output, self-similarity positive, + /// antipodal negative, and a non-trivial gap between the two. + #[rstest::rstest] + #[case::bits1(TQBits::Bits1)] + #[case::bits2(TQBits::Bits2)] + #[case::bits4(TQBits::Bits4)] + fn score_precomputed_dispatches_all_bit_widths(#[case] bits: TQBits) { + let dim = 512; + + for &distance in &[DistanceType::Dot, DistanceType::Cosine] { + let mut rng = StdRng::seed_from_u64(0xD15_DA7C4); // same across distances → stable + let tq = make_tq(dim, bits, distance); + let mut buf = vec![0.0f64; tq.padded_dim]; + + let raw = random_vector(dim, &mut rng); + let v = match distance { + DistanceType::Cosine => normalize_vector(&raw), + _ => raw, + }; + let neg_v: Vec = v.iter().map(|&x| -x).collect(); + + let v_q = tq.quantize(&v, &mut buf); + let neg_q = tq.quantize(&neg_v, &mut buf); + + let self_score = asymmetric_score_helper(&tq, &v, &v_q); + let anti_score = asymmetric_score_helper(&tq, &v, &neg_q); + + assert!( + self_score.is_finite() && anti_score.is_finite(), + "non-finite score for {bits:?}/{distance:?}: self={self_score}, anti={anti_score}", + ); + assert!( + self_score > 0.0, + "self-similarity should be positive for {bits:?}/{distance:?}, got {self_score}", + ); + assert!( + anti_score < 0.0, + "antipodal score should be negative for {bits:?}/{distance:?}, got {anti_score}", + ); + // Gap must be large relative to the score magnitudes — guards + // against a constant-output or sign-swapped dispatch bug. + let gap = self_score - anti_score; + let ref_mag = self_score.abs().max(anti_score.abs()); + assert!( + gap > ref_mag, + "score spread too small for {bits:?}/{distance:?}: self={self_score}, anti={anti_score}", + ); + } + } } diff --git a/lib/quantization/src/turboquant/simd/mod.rs b/lib/quantization/src/turboquant/simd/mod.rs new file mode 100644 index 0000000000..92e7ae8958 --- /dev/null +++ b/lib/quantization/src/turboquant/simd/mod.rs @@ -0,0 +1,104 @@ +//! SIMD query-encoding + dot-product routines, one submodule per bit-width. +//! +//! Every `query{N}bit` submodule exposes two public entry points: +//! +//! * [`Query{N}bitSimd`](query4bit::Query4bitSimd) — a rotation-applied query +//! precomputed for fast asymmetric scoring (original-query × PQ-vector). +//! `dotprod(vector)` dispatches at runtime to the best SIMD backend available +//! on the host CPU. +//! * [`score_{N}bit_internal`](query4bit::score_4bit_internal) — dot product of +//! two already-encoded PQ vectors (symmetric scoring), same runtime dispatch. +//! +//! Available SIMD backends per bit-width: +//! +//! | Bits | x86_64 | aarch64 | +//! |------|-----------------------------------------------|----------------------| +//! | 1 | AVX-512 VPOPCNTDQ, AVX2, SSE4.1+SSSE3 | NEON | +//! | 2 | AVX-512 VNNI, AVX2, SSE4.1+SSSE3 | NEON + SDOT, NEON | +//! | 4 | AVX-512 VNNI, AVX2, SSE4.1+SSSE3 | NEON + SDOT, NEON | +//! +//! On any other target the scalar reference kernels in each module take over. + +pub mod query1bit; +pub mod query2bit; +pub mod query4bit; + +// Re-exports below include the runtime-dispatching entry points used by the +// crate's scoring paths (`Query{N}bitSimd`, `score_{N}bit_internal`) plus +// scalar-reference and arch-specific kernels the benchmarks at +// `benches/turbo_simd.rs` target directly. Every symbol here is consumed +// either by `turboquant::quantization` inside the crate or by benches/ +// outside — narrowing them to `pub(crate)` would break the bench build. +#[cfg(all(target_arch = "aarch64", target_feature = "neon"))] +pub use query1bit::score_1bit_internal_neon; +pub use query1bit::{Query1bitSimd, score_1bit_internal, score_1bit_internal_scalar}; +#[cfg(target_arch = "x86_64")] +pub use query1bit::{ + score_1bit_internal_avx2, score_1bit_internal_avx512_vpopcntdq, score_1bit_internal_sse, +}; +pub use query2bit::{Query2bitSimd, score_2bit_internal, score_2bit_internal_scalar}; +#[cfg(target_arch = "x86_64")] +pub use query2bit::{ + score_2bit_internal_avx2, score_2bit_internal_avx512_vnni, score_2bit_internal_sse, +}; +#[cfg(all(target_arch = "aarch64", target_feature = "neon"))] +pub use query2bit::{score_2bit_internal_neon, score_2bit_internal_neon_sdot}; +pub use query4bit::{Query4bitSimd, score_4bit_internal, score_4bit_internal_scalar}; +#[cfg(target_arch = "x86_64")] +pub use query4bit::{ + score_4bit_internal_avx2, score_4bit_internal_avx512_vnni, score_4bit_internal_sse, +}; +#[cfg(all(target_arch = "aarch64", target_feature = "neon"))] +pub use query4bit::{score_4bit_internal_neon, score_4bit_internal_neon_sdot}; + +/// Test-only helpers shared by every `query{N}bit` submodule. +/// +/// Per-bit-width specifics (`PARITY_DIMS`, `random_inputs`) still live in +/// their own module's `shared` — this holds only the genuinely +/// bit-width-agnostic pieces. +#[cfg(test)] +mod shared { + use common::bitpacking::BitWriter; + use rand::RngExt; + use rand::prelude::StdRng; + use rand_distr::{Distribution, StandardNormal}; + + /// Uniformly random bytes — used directly by 1-bit tests (where bytes *are* + /// the packed form) and indirectly via [`pack_codes`] for wider widths. + pub fn random_bytes(rng: &mut StdRng, len: usize) -> Vec { + (0..len).map(|_| rng.random_range(0..=u8::MAX)).collect() + } + + /// `len` samples from N(0, 1). + pub fn sample_normal_vec(rng: &mut StdRng, len: usize) -> Vec { + (0..len).map(|_| StandardNormal.sample(rng)).collect() + } + + /// Map each raw float to the index of its nearest centroid. + pub fn encode_to_nearest_centroid(centroids: &[f32], raw: &[f32]) -> Vec { + raw.iter() + .map(|&v| { + centroids + .iter() + .enumerate() + .min_by(|a, b| (a.1 - v).abs().partial_cmp(&(b.1 - v).abs()).unwrap()) + .map(|(k, _)| k as u8) + .unwrap() + }) + .collect() + } + + /// Pack `indices` into bytes with `bits` bits per code, LSB-first — same + /// layout [`crate::turboquant::encoding::TurboQuantizer::pack_vector`] + /// produces (both go through [`BitWriter`]). Caller guarantees every + /// index fits in `bits` bits and `indices.len() * bits` is a multiple of 8. + pub fn pack_codes(indices: &[u8], bits: u8) -> Vec { + let mut out = Vec::with_capacity((indices.len() * bits as usize).div_ceil(8)); + let mut writer = BitWriter::new(&mut out); + for &idx in indices { + writer.write(idx, bits); + } + writer.finish(); + out + } +} diff --git a/lib/quantization/src/turboquant/simd/query1bit/arm.rs b/lib/quantization/src/turboquant/simd/query1bit/arm.rs new file mode 100644 index 0000000000..bed5e7ff31 --- /dev/null +++ b/lib/quantization/src/turboquant/simd/query1bit/arm.rs @@ -0,0 +1,228 @@ +//! NEON popcount path for [`super::score_1bit_internal`]. +//! +//! `vcntq_u8` returns per-byte popcounts (each ≤ 8), which we widen through +//! `vpaddlq_u8` (u8 → u16) and `vpadalq_u16` (accumulate 8 × u16 into 4 × u32 +//! pairwise sums). Per 16-byte chunk each u32 lane gains at most +//! `2 · 16 = 32`, so `u32::MAX / 32 ≈ 134 M` chunks (~2 GB) before the +//! accumulator lanes could overflow — far beyond any realistic vector. + +/// NEON implementation of [`super::score_1bit_internal`]. +/// +/// # Safety +/// CPU must support the `neon` feature (always true on aarch64). +#[target_feature(enable = "neon")] +pub unsafe fn score_1bit_internal_neon(a: &[u8], b: &[u8]) -> f32 { + use core::arch::aarch64::*; + + assert_eq!( + a.len(), + b.len(), + "score_1bit_internal_neon: vector length mismatch ({} vs {})", + a.len(), + b.len(), + ); + + unsafe { + let mut acc = vdupq_n_u32(0); + let chunks = a.len() / 16; + for i in 0..chunks { + let va = vld1q_u8(a.as_ptr().add(i * 16)); + let vb = vld1q_u8(b.as_ptr().add(i * 16)); + let x = veorq_u8(va, vb); + let cnt8 = vcntq_u8(x); + let cnt16 = vpaddlq_u8(cnt8); + acc = vpadalq_u16(acc, cnt16); + } + + let mut popcnt = u64::from(vaddvq_u32(acc)); + + let tail_start = chunks * 16; + for i in tail_start..a.len() { + popcnt += u64::from((a[i] ^ b[i]).count_ones()); + } + + super::popcount_to_score(a.len(), popcnt) + } +} + +impl super::Query1bitSimd { + /// NEON implementation of [`super::Query1bitSimd::dotprod_raw`]. + /// + /// Per block: one 16-byte data load + `BITS` plane loads; each plane's + /// `vandq_u8 · vcntq_u8` pair is pair-added through `vpaddlq_u8` / + /// `vpadalq_u16` into a dedicated `uint32x4_t` accumulator. `BITS` + /// accumulators live in registers (≤ 16 of 32 vregs available), so the + /// inner loop is purely ALU. + /// + /// # Safety + /// CPU must support the `neon` feature (always true on aarch64). + #[target_feature(enable = "neon")] + pub unsafe fn dotprod_raw_neon(&self, vector: &[u8]) -> i64 { + use core::arch::aarch64::*; + + unsafe { + let mut acc: [uint32x4_t; BITS] = core::array::from_fn(|_| vdupq_n_u32(0)); + + // Main loop: full 128-dim blocks read directly from `vector`. + for block_idx in 0..self.num_full_blocks() { + let data = vld1q_u8(vector.as_ptr().add(block_idx * super::BLOCK_BYTES)); + let block_base = block_idx * BITS * super::BLOCK_BYTES; + for (b, acc_b) in acc.iter_mut().enumerate() { + let plane = vld1q_u8( + self.planes + .as_ptr() + .add(block_base + b * super::BLOCK_BYTES), + ); + let cnt = vcntq_u8(vandq_u8(data, plane)); + let cnt16 = vpaddlq_u8(cnt); + *acc_b = vpadalq_u16(*acc_b, cnt16); + } + } + + // Partial tail block via zero-padded stack buffer — same SIMD + // kernel, data bytes beyond `tail_bytes` are zero so they can't + // contribute to any plane's AND-popcount. + if let Some((buf, block_idx)) = self.tail_block_scratch(vector) { + let data = vld1q_u8(buf.as_ptr()); + let block_base = block_idx * BITS * super::BLOCK_BYTES; + for (b, acc_b) in acc.iter_mut().enumerate() { + let plane = vld1q_u8( + self.planes + .as_ptr() + .add(block_base + b * super::BLOCK_BYTES), + ); + let cnt = vcntq_u8(vandq_u8(data, plane)); + let cnt16 = vpaddlq_u8(cnt); + *acc_b = vpadalq_u16(*acc_b, cnt16); + } + } + + let mut v_dot_q: i64 = 0; + for (b, acc_b) in acc.iter().enumerate() { + let popcnt = u64::from(vaddvq_u32(*acc_b)); + let w_b: i64 = if b == BITS - 1 { + -(1i64 << (BITS - 1)) + } else { + 1i64 << b + }; + v_dot_q += w_b * popcnt as i64; + } + v_dot_q + } + } +} + +#[cfg(test)] +mod tests { + use rand::SeedableRng as _; + use rand::prelude::StdRng; + + use super::super::super::shared::random_bytes; + use super::super::shared::PARITY_BYTE_LENS; + use super::super::{score_1bit_internal_neon, score_1bit_internal_scalar}; + + #[test] + fn test_score_neon_matches_scalar() { + let mut rng = StdRng::seed_from_u64(7); + for &byte_len in PARITY_BYTE_LENS { + let a = random_bytes(&mut rng, byte_len); + let b = random_bytes(&mut rng, byte_len); + let scalar = score_1bit_internal_scalar(&a, &b); + let neon = unsafe { score_1bit_internal_neon(&a, &b) }; + assert_eq!( + scalar.to_bits(), + neon.to_bits(), + "scalar {scalar} != neon {neon} at byte_len {byte_len}", + ); + } + } + + /// Overflow safety: at 64 KiB (= 524 288 bits) with `a = all 0xFF`, + /// `b = all 0x00`, every bit disagrees → `popcnt = n_bits`. Scalar is + /// the reference (u64 accumulator throughout, can't overflow). NEON + /// must match exactly — a mismatch would indicate an intermediate + /// `u32`/`u16` lane saturated. + #[test] + fn test_score_neon_overflow_safety_64k() { + let byte_len = 65_536 / 8; + let a = vec![0xFF_u8; byte_len]; + let b = vec![0x00_u8; byte_len]; + + let scalar = score_1bit_internal_scalar(&a, &b); + let neon = unsafe { score_1bit_internal_neon(&a, &b) }; + assert_eq!( + scalar.to_bits(), + neon.to_bits(), + "neon disagrees at byte_len={byte_len}: scalar={scalar} neon={neon}", + ); + + // Also sanity-check the known closed-form: all bits disagree → + // sign_sum = −n_bits → score = −c² · n_bits. + let expected = -super::super::CENTROID_SQ * (byte_len * 8) as f32; + assert!((scalar - expected).abs() / expected.abs() < 1e-6); + } + + /// Parity for `Query1bitSimd::dotprod_raw_neon` vs the scalar kernel, + /// at a few `BITS` values and dims. Integer result must match bit-exactly. + #[test] + fn test_query_dotprod_neon_matches_scalar() { + use rand_distr::{Distribution, StandardNormal}; + + use super::super::Query1bitSimd; + + fn check(dim: usize, seed: u64) { + let mut rng = StdRng::seed_from_u64(seed); + let query: Vec = (0..dim).map(|_| StandardNormal.sample(&mut rng)).collect(); + let data = random_bytes(&mut rng, dim / 8); + let q = Query1bitSimd::::new(&query); + let scalar = q.dotprod_raw(&data); + let neon = unsafe { q.dotprod_raw_neon(&data) }; + assert_eq!( + scalar, neon, + "BITS={BITS} dim={dim}: scalar={scalar} neon={neon}" + ); + } + + // Corner-case dims exercising every tail size the 1-bit pipeline + // can produce (tail ∈ {0, 8, 16, …, 120}, always a multiple of 8): + // • 128, 256, 512, 1024, 2048 — full blocks, no tail. + // • 8, 64, 120 — zero blocks + tail-only scoring paths. + // • 136, 1032, 2040 — realistic matryoshka slices with tails. + // • 640, 768, 896 — Gemma/BGE-style matryoshka dims. + for &dim in &[ + 8usize, 64, 120, 128, 136, 256, 512, 640, 768, 896, 1024, 1032, 2040, 2048, + ] { + check::<8>(dim, 0xCAFE); + check::<10>(dim, 0xBEEF); + check::<12>(dim, 0xDEAD); + } + } + + /// Overflow safety for `dotprod_raw_neon` at dim=64K and max BITS + /// (quantization constants stressed to the extreme): all-1 data vs a + /// query scaled to saturate the signed range. Scalar is u64-accumulator + /// reference; NEON u32 per-plane accumulators must match exactly. + #[test] + fn test_query_dotprod_neon_overflow_safety_64k() { + use super::super::Query1bitSimd; + + let dim = 65_536; + // Query = all +1.0 float → maps to +max signed int in every lane. + let query = vec![1.0_f32; dim]; + let data = vec![0xFFu8; dim / 8]; + + let q8 = Query1bitSimd::<8>::new(&query); + assert_eq!( + q8.dotprod_raw(&data), + unsafe { q8.dotprod_raw_neon(&data) }, + "BITS=8 dim={dim}", + ); + + let q16 = Query1bitSimd::<16>::new(&query); + assert_eq!( + q16.dotprod_raw(&data), + unsafe { q16.dotprod_raw_neon(&data) }, + "BITS=16 dim={dim}", + ); + } +} diff --git a/lib/quantization/src/turboquant/simd/query1bit/mod.rs b/lib/quantization/src/turboquant/simd/query1bit/mod.rs new file mode 100644 index 0000000000..9ab52f5b09 --- /dev/null +++ b/lib/quantization/src/turboquant/simd/query1bit/mod.rs @@ -0,0 +1,472 @@ +//! 1-bit product-quantization scoring. +//! +//! The 1-bit codebook is `[-c, +c]` (see `CENTROIDS_1BIT`), so each code is a +//! single sign bit. Input vectors pack 8 lanes per byte. +//! +//! Two scoring paths live here: +//! +//! * [`score_1bit_internal`] — dot between two packed vectors. Collapses to +//! `c² · (n_match − n_mismatch)`, a plain `XOR + popcount`, so the per-arch +//! work is just picking the fastest popcount primitive (AVX-512 VPOPCNTDQ, +//! AVX2 / SSE `pshufb`-nibble lookup, NEON `vcntq_u8`). +//! +//! * [`Query1bitSimd`] — asymmetric scoring of an original query against +//! packed data. The query is quantized to `BITS`-bit signed integers +//! (default 8), bit-plane-transposed into a block-interleaved layout, and +//! scored per block as `Σ_b w_b · popcount(data AND plane_b)` where +//! `w_b = 2^b` for b < BITS−1 and `−2^(BITS−1)` for the sign plane. + +/// `|c|` for the 1-bit codebook — Lloyd-Max on N(0, 1) gives `sqrt(2/π)`. +/// Kept in sync with `CENTROIDS_1BIT` in `lloyd_max.rs` by a test below. +const CENTROID_ABS: f32 = 0.797_884_6; + +/// `c²` — converts the signed-agreement count into a centroid dot product: +/// each lane contributes `(±c)·(±c) = ±c²`, summing to `c² · sign_sum`. +const CENTROID_SQ: f32 = CENTROID_ABS * CENTROID_ABS; + +/// Block size for the bit-plane interleave: 16 packed bytes = 128 data dims. +/// Matches an XMM register; AVX2 processes 2 blocks per iter, AVX-512 takes 4. +const BLOCK_BYTES: usize = 16; + +/// Encoded query for 1-bit PQ scoring with `BITS`-bit signed query +/// quantization (default 8). +/// +/// # Encoding +/// Query is quantized to signed `BITS`-bit integers `q ∈ [−2^(BITS−1), +/// 2^(BITS−1) − 1]` in two's complement, then bit-plane-transposed and +/// block-interleaved: for each group of 128 dims (16 packed data bytes), +/// `BITS` × 16 consecutive bytes hold plane 0's 16 bytes, then plane 1's +/// 16 bytes, ..., then plane `BITS−1`'s 16 bytes. This lets scoring read +/// data once per block (16-byte SIMD load) and consume `BITS` contiguous +/// plane chunks in a tight inner loop. +/// +/// # Scoring +/// The signed dot product is +/// ```text +/// signed_dot = 2 · v_dot_q − Σ q_signed, +/// v_dot_q = Σ_b w_b · popcount(v_block AND plane_b), +/// w_b = 2^b for b ∈ 0..BITS−1, w_{BITS−1} = −2^(BITS−1). +/// ``` +/// The `Σ q_signed` correction is query-side (stored in `sum_q_signed`), +/// so no per-data-vector precomputation is needed. The final float output +/// is `(c / q_scale) · signed_dot`, folded into `postprocess_scale`. +pub struct Query1bitSimd { + /// Block-interleaved bit-planes. Length is `total_blocks · BITS · 16` + /// bytes, where `total_blocks = num_full_blocks + (tail_bytes > 0)`. + /// When a partial tail block is present it sits at the end, with bit-plane + /// bytes beyond `tail_bytes` zero-padded so `data AND plane = 0` for the + /// padding lanes. Same SIMD kernel handles full and partial blocks once + /// the data side is zero-padded into a 16-byte stack buffer. + planes: Vec, + /// Number of **full** 128-dim blocks (excludes the partial tail). + num_full_blocks: usize, + /// Bytes in the trailing partial block — `0` if dim is block-aligned, + /// `1..=15` otherwise (`tail_bytes · 8 = tail_dims`). + tail_bytes: u8, + /// `c / q_scale` — single scalar reconstruction factor. + postprocess_scale: f32, + /// `Σ q_signed` over **all** dims (full blocks + tail). + sum_q_signed: i64, +} + +impl Query1bitSimd { + /// Query dim must be a multiple of 8 (the 1-bit packing width). + /// Matryoshka-friendly — any such dim is accepted, with a tail of up to + /// 15 packed bytes (120 dims) handled by the same SIMD kernel as full + /// blocks via a zero-padded 16-byte scratch buffer. + pub fn new(data: &[f32]) -> Self { + assert!( + (2..=16).contains(&BITS), + "Query1bitSimd: BITS must be in [2, 16], got {BITS}", + ); + assert!( + data.len().is_multiple_of(8), + "Query1bitSimd: dim must be a multiple of 8 (got {})", + data.len(), + ); + + let q_abs_max_int = (1i64 << (BITS - 1)) - 1; + let q_abs_max = data + .iter() + .copied() + .map(f32::abs) + .fold(0.0_f32, f32::max) + .max(f32::EPSILON); + let q_scale = q_abs_max_int as f32 / q_abs_max; + let clamp_hi = q_abs_max_int as f32; + let clamp_lo = -clamp_hi; + + let encode = + |value: f32| -> i64 { (value * q_scale).round().clamp(clamp_lo, clamp_hi) as i64 }; + + let num_full_blocks = data.len() / (8 * BLOCK_BYTES); + let full_dims = num_full_blocks * 8 * BLOCK_BYTES; + let tail_dims = data.len() - full_dims; + debug_assert!(tail_dims < 8 * BLOCK_BYTES && tail_dims.is_multiple_of(8)); + let tail_bytes = tail_dims / 8; + let has_tail = tail_bytes > 0; + let total_blocks = num_full_blocks + usize::from(has_tail); + + let mut planes = vec![0u8; total_blocks * BITS * BLOCK_BYTES]; + let mut sum_q_signed: i64 = 0; + let bits_mask = (1u64 << BITS) - 1; + + // Helper to deposit a single dim's bits into the bit-plane layout. + let mut deposit = |q: i64, block_idx: usize, byte_in_block: usize, bit_in_byte: usize| { + let q_bits = (q as u64) & bits_mask; + let block_base = block_idx * BITS * BLOCK_BYTES; + for b in 0..BITS { + let bit = ((q_bits >> b) & 1) as u8; + planes[block_base + b * BLOCK_BYTES + byte_in_block] |= bit << bit_in_byte; + } + }; + + for block_idx in 0..num_full_blocks { + for byte_in_block in 0..BLOCK_BYTES { + for bit_in_byte in 0..8 { + let dim = block_idx * 8 * BLOCK_BYTES + byte_in_block * 8 + bit_in_byte; + let q = encode(data[dim]); + sum_q_signed += q; + deposit(q, block_idx, byte_in_block, bit_in_byte); + } + } + } + + // Partial tail block: same bit-plane encoding, padding bytes stay zero. + if has_tail { + for i in 0..tail_dims { + let q = encode(data[full_dims + i]); + sum_q_signed += q; + deposit(q, num_full_blocks, i / 8, i % 8); + } + } + + Self { + planes, + num_full_blocks, + tail_bytes: tail_bytes as u8, + postprocess_scale: CENTROID_ABS / q_scale, + sum_q_signed, + } + } + + /// Score the encoded query against a PQ-encoded `vector` (8 lanes / byte). + /// `vector.len()` must equal the original query `dim / 8` — full blocks + /// first, then up to 15 tail bytes. + pub fn dotprod(&self, vector: &[u8]) -> f32 { + let v_dot_q = self.dotprod_raw_best(vector); + let signed_dot = 2 * v_dot_q - self.sum_q_signed; + self.postprocess_scale * signed_dot as f32 + } + + #[inline] + fn dotprod_raw_best(&self, vector: &[u8]) -> i64 { + #[cfg(target_arch = "x86_64")] + { + if std::is_x86_feature_detected!("avx512vl") + && std::is_x86_feature_detected!("avx512vpopcntdq") + { + return unsafe { self.dotprod_raw_avx512_vpopcntdq(vector) }; + } + if std::is_x86_feature_detected!("sse4.1") && std::is_x86_feature_detected!("ssse3") { + return unsafe { self.dotprod_raw_sse(vector) }; + } + } + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + { + return unsafe { self.dotprod_raw_neon(vector) }; + } + #[allow(unreachable_code)] + self.dotprod_raw(vector) + } + + /// Integer kernel: `Σ v_j · q_signed[j]` across all dims — full blocks + /// decoded from the bit-plane AND-popcount form, plus scalar tail. + /// Reference implementation — SIMD variants in [`arm`] / [`x64`] must + /// match this bit-exactly. + pub fn dotprod_raw(&self, vector: &[u8]) -> i64 { + let mut v_dot_q: i64 = 0; + for block_idx in 0..self.num_full_blocks { + let data_block = &vector[block_idx * BLOCK_BYTES..(block_idx + 1) * BLOCK_BYTES]; + v_dot_q += self.score_block_scalar(data_block, block_idx); + } + if let Some((buf, block_idx)) = self.tail_block_scratch(vector) { + v_dot_q += self.score_block_scalar(&buf, block_idx); + } + v_dot_q + } + + /// Score a single 16-byte data block against the plane entries for + /// `block_idx`. Scalar reference used by [`Self::dotprod_raw`] and by + /// the shared tail-copy helpers in arch modules. + #[inline] + fn score_block_scalar(&self, data_block: &[u8], block_idx: usize) -> i64 { + let plane_base = block_idx * BITS * BLOCK_BYTES; + let mut v_dot_q: i64 = 0; + for b in 0..BITS { + let plane = &self.planes[plane_base + b * BLOCK_BYTES..][..BLOCK_BYTES]; + let mut c: u32 = 0; + for i in 0..BLOCK_BYTES { + c += (data_block[i] & plane[i]).count_ones(); + } + let w_b: i64 = if b == BITS - 1 { + -(1i64 << (BITS - 1)) + } else { + 1i64 << b + }; + v_dot_q += w_b * i64::from(c); + } + v_dot_q + } + + /// Copy `tail_bytes` bytes from `vector`'s trailing partial block into a + /// zero-padded 16-byte scratch buffer. Returns `None` if the query has + /// no tail; otherwise returns `Some((buf, block_idx))` — the block index + /// is where the SIMD backend should read plane bytes from. + #[inline] + pub(super) fn tail_block_scratch(&self, vector: &[u8]) -> Option<([u8; BLOCK_BYTES], usize)> { + if self.tail_bytes == 0 { + return None; + } + let mut buf = [0u8; BLOCK_BYTES]; + let tail_start = self.num_full_blocks * BLOCK_BYTES; + let tail_len = self.tail_bytes as usize; + buf[..tail_len].copy_from_slice(&vector[tail_start..tail_start + tail_len]); + Some((buf, self.num_full_blocks)) + } + + /// Number of full 128-dim blocks the SIMD main-loop should iterate over. + /// Exposed to arch backends so they don't depend on the internal field name. + #[inline] + pub(super) fn num_full_blocks(&self) -> usize { + self.num_full_blocks + } +} + +/// Dot product between two already-encoded 1-bit PQ vectors. +/// +/// Both `a` and `b` are packed 8 lanes per byte. Returns `c² · sign_sum` +/// where `sign_sum = n_bits − 2 · popcount(a ⊕ b)`: bits that agree +/// contribute `+c²`, bits that disagree contribute `−c²`. +/// +/// Dispatches at runtime to the fastest popcount backend available on the +/// host CPU; falls back to [`score_1bit_internal_scalar`] otherwise. +/// +/// # Panics +/// Panics if `a` and `b` have different lengths. +pub fn score_1bit_internal(a: &[u8], b: &[u8]) -> f32 { + assert_eq!( + a.len(), + b.len(), + "score_1bit_internal: vector length mismatch ({} vs {})", + a.len(), + b.len(), + ); + + #[cfg(target_arch = "x86_64")] + { + if std::is_x86_feature_detected!("avx512f") + && std::is_x86_feature_detected!("avx512vpopcntdq") + { + return unsafe { x64::score_1bit_internal_avx512_vpopcntdq(a, b) }; + } + if std::is_x86_feature_detected!("avx2") { + return unsafe { x64::score_1bit_internal_avx2(a, b) }; + } + if std::is_x86_feature_detected!("sse4.1") && std::is_x86_feature_detected!("ssse3") { + return unsafe { x64::score_1bit_internal_sse(a, b) }; + } + } + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + { + return unsafe { arm::score_1bit_internal_neon(a, b) }; + } + #[allow(unreachable_code)] + score_1bit_internal_scalar(a, b) +} + +/// Scalar reference implementation of [`score_1bit_internal`] — byte-wise +/// `XOR + count_ones`. Same result as the public function; used as the +/// parity baseline for SIMD tests and the ultimate fallback on +/// architectures without a SIMD variant. +pub fn score_1bit_internal_scalar(a: &[u8], b: &[u8]) -> f32 { + let mut popcnt: u64 = 0; + for (&ba, &bb) in a.iter().zip(b.iter()) { + popcnt += u64::from((ba ^ bb).count_ones()); + } + popcount_to_score(a.len(), popcnt) +} + +/// Convert `popcount(a ⊕ b)` over `byte_len` bytes into the centroid dot +/// product. Shared across scalar and SIMD paths so every backend agrees on +/// the final rounding. +#[inline] +pub(super) fn popcount_to_score(byte_len: usize, popcnt: u64) -> f32 { + let total_bits = (byte_len as i64) * 8; + let sign_sum = total_bits - 2 * (popcnt as i64); + CENTROID_SQ * sign_sum as f32 +} + +#[cfg(all(target_arch = "aarch64", target_feature = "neon"))] +mod arm; +#[cfg(target_arch = "x86_64")] +mod x64; + +#[cfg(all(target_arch = "aarch64", target_feature = "neon"))] +pub use arm::score_1bit_internal_neon; +#[cfg(target_arch = "x86_64")] +pub use x64::{ + score_1bit_internal_avx2, score_1bit_internal_avx512_vpopcntdq, score_1bit_internal_sse, +}; + +#[cfg(test)] +pub(super) mod shared { + /// Byte lengths used by parity tests. Cover: below SSE width, a single + /// SSE chunk, an AVX2 chunk, an AVX-512 chunk, two AVX-512 chunks + tail. + pub const PARITY_BYTE_LENS: &[usize] = &[1, 7, 8, 15, 16, 31, 32, 63, 64, 127, 128, 257, 513]; +} + +/// Accuracy / precision tests for `score_1bit_internal` and `Query1bitSimd`. +/// Per-arch SIMD parity tests live in the `arm` / `x64` submodules. +#[cfg(test)] +mod tests { + use rand::SeedableRng as _; + use rand::prelude::StdRng; + + use super::super::shared::random_bytes; + use super::*; + use crate::turboquant::TQBits; + + /// `CENTROID_ABS` must match the magnitude in `CENTROIDS_1BIT`. + #[test] + fn test_codebook_matches_lloyd_max() { + let centroids = TQBits::Bits1.get_centroids(); + assert_eq!(centroids.len(), 2); + assert!((centroids[0] + CENTROID_ABS).abs() < 1e-6); + assert!((centroids[1] - CENTROID_ABS).abs() < 1e-6); + } + + /// Scalar result equals the naive sum-of-centroid-products. + #[test] + fn test_scalar_matches_centroid_product() { + let mut rng = StdRng::seed_from_u64(0xC0FFEE); + for &byte_len in &[1usize, 4, 8, 9, 16, 32, 128, 257] { + let a = random_bytes(&mut rng, byte_len); + let b = random_bytes(&mut rng, byte_len); + + let mut expected = 0.0_f32; + for byte_idx in 0..byte_len { + for bit in 0..8 { + let sa = if (a[byte_idx] >> bit) & 1 == 1 { + CENTROID_ABS + } else { + -CENTROID_ABS + }; + let sb = if (b[byte_idx] >> bit) & 1 == 1 { + CENTROID_ABS + } else { + -CENTROID_ABS + }; + expected += sa * sb; + } + } + let got = score_1bit_internal_scalar(&a, &b); + assert!( + (expected - got).abs() < 1e-3, + "byte_len={byte_len} expected {expected} got {got}", + ); + } + } + + /// Dispatched result matches scalar bit-for-bit. + #[test] + fn test_dispatch_matches_scalar() { + let mut rng = StdRng::seed_from_u64(42); + for &byte_len in super::shared::PARITY_BYTE_LENS { + let a = random_bytes(&mut rng, byte_len); + let b = random_bytes(&mut rng, byte_len); + let s = score_1bit_internal_scalar(&a, &b); + let d = score_1bit_internal(&a, &b); + assert_eq!(s.to_bits(), d.to_bits(), "byte_len={byte_len}"); + } + } + + /// score(a, a) == +c² · n_bits, score(a, ¬a) == −c² · n_bits. + #[test] + fn test_score_extremes() { + let mut rng = StdRng::seed_from_u64(1); + let byte_len = 64; + let a = random_bytes(&mut rng, byte_len); + let not_a: Vec = a.iter().map(|&x| !x).collect(); + + let n_bits = (byte_len * 8) as f32; + let max = CENTROID_SQ * n_bits; + + assert!((score_1bit_internal(&a, &a) - max).abs() < 1e-3); + assert!((score_1bit_internal(&a, ¬_a) + max).abs() < 1e-3); + } + + /// Round-trip test: `Query1bitSimd::dotprod(data)` should approximate + /// the exact centroid-dot-product `Σ query_i · sign(v_i) · c`, modulo + /// query quantization noise. Uses BITS=8 (default) and BITS=12 to + /// sanity-check that widening actually reduces error. + /// + /// Parameterized over matryoshka-style corner-case dims — every case + /// exercises the tail path for BITS=8 and BITS=12 independently. + #[rstest::rstest] + #[case::full_blocks(1024)] + #[case::tail_only(120)] + #[case::block_plus_small_tail(136)] + #[case::block_plus_max_tail(1144)] // 8 blocks + 120 tail dims + #[case::matryoshka_640(640)] + #[case::matryoshka_768(768)] + #[case::matryoshka_896(896)] + fn test_query_dotprod_matches_reference(#[case] dim: usize) { + use rand_distr::{Distribution, StandardNormal}; + + let byte_len = dim / 8; + let mut rng = StdRng::seed_from_u64(1234); + + let query: Vec = (0..dim).map(|_| StandardNormal.sample(&mut rng)).collect(); + let data: Vec = random_bytes(&mut rng, byte_len); + + // Reference: decode data bits to ±c and compute plain float dot. + let mut expected = 0.0_f32; + for (i, &q_i) in query.iter().enumerate() { + let bit = (data[i / 8] >> (i % 8)) & 1; + let sign = if bit == 1 { + CENTROID_ABS + } else { + -CENTROID_ABS + }; + expected += q_i * sign; + } + + let q8 = Query1bitSimd::<8>::new(&query); + let got8 = q8.dotprod(&data); + let q12 = Query1bitSimd::<12>::new(&query); + let got12 = q12.dotprod(&data); + + // Compare against `|expected| + sqrt(dim)` so small dims (where + // `expected` may be near zero by chance) don't artificially blow up + // the relative error. + let scale = expected.abs().max((dim as f32).sqrt()); + let rel_err_8 = (got8 - expected).abs() / scale; + let rel_err_12 = (got12 - expected).abs() / scale; + + assert!( + rel_err_8 < 2e-2, + "dim={dim} BITS=8 rel_err={rel_err_8} (got {got8} vs {expected})" + ); + assert!( + rel_err_12 < 2e-3, + "dim={dim} BITS=12 rel_err={rel_err_12} (got {got12} vs {expected})" + ); + // Widening must shrink error (strictly, with overwhelming probability + // — tightens to `<=` here to be robust against tiny-dim rng noise). + assert!( + rel_err_12 <= rel_err_8, + "dim={dim} BITS=12 err {rel_err_12} should be ≤ BITS=8 err {rel_err_8}", + ); + } +} diff --git a/lib/quantization/src/turboquant/simd/query1bit/x64.rs b/lib/quantization/src/turboquant/simd/query1bit/x64.rs new file mode 100644 index 0000000000..1967a24567 --- /dev/null +++ b/lib/quantization/src/turboquant/simd/query1bit/x64.rs @@ -0,0 +1,502 @@ +//! x86_64 popcount paths for [`super::score_1bit_internal`]. +//! +//! Three backends, picked at runtime by the dispatcher in `mod.rs`: +//! +//! * **SSE4.1 / SSSE3** — `pshufb`-based nibble-lookup popcount (Muła): split +//! each byte into its two nibbles, `pshufb` a 16-entry popcount table, +//! sum. `psadbw` against zero horizontally sums 8 bytes into a u16 lane, +//! which we accumulate into a u64 pair. +//! +//! * **AVX2** — same Muła trick on 32-byte YMM registers. +//! +//! * **AVX-512 VPOPCNTDQ** — `vpopcntq` in hardware: 8 × u64 popcounts per +//! instruction, summed into a 512-bit u64 accumulator and reduced with +//! `_mm512_reduce_add_epi64`. +//! +//! All three use a u64 accumulator so no intermediate can saturate at any +//! reasonable vector size (see `test_score_*_overflow_safety_64k`). + +/// 16-byte popcount-of-nibble lookup table. Index = nibble value, +/// value = number of 1-bits. Broadcast to YMM/ZMM as needed. +const NIBBLE_POPCNT: [i8; 16] = [0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4]; + +/// Raw popcount of `a ⊕ b` using 16-byte SSE pshufb-nibble-lookup chunks +/// plus a scalar byte tail. Shared between [`score_1bit_internal_sse`] and +/// the tail path of [`score_1bit_internal_avx512_vpopcntdq`]. +/// +/// # Safety +/// `a.len() == b.len()`. CPU must support `ssse3` and `sse4.1`. +#[inline] +#[target_feature(enable = "sse4.1,ssse3")] +unsafe fn popcount_sse(a: &[u8], b: &[u8]) -> u64 { + use core::arch::x86_64::*; + + unsafe { + let lookup = _mm_loadu_si128(NIBBLE_POPCNT.as_ptr().cast::<__m128i>()); + let low_mask = _mm_set1_epi8(0x0F); + let zero = _mm_setzero_si128(); + let mut acc = _mm_setzero_si128(); + + let chunks = a.len() / 16; + for i in 0..chunks { + let va = _mm_loadu_si128(a.as_ptr().add(i * 16).cast::<__m128i>()); + let vb = _mm_loadu_si128(b.as_ptr().add(i * 16).cast::<__m128i>()); + let x = _mm_xor_si128(va, vb); + + let lo = _mm_and_si128(x, low_mask); + let hi = _mm_and_si128(_mm_srli_epi16(x, 4), low_mask); + let cnt_lo = _mm_shuffle_epi8(lookup, lo); + let cnt_hi = _mm_shuffle_epi8(lookup, hi); + // Per-byte popcount ≤ 8; two halves summed ≤ 16 per u8 — fits u8. + let cnt = _mm_add_epi8(cnt_lo, cnt_hi); + + // `psadbw(cnt, 0)` horizontally sums 8 bytes into each u64 lane + // (max 8 · 16 = 128 per lane per chunk — zero overflow risk). + acc = _mm_add_epi64(acc, _mm_sad_epu8(cnt, zero)); + } + + let lo = _mm_cvtsi128_si64(acc) as u64; + let hi = _mm_cvtsi128_si64(_mm_unpackhi_epi64(acc, acc)) as u64; + let mut popcnt = lo + hi; + + let tail_start = chunks * 16; + for i in tail_start..a.len() { + popcnt += u64::from((a[i] ^ b[i]).count_ones()); + } + popcnt + } +} + +/// SSE4.1 + SSSE3 implementation of [`super::score_1bit_internal`]. +/// +/// # Safety +/// CPU must support `ssse3` and `sse4.1`. +#[target_feature(enable = "sse4.1,ssse3")] +pub unsafe fn score_1bit_internal_sse(a: &[u8], b: &[u8]) -> f32 { + assert_eq!( + a.len(), + b.len(), + "score_1bit_internal_sse: vector length mismatch ({} vs {})", + a.len(), + b.len(), + ); + super::popcount_to_score(a.len(), unsafe { popcount_sse(a, b) }) +} + +/// AVX2 implementation of [`super::score_1bit_internal`]. +/// +/// Tail after the 32-byte bulk loop (up to 31 bytes) is routed through +/// [`popcount_sse`] — at most 1 SSE chunk + scalar bytes, still cheaper +/// than a 31-iteration scalar loop on short vectors. +/// +/// # Safety +/// CPU must support `avx2`, `ssse3`, and `sse4.1`. +#[target_feature(enable = "avx2,sse4.1,ssse3")] +pub unsafe fn score_1bit_internal_avx2(a: &[u8], b: &[u8]) -> f32 { + use core::arch::x86_64::*; + + assert_eq!( + a.len(), + b.len(), + "score_1bit_internal_avx2: vector length mismatch ({} vs {})", + a.len(), + b.len(), + ); + + unsafe { + // Broadcast the 16-byte lookup into both 128-bit halves — pshufb + // operates per-lane, so each half needs its own copy of the table. + let lookup_half = _mm_loadu_si128(NIBBLE_POPCNT.as_ptr().cast::<__m128i>()); + let lookup = _mm256_set_m128i(lookup_half, lookup_half); + let low_mask = _mm256_set1_epi8(0x0F); + let zero = _mm256_setzero_si256(); + let mut acc = _mm256_setzero_si256(); + + let chunks = a.len() / 32; + for i in 0..chunks { + let va = _mm256_loadu_si256(a.as_ptr().add(i * 32).cast::<__m256i>()); + let vb = _mm256_loadu_si256(b.as_ptr().add(i * 32).cast::<__m256i>()); + let x = _mm256_xor_si256(va, vb); + + let lo = _mm256_and_si256(x, low_mask); + let hi = _mm256_and_si256(_mm256_srli_epi16(x, 4), low_mask); + let cnt_lo = _mm256_shuffle_epi8(lookup, lo); + let cnt_hi = _mm256_shuffle_epi8(lookup, hi); + let cnt = _mm256_add_epi8(cnt_lo, cnt_hi); + + let sum64 = _mm256_sad_epu8(cnt, zero); + acc = _mm256_add_epi64(acc, sum64); + } + + // Reduce 4 × u64 → scalar u64 via two 128-bit halves. + let lo128 = _mm256_castsi256_si128(acc); + let hi128 = _mm256_extracti128_si256(acc, 1); + let sum128 = _mm_add_epi64(lo128, hi128); + let lo = _mm_cvtsi128_si64(sum128) as u64; + let hi = _mm_cvtsi128_si64(_mm_unpackhi_epi64(sum128, sum128)) as u64; + let mut popcnt = lo + hi; + + let tail_start = chunks * 32; + popcnt += popcount_sse(&a[tail_start..], &b[tail_start..]); + + super::popcount_to_score(a.len(), popcnt) + } +} + +/// AVX-512 VPOPCNTDQ implementation of [`super::score_1bit_internal`]. +/// +/// Tail after the 64-byte bulk loop (up to 63 bytes) is handled via the +/// [`popcount_sse`] helper — 3 SSE chunks + scalar bytes is ~10× cheaper +/// than a 63-iteration scalar loop when the tail is non-trivial. +/// +/// # Safety +/// CPU must support `avx512f`, `avx512vpopcntdq`, `ssse3`, and `sse4.1`. +#[target_feature(enable = "avx512f,avx512vpopcntdq,sse4.1,ssse3")] +pub unsafe fn score_1bit_internal_avx512_vpopcntdq(a: &[u8], b: &[u8]) -> f32 { + use core::arch::x86_64::*; + + assert_eq!( + a.len(), + b.len(), + "score_1bit_internal_avx512_vpopcntdq: vector length mismatch ({} vs {})", + a.len(), + b.len(), + ); + + unsafe { + let mut acc = _mm512_setzero_si512(); + let chunks = a.len() / 64; + for i in 0..chunks { + let va = _mm512_loadu_si512(a.as_ptr().add(i * 64).cast::<__m512i>()); + let vb = _mm512_loadu_si512(b.as_ptr().add(i * 64).cast::<__m512i>()); + let x = _mm512_xor_si512(va, vb); + let cnt = _mm512_popcnt_epi64(x); + acc = _mm512_add_epi64(acc, cnt); + } + let mut popcnt = _mm512_reduce_add_epi64(acc) as u64; + + let tail_start = chunks * 64; + popcnt += popcount_sse(&a[tail_start..], &b[tail_start..]); + + super::popcount_to_score(a.len(), popcnt) + } +} + +impl super::Query1bitSimd { + /// SSE4.1 + SSSE3 implementation of + /// [`super::Query1bitSimd::dotprod_raw`]. + /// + /// Per block: load 16-byte data chunk, then for each of `BITS` planes + /// `pshufb`-nibble-lookup popcount of `AND` reduced via `psadbw` into + /// u64 pair; accumulate in `[__m128i; BITS]` regs (one per plane). + /// + /// # Safety + /// CPU must support `ssse3` and `sse4.1`. + #[target_feature(enable = "sse4.1,ssse3")] + pub unsafe fn dotprod_raw_sse(&self, vector: &[u8]) -> i64 { + use core::arch::x86_64::*; + + unsafe { + let lookup = _mm_loadu_si128(NIBBLE_POPCNT.as_ptr().cast::<__m128i>()); + let low_mask = _mm_set1_epi8(0x0F); + let zero = _mm_setzero_si128(); + let mut acc: [__m128i; BITS] = core::array::from_fn(|_| _mm_setzero_si128()); + + // Main loop: full blocks from the vector directly. + for block_idx in 0..self.num_full_blocks() { + let data = _mm_loadu_si128( + vector + .as_ptr() + .add(block_idx * super::BLOCK_BYTES) + .cast::<__m128i>(), + ); + let block_base = block_idx * BITS * super::BLOCK_BYTES; + for (b, acc_b) in acc.iter_mut().enumerate() { + let plane = _mm_loadu_si128( + self.planes + .as_ptr() + .add(block_base + b * super::BLOCK_BYTES) + .cast::<__m128i>(), + ); + let x = _mm_and_si128(data, plane); + let lo = _mm_and_si128(x, low_mask); + let hi = _mm_and_si128(_mm_srli_epi16(x, 4), low_mask); + let cnt_lo = _mm_shuffle_epi8(lookup, lo); + let cnt_hi = _mm_shuffle_epi8(lookup, hi); + let cnt = _mm_add_epi8(cnt_lo, cnt_hi); + *acc_b = _mm_add_epi64(*acc_b, _mm_sad_epu8(cnt, zero)); + } + } + + // Partial tail block via zero-padded stack buffer. + if let Some((buf, block_idx)) = self.tail_block_scratch(vector) { + let data = _mm_loadu_si128(buf.as_ptr().cast::<__m128i>()); + let block_base = block_idx * BITS * super::BLOCK_BYTES; + for (b, acc_b) in acc.iter_mut().enumerate() { + let plane = _mm_loadu_si128( + self.planes + .as_ptr() + .add(block_base + b * super::BLOCK_BYTES) + .cast::<__m128i>(), + ); + let x = _mm_and_si128(data, plane); + let lo = _mm_and_si128(x, low_mask); + let hi = _mm_and_si128(_mm_srli_epi16(x, 4), low_mask); + let cnt_lo = _mm_shuffle_epi8(lookup, lo); + let cnt_hi = _mm_shuffle_epi8(lookup, hi); + let cnt = _mm_add_epi8(cnt_lo, cnt_hi); + *acc_b = _mm_add_epi64(*acc_b, _mm_sad_epu8(cnt, zero)); + } + } + + reduce_planes::(&acc) + } + } + + /// AVX-512 VPOPCNTDQ (on XMM via AVX-512VL) implementation of + /// [`super::Query1bitSimd::dotprod_raw`]. + /// + /// Replaces the Muła nibble-lookup with hardware `_mm_popcnt_epi64` + /// (one instruction per block per plane). Block stays at 16 bytes + /// since the interleave layout keeps plane chunks at 16-byte granularity. + /// + /// # Safety + /// CPU must support `avx512vl` and `avx512vpopcntdq` (and thus SSE2 for + /// the XMM load/store pairs). + #[target_feature(enable = "avx512vl,avx512vpopcntdq")] + pub unsafe fn dotprod_raw_avx512_vpopcntdq(&self, vector: &[u8]) -> i64 { + use core::arch::x86_64::*; + + unsafe { + let mut acc: [__m128i; BITS] = core::array::from_fn(|_| _mm_setzero_si128()); + + for block_idx in 0..self.num_full_blocks() { + let data = _mm_loadu_si128( + vector + .as_ptr() + .add(block_idx * super::BLOCK_BYTES) + .cast::<__m128i>(), + ); + let block_base = block_idx * BITS * super::BLOCK_BYTES; + for (b, acc_b) in acc.iter_mut().enumerate() { + let plane = _mm_loadu_si128( + self.planes + .as_ptr() + .add(block_base + b * super::BLOCK_BYTES) + .cast::<__m128i>(), + ); + let cnt = _mm_popcnt_epi64(_mm_and_si128(data, plane)); + *acc_b = _mm_add_epi64(*acc_b, cnt); + } + } + + // Partial tail block via zero-padded stack buffer. + if let Some((buf, block_idx)) = self.tail_block_scratch(vector) { + let data = _mm_loadu_si128(buf.as_ptr().cast::<__m128i>()); + let block_base = block_idx * BITS * super::BLOCK_BYTES; + for (b, acc_b) in acc.iter_mut().enumerate() { + let plane = _mm_loadu_si128( + self.planes + .as_ptr() + .add(block_base + b * super::BLOCK_BYTES) + .cast::<__m128i>(), + ); + let cnt = _mm_popcnt_epi64(_mm_and_si128(data, plane)); + *acc_b = _mm_add_epi64(*acc_b, cnt); + } + } + + reduce_planes::(&acc) + } + } +} + +/// Reduce `[__m128i; BITS]` plane accumulators (each holding 2 × u64 +/// popcount lanes) into the weighted `v_dot_q` integer sum. +/// +/// # Safety +/// Caller must have enabled at least SSE2 (true of every caller here). +#[inline] +#[target_feature(enable = "sse2")] +unsafe fn reduce_planes(acc: &[core::arch::x86_64::__m128i; BITS]) -> i64 { + use core::arch::x86_64::*; + + let mut v_dot_q: i64 = 0; + for (b, acc_b) in acc.iter().enumerate() { + let lo = _mm_cvtsi128_si64(*acc_b) as u64; + let hi = _mm_cvtsi128_si64(_mm_unpackhi_epi64(*acc_b, *acc_b)) as u64; + let popcnt = lo + hi; + let w_b: i64 = if b == BITS - 1 { + -(1i64 << (BITS - 1)) + } else { + 1i64 << b + }; + v_dot_q += w_b * popcnt as i64; + } + v_dot_q +} + +#[cfg(test)] +mod tests { + use rand::SeedableRng as _; + use rand::prelude::StdRng; + + use super::super::super::shared::random_bytes; + use super::super::score_1bit_internal_scalar; + use super::super::shared::PARITY_BYTE_LENS; + use super::*; + + #[test] + fn test_score_sse_matches_scalar() { + if !std::is_x86_feature_detected!("ssse3") || !std::is_x86_feature_detected!("sse4.1") { + return; + } + let mut rng = StdRng::seed_from_u64(7); + for &byte_len in PARITY_BYTE_LENS { + let a = random_bytes(&mut rng, byte_len); + let b = random_bytes(&mut rng, byte_len); + let scalar = score_1bit_internal_scalar(&a, &b); + let got = unsafe { score_1bit_internal_sse(&a, &b) }; + assert_eq!( + scalar.to_bits(), + got.to_bits(), + "sse mismatch at byte_len={byte_len}", + ); + } + } + + #[test] + fn test_score_avx2_matches_scalar() { + if !std::is_x86_feature_detected!("avx2") { + return; + } + let mut rng = StdRng::seed_from_u64(7); + for &byte_len in PARITY_BYTE_LENS { + let a = random_bytes(&mut rng, byte_len); + let b = random_bytes(&mut rng, byte_len); + let scalar = score_1bit_internal_scalar(&a, &b); + let got = unsafe { score_1bit_internal_avx2(&a, &b) }; + assert_eq!( + scalar.to_bits(), + got.to_bits(), + "avx2 mismatch at byte_len={byte_len}", + ); + } + } + + #[test] + fn test_score_avx512_vpopcntdq_matches_scalar() { + if !(std::is_x86_feature_detected!("avx512f") + && std::is_x86_feature_detected!("avx512vpopcntdq")) + { + return; + } + let mut rng = StdRng::seed_from_u64(7); + for &byte_len in PARITY_BYTE_LENS { + let a = random_bytes(&mut rng, byte_len); + let b = random_bytes(&mut rng, byte_len); + let scalar = score_1bit_internal_scalar(&a, &b); + let got = unsafe { score_1bit_internal_avx512_vpopcntdq(&a, &b) }; + assert_eq!( + scalar.to_bits(), + got.to_bits(), + "avx512 mismatch at byte_len={byte_len}", + ); + } + } + + /// Overflow safety at 64 KiB (524 288 bits) with `a = all 0xFF`, + /// `b = all 0x00` — every bit disagrees, so `popcnt = n_bits`. Each + /// SIMD path must match scalar (u64 throughout) exactly; a mismatch + /// would mean an intermediate `u16`/`u32`/u64 lane overflowed. + #[test] + fn test_score_overflow_safety_64k() { + let byte_len = 65_536 / 8; + let a = vec![0xFF_u8; byte_len]; + let b = vec![0x00_u8; byte_len]; + let scalar = score_1bit_internal_scalar(&a, &b); + + unsafe { + if std::is_x86_feature_detected!("ssse3") && std::is_x86_feature_detected!("sse4.1") { + let sse = score_1bit_internal_sse(&a, &b); + assert_eq!(scalar.to_bits(), sse.to_bits(), "sse overflow at 64k"); + } + if std::is_x86_feature_detected!("avx2") { + let avx2 = score_1bit_internal_avx2(&a, &b); + assert_eq!(scalar.to_bits(), avx2.to_bits(), "avx2 overflow at 64k"); + } + if std::is_x86_feature_detected!("avx512f") + && std::is_x86_feature_detected!("avx512vpopcntdq") + { + let avx512 = score_1bit_internal_avx512_vpopcntdq(&a, &b); + assert_eq!(scalar.to_bits(), avx512.to_bits(), "avx512 overflow at 64k"); + } + } + } + + /// Parity of `Query1bitSimd::dotprod_raw_{sse, avx512_vpopcntdq}` vs + /// the scalar kernel across several BITS values and dims. + #[test] + fn test_query_dotprod_x86_matches_scalar() { + use rand_distr::{Distribution, StandardNormal}; + + use super::super::Query1bitSimd; + + fn check(dim: usize, seed: u64) { + let mut rng = StdRng::seed_from_u64(seed); + let query: Vec = (0..dim).map(|_| StandardNormal.sample(&mut rng)).collect(); + let data = random_bytes(&mut rng, dim / 8); + let q = Query1bitSimd::::new(&query); + let scalar = q.dotprod_raw(&data); + + if std::is_x86_feature_detected!("ssse3") && std::is_x86_feature_detected!("sse4.1") { + let sse = unsafe { q.dotprod_raw_sse(&data) }; + assert_eq!(scalar, sse, "BITS={BITS} dim={dim}: sse mismatch"); + } + if std::is_x86_feature_detected!("avx512vl") + && std::is_x86_feature_detected!("avx512vpopcntdq") + { + let avx512 = unsafe { q.dotprod_raw_avx512_vpopcntdq(&data) }; + assert_eq!(scalar, avx512, "BITS={BITS} dim={dim}: avx512 mismatch"); + } + } + + for &dim in &[128usize, 256, 384, 512, 1024, 2048] { + check::<8>(dim, 0xCAFE); + check::<10>(dim, 0xBEEF); + check::<12>(dim, 0xDEAD); + } + } + + /// Overflow safety at dim=64K with max-magnitude query against all-1 + /// data. Each SIMD path (when available on the CPU) must match scalar + /// exactly; a mismatch would mean an intermediate `u32` per-plane + /// accumulator (or the u64 lane in the VPOPCNTDQ variant) saturated. + #[test] + fn test_query_dotprod_x86_overflow_safety_64k() { + use super::super::Query1bitSimd; + + let dim = 65_536; + let query = vec![1.0_f32; dim]; + let data = vec![0xFFu8; dim / 8]; + + fn check(query: &[f32], data: &[u8]) { + let q = Query1bitSimd::::new(query); + let scalar = q.dotprod_raw(data); + + if std::is_x86_feature_detected!("ssse3") && std::is_x86_feature_detected!("sse4.1") { + let sse = unsafe { q.dotprod_raw_sse(data) }; + assert_eq!(scalar, sse, "BITS={BITS} sse overflow at 64k"); + } + if std::is_x86_feature_detected!("avx512vl") + && std::is_x86_feature_detected!("avx512vpopcntdq") + { + let avx512 = unsafe { q.dotprod_raw_avx512_vpopcntdq(data) }; + assert_eq!(scalar, avx512, "BITS={BITS} avx512 overflow at 64k"); + } + } + + check::<8>(&query, &data); + check::<16>(&query, &data); + } +} diff --git a/lib/quantization/src/turboquant/simd/query2bit/arm.rs b/lib/quantization/src/turboquant/simd/query2bit/arm.rs new file mode 100644 index 0000000000..affc47690a --- /dev/null +++ b/lib/quantization/src/turboquant/simd/query2bit/arm.rs @@ -0,0 +1,420 @@ +//! NEON SIMD paths for [`Query2bitSimd`] on aarch64. +//! +//! Codebook storage mirrors [`super::query4bit::arm`] — full signed `i8` with +//! no offset — so `vmull_s8` / `sdot` run as true signed-signed multiplies. +//! +//! # Unpack trick +//! Each packed data byte holds 4 × 2-bit codes; consecutive pairs (`c0c1`, +//! `c2c3`) are exactly the low / high nibble of the byte. A nibble takes +//! 16 values → one-to-one with a 16-entry `vqtbl1_s8` table. We precompute +//! two 16-byte tables: +//! +//! * `PAIR_TABLE_EVEN[pair] = CODEBOOK_I8[pair & 0b11]` — "even code" +//! * `PAIR_TABLE_ODD[pair] = CODEBOOK_I8[(pair >> 2) & 0b11]` — "odd code" +//! +//! For 4 packed data bytes (= 16 codes = one query chunk): +//! 1. Split low/high nibbles of each byte (4 + 4 = 8 nibbles). +//! 2. Interleave them so lane `2i` = low nib of byte i, lane `2i+1` = high nib. +//! 3. Two `vqtbl1_s8` lookups yield 8 + 8 centroid bytes. +//! 4. `vzip1_s8` + `vzip2_s8` + `vcombine_s8` → 16 centroids in natural +//! dim order. +//! +//! Same pipeline as 4-bit from that point on (i8 × i8 → i16 + `vpadalq_s16`). + +use super::{CODEBOOK_I8, CODEBOOK_SCALE, QUERY_HIGH_COEF, Query2bitSimd}; + +/// `PAIR_TABLE_EVEN[nibble]` = `CODEBOOK_I8[nibble & 0b11]`. +const PAIR_TABLE_EVEN: [i8; 16] = { + let mut tbl = [0_i8; 16]; + let mut k = 0; + while k < 16 { + tbl[k] = CODEBOOK_I8[k & 0b11]; + k += 1; + } + tbl +}; + +/// `PAIR_TABLE_ODD[nibble]` = `CODEBOOK_I8[(nibble >> 2) & 0b11]`. +const PAIR_TABLE_ODD: [i8; 16] = { + let mut tbl = [0_i8; 16]; + let mut k = 0; + while k < 16 { + tbl[k] = CODEBOOK_I8[(k >> 2) & 0b11]; + k += 1; + } + tbl +}; + +/// Unpack 4 packed data bytes into a natural-order `int8x16_t` of 16 +/// centroid i8 values (one per code). +#[inline] +#[target_feature(enable = "neon")] +unsafe fn unpack_16_codes(bytes4: *const u8) -> core::arch::aarch64::int8x16_t { + use core::arch::aarch64::*; + unsafe { + let data = vreinterpret_u8_u32(vld1_dup_u32(bytes4.cast::())); + // data is uint8x8 = [b0, b1, b2, b3, b0, b1, b2, b3]; we only use low 4. + + let nibble_mask = vdup_n_u8(0x0F); + let lo_nibs = vand_u8(data, nibble_mask); // [b0&F, b1&F, b2&F, b3&F, ...] + let hi_nibs = vshr_n_u8(data, 4); + + // Interleave: [lo(b0), hi(b0), lo(b1), hi(b1), lo(b2), hi(b2), lo(b3), hi(b3)] + // `vzip1_u8` on the 8-lane inputs takes lanes [0..4] of each, zipped. + let pair_indices = vzip1_u8(lo_nibs, hi_nibs); + + let t_even = vld1q_s8(PAIR_TABLE_EVEN.as_ptr()); + let t_odd = vld1q_s8(PAIR_TABLE_ODD.as_ptr()); + + let c_even = vqtbl1_s8(t_even, pair_indices); // int8x8 + let c_odd = vqtbl1_s8(t_odd, pair_indices); // int8x8 + + // Natural order: [c0, c1, c2, c3, ..., c15] + let c_lo = vzip1_s8(c_even, c_odd); + let c_hi = vzip2_s8(c_even, c_odd); + vcombine_s8(c_lo, c_hi) + } +} + +impl Query2bitSimd { + /// NEON implementation of [`Query2bitSimd::dotprod_raw`]. + /// + /// # Safety + /// CPU must support the `neon` feature (always true on aarch64). + #[target_feature(enable = "neon")] + pub unsafe fn dotprod_raw_neon(&self, vector: &[u8]) -> i64 { + use core::arch::aarch64::*; + + assert_eq!( + vector.len(), + self.expected_vector_bytes(), + "Query2bitSimd::dotprod_raw_neon: vector length mismatch ({} vs expected {})", + vector.len(), + self.expected_vector_bytes(), + ); + + unsafe { + let mut acc_low = vdupq_n_s32(0); + let mut acc_high = vdupq_n_s32(0); + + for (chunk_idx, [low, high]) in self.query_data.iter().enumerate() { + let c = unpack_16_codes(vector.as_ptr().add(chunk_idx * 4)); + + let q_low = vld1q_s8(low.as_ptr()); + let q_high = vld1q_s8(high.as_ptr()); + + let prod_low_lo = vmull_s8(vget_low_s8(q_low), vget_low_s8(c)); + let prod_low_hi = vmull_high_s8(q_low, c); + let prod_high_lo = vmull_s8(vget_low_s8(q_high), vget_low_s8(c)); + let prod_high_hi = vmull_high_s8(q_high, c); + + acc_low = vpadalq_s16(acc_low, prod_low_lo); + acc_low = vpadalq_s16(acc_low, prod_low_hi); + acc_high = vpadalq_s16(acc_high, prod_high_lo); + acc_high = vpadalq_s16(acc_high, prod_high_hi); + } + + let full = + i64::from(vaddvq_s32(acc_low)) + QUERY_HIGH_COEF * i64::from(vaddvq_s32(acc_high)); + full + self.dotprod_raw_tail(vector) + } + } + + /// ARMv8.2-A Dot Product variant. Uses SDOT to sum 4 × i8 × i8 products + /// per i32 lane per instruction, with a 2× unroll (2 chunks per iter) + + /// a 1-chunk tail for odd chunk counts (allows `dim % 16 == 0`, matching + /// the scalar / plain-NEON contract). + /// + /// # Safety + /// CPU must support `neon` and `dotprod`. + #[target_feature(enable = "neon,dotprod")] + pub unsafe fn dotprod_raw_neon_sdot(&self, vector: &[u8]) -> i64 { + use core::arch::aarch64::*; + + assert_eq!( + vector.len(), + self.expected_vector_bytes(), + "Query2bitSimd::dotprod_raw_neon_sdot: vector length mismatch ({} vs expected {})", + vector.len(), + self.expected_vector_bytes(), + ); + + unsafe { + let mut acc_low_0 = vdupq_n_s32(0); + let mut acc_low_1 = vdupq_n_s32(0); + let mut acc_high_0 = vdupq_n_s32(0); + let mut acc_high_1 = vdupq_n_s32(0); + + let n_pairs = self.query_data.len() / 2; + let data_ptr = vector.as_ptr(); + for i in 0..n_pairs { + let [low_0, high_0] = &self.query_data[2 * i]; + let [low_1, high_1] = &self.query_data[2 * i + 1]; + + let c_0 = unpack_16_codes(data_ptr.add(8 * i)); + let c_1 = unpack_16_codes(data_ptr.add(8 * i + 4)); + + let q_low_0 = vld1q_s8(low_0.as_ptr()); + let q_high_0 = vld1q_s8(high_0.as_ptr()); + let q_low_1 = vld1q_s8(low_1.as_ptr()); + let q_high_1 = vld1q_s8(high_1.as_ptr()); + + core::arch::asm!( + "sdot {acc:v}.4s, {a:v}.16b, {b:v}.16b", + acc = inout(vreg) acc_low_0, + a = in(vreg) q_low_0, + b = in(vreg) c_0, + options(pure, nomem, nostack, preserves_flags), + ); + core::arch::asm!( + "sdot {acc:v}.4s, {a:v}.16b, {b:v}.16b", + acc = inout(vreg) acc_high_0, + a = in(vreg) q_high_0, + b = in(vreg) c_0, + options(pure, nomem, nostack, preserves_flags), + ); + core::arch::asm!( + "sdot {acc:v}.4s, {a:v}.16b, {b:v}.16b", + acc = inout(vreg) acc_low_1, + a = in(vreg) q_low_1, + b = in(vreg) c_1, + options(pure, nomem, nostack, preserves_flags), + ); + core::arch::asm!( + "sdot {acc:v}.4s, {a:v}.16b, {b:v}.16b", + acc = inout(vreg) acc_high_1, + a = in(vreg) q_high_1, + b = in(vreg) c_1, + options(pure, nomem, nostack, preserves_flags), + ); + } + + // Tail: odd chunk count → one extra chunk via single SDOT per half. + if self.query_data.len() % 2 == 1 { + let tail = 2 * n_pairs; + let [low_t, high_t] = &self.query_data[tail]; + let c_t = unpack_16_codes(data_ptr.add(4 * tail)); + let q_low_t = vld1q_s8(low_t.as_ptr()); + let q_high_t = vld1q_s8(high_t.as_ptr()); + + core::arch::asm!( + "sdot {acc:v}.4s, {a:v}.16b, {b:v}.16b", + acc = inout(vreg) acc_low_0, + a = in(vreg) q_low_t, + b = in(vreg) c_t, + options(pure, nomem, nostack, preserves_flags), + ); + core::arch::asm!( + "sdot {acc:v}.4s, {a:v}.16b, {b:v}.16b", + acc = inout(vreg) acc_high_0, + a = in(vreg) q_high_t, + b = in(vreg) c_t, + options(pure, nomem, nostack, preserves_flags), + ); + } + + let acc_low = vaddq_s32(acc_low_0, acc_low_1); + let acc_high = vaddq_s32(acc_high_0, acc_high_1); + let full = + i64::from(vaddvq_s32(acc_low)) + QUERY_HIGH_COEF * i64::from(vaddvq_s32(acc_high)); + full + self.dotprod_raw_tail(vector) + } + } +} + +/// NEON implementation of [`super::score_2bit_internal`] — vector × vector +/// centroid dot, both sides unpacked via the pair-table trick. +/// +/// # Safety +/// CPU must support `neon`. +#[target_feature(enable = "neon")] +pub unsafe fn score_2bit_internal_neon(a: &[u8], b: &[u8]) -> f32 { + use core::arch::aarch64::*; + + assert_eq!( + a.len(), + b.len(), + "score_2bit_internal_neon: vector length mismatch ({} vs {})", + a.len(), + b.len(), + ); + + unsafe { + let mut acc = vdupq_n_s32(0); + + let n_full = a.len() / 4; + for i in 0..n_full { + let c_a = unpack_16_codes(a.as_ptr().add(i * 4)); + let c_b = unpack_16_codes(b.as_ptr().add(i * 4)); + + let prod_lo = vmull_s8(vget_low_s8(c_a), vget_low_s8(c_b)); + let prod_hi = vmull_high_s8(c_a, c_b); + acc = vpadalq_s16(acc, prod_lo); + acc = vpadalq_s16(acc, prod_hi); + } + + let simd_bytes = n_full * 4; + let acc_i64 = i64::from(vaddvq_s32(acc)) + + super::score_2bit_internal_integer(&a[simd_bytes..], &b[simd_bytes..]); + acc_i64 as f32 / (CODEBOOK_SCALE * CODEBOOK_SCALE) + } +} + +/// SDOT variant of [`score_2bit_internal_neon`]. 2× chunk unroll mirrors +/// [`Query2bitSimd::dotprod_raw_neon_sdot`]. +/// +/// # Safety +/// CPU must support `neon` and `dotprod`. +#[target_feature(enable = "neon,dotprod")] +pub unsafe fn score_2bit_internal_neon_sdot(a: &[u8], b: &[u8]) -> f32 { + use core::arch::aarch64::*; + + assert_eq!( + a.len(), + b.len(), + "score_2bit_internal_neon_sdot: vector length mismatch ({} vs {})", + a.len(), + b.len(), + ); + unsafe { + let mut acc_0 = vdupq_n_s32(0); + let mut acc_1 = vdupq_n_s32(0); + + let n_pairs = a.len() / 8; + for i in 0..n_pairs { + let c_a_0 = unpack_16_codes(a.as_ptr().add(8 * i)); + let c_a_1 = unpack_16_codes(a.as_ptr().add(8 * i + 4)); + let c_b_0 = unpack_16_codes(b.as_ptr().add(8 * i)); + let c_b_1 = unpack_16_codes(b.as_ptr().add(8 * i + 4)); + + core::arch::asm!( + "sdot {acc:v}.4s, {a:v}.16b, {b:v}.16b", + acc = inout(vreg) acc_0, + a = in(vreg) c_a_0, + b = in(vreg) c_b_0, + options(pure, nomem, nostack, preserves_flags), + ); + core::arch::asm!( + "sdot {acc:v}.4s, {a:v}.16b, {b:v}.16b", + acc = inout(vreg) acc_1, + a = in(vreg) c_a_1, + b = in(vreg) c_b_1, + options(pure, nomem, nostack, preserves_flags), + ); + } + + // Odd 4-byte chunk leftover (a.len() % 8 ∈ {4..7}): run via vmull. + let mut acc_tail = vdupq_n_s32(0); + let mut offset = n_pairs * 8; + if a.len() - offset >= 4 { + let c_a = unpack_16_codes(a.as_ptr().add(offset)); + let c_b = unpack_16_codes(b.as_ptr().add(offset)); + let prod_lo = vmull_s8(vget_low_s8(c_a), vget_low_s8(c_b)); + let prod_hi = vmull_high_s8(c_a, c_b); + acc_tail = vpadalq_s16(acc_tail, prod_lo); + acc_tail = vpadalq_s16(acc_tail, prod_hi); + offset += 4; + } + + let acc = vaddq_s32(vaddq_s32(acc_0, acc_1), acc_tail); + let acc_i64 = i64::from(vaddvq_s32(acc)) + + super::score_2bit_internal_integer(&a[offset..], &b[offset..]); + acc_i64 as f32 / (CODEBOOK_SCALE * CODEBOOK_SCALE) + } +} + +#[cfg(test)] +mod tests { + use rand::SeedableRng as _; + use rand::prelude::StdRng; + + use super::super::super::shared::pack_codes; + use super::super::shared::{PARITY_DIMS, random_inputs}; + use super::super::{Query2bitSimd, score_2bit_internal_scalar}; + use super::{score_2bit_internal_neon, score_2bit_internal_neon_sdot}; + + #[test] + fn test_neon_matches_scalar() { + let mut rng = StdRng::seed_from_u64(7); + for &dim in PARITY_DIMS { + let (simd_query, vector) = random_inputs(&mut rng, dim); + let scalar = simd_query.dotprod_raw(&vector); + let neon = unsafe { simd_query.dotprod_raw_neon(&vector) }; + assert_eq!(scalar, neon, "scalar {scalar} != neon {neon} at dim {dim}"); + } + } + + #[test] + fn test_neon_sdot_matches_scalar() { + if !std::arch::is_aarch64_feature_detected!("dotprod") { + return; + } + let mut rng = StdRng::seed_from_u64(7); + for &dim in PARITY_DIMS { + let (simd_query, vector) = random_inputs(&mut rng, dim); + let scalar = simd_query.dotprod_raw(&vector); + let sdot = unsafe { simd_query.dotprod_raw_neon_sdot(&vector) }; + assert_eq!(scalar, sdot, "scalar {scalar} != sdot {sdot} at dim {dim}"); + } + } + + /// Saturation safety for query-vs-vector at dim=64K: query maxed out, + /// every lane pointing at CODEBOOK_I8[3] = +127. + #[test] + fn test_saturation_safety_64k() { + let dim = 65_536; + let query = vec![1.0_f32; dim]; + let indices: Vec = vec![3; dim]; // +127 centroid + let vector = pack_codes(&indices, 2); + + let q = Query2bitSimd::new(&query); + let scalar = q.dotprod_raw(&vector); + + unsafe { + let neon = q.dotprod_raw_neon(&vector); + assert_eq!(scalar, neon, "neon disagrees at dim={dim}"); + + if std::arch::is_aarch64_feature_detected!("dotprod") { + let sdot = q.dotprod_raw_neon_sdot(&vector); + assert_eq!(scalar, sdot, "sdot disagrees at dim={dim}"); + } + } + } + + #[test] + fn test_score_neon_matches_scalar() { + let mut rng = StdRng::seed_from_u64(7); + for &dim in PARITY_DIMS { + let (_, vec_a) = random_inputs(&mut rng, dim); + let (_, vec_b) = random_inputs(&mut rng, dim); + let scalar = score_2bit_internal_scalar(&vec_a, &vec_b); + + let neon = unsafe { score_2bit_internal_neon(&vec_a, &vec_b) }; + assert_eq!(scalar, neon, "score neon parity fail at dim {dim}"); + + if std::arch::is_aarch64_feature_detected!("dotprod") { + let sdot = unsafe { score_2bit_internal_neon_sdot(&vec_a, &vec_b) }; + assert_eq!(scalar, sdot, "score sdot parity fail at dim {dim}"); + } + } + } + + #[test] + fn test_score_saturation_safety_64k() { + let dim = 65_536; + let indices: Vec = vec![3; dim]; + let vec_a = pack_codes(&indices, 2); + let vec_b = pack_codes(&indices, 2); + + let scalar = score_2bit_internal_scalar(&vec_a, &vec_b); + unsafe { + let neon = score_2bit_internal_neon(&vec_a, &vec_b); + assert_eq!(scalar, neon, "score neon overflow at 64k"); + + if std::arch::is_aarch64_feature_detected!("dotprod") { + let sdot = score_2bit_internal_neon_sdot(&vec_a, &vec_b); + assert_eq!(scalar, sdot, "score sdot overflow at 64k"); + } + } + } +} diff --git a/lib/quantization/src/turboquant/simd/query2bit/mod.rs b/lib/quantization/src/turboquant/simd/query2bit/mod.rs new file mode 100644 index 0000000000..67ffcbfe56 --- /dev/null +++ b/lib/quantization/src/turboquant/simd/query2bit/mod.rs @@ -0,0 +1,515 @@ +//! 2-bit product-quantization scoring. +//! +//! Mirrors [`super::query4bit`]: query is quantized to two signed 7-bit +//! halves (combined via `q_signed = QUERY_HIGH_COEF · high + low`), the +//! codebook is the arch-native storage form (`CODEBOOK_I8` on aarch64, +//! `CODEBOOK_U8` on x86_64), and scoring uses bias-corrected integer +//! accumulation. The only real differences from 4-bit are: +//! +//! 1. 4 centroids instead of 16 — `CENTROIDS_2BIT` from `lloyd_max`. +//! 2. 4 codes packed per byte (2 bits each) instead of 2 nibbles per byte. +//! 3. The SIMD unpack uses a **pair-table** trick: a nibble of the packed +//! data byte encodes a pair of 2-bit codes (16 possible combinations), +//! which maps one-to-one to a 16-entry `vqtbl1q_s8` / `pshufb` table. +//! Two such lookups (even / odd centroid of each pair), zipped, give a +//! natural-order `int8x16` of 16 centroid bytes per 4 packed data bytes. + +/// `max|c|` over `CENTROIDS_2BIT` — the extreme centroid magnitude. +const CODEBOOK_ABS_MAX: f32 = 1.510; + +/// Signed `i8` codebook for aarch64: `c_scale = 127 / max|c|`, no offset. +#[cfg(all(target_arch = "aarch64", target_feature = "neon"))] +const CODEBOOK_I8: [i8; 4] = [-127, -38, 38, 127]; + +#[cfg(all(target_arch = "aarch64", target_feature = "neon"))] +const CODEBOOK_SCALE: f32 = 127.0 / CODEBOOK_ABS_MAX; + +#[cfg(all(target_arch = "aarch64", target_feature = "neon"))] +const CODEBOOK_OFFSET: i64 = 0; + +#[cfg(all(target_arch = "aarch64", target_feature = "neon"))] +const QUERY_ABS_MAX: f32 = 32639.0; + +#[cfg(all(target_arch = "aarch64", target_feature = "neon"))] +const QUERY_HIGH_COEF: i64 = 256; + +/// Unsigned `u8` codebook for x86_64: `c_u = c_signed + 128`. +#[cfg(target_arch = "x86_64")] +const CODEBOOK_U8: [u8; 4] = [0, 90, 166, 255]; + +#[cfg(target_arch = "x86_64")] +const CODEBOOK_OFFSET: i64 = 128; + +#[cfg(target_arch = "x86_64")] +const CODEBOOK_SCALE: f32 = 128.0 / CODEBOOK_ABS_MAX; + +#[cfg(target_arch = "x86_64")] +const QUERY_ABS_MAX: f32 = 8127.0; + +#[cfg(target_arch = "x86_64")] +const QUERY_HIGH_COEF: i64 = 128; + +// Fallback for architectures with neither NEON nor x86_64. Matches the +// x86_64 scheme so the scalar reference produces the same numeric result. +#[cfg(not(any( + all(target_arch = "aarch64", target_feature = "neon"), + target_arch = "x86_64", +)))] +const CODEBOOK_U8: [u8; 4] = [0, 90, 166, 255]; +#[cfg(not(any( + all(target_arch = "aarch64", target_feature = "neon"), + target_arch = "x86_64", +)))] +const CODEBOOK_OFFSET: i64 = 128; +#[cfg(not(any( + all(target_arch = "aarch64", target_feature = "neon"), + target_arch = "x86_64", +)))] +const CODEBOOK_SCALE: f32 = 128.0 / CODEBOOK_ABS_MAX; +#[cfg(not(any( + all(target_arch = "aarch64", target_feature = "neon"), + target_arch = "x86_64", +)))] +const QUERY_ABS_MAX: f32 = 8127.0; +#[cfg(not(any( + all(target_arch = "aarch64", target_feature = "neon"), + target_arch = "x86_64", +)))] +const QUERY_HIGH_COEF: i64 = 128; + +/// Read the codebook value at `idx` (a 2-bit code in `0..=3`) in arch-native +/// storage form. On aarch64 that is the signed `CODEBOOK_I8`; on x86_64 +/// (and the fallback) the unsigned `CODEBOOK_U8` — the `+OFFSET` shift is +/// unwound later by the query-side `bias_correction`. +#[inline] +fn codebook_value_i64(idx: u8) -> i64 { + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + { + i64::from(CODEBOOK_I8[idx as usize]) + } + #[cfg(not(all(target_arch = "aarch64", target_feature = "neon")))] + { + i64::from(CODEBOOK_U8[idx as usize]) + } +} + +/// Read the codebook value at `idx` as a **true signed** integer, regardless +/// of which arch's storage is active. Used by vector-vs-vector scoring. +#[inline] +fn codebook_signed_i64(idx: u8) -> i64 { + codebook_value_i64(idx) - CODEBOOK_OFFSET +} + +/// Encoded query for asymmetric 2-bit PQ scoring. +/// +/// # Encoding +/// The f32 query is quantized to signed integers +/// `q_signed ∈ [−QUERY_ABS_MAX, QUERY_ABS_MAX]` and split into two i8 halves +/// combined as `q_signed = QUERY_HIGH_COEF · high + low` (`K = 256` on +/// aarch64 for full-range i8 halves, `K = 128` on x86_64 to keep the x86 +/// `maddubs` pair sum inside i16). Storage is 16-dim chunks of `[low, high]` +/// plus a scalar-handled tail of up to 12 dims (`dim % 4 == 0`). +/// +/// # Scoring +/// `dotprod_raw = Σ_j q_signed[j] · c_raw[v[j]]` accumulated from the SIMD +/// kernel's chunk pass plus the scalar tail. The float result is +/// `postprocess_scale · (dot_raw − bias_correction)`, where `bias_correction` +/// absorbs the `+OFFSET` shift in the x86 unsigned-codebook layout and is 0 +/// on aarch64 (signed codebook, no shift). +pub struct Query2bitSimd { + /// Full 16-dim chunks of the query — each covers 16 dims → 4 packed data bytes. + query_data: Vec<[[i8; 16]; 2]>, + /// Trailing dims that didn't fill a 16-dim chunk — up to 12 (since + /// `dim % 4 == 0`: tail is one of 0, 4, 8, 12 dims). + tail_low: [i8; 12], + tail_high: [i8; 12], + /// Number of meaningful entries in the tail arrays (`0..=12`, multiple of 4). + tail_dims: u8, + /// `1 / (q_scale · c_scale)` — prefactor from integer to float dot product. + postprocess_scale: f32, + /// `CODEBOOK_OFFSET · Σ q_signed[j]` — subtracted from `dot_raw` to + /// recover the true signed dot. `0` on aarch64 (signed codebook). + bias_correction: i64, +} + +impl Query2bitSimd { + /// Query dim must be a multiple of 4 (the 2-bit packing width: four codes + /// per byte). Dims that don't fill a 16-dim chunk produce up to a 12-dim + /// tail handled scalar-wise in every SIMD path — Matryoshka-friendly. + pub fn new(data: &[f32]) -> Self { + assert!( + data.len().is_multiple_of(4), + "Query2bitSimd requires query dim to be a multiple of 4 (got {})", + data.len(), + ); + + let q_abs_max = data + .iter() + .copied() + .map(f32::abs) + .fold(0.0_f32, f32::max) + .max(f32::EPSILON); + let q_scale = QUERY_ABS_MAX / q_abs_max; + + let k = QUERY_HIGH_COEF as i32; + let half_k = k / 2; + let clamp_hi = QUERY_ABS_MAX; + let clamp_lo = -QUERY_ABS_MAX; + + let encode = |value: f32| -> (i8, i8, i64) { + let q_signed = (value * q_scale).round().clamp(clamp_lo, clamp_hi) as i32; + let l_mod = q_signed.rem_euclid(k); + let l = if l_mod >= half_k { l_mod - k } else { l_mod } as i8; + let h = ((q_signed - i32::from(l)) / k) as i8; + (l, h, i64::from(q_signed)) + }; + + let num_chunks = data.len() / 16; + let full_dims = num_chunks * 16; + let tail_dims = data.len() - full_dims; + debug_assert!(tail_dims < 16 && tail_dims.is_multiple_of(4)); + + let mut query_data: Vec<[[i8; 16]; 2]> = Vec::with_capacity(num_chunks); + let mut sum_q_signed: i64 = 0; + for chunk_idx in 0..num_chunks { + let mut low = [0_i8; 16]; + let mut high = [0_i8; 16]; + for i in 0..16 { + let (l, h, q) = encode(data[chunk_idx * 16 + i]); + low[i] = l; + high[i] = h; + sum_q_signed += q; + } + query_data.push([low, high]); + } + + let mut tail_low = [0_i8; 12]; + let mut tail_high = [0_i8; 12]; + for i in 0..tail_dims { + let (l, h, q) = encode(data[full_dims + i]); + tail_low[i] = l; + tail_high[i] = h; + sum_q_signed += q; + } + + Self { + query_data, + tail_low, + tail_high, + tail_dims: tail_dims as u8, + postprocess_scale: 1.0 / (q_scale * CODEBOOK_SCALE), + bias_correction: CODEBOOK_OFFSET * sum_q_signed, + } + } + + /// Number of `vector` bytes the encoded query expects: 4 bytes per full + /// 16-dim chunk plus the packed tail (four 2-bit codes per byte). + #[inline] + pub(super) fn expected_vector_bytes(&self) -> usize { + self.query_data.len() * 4 + (self.tail_dims as usize).div_ceil(4) + } + + /// Score the encoded query against a 2-bit PQ-encoded `vector` + /// (four centroid indices per byte; bits `[2k..2k+2]` hold the code for + /// lane `k ∈ 0..=3`). `vector.len()` must equal `ceil(dim * 2 / 8)` — + /// full chunks first, then the tail bytes covering `tail_dims`. + /// + /// Dispatches at runtime to the best SIMD backend available on the host + /// CPU (AVX-512 VNNI → AVX2 → SSE → NEON + SDOT → NEON → scalar). + pub fn dotprod(&self, vector: &[u8]) -> f32 { + let dot_raw = self.dotprod_raw_best(vector); + self.postprocess_scale * (dot_raw - self.bias_correction) as f32 + } + + #[inline] + fn dotprod_raw_best(&self, vector: &[u8]) -> i64 { + #[cfg(target_arch = "x86_64")] + { + if std::is_x86_feature_detected!("avx512f") + && std::is_x86_feature_detected!("avx512bw") + && std::is_x86_feature_detected!("avx512vnni") + { + return unsafe { self.dotprod_raw_avx512_vnni(vector) }; + } + if std::is_x86_feature_detected!("avx2") { + return unsafe { self.dotprod_raw_avx2(vector) }; + } + if std::is_x86_feature_detected!("sse4.1") && std::is_x86_feature_detected!("ssse3") { + return unsafe { self.dotprod_raw_sse(vector) }; + } + } + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + { + if std::arch::is_aarch64_feature_detected!("dotprod") { + return unsafe { self.dotprod_raw_neon_sdot(vector) }; + } + return unsafe { self.dotprod_raw_neon(vector) }; + } + #[allow(unreachable_code)] + self.dotprod_raw(vector) + } + + /// Compute `Σ q_signed[j] · c_raw[v[j]]` across all dims — full chunks + /// plus tail. Returns `acc_low + QUERY_HIGH_COEF · acc_high`. + /// + /// `vector` is PQ-encoded with four 2-bit codes packed per byte: + /// bits `[2k..2k+2]` for `k ∈ 0..=3` hold codes `0..=3` of the byte, + /// in low-to-high bit order. + pub fn dotprod_raw(&self, vector: &[u8]) -> i64 { + assert_eq!( + vector.len(), + self.expected_vector_bytes(), + "Query2bitSimd::dotprod_raw: vector length mismatch ({} vs expected {})", + vector.len(), + self.expected_vector_bytes(), + ); + let mut acc_low: i64 = 0; + let mut acc_high: i64 = 0; + for (chunk_idx, [low, high]) in self.query_data.iter().enumerate() { + let v = &vector[chunk_idx * 4..(chunk_idx + 1) * 4]; + for i in 0..16 { + let byte = v[i / 4]; + let shift = 2 * (i % 4); + let idx = (byte >> shift) & 0x03; + let c = codebook_value_i64(idx); + acc_low += i64::from(low[i]) * c; + acc_high += i64::from(high[i]) * c; + } + } + acc_low + QUERY_HIGH_COEF * acc_high + self.dotprod_raw_tail(vector) + } + + /// Scalar contribution from the trailing `tail_dims` query entries. + /// Shared by every SIMD backend so they only implement the full-chunk + /// loop and forward the tail here. + #[inline] + pub(super) fn dotprod_raw_tail(&self, vector: &[u8]) -> i64 { + if self.tail_dims == 0 { + return 0; + } + let tail_byte_start = self.query_data.len() * 4; + let mut acc_low: i64 = 0; + let mut acc_high: i64 = 0; + for i in 0..self.tail_dims as usize { + let byte = vector[tail_byte_start + i / 4]; + let shift = 2 * (i % 4); + let idx = (byte >> shift) & 0x03; + let c = codebook_value_i64(idx); + acc_low += i64::from(self.tail_low[i]) * c; + acc_high += i64::from(self.tail_high[i]) * c; + } + acc_low + QUERY_HIGH_COEF * acc_high + } +} + +/// Dot product between two already-encoded 2-bit PQ vectors. Any byte length +/// is accepted — bytes beyond the last SIMD chunk are folded in scalar-wise. +/// +/// # Panics +/// Panics if the two vectors have different lengths. +pub fn score_2bit_internal(a: &[u8], b: &[u8]) -> f32 { + assert_eq!( + a.len(), + b.len(), + "score_2bit_internal: vector length mismatch ({} vs {})", + a.len(), + b.len(), + ); + + #[cfg(target_arch = "x86_64")] + { + if std::is_x86_feature_detected!("avx512f") + && std::is_x86_feature_detected!("avx512bw") + && std::is_x86_feature_detected!("avx512vnni") + { + return unsafe { x64::score_2bit_internal_avx512_vnni(a, b) }; + } + if std::is_x86_feature_detected!("avx2") { + return unsafe { x64::score_2bit_internal_avx2(a, b) }; + } + if std::is_x86_feature_detected!("sse4.1") && std::is_x86_feature_detected!("ssse3") { + return unsafe { x64::score_2bit_internal_sse(a, b) }; + } + } + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + { + if std::arch::is_aarch64_feature_detected!("dotprod") { + return unsafe { arm::score_2bit_internal_neon_sdot(a, b) }; + } + return unsafe { arm::score_2bit_internal_neon(a, b) }; + } + #[allow(unreachable_code)] + score_2bit_internal_scalar(a, b) +} + +/// Scalar reference for [`score_2bit_internal`] — see the 4-bit counterpart +/// for the design rationale. +pub fn score_2bit_internal_scalar(a: &[u8], b: &[u8]) -> f32 { + score_2bit_internal_integer(a, b) as f32 / (CODEBOOK_SCALE * CODEBOOK_SCALE) +} + +/// Integer-only scalar kernel — used by SIMD paths to fold in bytes that +/// didn't fit a full SIMD chunk, and by [`score_2bit_internal_scalar`] as its +/// inner loop. +#[inline] +pub(super) fn score_2bit_internal_integer(a: &[u8], b: &[u8]) -> i64 { + let mut acc: i64 = 0; + for (&byte_a, &byte_b) in a.iter().zip(b.iter()) { + for k in 0..4 { + let shift = 2 * k; + let a_k = (byte_a >> shift) & 0x03; + let b_k = (byte_b >> shift) & 0x03; + acc += codebook_signed_i64(a_k) * codebook_signed_i64(b_k); + } + } + acc +} + +#[cfg(all(target_arch = "aarch64", target_feature = "neon"))] +mod arm; + +#[cfg(target_arch = "x86_64")] +mod x64; + +#[cfg(all(target_arch = "aarch64", target_feature = "neon"))] +pub use arm::{score_2bit_internal_neon, score_2bit_internal_neon_sdot}; +#[cfg(target_arch = "x86_64")] +pub use x64::{score_2bit_internal_avx2, score_2bit_internal_avx512_vnni, score_2bit_internal_sse}; + +/// 2-bit-specific test helpers. Bit-width-agnostic helpers (`pack_codes`, +/// `sample_normal_vec`, `encode_to_nearest_centroid`) live in +/// [`super::super::shared`]. +#[cfg(test)] +pub(super) mod shared { + use rand::prelude::StdRng; + use rand::seq::SliceRandom; + + use super::super::shared::{pack_codes, sample_normal_vec}; + use super::Query2bitSimd; + + /// Corner-case dims covering every tail size the 2-bit pipeline can + /// produce (tail is 0, 4, 8, or 12 dims since `dim % 4 == 0`): + /// • `16, 64, 128, 256, 1024, 2048` — full chunks, no tail. + /// • `48` — 3 chunks (odd for SDOT/AVX2/AVX-512 unrolls), no tail. + /// • `20, 28, 44, 60, 1028, 2044` — full chunks + 4/12/12/12/4/12-dim tail. + /// • `268` — 16 chunks + 12-dim tail (realistic matryoshka). + pub const PARITY_DIMS: &[usize] = &[ + 16, 20, 28, 32, 44, 48, 60, 64, 128, 256, 268, 1024, 1028, 2044, 2048, + ]; + + /// Parity-test helper: query ~ N(0, 1), balanced index distribution. + pub fn random_inputs(rng: &mut StdRng, dim: usize) -> (Query2bitSimd, Vec) { + let query = sample_normal_vec(rng, dim); + let mut indices: Vec = (0..dim).map(|i| (i % 4) as u8).collect(); + indices.shuffle(rng); + (Query2bitSimd::new(&query), pack_codes(&indices, 2)) + } +} + +/// Accuracy / precision tests for `Query2bitSimd` and `score_2bit_internal`. +/// Per-arch SIMD parity tests live in the `arm` / `x64` submodules. +#[cfg(test)] +mod tests { + use rand::SeedableRng as _; + use rand::prelude::StdRng; + + use super::super::shared::{encode_to_nearest_centroid, pack_codes, sample_normal_vec}; + use super::{CODEBOOK_ABS_MAX, Query2bitSimd, score_2bit_internal_scalar}; + use crate::turboquant::TQBits; + + #[test] + fn test_codebook_matches_lloyd_max() { + let centroids = TQBits::Bits2.get_centroids(); + assert_eq!(centroids.len(), 4); + + let c_abs_max = centroids + .iter() + .copied() + .map(f32::abs) + .fold(0.0_f32, f32::max); + assert!( + (CODEBOOK_ABS_MAX - c_abs_max).abs() < 1e-6, + "CODEBOOK_ABS_MAX ({CODEBOOK_ABS_MAX}) != max|CENTROIDS_2BIT| ({c_abs_max})" + ); + + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + { + let c_scale = 127.0 / c_abs_max; + let quantized: [i8; 4] = std::array::from_fn(|k| { + (centroids[k] * c_scale).round().clamp(-127.0, 127.0) as i8 + }); + assert_eq!(quantized, super::CODEBOOK_I8); + } + #[cfg(not(all(target_arch = "aarch64", target_feature = "neon")))] + { + let c_scale = 128.0 / c_abs_max; + let offset = super::CODEBOOK_OFFSET as i32; + let quantized: [u8; 4] = std::array::from_fn(|k| { + let signed = (centroids[k] * c_scale).round().clamp(-128.0, 127.0) as i32; + (signed + offset) as u8 + }); + assert_eq!(quantized, super::CODEBOOK_U8); + } + } + + #[rstest::rstest] + #[case::full_chunks(256)] + #[case::small_tail(20)] + #[case::max_tail(28)] + #[case::odd_chunks_only(48)] + #[case::odd_chunks_plus_tail(60)] + #[case::matryoshka(268)] + #[case::large_with_tail(2044)] + fn test_dotprod_matches_float(#[case] dim: usize) { + let mut rng = StdRng::seed_from_u64(42); + let n_trials = 64; + + let centroids = TQBits::Bits2.get_centroids(); + + for _ in 0..n_trials { + let query = sample_normal_vec(&mut rng, dim); + let v_raw = sample_normal_vec(&mut rng, dim); + let indices = encode_to_nearest_centroid(centroids, &v_raw); + let v_pq: Vec = indices.iter().map(|&k| centroids[k as usize]).collect(); + + let pq_dot: f32 = query.iter().zip(v_pq.iter()).map(|(a, b)| a * b).sum(); + let simd_dot = Query2bitSimd::new(&query).dotprod(&pack_codes(&indices, 2)); + + // SIMD-added error scales like √dim · σ_q · ε_c. Scale tolerance + // with √dim so large-dim trials don't falsely fail on 3σ tails. + let tol = (0.5_f32).max(0.03 * (dim as f32).sqrt()); + assert!( + (pq_dot - simd_dot).abs() < tol, + "dim={dim}: simd_dot {simd_dot} too far from ideal PQ dot {pq_dot} (tol={tol})", + ); + } + } + + /// `score_2bit_internal_scalar(a, b)` ≈ `Σ centroid(a_k) · centroid(b_k)`. + #[test] + fn test_score_2bit_internal_matches_centroid_product() { + let mut rng = StdRng::seed_from_u64(0xBAD); + let dim = 256; + let centroids = TQBits::Bits2.get_centroids(); + let n_trials = 16; + + for _ in 0..n_trials { + let raw_a = sample_normal_vec(&mut rng, dim); + let raw_b = sample_normal_vec(&mut rng, dim); + let idx_a = encode_to_nearest_centroid(centroids, &raw_a); + let idx_b = encode_to_nearest_centroid(centroids, &raw_b); + + let expected: f32 = idx_a + .iter() + .zip(idx_b.iter()) + .map(|(&a, &b)| centroids[a as usize] * centroids[b as usize]) + .sum(); + let got = score_2bit_internal_scalar(&pack_codes(&idx_a, 2), &pack_codes(&idx_b, 2)); + + assert!( + (expected - got).abs() < 0.5, + "scalar score {got} too far from centroid product {expected}", + ); + } + } +} diff --git a/lib/quantization/src/turboquant/simd/query2bit/x64.rs b/lib/quantization/src/turboquant/simd/query2bit/x64.rs new file mode 100644 index 0000000000..51b60df20e --- /dev/null +++ b/lib/quantization/src/turboquant/simd/query2bit/x64.rs @@ -0,0 +1,646 @@ +//! x86_64 SIMD paths for [`Query2bitSimd`] and [`super::score_2bit_internal`]. +//! +//! Storage / encoding mirror [`super::query4bit::x64`]: unsigned `CODEBOOK_U8` +//! consumed as the `u8` operand of `maddubs` / `VPDPBUSD`, with the query +//! quantized to 7-bit signed halves (K=128) so the `u8 × i8 → i16` pair-sum +//! never saturates. The only new machinery is the 2-bit → centroid unpack, +//! which uses a pair-table trick analogous to [`super::arm`]: +//! +//! * `PAIR_TABLE_EVEN_U8[nibble] = CODEBOOK_U8[nibble & 0b11]` +//! * `PAIR_TABLE_ODD_U8[nibble] = CODEBOOK_U8[(nibble >> 2) & 0b11]` +//! +//! Per 4 packed data bytes: split low/high nibbles, `pshufb` both pair tables +//! with the nibble indices, interleave via `punpcklbw` → 16 centroid bytes +//! in natural dim order. From there on the pipeline is identical to 4-bit. + +use super::{CODEBOOK_SCALE, CODEBOOK_U8, QUERY_HIGH_COEF, Query2bitSimd}; + +/// `PAIR_TABLE_EVEN_U8[nibble]` = `CODEBOOK_U8[nibble & 0b11]`. +const PAIR_TABLE_EVEN_U8: [u8; 16] = { + let mut tbl = [0_u8; 16]; + let mut k = 0; + while k < 16 { + tbl[k] = CODEBOOK_U8[k & 0b11]; + k += 1; + } + tbl +}; + +/// `PAIR_TABLE_ODD_U8[nibble]` = `CODEBOOK_U8[(nibble >> 2) & 0b11]`. +const PAIR_TABLE_ODD_U8: [u8; 16] = { + let mut tbl = [0_u8; 16]; + let mut k = 0; + while k < 16 { + tbl[k] = CODEBOOK_U8[(k >> 2) & 0b11]; + k += 1; + } + tbl +}; + +/// Unpack 4 packed data bytes into a natural-order `__m128i` of 16 centroid +/// `u8` values (one per code). +#[inline] +#[target_feature(enable = "sse4.1,ssse3")] +unsafe fn unpack_16_codes_sse(bytes4: *const u8) -> core::arch::x86_64::__m128i { + use core::arch::x86_64::*; + unsafe { + // Load 4 bytes into low 32 bits. + let data = _mm_cvtsi32_si128(bytes4.cast::().read_unaligned()); + + let low_mask = _mm_set1_epi8(0x0F); + let lo_nibs = _mm_and_si128(data, low_mask); + let hi_nibs = _mm_and_si128(_mm_srli_epi16(data, 4), low_mask); + + // Interleave low/high nibbles of the first 4 lanes: + // [lo(b0), hi(b0), lo(b1), hi(b1), lo(b2), hi(b2), lo(b3), hi(b3), …] + let pair_indices = _mm_unpacklo_epi8(lo_nibs, hi_nibs); + + let t_even = _mm_loadu_si128(PAIR_TABLE_EVEN_U8.as_ptr().cast::<__m128i>()); + let t_odd = _mm_loadu_si128(PAIR_TABLE_ODD_U8.as_ptr().cast::<__m128i>()); + let c_even = _mm_shuffle_epi8(t_even, pair_indices); + let c_odd = _mm_shuffle_epi8(t_odd, pair_indices); + + // Natural dim order: [c0, c1, c2, c3, c4, ..., c15]. + _mm_unpacklo_epi8(c_even, c_odd) + } +} + +#[target_feature(enable = "sse2")] +unsafe fn hsum_i32_sse(v: core::arch::x86_64::__m128i) -> i32 { + use core::arch::x86_64::*; + let v = _mm_add_epi32(v, _mm_shuffle_epi32(v, 0x4E)); + let v = _mm_add_epi32(v, _mm_shuffle_epi32(v, 0xB1)); + _mm_cvtsi128_si32(v) +} + +impl Query2bitSimd { + /// SSE4.1 + SSSE3 implementation of [`Query2bitSimd::dotprod_raw`]. + /// + /// # Safety + /// CPU must support `ssse3` and `sse4.1`. + #[target_feature(enable = "sse4.1,ssse3")] + pub unsafe fn dotprod_raw_sse(&self, vector: &[u8]) -> i64 { + use core::arch::x86_64::*; + + assert_eq!( + vector.len(), + self.expected_vector_bytes(), + "Query2bitSimd::dotprod_raw_sse: vector length mismatch ({} vs expected {})", + vector.len(), + self.expected_vector_bytes(), + ); + + unsafe { + let mut acc_low = _mm_setzero_si128(); + let mut acc_high = _mm_setzero_si128(); + let ones = _mm_set1_epi16(1); + + for (chunk_idx, [low, high]) in self.query_data.iter().enumerate() { + let c = unpack_16_codes_sse(vector.as_ptr().add(chunk_idx * 4)); + + let q_low = _mm_loadu_si128(low.as_ptr().cast::<__m128i>()); + let q_high = _mm_loadu_si128(high.as_ptr().cast::<__m128i>()); + + let prod_low = _mm_maddubs_epi16(c, q_low); + let prod_high = _mm_maddubs_epi16(c, q_high); + + acc_low = _mm_add_epi32(acc_low, _mm_madd_epi16(prod_low, ones)); + acc_high = _mm_add_epi32(acc_high, _mm_madd_epi16(prod_high, ones)); + } + + let sum_low = i64::from(hsum_i32_sse(acc_low)); + let sum_high = i64::from(hsum_i32_sse(acc_high)); + sum_low + QUERY_HIGH_COEF * sum_high + self.dotprod_raw_tail(vector) + } + } + + /// AVX2 implementation. Built on top of the SSE unpack (`_mm_shuffle_epi8` + /// stays 128-bit-lane-scoped on AVX2, so doubling up to YMM for 8 bytes of + /// data at once requires extra lane-management that costs more than the + /// unroll saves). We call the SSE unpack twice per iteration and pair the + /// `maddubs` / `madd_epi16` paths on 256-bit vectors where they're cheap. + /// + /// # Safety + /// CPU must support `avx2`, `ssse3` and `sse4.1`. + #[target_feature(enable = "avx2,sse4.1,ssse3")] + pub unsafe fn dotprod_raw_avx2(&self, vector: &[u8]) -> i64 { + use core::arch::x86_64::*; + + assert_eq!( + vector.len(), + self.expected_vector_bytes(), + "Query2bitSimd::dotprod_raw_avx2: vector length mismatch ({} vs expected {})", + vector.len(), + self.expected_vector_bytes(), + ); + + unsafe { + let ones = _mm256_set1_epi16(1); + let mut acc_low = _mm256_setzero_si256(); + let mut acc_high = _mm256_setzero_si256(); + + let n_chunks = self.query_data.len(); + let n_pairs = n_chunks / 2; + + // 2× unroll: fold two SSE chunks into one YMM accumulation per iter. + for p in 0..n_pairs { + let [qa_lo, qa_hi] = &self.query_data[2 * p]; + let [qb_lo, qb_hi] = &self.query_data[2 * p + 1]; + let c_a = unpack_16_codes_sse(vector.as_ptr().add(4 * (2 * p))); + let c_b = unpack_16_codes_sse(vector.as_ptr().add(4 * (2 * p + 1))); + let c = _mm256_set_m128i(c_b, c_a); + + let q_lo_a = _mm_loadu_si128(qa_lo.as_ptr().cast::<__m128i>()); + let q_lo_b = _mm_loadu_si128(qb_lo.as_ptr().cast::<__m128i>()); + let q_hi_a = _mm_loadu_si128(qa_hi.as_ptr().cast::<__m128i>()); + let q_hi_b = _mm_loadu_si128(qb_hi.as_ptr().cast::<__m128i>()); + let q_low = _mm256_set_m128i(q_lo_b, q_lo_a); + let q_high = _mm256_set_m128i(q_hi_b, q_hi_a); + + let prod_low = _mm256_maddubs_epi16(c, q_low); + let prod_high = _mm256_maddubs_epi16(c, q_high); + acc_low = _mm256_add_epi32(acc_low, _mm256_madd_epi16(prod_low, ones)); + acc_high = _mm256_add_epi32(acc_high, _mm256_madd_epi16(prod_high, ones)); + } + + // Fold YMM accumulators into XMM. + let mut sum_low_sse = _mm_add_epi32( + _mm256_castsi256_si128(acc_low), + _mm256_extracti128_si256(acc_low, 1), + ); + let mut sum_high_sse = _mm_add_epi32( + _mm256_castsi256_si128(acc_high), + _mm256_extracti128_si256(acc_high, 1), + ); + + // Tail: odd chunk count → one extra chunk via SSE. + if n_chunks % 2 == 1 { + let idx = 2 * n_pairs; + let [q_lo_t, q_hi_t] = &self.query_data[idx]; + let c = unpack_16_codes_sse(vector.as_ptr().add(4 * idx)); + let q_low = _mm_loadu_si128(q_lo_t.as_ptr().cast::<__m128i>()); + let q_high = _mm_loadu_si128(q_hi_t.as_ptr().cast::<__m128i>()); + let prod_low = _mm_maddubs_epi16(c, q_low); + let prod_high = _mm_maddubs_epi16(c, q_high); + let ones = _mm_set1_epi16(1); + sum_low_sse = _mm_add_epi32(sum_low_sse, _mm_madd_epi16(prod_low, ones)); + sum_high_sse = _mm_add_epi32(sum_high_sse, _mm_madd_epi16(prod_high, ones)); + } + + let sum_low = i64::from(hsum_i32_sse(sum_low_sse)); + let sum_high = i64::from(hsum_i32_sse(sum_high_sse)); + sum_low + QUERY_HIGH_COEF * sum_high + self.dotprod_raw_tail(vector) + } + } + + /// AVX-512 + VNNI implementation — uses `VPDPBUSD` on 512-bit ZMM for + /// fused `u8 × i8 → i32` MAC. Processes 4 chunks (64 codes) per iter. + /// + /// Tail handling falls back to SSE `maddubs + madd_epi16` rather than the + /// narrower 128/256-bit VPDPBUSD variants, because those need `avx512vl`. + /// + /// # Safety + /// CPU must support `avx512f`, `avx512bw`, `avx512vnni`, `ssse3`, `sse4.1`. + #[target_feature(enable = "avx512f,avx512bw,avx512vnni,sse4.1,ssse3")] + pub unsafe fn dotprod_raw_avx512_vnni(&self, vector: &[u8]) -> i64 { + use core::arch::x86_64::*; + + assert_eq!( + vector.len(), + self.expected_vector_bytes(), + "Query2bitSimd::dotprod_raw_avx512_vnni: vector length mismatch ({} vs expected {})", + vector.len(), + self.expected_vector_bytes(), + ); + + unsafe { + let mut acc_low = _mm512_setzero_si512(); + let mut acc_high = _mm512_setzero_si512(); + + let n_chunks = self.query_data.len(); + let n_quads = n_chunks / 4; + + for q in 0..n_quads { + let base = 4 * q; + + // Unpack 4 × 16 centroids from 16 packed data bytes. + let c_0 = unpack_16_codes_sse(vector.as_ptr().add(4 * base)); + let c_1 = unpack_16_codes_sse(vector.as_ptr().add(4 * (base + 1))); + let c_2 = unpack_16_codes_sse(vector.as_ptr().add(4 * (base + 2))); + let c_3 = unpack_16_codes_sse(vector.as_ptr().add(4 * (base + 3))); + let c_ab = _mm256_set_m128i(c_1, c_0); + let c_cd = _mm256_set_m128i(c_3, c_2); + let c = _mm512_inserti64x4(_mm512_castsi256_si512(c_ab), c_cd, 1); + + // Load 4 × 16 query-low i8 and 4 × 16 query-high i8 into ZMMs. + let q_lo_0 = _mm_loadu_si128(self.query_data[base][0].as_ptr().cast::<__m128i>()); + let q_lo_1 = + _mm_loadu_si128(self.query_data[base + 1][0].as_ptr().cast::<__m128i>()); + let q_lo_2 = + _mm_loadu_si128(self.query_data[base + 2][0].as_ptr().cast::<__m128i>()); + let q_lo_3 = + _mm_loadu_si128(self.query_data[base + 3][0].as_ptr().cast::<__m128i>()); + let q_lo_ab = _mm256_set_m128i(q_lo_1, q_lo_0); + let q_lo_cd = _mm256_set_m128i(q_lo_3, q_lo_2); + let q_low = _mm512_inserti64x4(_mm512_castsi256_si512(q_lo_ab), q_lo_cd, 1); + + let q_hi_0 = _mm_loadu_si128(self.query_data[base][1].as_ptr().cast::<__m128i>()); + let q_hi_1 = + _mm_loadu_si128(self.query_data[base + 1][1].as_ptr().cast::<__m128i>()); + let q_hi_2 = + _mm_loadu_si128(self.query_data[base + 2][1].as_ptr().cast::<__m128i>()); + let q_hi_3 = + _mm_loadu_si128(self.query_data[base + 3][1].as_ptr().cast::<__m128i>()); + let q_hi_ab = _mm256_set_m128i(q_hi_1, q_hi_0); + let q_hi_cd = _mm256_set_m128i(q_hi_3, q_hi_2); + let q_high = _mm512_inserti64x4(_mm512_castsi256_si512(q_hi_ab), q_hi_cd, 1); + + acc_low = _mm512_dpbusd_epi32(acc_low, c, q_low); + acc_high = _mm512_dpbusd_epi32(acc_high, c, q_high); + } + + let mut total_low = i64::from(_mm512_reduce_add_epi32(acc_low)); + let mut total_high = i64::from(_mm512_reduce_add_epi32(acc_high)); + + // Tail (0..3 chunks) via plain SSE `maddubs + madd_epi16` — + // avoids pulling in avx512vl for the narrow VPDPBUSD variants. + let tail_start = n_quads * 4; + let ones = _mm_set1_epi16(1); + for p in tail_start..n_chunks { + let [q_lo_chunk, q_hi_chunk] = &self.query_data[p]; + let c = unpack_16_codes_sse(vector.as_ptr().add(4 * p)); + let q_low = _mm_loadu_si128(q_lo_chunk.as_ptr().cast::<__m128i>()); + let q_high = _mm_loadu_si128(q_hi_chunk.as_ptr().cast::<__m128i>()); + let prod_low = _mm_maddubs_epi16(c, q_low); + let prod_high = _mm_maddubs_epi16(c, q_high); + total_low += i64::from(hsum_i32_sse(_mm_madd_epi16(prod_low, ones))); + total_high += i64::from(hsum_i32_sse(_mm_madd_epi16(prod_high, ones))); + } + + total_low + QUERY_HIGH_COEF * total_high + self.dotprod_raw_tail(vector) + } + } +} + +// ------------------------------------------------------------------ +// score_2bit_internal — both operands signed, same widen-to-i16 +// pattern as query4bit's score_4bit_internal. XOR 0x80 converts the +// CODEBOOK_U8 pair-table bytes back to signed i8 in-register. +// ------------------------------------------------------------------ + +/// SSE4.1 + SSSE3 implementation of [`super::score_2bit_internal`]. +/// +/// # Safety +/// CPU must support `ssse3` and `sse4.1`. +#[target_feature(enable = "sse4.1,ssse3")] +pub unsafe fn score_2bit_internal_sse(a: &[u8], b: &[u8]) -> f32 { + use core::arch::x86_64::*; + + assert_eq!( + a.len(), + b.len(), + "score_2bit_internal_sse: vector length mismatch ({} vs {})", + a.len(), + b.len(), + ); + + unsafe { + let shift = _mm_set1_epi8(-128_i8); + let mut acc = _mm_setzero_si128(); + let n_full = a.len() / 4; + for i in 0..n_full { + let c_a_u8 = unpack_16_codes_sse(a.as_ptr().add(i * 4)); + let c_b_u8 = unpack_16_codes_sse(b.as_ptr().add(i * 4)); + let c_a = _mm_xor_si128(c_a_u8, shift); + let c_b = _mm_xor_si128(c_b_u8, shift); + + let c_a_lo = _mm_cvtepi8_epi16(c_a); + let c_a_hi = _mm_cvtepi8_epi16(_mm_srli_si128(c_a, 8)); + let c_b_lo = _mm_cvtepi8_epi16(c_b); + let c_b_hi = _mm_cvtepi8_epi16(_mm_srli_si128(c_b, 8)); + + let prod_lo = _mm_madd_epi16(c_a_lo, c_b_lo); + let prod_hi = _mm_madd_epi16(c_a_hi, c_b_hi); + acc = _mm_add_epi32(acc, _mm_add_epi32(prod_lo, prod_hi)); + } + let simd_bytes = n_full * 4; + let acc_i64 = i64::from(hsum_i32_sse(acc)) + + super::score_2bit_internal_integer(&a[simd_bytes..], &b[simd_bytes..]); + acc_i64 as f32 / (CODEBOOK_SCALE * CODEBOOK_SCALE) + } +} + +/// AVX2 implementation of [`super::score_2bit_internal`]. +/// +/// # Safety +/// CPU must support `avx2`, `ssse3`, `sse4.1`. +#[target_feature(enable = "avx2,sse4.1,ssse3")] +pub unsafe fn score_2bit_internal_avx2(a: &[u8], b: &[u8]) -> f32 { + use core::arch::x86_64::*; + + assert_eq!(a.len(), b.len()); + + unsafe { + let shift256 = _mm256_set1_epi8(-128_i8); + let mut acc = _mm256_setzero_si256(); + + let n_chunks = a.len() / 4; + let n_pairs = n_chunks / 2; + + for p in 0..n_pairs { + let off_a0 = 4 * (2 * p); + let off_a1 = 4 * (2 * p + 1); + let c_a = _mm256_set_m128i( + unpack_16_codes_sse(a.as_ptr().add(off_a1)), + unpack_16_codes_sse(a.as_ptr().add(off_a0)), + ); + let c_b = _mm256_set_m128i( + unpack_16_codes_sse(b.as_ptr().add(off_a1)), + unpack_16_codes_sse(b.as_ptr().add(off_a0)), + ); + let c_a = _mm256_xor_si256(c_a, shift256); + let c_b = _mm256_xor_si256(c_b, shift256); + + let c_a_lo = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(c_a)); + let c_a_hi = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(c_a, 1)); + let c_b_lo = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(c_b)); + let c_b_hi = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(c_b, 1)); + + let prod_lo = _mm256_madd_epi16(c_a_lo, c_b_lo); + let prod_hi = _mm256_madd_epi16(c_a_hi, c_b_hi); + acc = _mm256_add_epi32(acc, _mm256_add_epi32(prod_lo, prod_hi)); + } + + let mut acc_sse = _mm_add_epi32( + _mm256_castsi256_si128(acc), + _mm256_extracti128_si256(acc, 1), + ); + if n_chunks % 2 == 1 { + let off = 4 * (2 * n_pairs); + let shift = _mm_set1_epi8(-128_i8); + let c_a = _mm_xor_si128(unpack_16_codes_sse(a.as_ptr().add(off)), shift); + let c_b = _mm_xor_si128(unpack_16_codes_sse(b.as_ptr().add(off)), shift); + let c_a_lo = _mm_cvtepi8_epi16(c_a); + let c_a_hi = _mm_cvtepi8_epi16(_mm_srli_si128(c_a, 8)); + let c_b_lo = _mm_cvtepi8_epi16(c_b); + let c_b_hi = _mm_cvtepi8_epi16(_mm_srli_si128(c_b, 8)); + acc_sse = _mm_add_epi32(acc_sse, _mm_madd_epi16(c_a_lo, c_b_lo)); + acc_sse = _mm_add_epi32(acc_sse, _mm_madd_epi16(c_a_hi, c_b_hi)); + } + let simd_bytes = n_chunks * 4; + let acc_i64 = i64::from(hsum_i32_sse(acc_sse)) + + super::score_2bit_internal_integer(&a[simd_bytes..], &b[simd_bytes..]); + acc_i64 as f32 / (CODEBOOK_SCALE * CODEBOOK_SCALE) + } +} + +/// AVX-512 + VNNI implementation — uses `VPDPWSSD` for fused i16×i16 MAC. +/// +/// # Safety +/// CPU must support `avx512f`, `avx512bw`, `avx512vnni`, `ssse3`, `sse4.1`. +#[target_feature(enable = "avx512f,avx512bw,avx512vnni,sse4.1,ssse3")] +pub unsafe fn score_2bit_internal_avx512_vnni(a: &[u8], b: &[u8]) -> f32 { + use core::arch::x86_64::*; + + assert_eq!(a.len(), b.len()); + + unsafe { + // Full-width 512-bit VNNI. Process 4 chunks (64 codes) per iter. + let shift512 = _mm512_set1_epi8(-128_i8); + let mut acc = _mm512_setzero_si512(); + + let n_chunks = a.len() / 4; + let n_quads = n_chunks / 4; + + for q in 0..n_quads { + let base = 4 * q; + let c_a_0 = unpack_16_codes_sse(a.as_ptr().add(4 * base)); + let c_a_1 = unpack_16_codes_sse(a.as_ptr().add(4 * (base + 1))); + let c_a_2 = unpack_16_codes_sse(a.as_ptr().add(4 * (base + 2))); + let c_a_3 = unpack_16_codes_sse(a.as_ptr().add(4 * (base + 3))); + let c_b_0 = unpack_16_codes_sse(b.as_ptr().add(4 * base)); + let c_b_1 = unpack_16_codes_sse(b.as_ptr().add(4 * (base + 1))); + let c_b_2 = unpack_16_codes_sse(b.as_ptr().add(4 * (base + 2))); + let c_b_3 = unpack_16_codes_sse(b.as_ptr().add(4 * (base + 3))); + + let c_a_ab = _mm256_set_m128i(c_a_1, c_a_0); + let c_a_cd = _mm256_set_m128i(c_a_3, c_a_2); + let c_a = _mm512_inserti64x4(_mm512_castsi256_si512(c_a_ab), c_a_cd, 1); + let c_b_ab = _mm256_set_m128i(c_b_1, c_b_0); + let c_b_cd = _mm256_set_m128i(c_b_3, c_b_2); + let c_b = _mm512_inserti64x4(_mm512_castsi256_si512(c_b_ab), c_b_cd, 1); + + // Unwind the `+128` shift to recover signed i8. + let c_a = _mm512_xor_si512(c_a, shift512); + let c_b = _mm512_xor_si512(c_b, shift512); + + // Widen i8 → i16 in both halves of each ZMM. + let c_a_lo = _mm512_cvtepi8_epi16(_mm512_castsi512_si256(c_a)); + let c_a_hi = _mm512_cvtepi8_epi16(_mm512_extracti64x4_epi64(c_a, 1)); + let c_b_lo = _mm512_cvtepi8_epi16(_mm512_castsi512_si256(c_b)); + let c_b_hi = _mm512_cvtepi8_epi16(_mm512_extracti64x4_epi64(c_b, 1)); + + acc = _mm512_dpwssd_epi32(acc, c_a_lo, c_b_lo); + acc = _mm512_dpwssd_epi32(acc, c_a_hi, c_b_hi); + } + + let mut acc_total = i64::from(_mm512_reduce_add_epi32(acc)); + + // Tail (0..3 chunks) via SSE `madd_epi16` — no avx512vl needed. + let tail_start = n_quads * 4; + let shift_128 = _mm_set1_epi8(-128_i8); + for p in tail_start..n_chunks { + let c_a = _mm_xor_si128(unpack_16_codes_sse(a.as_ptr().add(4 * p)), shift_128); + let c_b = _mm_xor_si128(unpack_16_codes_sse(b.as_ptr().add(4 * p)), shift_128); + let c_a_lo = _mm_cvtepi8_epi16(c_a); + let c_a_hi = _mm_cvtepi8_epi16(_mm_srli_si128(c_a, 8)); + let c_b_lo = _mm_cvtepi8_epi16(c_b); + let c_b_hi = _mm_cvtepi8_epi16(_mm_srli_si128(c_b, 8)); + let prod = _mm_add_epi32( + _mm_madd_epi16(c_a_lo, c_b_lo), + _mm_madd_epi16(c_a_hi, c_b_hi), + ); + acc_total += i64::from(hsum_i32_sse(prod)); + } + + // Scalar tail for any bytes past the last SIMD chunk (a.len() % 4). + let simd_bytes = n_chunks * 4; + acc_total += super::score_2bit_internal_integer(&a[simd_bytes..], &b[simd_bytes..]); + acc_total as f32 / (CODEBOOK_SCALE * CODEBOOK_SCALE) + } +} + +#[cfg(test)] +mod tests { + use rand::SeedableRng as _; + use rand::prelude::StdRng; + + use super::super::super::shared::pack_codes; + use super::super::shared::{PARITY_DIMS, random_inputs}; + use super::super::{Query2bitSimd, score_2bit_internal_scalar}; + use super::*; + + #[test] + fn test_sse_matches_scalar() { + if !std::is_x86_feature_detected!("ssse3") || !std::is_x86_feature_detected!("sse4.1") { + return; + } + let mut rng = StdRng::seed_from_u64(7); + for &dim in PARITY_DIMS { + let (simd_query, vector) = random_inputs(&mut rng, dim); + let scalar = simd_query.dotprod_raw(&vector); + let got = unsafe { simd_query.dotprod_raw_sse(&vector) }; + assert_eq!(scalar, got, "sse mismatch at dim {dim}"); + } + } + + #[test] + fn test_avx2_matches_scalar() { + if !std::is_x86_feature_detected!("avx2") { + return; + } + let mut rng = StdRng::seed_from_u64(7); + for &dim in PARITY_DIMS { + let (simd_query, vector) = random_inputs(&mut rng, dim); + let scalar = simd_query.dotprod_raw(&vector); + let got = unsafe { simd_query.dotprod_raw_avx2(&vector) }; + assert_eq!(scalar, got, "avx2 mismatch at dim {dim}"); + } + } + + #[test] + fn test_avx512_vnni_matches_scalar() { + if !(std::is_x86_feature_detected!("avx512f") + && std::is_x86_feature_detected!("avx512bw") + && std::is_x86_feature_detected!("avx512vnni")) + { + return; + } + let mut rng = StdRng::seed_from_u64(7); + for &dim in PARITY_DIMS { + let (simd_query, vector) = random_inputs(&mut rng, dim); + let scalar = simd_query.dotprod_raw(&vector); + let got = unsafe { simd_query.dotprod_raw_avx512_vnni(&vector) }; + assert_eq!(scalar, got, "avx512 mismatch at dim {dim}"); + } + } + + #[test] + fn test_saturation_safety_64k() { + let dim = 65_536; + let query = vec![1.0_f32; dim]; + let indices: Vec = vec![3; dim]; // max-magnitude centroid + let vector = pack_codes(&indices, 2); + + let q = Query2bitSimd::new(&query); + let scalar = q.dotprod_raw(&vector); + + unsafe { + if std::is_x86_feature_detected!("ssse3") && std::is_x86_feature_detected!("sse4.1") { + assert_eq!(scalar, q.dotprod_raw_sse(&vector)); + } + if std::is_x86_feature_detected!("avx2") { + assert_eq!(scalar, q.dotprod_raw_avx2(&vector)); + } + if std::is_x86_feature_detected!("avx512f") + && std::is_x86_feature_detected!("avx512bw") + && std::is_x86_feature_detected!("avx512vnni") + { + assert_eq!(scalar, q.dotprod_raw_avx512_vnni(&vector)); + } + } + } + + #[test] + fn test_score_sse_matches_scalar() { + if !std::is_x86_feature_detected!("ssse3") || !std::is_x86_feature_detected!("sse4.1") { + return; + } + let mut rng = StdRng::seed_from_u64(7); + for &dim in PARITY_DIMS { + let (_, vec_a) = random_inputs(&mut rng, dim); + let (_, vec_b) = random_inputs(&mut rng, dim); + let scalar = score_2bit_internal_scalar(&vec_a, &vec_b); + let sse = unsafe { score_2bit_internal_sse(&vec_a, &vec_b) }; + assert_eq!( + scalar.to_bits(), + sse.to_bits(), + "score sse mismatch at dim {dim}", + ); + } + } + + #[test] + fn test_score_avx2_matches_scalar() { + if !std::is_x86_feature_detected!("avx2") { + return; + } + let mut rng = StdRng::seed_from_u64(7); + for &dim in PARITY_DIMS { + let (_, vec_a) = random_inputs(&mut rng, dim); + let (_, vec_b) = random_inputs(&mut rng, dim); + let scalar = score_2bit_internal_scalar(&vec_a, &vec_b); + let got = unsafe { score_2bit_internal_avx2(&vec_a, &vec_b) }; + assert_eq!( + scalar.to_bits(), + got.to_bits(), + "score avx2 mismatch at dim {dim}", + ); + } + } + + #[test] + fn test_score_avx512_vnni_matches_scalar() { + if !(std::is_x86_feature_detected!("avx512f") + && std::is_x86_feature_detected!("avx512bw") + && std::is_x86_feature_detected!("avx512vnni")) + { + return; + } + let mut rng = StdRng::seed_from_u64(7); + for &dim in PARITY_DIMS { + let (_, vec_a) = random_inputs(&mut rng, dim); + let (_, vec_b) = random_inputs(&mut rng, dim); + let scalar = score_2bit_internal_scalar(&vec_a, &vec_b); + let got = unsafe { score_2bit_internal_avx512_vnni(&vec_a, &vec_b) }; + assert_eq!( + scalar.to_bits(), + got.to_bits(), + "score avx512 mismatch at dim {dim}", + ); + } + } + + #[test] + fn test_score_saturation_safety_64k() { + let dim = 65_536; + let indices: Vec = vec![3; dim]; + let vec_a = pack_codes(&indices, 2); + let vec_b = pack_codes(&indices, 2); + let scalar = score_2bit_internal_scalar(&vec_a, &vec_b); + + unsafe { + if std::is_x86_feature_detected!("ssse3") && std::is_x86_feature_detected!("sse4.1") { + assert_eq!( + scalar.to_bits(), + score_2bit_internal_sse(&vec_a, &vec_b).to_bits() + ); + } + if std::is_x86_feature_detected!("avx2") { + assert_eq!( + scalar.to_bits(), + score_2bit_internal_avx2(&vec_a, &vec_b).to_bits() + ); + } + if std::is_x86_feature_detected!("avx512f") + && std::is_x86_feature_detected!("avx512bw") + && std::is_x86_feature_detected!("avx512vnni") + { + assert_eq!( + scalar.to_bits(), + score_2bit_internal_avx512_vnni(&vec_a, &vec_b).to_bits(), + ); + } + } + } +} diff --git a/lib/quantization/src/turboquant/simd/query4bit/arm.rs b/lib/quantization/src/turboquant/simd/query4bit/arm.rs new file mode 100644 index 0000000000..d4c42d1f17 --- /dev/null +++ b/lib/quantization/src/turboquant/simd/query4bit/arm.rs @@ -0,0 +1,467 @@ +//! NEON SIMD paths for [`Query4bitSimd`] on aarch64. +//! +//! The codebook is stored as signed i8 here (`CODEBOOK_I8`), so `vmull_s8` and +//! `sdot` operate on true i8×i8 products — no bias correction. The +//! combining coefficient `QUERY_HIGH_COEF = 256` pairs with full-range i8 +//! query halves, yielding ~15.9-bit query precision. + +use super::{CODEBOOK_I8, CODEBOOK_SCALE, QUERY_HIGH_COEF, Query4bitSimd}; + +impl Query4bitSimd { + /// ARM NEON implementation of [`Query4bitSimd::dotprod_raw`]. + /// + /// # Safety + /// CPU must support the `neon` feature (always true on aarch64). + #[target_feature(enable = "neon")] + pub unsafe fn dotprod_raw_neon(&self, vector: &[u8]) -> i64 { + use core::arch::aarch64::*; + + assert_eq!( + vector.len(), + self.expected_vector_bytes(), + "Query4bitSimd::dotprod_raw_neon: vector length mismatch ({} vs expected {})", + vector.len(), + self.expected_vector_bytes(), + ); + + unsafe { + let codebook = vld1q_s8(CODEBOOK_I8.as_ptr()); + let mut acc_low = vdupq_n_s32(0); + let mut acc_high = vdupq_n_s32(0); + + let nibble_mask = vdup_n_u8(0x0F); + for (chunk_idx, [low, high]) in self.query_data.iter().enumerate() { + let v_packed = vld1_u8(vector.as_ptr().add(chunk_idx * 8)); + let v_lo = vand_u8(v_packed, nibble_mask); + let v_hi = vshr_n_u8(v_packed, 4); + let v = vcombine_u8(vzip1_u8(v_lo, v_hi), vzip2_u8(v_lo, v_hi)); + + let c = vqtbl1q_s8(codebook, v); + + let q_low = vld1q_s8(low.as_ptr()); + let q_high = vld1q_s8(high.as_ptr()); + + let prod_low_lo = vmull_s8(vget_low_s8(q_low), vget_low_s8(c)); + let prod_low_hi = vmull_high_s8(q_low, c); + let prod_high_lo = vmull_s8(vget_low_s8(q_high), vget_low_s8(c)); + let prod_high_hi = vmull_high_s8(q_high, c); + + acc_low = vpadalq_s16(acc_low, prod_low_lo); + acc_low = vpadalq_s16(acc_low, prod_low_hi); + acc_high = vpadalq_s16(acc_high, prod_high_lo); + acc_high = vpadalq_s16(acc_high, prod_high_hi); + } + + // Tail: one extra chunk via NEON on a zero-padded 8-byte scratch. + // Data bytes beyond `tail_dims / 2` are zero, and `tail_low / + // tail_high` slots beyond `tail_dims` are zero — their products + // contribute `0` to the final sum. + if let Some(buf) = self.tail_chunk_scratch(vector) { + let v_packed = vld1_u8(buf.as_ptr()); + let v_lo = vand_u8(v_packed, nibble_mask); + let v_hi = vshr_n_u8(v_packed, 4); + let v = vcombine_u8(vzip1_u8(v_lo, v_hi), vzip2_u8(v_lo, v_hi)); + let c = vqtbl1q_s8(codebook, v); + + let q_low = vld1q_s8(self.tail_low.as_ptr()); + let q_high = vld1q_s8(self.tail_high.as_ptr()); + + let prod_low_lo = vmull_s8(vget_low_s8(q_low), vget_low_s8(c)); + let prod_low_hi = vmull_high_s8(q_low, c); + let prod_high_lo = vmull_s8(vget_low_s8(q_high), vget_low_s8(c)); + let prod_high_hi = vmull_high_s8(q_high, c); + + acc_low = vpadalq_s16(acc_low, prod_low_lo); + acc_low = vpadalq_s16(acc_low, prod_low_hi); + acc_high = vpadalq_s16(acc_high, prod_high_lo); + acc_high = vpadalq_s16(acc_high, prod_high_hi); + } + + i64::from(vaddvq_s32(acc_low)) + QUERY_HIGH_COEF * i64::from(vaddvq_s32(acc_high)) + } + } + + /// ARMv8.2-A Dot Product variant. Uses `SDOT` to sum four i8×i8 products + /// per i32 lane per instruction, emitted via inline asm because + /// `vdotq_s32` is still unstable (rust-lang/rust#117224). + /// + /// 2× unrolled: two chunks per iteration with four independent `i32×4` + /// accumulators (two for low, two for high) break the single dependency + /// chain of a naive implementation. On Apple M-series SDOT has ~3-cycle + /// latency at 4/cycle throughput, and four parallel chains lift + /// throughput ~7–20% over a 1× version (larger dims benefit more — + /// latency dominates there). + /// + /// # Safety + /// CPU must support `neon` and `dotprod`. + #[target_feature(enable = "neon,dotprod")] + pub unsafe fn dotprod_raw_neon_sdot(&self, vector: &[u8]) -> i64 { + use core::arch::aarch64::*; + + assert_eq!( + vector.len(), + self.expected_vector_bytes(), + "Query4bitSimd::dotprod_raw_neon_sdot: vector length mismatch ({} vs expected {})", + vector.len(), + self.expected_vector_bytes(), + ); + + unsafe { + let codebook = vld1q_s8(CODEBOOK_I8.as_ptr()); + let mut acc_low_0 = vdupq_n_s32(0); + let mut acc_low_1 = vdupq_n_s32(0); + let mut acc_high_0 = vdupq_n_s32(0); + let mut acc_high_1 = vdupq_n_s32(0); + let nibble_mask_q = vdupq_n_u8(0x0F); + + let chunks = self.query_data.as_slice(); + let n_pairs = chunks.len() / 2; + + for i in 0..n_pairs { + let [low_0, high_0] = chunks[2 * i]; + let [low_1, high_1] = chunks[2 * i + 1]; + + // One 16-byte load covers both chunks' 8 packed bytes each. + let v_packed = vld1q_u8(vector.as_ptr().add(16 * i)); + let v_lo = vandq_u8(v_packed, nibble_mask_q); + let v_hi = vshrq_n_u8(v_packed, 4); + let v_0 = vzip1q_u8(v_lo, v_hi); + let v_1 = vzip2q_u8(v_lo, v_hi); + let c_0 = vqtbl1q_s8(codebook, v_0); + let c_1 = vqtbl1q_s8(codebook, v_1); + + let q_low_0 = vld1q_s8(low_0.as_ptr()); + let q_high_0 = vld1q_s8(high_0.as_ptr()); + let q_low_1 = vld1q_s8(low_1.as_ptr()); + let q_high_1 = vld1q_s8(high_1.as_ptr()); + + // Four independent SDOT dependency chains — 2× unroll for ILP. + core::arch::asm!( + "sdot {acc:v}.4s, {a:v}.16b, {b:v}.16b", + acc = inout(vreg) acc_low_0, + a = in(vreg) q_low_0, + b = in(vreg) c_0, + options(pure, nomem, nostack, preserves_flags), + ); + core::arch::asm!( + "sdot {acc:v}.4s, {a:v}.16b, {b:v}.16b", + acc = inout(vreg) acc_low_1, + a = in(vreg) q_low_1, + b = in(vreg) c_1, + options(pure, nomem, nostack, preserves_flags), + ); + core::arch::asm!( + "sdot {acc:v}.4s, {a:v}.16b, {b:v}.16b", + acc = inout(vreg) acc_high_0, + a = in(vreg) q_high_0, + b = in(vreg) c_0, + options(pure, nomem, nostack, preserves_flags), + ); + core::arch::asm!( + "sdot {acc:v}.4s, {a:v}.16b, {b:v}.16b", + acc = inout(vreg) acc_high_1, + a = in(vreg) q_high_1, + b = in(vreg) c_1, + options(pure, nomem, nostack, preserves_flags), + ); + } + + // Odd leftover chunk (after the paired loop) — single-chunk SDOT + // reusing acc_*_0 accumulators. Needed when `query_data.len()` is + // odd, i.e. `dim ∈ (16·n, 16·n + 14]` rather than a multiple of 32. + if chunks.len() % 2 == 1 { + let tail_chunk = n_pairs * 2; + let [low_t, high_t] = chunks[tail_chunk]; + let v_packed = vld1_u8(vector.as_ptr().add(8 * tail_chunk)); + let v_lo = vand_u8(v_packed, vdup_n_u8(0x0F)); + let v_hi = vshr_n_u8(v_packed, 4); + let v = vcombine_u8(vzip1_u8(v_lo, v_hi), vzip2_u8(v_lo, v_hi)); + let c = vqtbl1q_s8(codebook, v); + let q_low_t = vld1q_s8(low_t.as_ptr()); + let q_high_t = vld1q_s8(high_t.as_ptr()); + + core::arch::asm!( + "sdot {acc:v}.4s, {a:v}.16b, {b:v}.16b", + acc = inout(vreg) acc_low_0, + a = in(vreg) q_low_t, + b = in(vreg) c, + options(pure, nomem, nostack, preserves_flags), + ); + core::arch::asm!( + "sdot {acc:v}.4s, {a:v}.16b, {b:v}.16b", + acc = inout(vreg) acc_high_0, + a = in(vreg) q_high_t, + b = in(vreg) c, + options(pure, nomem, nostack, preserves_flags), + ); + } + + // Tail: one extra chunk via single-SDOT on a zero-padded scratch. + if let Some(buf) = self.tail_chunk_scratch(vector) { + let v_packed = vld1_u8(buf.as_ptr()); + let v_lo = vand_u8(v_packed, vdup_n_u8(0x0F)); + let v_hi = vshr_n_u8(v_packed, 4); + let v = vcombine_u8(vzip1_u8(v_lo, v_hi), vzip2_u8(v_lo, v_hi)); + let c = vqtbl1q_s8(codebook, v); + let q_low_t = vld1q_s8(self.tail_low.as_ptr()); + let q_high_t = vld1q_s8(self.tail_high.as_ptr()); + + core::arch::asm!( + "sdot {acc:v}.4s, {a:v}.16b, {b:v}.16b", + acc = inout(vreg) acc_low_0, + a = in(vreg) q_low_t, + b = in(vreg) c, + options(pure, nomem, nostack, preserves_flags), + ); + core::arch::asm!( + "sdot {acc:v}.4s, {a:v}.16b, {b:v}.16b", + acc = inout(vreg) acc_high_0, + a = in(vreg) q_high_t, + b = in(vreg) c, + options(pure, nomem, nostack, preserves_flags), + ); + } + + let acc_low = vaddq_s32(acc_low_0, acc_low_1); + let acc_high = vaddq_s32(acc_high_0, acc_high_1); + i64::from(vaddvq_s32(acc_low)) + QUERY_HIGH_COEF * i64::from(vaddvq_s32(acc_high)) + } + } +} + +/// NEON implementation of [`super::score_4bit_internal`]. Both PQ-encoded +/// vectors are unpacked into i8 codebook values via `vqtbl1q_s8`, multiplied +/// signed-signed via `vmull_s8` (16 256-max product fits i16), and pair-added +/// into an i32×4 accumulator. Reconstruction divides by `c_scale²`. +/// +/// # Safety +/// CPU must support the `neon` feature (always true on aarch64). +#[target_feature(enable = "neon")] +pub unsafe fn score_4bit_internal_neon(a: &[u8], b: &[u8]) -> f32 { + use core::arch::aarch64::*; + + assert_eq!( + a.len(), + b.len(), + "score_4bit_internal_neon: vector length mismatch ({} vs {})", + a.len(), + b.len(), + ); + + unsafe { + let codebook = vld1q_s8(CODEBOOK_I8.as_ptr()); + let mut acc = vdupq_n_s32(0); + let nibble_mask = vdup_n_u8(0x0F); + + let n_full = a.len() / 8; + for i in 0..n_full { + let va = vld1_u8(a.as_ptr().add(i * 8)); + let va_lo = vand_u8(va, nibble_mask); + let va_hi = vshr_n_u8(va, 4); + let a_idx = vcombine_u8(vzip1_u8(va_lo, va_hi), vzip2_u8(va_lo, va_hi)); + + let vb = vld1_u8(b.as_ptr().add(i * 8)); + let vb_lo = vand_u8(vb, nibble_mask); + let vb_hi = vshr_n_u8(vb, 4); + let b_idx = vcombine_u8(vzip1_u8(vb_lo, vb_hi), vzip2_u8(vb_lo, vb_hi)); + + let c_a = vqtbl1q_s8(codebook, a_idx); + let c_b = vqtbl1q_s8(codebook, b_idx); + + let prod_lo = vmull_s8(vget_low_s8(c_a), vget_low_s8(c_b)); + let prod_hi = vmull_high_s8(c_a, c_b); + + acc = vpadalq_s16(acc, prod_lo); + acc = vpadalq_s16(acc, prod_hi); + } + + let simd_bytes = n_full * 8; + let acc_i64 = i64::from(vaddvq_s32(acc)) + + super::score_4bit_internal_integer(&a[simd_bytes..], &b[simd_bytes..]); + acc_i64 as f32 / (CODEBOOK_SCALE * CODEBOOK_SCALE) + } +} + +/// SDOT variant of [`score_4bit_internal_neon`]. Two independent `i32×4` +/// accumulators consume two chunks (32 elements) per iteration, following +/// the same 2× unroll pattern as [`Query4bitSimd::dotprod_raw_neon_sdot`]. +/// +/// # Safety +/// CPU must support `neon` and `dotprod`. +#[target_feature(enable = "neon,dotprod")] +pub unsafe fn score_4bit_internal_neon_sdot(a: &[u8], b: &[u8]) -> f32 { + use core::arch::aarch64::*; + + assert_eq!( + a.len(), + b.len(), + "score_4bit_internal_neon_sdot: vector length mismatch ({} vs {})", + a.len(), + b.len(), + ); + + unsafe { + let codebook = vld1q_s8(CODEBOOK_I8.as_ptr()); + let mut acc_0 = vdupq_n_s32(0); + let mut acc_1 = vdupq_n_s32(0); + let nibble_mask_q = vdupq_n_u8(0x0F); + + let n_pairs = a.len() / 16; + + for i in 0..n_pairs { + let va = vld1q_u8(a.as_ptr().add(16 * i)); + let va_lo = vandq_u8(va, nibble_mask_q); + let va_hi = vshrq_n_u8(va, 4); + let a_idx_0 = vzip1q_u8(va_lo, va_hi); + let a_idx_1 = vzip2q_u8(va_lo, va_hi); + + let vb = vld1q_u8(b.as_ptr().add(16 * i)); + let vb_lo = vandq_u8(vb, nibble_mask_q); + let vb_hi = vshrq_n_u8(vb, 4); + let b_idx_0 = vzip1q_u8(vb_lo, vb_hi); + let b_idx_1 = vzip2q_u8(vb_lo, vb_hi); + + let c_a_0 = vqtbl1q_s8(codebook, a_idx_0); + let c_a_1 = vqtbl1q_s8(codebook, a_idx_1); + let c_b_0 = vqtbl1q_s8(codebook, b_idx_0); + let c_b_1 = vqtbl1q_s8(codebook, b_idx_1); + + core::arch::asm!( + "sdot {acc:v}.4s, {a:v}.16b, {b:v}.16b", + acc = inout(vreg) acc_0, + a = in(vreg) c_a_0, + b = in(vreg) c_b_0, + options(pure, nomem, nostack, preserves_flags), + ); + core::arch::asm!( + "sdot {acc:v}.4s, {a:v}.16b, {b:v}.16b", + acc = inout(vreg) acc_1, + a = in(vreg) c_a_1, + b = in(vreg) c_b_1, + options(pure, nomem, nostack, preserves_flags), + ); + } + + let acc = vaddq_s32(acc_0, acc_1); + let simd_bytes = n_pairs * 16; + let acc_i64 = i64::from(vaddvq_s32(acc)) + + super::score_4bit_internal_integer(&a[simd_bytes..], &b[simd_bytes..]); + acc_i64 as f32 / (CODEBOOK_SCALE * CODEBOOK_SCALE) + } +} + +#[cfg(test)] +mod tests { + use rand::SeedableRng as _; + use rand::prelude::StdRng; + + use super::super::super::shared::pack_codes; + use super::super::shared::{PARITY_DIMS, random_inputs}; + use super::super::{Query4bitSimd, score_4bit_internal_scalar}; + use super::{score_4bit_internal_neon, score_4bit_internal_neon_sdot}; + + #[test] + fn test_neon_matches_scalar() { + let mut rng = StdRng::seed_from_u64(7); + for &dim in PARITY_DIMS { + let (simd_query, vector) = random_inputs(&mut rng, dim); + let scalar = simd_query.dotprod_raw(&vector); + let neon = unsafe { simd_query.dotprod_raw_neon(&vector) }; + assert_eq!(scalar, neon, "scalar {scalar} != neon {neon} at dim {dim}"); + } + } + + #[test] + fn test_neon_sdot_matches_scalar() { + if !std::arch::is_aarch64_feature_detected!("dotprod") { + return; + } + let mut rng = StdRng::seed_from_u64(7); + for &dim in PARITY_DIMS { + let (simd_query, vector) = random_inputs(&mut rng, dim); + let scalar = simd_query.dotprod_raw(&vector); + let sdot = unsafe { simd_query.dotprod_raw_neon_sdot(&vector) }; + assert_eq!(scalar, sdot, "scalar {scalar} != sdot {sdot} at dim {dim}"); + } + } + + /// Single saturation-safety check at an extreme dim (64K) with the + /// worst-case combination: query maxed out and every lane of the vector + /// pointing at the extreme-magnitude codebook slot. Scalar is the + /// reference (i64 throughout, saturation-free by construction); each + /// SIMD path must match it exactly. A mismatch proves that some + /// intermediate integer saturated or overflowed. + #[test] + fn test_saturation_safety_64k() { + let dim = 65_536; + let query = vec![1.0_f32; dim]; + let indices: Vec = vec![15; dim]; // CODEBOOK_I8[15] = +127 + let vector = pack_codes(&indices, 4); + + let q = Query4bitSimd::new(&query); + let scalar = q.dotprod_raw(&vector); + + unsafe { + let neon = q.dotprod_raw_neon(&vector); + assert_eq!(scalar, neon, "neon disagrees at dim={dim}"); + + if std::arch::is_aarch64_feature_detected!("dotprod") { + let sdot = q.dotprod_raw_neon_sdot(&vector); + assert_eq!(scalar, sdot, "sdot disagrees at dim={dim}"); + } + } + } + + /// Parity: NEON `score_4bit_internal` variants must reproduce the scalar + /// reference bit-exactly. Integer accumulators are identical across + /// paths (no saturating intermediates at parity-test dims), so the f32 + /// outputs should match `assert_eq!` down to the last ulp. + #[test] + fn test_score_neon_matches_scalar() { + let mut rng = StdRng::seed_from_u64(7); + for &dim in PARITY_DIMS { + let (_, vec_a) = random_inputs(&mut rng, dim); + let (_, vec_b) = random_inputs(&mut rng, dim); + let scalar = score_4bit_internal_scalar(&vec_a, &vec_b); + + let neon = unsafe { score_4bit_internal_neon(&vec_a, &vec_b) }; + assert_eq!( + scalar, neon, + "score: scalar {scalar} != neon {neon} at dim {dim}" + ); + + if std::arch::is_aarch64_feature_detected!("dotprod") { + let sdot = unsafe { score_4bit_internal_neon_sdot(&vec_a, &vec_b) }; + assert_eq!( + scalar, sdot, + "score: scalar {scalar} != sdot {sdot} at dim {dim}" + ); + } + } + } + + /// Saturation-safety at 64K: both vectors every index 15 → every product + /// hits `CODEBOOK_I8[15]² = 127² = 16 129`. Total = 16 384 × 16 129 ≈ + /// 264 M which comfortably fits i32 (~8× headroom), but any intermediate + /// i16 accumulator (pre-vpadalq) maxes out at 16 129 < 32 767 so no + /// saturation. SDOT accumulates into i32 directly. + #[test] + fn test_score_saturation_safety_64k() { + let dim = 65_536; + let indices: Vec = vec![15; dim]; // CODEBOOK_I8[15] = +127 + let vec_a = pack_codes(&indices, 4); + let vec_b = pack_codes(&indices, 4); + + let scalar = score_4bit_internal_scalar(&vec_a, &vec_b); + + unsafe { + let neon = score_4bit_internal_neon(&vec_a, &vec_b); + assert_eq!(scalar, neon, "score neon disagrees at dim={dim}"); + + if std::arch::is_aarch64_feature_detected!("dotprod") { + let sdot = score_4bit_internal_neon_sdot(&vec_a, &vec_b); + assert_eq!(scalar, sdot, "score sdot disagrees at dim={dim}"); + } + } + } +} diff --git a/lib/quantization/src/turboquant/simd/query4bit/mod.rs b/lib/quantization/src/turboquant/simd/query4bit/mod.rs new file mode 100644 index 0000000000..131d6242a6 --- /dev/null +++ b/lib/quantization/src/turboquant/simd/query4bit/mod.rs @@ -0,0 +1,682 @@ +//! 4-bit product-quantization scoring. +//! +//! # Encoding strategy (arch-specific) +//! +//! Both architectures share the same storage layout for `query_data` +//! (`Vec<[[i8; 16]; 2]>`) and the same `dot_raw` → float reconstruction formula +//! (`postprocess_scale · (dot_raw − bias_correction)`), but the numerical +//! encoding of the codebook and query differs to squeeze the most precision out +//! of each SIMD instruction set: +//! +//! * **aarch64** — `vmull_s8` and `sdot` are true `i8 × i8 → i16/i32` signed +//! multiplies, so we can store the full `i8 ∈ [−127, 127]` codebook directly +//! with no offset. Query halves are full `i8 ∈ [−128, 127]` combined as +//! `q_signed = 256 · high + low`, giving ~15.9-bit query precision. The +//! reconstruction needs no bias correction. +//! +//! * **x86_64** — `_mm_maddubs_epi16` and `VPDPBUSD` consume one `u8` and one +//! `i8` operand. To carry full 8-bit codebook magnitude we feed it unsigned +//! `c_u ∈ [0, 255]` (shifted from signed by `+128`) and keep the query halves +//! narrower to stay under i16 pair-sum saturation +//! (`c_u ≤ 255, q ∈ [−64, 63]` → `|pair| ≤ 2·255·64 = 32 640 < 32 767` ✓). +//! Query halves are 7-bit signed combined as `q_signed = 128 · high + low`, +//! giving ~13.9-bit query precision. The shift contributes a per-query +//! bias `128 · Σ q_signed` that we subtract once in `dotprod`. +//! +//! Both codebooks derive from `CENTROIDS_4BIT` (Lloyd-Max on N(0,1)); see +//! `test_codebook_matches_lloyd_max` for the consistency check. + +/// `max|c|` over `CENTROIDS_4BIT` — the extreme centroid. Shared by both archs. +const CODEBOOK_ABS_MAX: f32 = 2.733; + +/// Full `i8` signed codebook for aarch64. `c_scale = 127 / max|c|` so the +/// extremes hit ±127. `c_signed[k] = CODEBOOK_I8[k]` directly — no offset. +#[cfg(all(target_arch = "aarch64", target_feature = "neon"))] +const CODEBOOK_I8: [i8; 16] = [ + -127, -96, -75, -58, -44, -31, -18, -6, 6, 18, 31, 44, 58, 75, 96, 127, +]; + +/// Codebook scale on aarch64: `c_scale = 127 / max|c|`. +#[cfg(all(target_arch = "aarch64", target_feature = "neon"))] +const CODEBOOK_SCALE: f32 = 127.0 / CODEBOOK_ABS_MAX; + +/// Aarch64 stores the codebook already signed, so no shift-recovery is needed. +/// Kept as a uniform symbol so `new()` / `dotprod` don't need cfg branches. +#[cfg(all(target_arch = "aarch64", target_feature = "neon"))] +const CODEBOOK_OFFSET: i64 = 0; + +/// Maximum signed-integer magnitude the query encoder targets. Derived so the +/// balanced `q_signed = K · high + low` split keeps both halves inside i8 (on +/// aarch64: |low|, |high| ≤ 128 with K=256; on x86_64: |low|, |high| ≤ 64 with +/// K=128 to satisfy maddubs saturation given the full u8 codebook). +#[cfg(all(target_arch = "aarch64", target_feature = "neon"))] +const QUERY_ABS_MAX: f32 = 32639.0; +#[cfg(all(target_arch = "aarch64", target_feature = "neon"))] +const QUERY_HIGH_COEF: i64 = 256; + +/// Full `u8` unsigned codebook for x86_64. `c_scale = 128 / max|c|` and +/// `c_u[k] = c_signed[k] + 128` puts the centroids into `[0, 255]`. +/// `maddubs` / `VPDPBUSD` consume this directly as their u8 operand. +#[cfg(target_arch = "x86_64")] +const CODEBOOK_U8: [u8; 16] = [ + 0, 31, 52, 69, 84, 97, 110, 122, 134, 146, 159, 172, 187, 204, 225, 255, +]; + +/// Codebook shift: `c_signed[k] = CODEBOOK_U8[k] − CODEBOOK_OFFSET`. Pair this +/// with `bias_correction = OFFSET · Σ q_signed` in `dotprod` to recover +/// `Σ q · c_signed` from the raw `Σ q · c_u`. +#[cfg(target_arch = "x86_64")] +const CODEBOOK_OFFSET: i64 = 128; + +/// Codebook scale on x86_64: `c_scale = 128 / max|c|`. +#[cfg(target_arch = "x86_64")] +const CODEBOOK_SCALE: f32 = 128.0 / CODEBOOK_ABS_MAX; + +/// Max q_signed on x86_64 — 7-bit signed half × 128 → roughly ±8127. Actual +/// symmetric cap of 8127 = 128·63 + 63 keeps both halves in `[−64, 63]`. +#[cfg(target_arch = "x86_64")] +const QUERY_ABS_MAX: f32 = 8127.0; +#[cfg(target_arch = "x86_64")] +const QUERY_HIGH_COEF: i64 = 128; + +// Fallback constants for architectures with neither NEON SIMD nor x86_64. +// They match the x86_64 scheme so the scalar reference path can run and +// produce numerically identical results to whichever arch is active. +#[cfg(not(any( + all(target_arch = "aarch64", target_feature = "neon"), + target_arch = "x86_64", +)))] +const CODEBOOK_U8: [u8; 16] = [ + 0, 31, 52, 69, 84, 97, 110, 122, 134, 146, 159, 172, 187, 204, 225, 255, +]; +#[cfg(not(any( + all(target_arch = "aarch64", target_feature = "neon"), + target_arch = "x86_64", +)))] +const CODEBOOK_OFFSET: i64 = 128; +#[cfg(not(any( + all(target_arch = "aarch64", target_feature = "neon"), + target_arch = "x86_64", +)))] +const CODEBOOK_SCALE: f32 = 128.0 / CODEBOOK_ABS_MAX; +#[cfg(not(any( + all(target_arch = "aarch64", target_feature = "neon"), + target_arch = "x86_64", +)))] +const QUERY_ABS_MAX: f32 = 8127.0; +#[cfg(not(any( + all(target_arch = "aarch64", target_feature = "neon"), + target_arch = "x86_64", +)))] +const QUERY_HIGH_COEF: i64 = 128; + +/// Read the codebook value at `idx` in its arch-native storage form as `i64`. +/// On aarch64 that's the signed `CODEBOOK_I8`; on x86_64 (and fallback) it's +/// the unsigned `CODEBOOK_U8` — the `+OFFSET` shift is unwound later by the +/// query-side `bias_correction`. Use this inside SIMD-adjacent code that +/// mirrors what the intrinsics actually see. +#[inline] +fn codebook_value_i64(idx: u8) -> i64 { + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + { + i64::from(CODEBOOK_I8[idx as usize]) + } + #[cfg(not(all(target_arch = "aarch64", target_feature = "neon")))] + { + i64::from(CODEBOOK_U8[idx as usize]) + } +} + +/// Read the codebook value at `idx` as a **true signed** integer, regardless +/// of which storage the current arch uses. Used by vector-vs-vector scoring +/// where both operands come from the codebook and there's no query-side +/// `bias_correction` to absorb the `+OFFSET` shift. +#[inline] +fn codebook_signed_i64(idx: u8) -> i64 { + codebook_value_i64(idx) - CODEBOOK_OFFSET +} + +/// Encoded query for asymmetric 4-bit PQ scoring. +/// +/// # Encoding +/// The f32 query is quantized to signed integers +/// `q_signed ∈ [−QUERY_ABS_MAX, QUERY_ABS_MAX]` and split into two i8 halves +/// combined as `q_signed = QUERY_HIGH_COEF · high + low` — see the module-level +/// docs for the per-arch values of `QUERY_HIGH_COEF` (256 on aarch64, 128 on +/// x86_64) and the reasoning behind the split. Storage is 16-dim chunks of +/// `[low, high]` plus a scalar-handled tail of up to 14 dims (`dim % 2 == 0`). +/// A matryoshka-trimmed model at any even dim fits without re-encoding. +/// +/// # Scoring +/// `dotprod_raw = Σ_j q_signed[j] · c_raw[v[j]]`. The float result is +/// `postprocess_scale · (dot_raw − bias_correction)`, where `bias_correction` +/// absorbs the `+OFFSET` shift in the x86 unsigned-codebook layout and is 0 +/// on aarch64 (signed codebook, no shift). +pub struct Query4bitSimd { + /// Full 16-dim chunks of the query, each stored as `[low, high]` i8 halves. + query_data: Vec<[[i8; 16]; 2]>, + /// Trailing dims that don't fill a 16-dim chunk — up to 14 (since + /// `dim % 2 == 0`). Arrays are sized to a full 16-lane SIMD register + /// (zero-padded beyond `tail_dims`) so arch backends can feed them + /// straight into one extra `maddubs` / `vmull_s8` iteration on top of a + /// zero-padded data chunk, replacing the scalar 14-iteration loop. + tail_low: [i8; 16], + tail_high: [i8; 16], + /// Number of meaningful entries in the tail arrays (`0..=14`, always even). + tail_dims: u8, + /// `1 / (q_scale · c_scale)` — prefactor from integer to float dot product. + postprocess_scale: f32, + /// `CODEBOOK_OFFSET · Σ q_signed[j]` — sums over **all** dims (full + /// chunks + tail). Subtracted from `dot_raw` to recover the true signed + /// dot product. `0` on aarch64, where the codebook is already signed. + bias_correction: i64, +} + +impl Query4bitSimd { + /// Query dim must be a multiple of 2 (the 4-bit packing width: two codes + /// per byte). Any such dim is accepted — dims that don't fill a full + /// 16-dim chunk produce up to a 14-dim tail handled scalar-wise in every + /// SIMD path. This makes `Query4bitSimd` Matryoshka-friendly: a model + /// trimmed to 640 / 768 / 896 / 1024 dims all work with the same storage. + pub fn new(data: &[f32]) -> Self { + assert!( + data.len().is_multiple_of(2), + "Query4bitSimd requires query dim to be a multiple of 2 (got {})", + data.len(), + ); + + let q_abs_max = data + .iter() + .copied() + .map(f32::abs) + .fold(0.0_f32, f32::max) + .max(f32::EPSILON); + let q_scale = QUERY_ABS_MAX / q_abs_max; + + let k = QUERY_HIGH_COEF as i32; + let half_k = k / 2; + let clamp_hi = QUERY_ABS_MAX; + let clamp_lo = -QUERY_ABS_MAX; + + // Balanced signed split, same math for full chunks and tail. + let encode = |value: f32| -> (i8, i8, i64) { + let q_signed = (value * q_scale).round().clamp(clamp_lo, clamp_hi) as i32; + let l_mod = q_signed.rem_euclid(k); + let l = if l_mod >= half_k { l_mod - k } else { l_mod } as i8; + let h = ((q_signed - i32::from(l)) / k) as i8; + (l, h, i64::from(q_signed)) + }; + + let num_chunks = data.len() / 16; + let full_dims = num_chunks * 16; + let tail_dims = data.len() - full_dims; + debug_assert!(tail_dims < 16 && tail_dims.is_multiple_of(2)); + + let mut query_data: Vec<[[i8; 16]; 2]> = Vec::with_capacity(num_chunks); + let mut sum_q_signed: i64 = 0; + for chunk_idx in 0..num_chunks { + let mut low = [0_i8; 16]; + let mut high = [0_i8; 16]; + for i in 0..16 { + let (l, h, q) = encode(data[chunk_idx * 16 + i]); + low[i] = l; + high[i] = h; + sum_q_signed += q; + } + query_data.push([low, high]); + } + + let mut tail_low = [0_i8; 16]; + let mut tail_high = [0_i8; 16]; + for i in 0..tail_dims { + let (l, h, q) = encode(data[full_dims + i]); + tail_low[i] = l; + tail_high[i] = h; + sum_q_signed += q; + } + + Self { + query_data, + tail_low, + tail_high, + tail_dims: tail_dims as u8, + postprocess_scale: 1.0 / (q_scale * CODEBOOK_SCALE), + bias_correction: CODEBOOK_OFFSET * sum_q_signed, + } + } + + /// Number of `vector` bytes the encoded query expects: 8 bytes per full + /// 16-dim chunk plus the packed tail (two 4-bit codes per byte). + #[inline] + pub(super) fn expected_vector_bytes(&self) -> usize { + self.query_data.len() * 8 + (self.tail_dims as usize).div_ceil(2) + } + + /// Score the encoded query against a 4-bit PQ-encoded `vector` + /// (two centroid indices per byte; low nibble = even lane, high nibble + /// = odd lane). `vector.len()` must equal `ceil(dim / 2)` — full chunks + /// first, then the tail bytes covering `tail_dims`. + /// + /// Dispatches at runtime to the best SIMD backend available on the host + /// CPU (AVX-512 VNNI → AVX2 → SSE → NEON + SDOT → NEON → scalar). + pub fn dotprod(&self, vector: &[u8]) -> f32 { + // No per-vector correction loop: `bias_correction` was baked in at `new()`. + let dot_raw = self.dotprod_raw_best(vector); + self.postprocess_scale * (dot_raw - self.bias_correction) as f32 + } + + #[inline] + fn dotprod_raw_best(&self, vector: &[u8]) -> i64 { + #[cfg(target_arch = "x86_64")] + { + if std::is_x86_feature_detected!("avx512f") + && std::is_x86_feature_detected!("avx512bw") + && std::is_x86_feature_detected!("avx512vnni") + { + return unsafe { self.dotprod_raw_avx512_vnni(vector) }; + } + if std::is_x86_feature_detected!("avx2") { + return unsafe { self.dotprod_raw_avx2(vector) }; + } + if std::is_x86_feature_detected!("sse4.1") && std::is_x86_feature_detected!("ssse3") { + return unsafe { self.dotprod_raw_sse(vector) }; + } + } + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + { + if std::arch::is_aarch64_feature_detected!("dotprod") { + return unsafe { self.dotprod_raw_neon_sdot(vector) }; + } + return unsafe { self.dotprod_raw_neon(vector) }; + } + #[allow(unreachable_code)] + self.dotprod_raw(vector) + } + + /// Compute `Σ q_signed[j] · c_raw[v[j]]` over all dims — both the + /// full-chunk section (SIMD-friendly 16 lanes per chunk) and the tail. + /// Returns `acc_low + QUERY_HIGH_COEF · acc_high`. + /// + /// `vector` is PQ-encoded with two 4-bit codebook indices packed per byte: + /// the low nibble is the index for the even lane (j = 2k), the high nibble + /// for the odd lane (j = 2k + 1). + pub fn dotprod_raw(&self, vector: &[u8]) -> i64 { + assert_eq!( + vector.len(), + self.expected_vector_bytes(), + "Query4bitSimd::dotprod_raw: vector length mismatch ({} vs expected {})", + vector.len(), + self.expected_vector_bytes(), + ); + let mut acc_low: i64 = 0; + let mut acc_high: i64 = 0; + for (chunk_idx, [low, high]) in self.query_data.iter().enumerate() { + let v = &vector[chunk_idx * 8..(chunk_idx + 1) * 8]; + for i in 0..16 { + let byte = v[i / 2]; + let idx = if i & 1 == 0 { byte & 0x0F } else { byte >> 4 }; + let c = codebook_value_i64(idx); + acc_low += i64::from(low[i]) * c; + acc_high += i64::from(high[i]) * c; + } + } + acc_low + QUERY_HIGH_COEF * acc_high + self.dotprod_raw_tail(vector) + } + + /// Scalar contribution from the `tail_dims` trailing query entries. + /// Used by the scalar [`Self::dotprod_raw`] reference. SIMD backends + /// have their own tail helpers that feed one zero-padded chunk into the + /// same kernel used for full chunks. + #[inline] + pub(super) fn dotprod_raw_tail(&self, vector: &[u8]) -> i64 { + if self.tail_dims == 0 { + return 0; + } + let tail_byte_start = self.query_data.len() * 8; + let mut acc_low: i64 = 0; + let mut acc_high: i64 = 0; + for i in 0..self.tail_dims as usize { + let byte = vector[tail_byte_start + i / 2]; + let idx = if i & 1 == 0 { byte & 0x0F } else { byte >> 4 }; + let c = codebook_value_i64(idx); + acc_low += i64::from(self.tail_low[i]) * c; + acc_high += i64::from(self.tail_high[i]) * c; + } + acc_low + QUERY_HIGH_COEF * acc_high + } + + /// Prepare an 8-byte zero-padded scratch buffer with the packed tail data, + /// suitable for feeding into a single SSE / NEON chunk kernel. + /// Returns `None` when there is no tail. The returned buffer has the + /// same nibble layout as any full chunk in `vector` — low nibble = even + /// lane, high nibble = odd lane. Unused lanes are zero: with + /// `tail_low[tail_dims..] = tail_high[tail_dims..] = 0` they contribute + /// nothing to `maddubs` / `vmull` products. + #[inline] + pub(super) fn tail_chunk_scratch(&self, vector: &[u8]) -> Option<[u8; 8]> { + if self.tail_dims == 0 { + return None; + } + let tail_byte_start = self.query_data.len() * 8; + let tail_bytes = (self.tail_dims as usize).div_ceil(2); + let mut buf = [0u8; 8]; + buf[..tail_bytes].copy_from_slice(&vector[tail_byte_start..tail_byte_start + tail_bytes]); + Some(buf) + } +} + +/// Dot product between two already-encoded 4-bit PQ vectors. Both `a` and +/// `b` are the packed-nibble format that [`Query4bitSimd::dotprod`] takes as +/// its `vector` argument — every byte holds two codebook indices (low nibble +/// = even lane, high nibble = odd lane). +/// +/// Computes `Σ c[a[j]] · c[b[j]]` in centroid-float space. Dispatches to the +/// fastest available SIMD implementation at runtime and falls back to +/// [`score_4bit_internal_scalar`] otherwise. Any byte length is accepted — +/// bytes that don't fill a full SIMD chunk are folded in scalar-wise so that +/// Matryoshka-style dim ∈ {2k: k ∈ ℕ} all work. +/// +/// # Panics +/// Panics if the two vectors have different lengths. +pub fn score_4bit_internal(a: &[u8], b: &[u8]) -> f32 { + assert_eq!( + a.len(), + b.len(), + "score_4bit_internal: vector length mismatch ({} vs {})", + a.len(), + b.len(), + ); + + #[cfg(target_arch = "x86_64")] + { + if std::is_x86_feature_detected!("avx512f") + && std::is_x86_feature_detected!("avx512bw") + && std::is_x86_feature_detected!("avx512vnni") + { + return unsafe { x64::score_4bit_internal_avx512_vnni(a, b) }; + } + if std::is_x86_feature_detected!("avx2") { + return unsafe { x64::score_4bit_internal_avx2(a, b) }; + } + if std::is_x86_feature_detected!("sse4.1") && std::is_x86_feature_detected!("ssse3") { + return unsafe { x64::score_4bit_internal_sse(a, b) }; + } + } + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + { + if std::arch::is_aarch64_feature_detected!("dotprod") { + return unsafe { arm::score_4bit_internal_neon_sdot(a, b) }; + } + return unsafe { arm::score_4bit_internal_neon(a, b) }; + } + #[allow(unreachable_code)] + score_4bit_internal_scalar(a, b) +} + +/// Scalar reference implementation of [`score_4bit_internal`]. Exposed as +/// the fallback on architectures without a SIMD variant, the `assert_eq!` +/// baseline for per-arch parity tests, and as a standalone bench target. +/// +/// Caller is responsible for checking the length preconditions — the public +/// [`score_4bit_internal`] enforces them before dispatching. +pub fn score_4bit_internal_scalar(a: &[u8], b: &[u8]) -> f32 { + // c_signed ≈ c_float · c_scale → c_signed_a · c_signed_b ≈ c_float_a · c_float_b · c_scale². + score_4bit_internal_integer(a, b) as f32 / (CODEBOOK_SCALE * CODEBOOK_SCALE) +} + +/// Integer-only scalar kernel shared by all backends — used by SIMD paths to +/// fold in any bytes that didn't fit a full SIMD chunk (and by +/// [`score_4bit_internal_scalar`] as its inner loop). +#[inline] +pub(super) fn score_4bit_internal_integer(a: &[u8], b: &[u8]) -> i64 { + let mut acc: i64 = 0; + for (&byte_a, &byte_b) in a.iter().zip(b.iter()) { + let a_lo = byte_a & 0x0F; + let a_hi = byte_a >> 4; + let b_lo = byte_b & 0x0F; + let b_hi = byte_b >> 4; + acc += codebook_signed_i64(a_lo) * codebook_signed_i64(b_lo); + acc += codebook_signed_i64(a_hi) * codebook_signed_i64(b_hi); + } + acc +} + +#[cfg(all(target_arch = "aarch64", target_feature = "neon"))] +mod arm; + +#[cfg(target_arch = "x86_64")] +mod x64; + +#[cfg(all(target_arch = "aarch64", target_feature = "neon"))] +pub use arm::{score_4bit_internal_neon, score_4bit_internal_neon_sdot}; +#[cfg(target_arch = "x86_64")] +pub use x64::{score_4bit_internal_avx2, score_4bit_internal_avx512_vnni, score_4bit_internal_sse}; + +/// 4-bit-specific test helpers. Bit-width-agnostic helpers (`pack_codes`, +/// `sample_normal_vec`, `encode_to_nearest_centroid`) live in +/// [`super::super::shared`]. +#[cfg(test)] +pub(super) mod shared { + use rand::prelude::StdRng; + use rand::seq::SliceRandom; + + use super::super::shared::{pack_codes, sample_normal_vec}; + use super::Query4bitSimd; + + /// Corner-case dims covering every tail size the 4-bit pipeline can + /// produce, for each SIMD backend: + /// • `16, 32, 128, 256, 1024, 2048` — full chunks, no tail (baseline). + /// • `48` — 3 chunks (odd SDOT/VNNI leftover), no tail. + /// • `18, 30, 46, 62, 1026, 2046` — full chunks + non-zero tail + /// (tail sizes 2, 14, 14, 14, 2, 14 dims respectively). + /// • `270` — 16 chunks + 14-dim tail (exercises largest tail at a + /// realistic matryoshka dim). + pub const PARITY_DIMS: &[usize] = &[ + 16, 18, 30, 32, 46, 48, 62, 128, 256, 270, 1024, 1026, 2046, 2048, + ]; + + /// Parity-test helper: query ~ N(0,1), balanced index distribution (each + /// centroid appears `dim/16` times, shuffled). Index distribution doesn't + /// affect scalar-vs-SIMD parity; we just need non-trivial data. + pub fn random_inputs(rng: &mut StdRng, dim: usize) -> (Query4bitSimd, Vec) { + let query = sample_normal_vec(rng, dim); + let mut indices: Vec = (0..dim).map(|i| (i % 16) as u8).collect(); + indices.shuffle(rng); + (Query4bitSimd::new(&query), pack_codes(&indices, 4)) + } +} + +/// Accuracy / precision tests for the public `Query4bitSimd` API. +/// +/// Per-arch SIMD parity tests and the saturation-safety test live in the +/// `arm` / `x64` submodules — they verify that each SIMD implementation +/// matches the scalar reference `dotprod_raw`, not the float ground truth. +#[cfg(test)] +mod tests { + // Anonymous `use _` brings the trait into scope for `StdRng::seed_from_u64` + // without introducing a name rustc then flags as unused. + use rand::SeedableRng as _; + use rand::prelude::StdRng; + + use super::super::shared::{encode_to_nearest_centroid, pack_codes, sample_normal_vec}; + use super::{CODEBOOK_ABS_MAX, Query4bitSimd}; + use crate::turboquant::TQBits; + + /// Whichever codebook representation the current arch uses (signed i8 on + /// aarch64, shifted u8 on x86_64), it must match what the runtime recipe + /// would produce from `CENTROIDS_4BIT`. + #[test] + fn test_codebook_matches_lloyd_max() { + let centroids = TQBits::Bits4.get_centroids(); + assert_eq!(centroids.len(), 16); + + let c_abs_max = centroids + .iter() + .copied() + .map(f32::abs) + .fold(0.0_f32, f32::max); + assert!( + (CODEBOOK_ABS_MAX - c_abs_max).abs() < 1e-6, + "CODEBOOK_ABS_MAX ({CODEBOOK_ABS_MAX}) != max|CENTROIDS_4BIT| ({c_abs_max})" + ); + + #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] + { + let c_scale = 127.0 / c_abs_max; + let quantized: [i8; 16] = std::array::from_fn(|k| { + (centroids[k] * c_scale).round().clamp(-127.0, 127.0) as i8 + }); + assert_eq!( + quantized, + super::CODEBOOK_I8, + "const CODEBOOK_I8 drifted from CENTROIDS_4BIT Lloyd-Max quantization", + ); + } + #[cfg(not(all(target_arch = "aarch64", target_feature = "neon")))] + { + let c_scale = 128.0 / c_abs_max; + let offset = super::CODEBOOK_OFFSET as i32; + let quantized: [u8; 16] = std::array::from_fn(|k| { + let signed = (centroids[k] * c_scale).round().clamp(-128.0, 127.0) as i32; + (signed + offset) as u8 + }); + assert_eq!( + quantized, + super::CODEBOOK_U8, + "const CODEBOOK_U8 drifted from CENTROIDS_4BIT Lloyd-Max quantization", + ); + } + } + + /// Reconstruction accuracy on realistic PQ inputs: query ∼ N(0,1), vector + /// drawn from N(0,1) then mapped to its nearest centroid. We compare + /// `simd.dotprod()` against the "ideal" PQ dot (sum of `q[j] · c[v[j]]` + /// with float-precision centroid lookup) — the error our SIMD path adds + /// over a hypothetical perfect-precision PQ should be tiny. + /// + /// Parameterized over matryoshka-style corner-case dims to exercise the + /// tail-handling logic end-to-end (not just bit-exact parity). + #[rstest::rstest] + #[case::full_chunks(256)] + #[case::small_tail(18)] + #[case::max_tail(30)] + #[case::odd_chunks_only(48)] + #[case::odd_chunks_plus_tail(62)] + #[case::matryoshka(270)] + #[case::large_with_tail(2046)] + fn test_dotprod_matches_float(#[case] dim: usize) { + let mut rng = StdRng::seed_from_u64(42); + let n_trials = 64; + + let centroids = TQBits::Bits4.get_centroids(); + + for _ in 0..n_trials { + let query = sample_normal_vec(&mut rng, dim); + let v_raw = sample_normal_vec(&mut rng, dim); + let indices = encode_to_nearest_centroid(centroids, &v_raw); + let v_pq: Vec = indices.iter().map(|&k| centroids[k as usize]).collect(); + + let pq_dot: f32 = query.iter().zip(v_pq.iter()).map(|(a, b)| a * b).sum(); + let simd_dot = Query4bitSimd::new(&query).dotprod(&pack_codes(&indices, 4)); + + // Error scales roughly like √dim · σ_q · ε_c. Allow a tolerance + // that is comfortably above the 3σ tail for dim up to ~2K. + let tol = (0.5_f32).max(0.03 * (dim as f32).sqrt()); + assert!( + (pq_dot - simd_dot).abs() < tol, + "dim={dim}: simd_dot {simd_dot} too far from ideal PQ dot {pq_dot} (tol={tol})", + ); + } + } + + /// Quantitative proof that our SIMD quantization is negligible next to PQ + /// centroid snapping: RMS error added by our encoding is at least 5× + /// smaller than the RMS error PQ itself introduces. If this invariant + /// ever flips, something in the quantization pipeline lost precision. + #[test] + fn test_simd_noise_below_pq_noise() { + let mut rng = StdRng::seed_from_u64(123); + let dim = 256; + let n_trials = 256; + + let centroids = TQBits::Bits4.get_centroids(); + + let mut sq_pq_noise = 0.0_f64; + let mut sq_simd_noise = 0.0_f64; + + for _ in 0..n_trials { + let query = sample_normal_vec(&mut rng, dim); + let v_raw = sample_normal_vec(&mut rng, dim); + let indices = encode_to_nearest_centroid(centroids, &v_raw); + let v_pq: Vec = indices.iter().map(|&k| centroids[k as usize]).collect(); + + let true_dot: f64 = query + .iter() + .zip(v_raw.iter()) + .map(|(a, b)| f64::from(*a) * f64::from(*b)) + .sum(); + let pq_dot: f64 = query + .iter() + .zip(v_pq.iter()) + .map(|(a, b)| f64::from(*a) * f64::from(*b)) + .sum(); + let simd_dot = f64::from(Query4bitSimd::new(&query).dotprod(&pack_codes(&indices, 4))); + + sq_pq_noise += (pq_dot - true_dot).powi(2); + sq_simd_noise += (simd_dot - pq_dot).powi(2); + } + + let rms_pq_noise = (sq_pq_noise / f64::from(n_trials)).sqrt(); + let rms_simd_noise = (sq_simd_noise / f64::from(n_trials)).sqrt(); + + // Print for easy comparison across encoding variants. + eprintln!( + "NOISE at dim={dim}: pq_rms={rms_pq_noise:.4} simd_rms={rms_simd_noise:.4} \ + ratio={:.2}×", + rms_pq_noise / rms_simd_noise, + ); + + assert!( + rms_simd_noise * 5.0 < rms_pq_noise, + "SIMD noise RMS {rms_simd_noise:.4} should be << PQ noise RMS \ + {rms_pq_noise:.4} (ratio {:.2}×)", + rms_pq_noise / rms_simd_noise, + ); + } + + /// `score_4bit_internal` should recover the pure centroid-space dot + /// product `Σ c[a[j]] · c[b[j]]` up to the i8 quantization step of the + /// codebook (≤ 1/c_scale ≈ 0.022 per centroid). For dim=256 the + /// cumulative RMS error stays well under 1.0. + #[test] + fn test_score_4bit_internal_matches_centroid_product() { + let mut rng = StdRng::seed_from_u64(7); + let centroids = TQBits::Bits4.get_centroids(); + let dim = 256; + let n_trials = 32; + + for _ in 0..n_trials { + let raw_a = sample_normal_vec(&mut rng, dim); + let raw_b = sample_normal_vec(&mut rng, dim); + let idx_a = encode_to_nearest_centroid(centroids, &raw_a); + let idx_b = encode_to_nearest_centroid(centroids, &raw_b); + + let truth: f64 = idx_a + .iter() + .zip(idx_b.iter()) + .map(|(&ia, &ib)| { + f64::from(centroids[ia as usize]) * f64::from(centroids[ib as usize]) + }) + .sum(); + let score = super::score_4bit_internal(&pack_codes(&idx_a, 4), &pack_codes(&idx_b, 4)); + + // Codebook quantization budget: Δc ≈ max|c|/127 ≈ 0.022, so each + // term has error ≲ 2·c·Δc ≲ 0.12; over d=256 independent terms + // the RMS error is ≲ √d · 0.12 ≈ 1.9. 2.0 is a loose 1σ-ish bound. + assert!( + (truth as f32 - score).abs() < 2.0, + "score {score} too far from centroid-product truth {truth}", + ); + } + } +} diff --git a/lib/quantization/src/turboquant/simd/query4bit/x64.rs b/lib/quantization/src/turboquant/simd/query4bit/x64.rs new file mode 100644 index 0000000000..ce7987b6e6 --- /dev/null +++ b/lib/quantization/src/turboquant/simd/query4bit/x64.rs @@ -0,0 +1,644 @@ +//! x86_64 SIMD paths for [`Query4bitSimd`]. +//! +//! The codebook is stored as unsigned u8 ∈ [0, 255] (`CODEBOOK_U8`), which is +//! exactly what `_mm_maddubs_epi16` / `VPDPBUSD` expect in their u8 operand +//! slot; the signed shift `c_signed = c_u − 128` is undone by +//! `Query4bitSimd`'s per-query `bias_correction`. Query halves are 7-bit +//! signed to keep the maddubs pair sum under i16 saturation: +//! c_u ≤ 255, q ∈ [−64, 63] → |pair| ≤ 2·255·64 = 32 640 < 32 767. +//! `QUERY_HIGH_COEF = 128` here — half the aarch64 value because the query +//! halves cover a 7-bit range. + +use super::{CODEBOOK_SCALE, CODEBOOK_U8, QUERY_HIGH_COEF, Query4bitSimd}; + +impl Query4bitSimd { + /// x86_64 SSE4.1 + SSSE3 implementation of [`Query4bitSimd::dotprod_raw`]. + /// + /// # Safety + /// CPU must support `ssse3` and `sse4.1`. + #[target_feature(enable = "sse4.1,ssse3")] + pub unsafe fn dotprod_raw_sse(&self, vector: &[u8]) -> i64 { + use core::arch::x86_64::*; + + assert_eq!( + vector.len(), + self.expected_vector_bytes(), + "Query4bitSimd::dotprod_raw_sse: vector length mismatch ({} vs expected {})", + vector.len(), + self.expected_vector_bytes(), + ); + + unsafe { + let codebook = _mm_loadu_si128(CODEBOOK_U8.as_ptr().cast::<__m128i>()); + let ones = _mm_set1_epi16(1); + let nibble_mask = _mm_set1_epi8(0x0F); + let mut acc_low = _mm_setzero_si128(); + let mut acc_high = _mm_setzero_si128(); + + for (chunk_idx, [low, high]) in self.query_data.iter().enumerate() { + let v_packed = + _mm_loadl_epi64(vector.as_ptr().add(chunk_idx * 8).cast::<__m128i>()); + let v_lo = _mm_and_si128(v_packed, nibble_mask); + let v_hi = _mm_and_si128(_mm_srli_epi16(v_packed, 4), nibble_mask); + let v = _mm_unpacklo_epi8(v_lo, v_hi); + let c_u = _mm_shuffle_epi8(codebook, v); + + let q_low = _mm_loadu_si128(low.as_ptr().cast::<__m128i>()); + let q_high = _mm_loadu_si128(high.as_ptr().cast::<__m128i>()); + + let prod_low = _mm_maddubs_epi16(c_u, q_low); + let prod_high = _mm_maddubs_epi16(c_u, q_high); + acc_low = _mm_add_epi32(acc_low, _mm_madd_epi16(prod_low, ones)); + acc_high = _mm_add_epi32(acc_high, _mm_madd_epi16(prod_high, ones)); + } + + // Tail: one extra SSE chunk on a zero-padded 8-byte scratch. + if let Some(buf) = self.tail_chunk_scratch(vector) { + let v_packed = _mm_loadl_epi64(buf.as_ptr().cast::<__m128i>()); + let v_lo = _mm_and_si128(v_packed, nibble_mask); + let v_hi = _mm_and_si128(_mm_srli_epi16(v_packed, 4), nibble_mask); + let v = _mm_unpacklo_epi8(v_lo, v_hi); + let c_u = _mm_shuffle_epi8(codebook, v); + let q_low = _mm_loadu_si128(self.tail_low.as_ptr().cast::<__m128i>()); + let q_high = _mm_loadu_si128(self.tail_high.as_ptr().cast::<__m128i>()); + let prod_low = _mm_maddubs_epi16(c_u, q_low); + let prod_high = _mm_maddubs_epi16(c_u, q_high); + acc_low = _mm_add_epi32(acc_low, _mm_madd_epi16(prod_low, ones)); + acc_high = _mm_add_epi32(acc_high, _mm_madd_epi16(prod_high, ones)); + } + + i64::from(hsum_i32_sse(acc_low)) + QUERY_HIGH_COEF * i64::from(hsum_i32_sse(acc_high)) + } + } + + /// x86_64 AVX2 implementation. `query_data` stores `[low, high]` as 32 + /// contiguous i8 bytes per chunk, so a single YMM load grabs both halves. + /// Codebook is broadcast to both 128-bit lanes; maddubs pairs the upper + /// lane with `high` and the lower lane with `low`, producing i32 sums + /// split cleanly by lane at the end. + /// + /// # Safety + /// CPU must support `avx2`. + #[target_feature(enable = "avx2")] + pub unsafe fn dotprod_raw_avx2(&self, vector: &[u8]) -> i64 { + use core::arch::x86_64::*; + + assert_eq!( + vector.len(), + self.expected_vector_bytes(), + "Query4bitSimd::dotprod_raw_avx2: vector length mismatch ({} vs expected {})", + vector.len(), + self.expected_vector_bytes(), + ); + + unsafe { + let codebook_128 = _mm_loadu_si128(CODEBOOK_U8.as_ptr().cast::<__m128i>()); + let codebook = _mm256_broadcastsi128_si256(codebook_128); + let ones = _mm256_set1_epi16(1); + let ones_128 = _mm_set1_epi16(1); + let nibble_mask = _mm_set1_epi8(0x0F); + let mut acc = _mm256_setzero_si256(); + + for (chunk_idx, chunk) in self.query_data.iter().enumerate() { + let low_high = _mm256_loadu_si256(chunk.as_ptr().cast::<__m256i>()); + + let v_packed = + _mm_loadl_epi64(vector.as_ptr().add(chunk_idx * 8).cast::<__m128i>()); + let v_lo = _mm_and_si128(v_packed, nibble_mask); + let v_hi = _mm_and_si128(_mm_srli_epi16(v_packed, 4), nibble_mask); + let v128 = _mm_unpacklo_epi8(v_lo, v_hi); + let v = _mm256_broadcastsi128_si256(v128); + let c = _mm256_shuffle_epi8(codebook, v); + + let prods = _mm256_maddubs_epi16(c, low_high); + acc = _mm256_add_epi32(acc, _mm256_madd_epi16(prods, ones)); + } + + let mut acc_low = _mm256_castsi256_si128(acc); + let mut acc_high = _mm256_extracti128_si256(acc, 1); + + // Tail: one extra SSE chunk on a zero-padded scratch — matches + // the SSE variant's post-loop kernel. + if let Some(buf) = self.tail_chunk_scratch(vector) { + let v_packed = _mm_loadl_epi64(buf.as_ptr().cast::<__m128i>()); + let v_lo = _mm_and_si128(v_packed, nibble_mask); + let v_hi = _mm_and_si128(_mm_srli_epi16(v_packed, 4), nibble_mask); + let v = _mm_unpacklo_epi8(v_lo, v_hi); + let c_u = _mm_shuffle_epi8(codebook_128, v); + let q_low = _mm_loadu_si128(self.tail_low.as_ptr().cast::<__m128i>()); + let q_high = _mm_loadu_si128(self.tail_high.as_ptr().cast::<__m128i>()); + let prod_low = _mm_maddubs_epi16(c_u, q_low); + let prod_high = _mm_maddubs_epi16(c_u, q_high); + acc_low = _mm_add_epi32(acc_low, _mm_madd_epi16(prod_low, ones_128)); + acc_high = _mm_add_epi32(acc_high, _mm_madd_epi16(prod_high, ones_128)); + } + + i64::from(hsum_i32_sse(acc_low)) + QUERY_HIGH_COEF * i64::from(hsum_i32_sse(acc_high)) + } + } + + /// AVX-512 VNNI (Ice Lake Xeon+, Zen 4+): 2 chunks per iteration. Two + /// consecutive `[low, high]` entries = 64 bytes = one ZMM load. `VPDPBUSD` + /// fuses the 4-wide u8×i8 dot with i32 accumulation; the ZMM layout puts + /// [low_a, high_a, low_b, high_b] into lanes 0..3. + /// + /// # Safety + /// CPU must support `avx512f`, `avx512bw`, and `avx512vnni`. + #[target_feature(enable = "avx512f,avx512bw,avx512vnni,sse4.1,ssse3")] + pub unsafe fn dotprod_raw_avx512_vnni(&self, vector: &[u8]) -> i64 { + use core::arch::x86_64::*; + + assert_eq!( + vector.len(), + self.expected_vector_bytes(), + "Query4bitSimd::dotprod_raw_avx512_vnni: vector length mismatch ({} vs expected {})", + vector.len(), + self.expected_vector_bytes(), + ); + + unsafe { + let codebook_128 = _mm_loadu_si128(CODEBOOK_U8.as_ptr().cast::<__m128i>()); + let codebook_512 = _mm512_broadcast_i32x4(codebook_128); + let nibble_mask_128 = _mm_set1_epi8(0x0F); + let ones_128 = _mm_set1_epi16(1); + let mut acc = _mm512_setzero_si512(); + + let chunks = self.query_data.as_slice(); + let n_pairs = chunks.len() / 2; + + for i in 0..n_pairs { + let pair_ptr = chunks.as_ptr().add(2 * i).cast::<__m512i>(); + let low_high_pair = _mm512_loadu_si512(pair_ptr); + + let v_packed_16 = _mm_loadu_si128(vector.as_ptr().add(16 * i).cast::<__m128i>()); + let v_lo = _mm_and_si128(v_packed_16, nibble_mask_128); + let v_hi = _mm_and_si128(_mm_srli_epi16(v_packed_16, 4), nibble_mask_128); + let v_chunk_a = _mm_unpacklo_epi8(v_lo, v_hi); + let v_chunk_b = _mm_unpackhi_epi8(v_lo, v_hi); + + let v_dup_a = _mm256_broadcastsi128_si256(v_chunk_a); + let v_dup_b = _mm256_broadcastsi128_si256(v_chunk_b); + let v_512 = _mm512_inserti64x4(_mm512_castsi256_si512(v_dup_a), v_dup_b, 1); + + let c_512 = _mm512_shuffle_epi8(codebook_512, v_512); + acc = _mm512_dpbusd_epi32(acc, c_512, low_high_pair); + } + + let acc_256_lo = _mm512_castsi512_si256(acc); + let acc_256_hi = _mm512_extracti64x4_epi64(acc, 1); + let lane_a_low = _mm256_castsi256_si128(acc_256_lo); + let lane_a_high = _mm256_extracti128_si256(acc_256_lo, 1); + let lane_b_low = _mm256_castsi256_si128(acc_256_hi); + let lane_b_high = _mm256_extracti128_si256(acc_256_hi, 1); + let mut sum_low_xmm = _mm_add_epi32(lane_a_low, lane_b_low); + let mut sum_high_xmm = _mm_add_epi32(lane_a_high, lane_b_high); + + // Odd leftover chunk via SSE-style `maddubs + madd_epi16` — 1 chunk + // is too narrow to benefit from VNNI here. Only reached when + // `query_data.len()` is odd (dim not a multiple of 32). + if chunks.len() % 2 == 1 { + let tail_chunk = 2 * n_pairs; + let [low, high] = chunks[tail_chunk]; + + let v_packed = + _mm_loadl_epi64(vector.as_ptr().add(tail_chunk * 8).cast::<__m128i>()); + let v_lo = _mm_and_si128(v_packed, nibble_mask_128); + let v_hi = _mm_and_si128(_mm_srli_epi16(v_packed, 4), nibble_mask_128); + let v = _mm_unpacklo_epi8(v_lo, v_hi); + let c_u = _mm_shuffle_epi8(codebook_128, v); + + let q_low = _mm_loadu_si128(low.as_ptr().cast::<__m128i>()); + let q_high = _mm_loadu_si128(high.as_ptr().cast::<__m128i>()); + let prod_low = _mm_maddubs_epi16(c_u, q_low); + let prod_high = _mm_maddubs_epi16(c_u, q_high); + sum_low_xmm = _mm_add_epi32(sum_low_xmm, _mm_madd_epi16(prod_low, ones_128)); + sum_high_xmm = _mm_add_epi32(sum_high_xmm, _mm_madd_epi16(prod_high, ones_128)); + } + + // Tail dims (< 14 remaining after the optional leftover chunk): + // one SSE chunk on a zero-padded scratch, same as the SSE variant. + if let Some(buf) = self.tail_chunk_scratch(vector) { + let v_packed = _mm_loadl_epi64(buf.as_ptr().cast::<__m128i>()); + let v_lo = _mm_and_si128(v_packed, nibble_mask_128); + let v_hi = _mm_and_si128(_mm_srli_epi16(v_packed, 4), nibble_mask_128); + let v = _mm_unpacklo_epi8(v_lo, v_hi); + let c_u = _mm_shuffle_epi8(codebook_128, v); + + let q_low = _mm_loadu_si128(self.tail_low.as_ptr().cast::<__m128i>()); + let q_high = _mm_loadu_si128(self.tail_high.as_ptr().cast::<__m128i>()); + let prod_low = _mm_maddubs_epi16(c_u, q_low); + let prod_high = _mm_maddubs_epi16(c_u, q_high); + sum_low_xmm = _mm_add_epi32(sum_low_xmm, _mm_madd_epi16(prod_low, ones_128)); + sum_high_xmm = _mm_add_epi32(sum_high_xmm, _mm_madd_epi16(prod_high, ones_128)); + } + + i64::from(hsum_i32_sse(sum_low_xmm)) + + QUERY_HIGH_COEF * i64::from(hsum_i32_sse(sum_high_xmm)) + } + } +} + +#[target_feature(enable = "sse2")] +unsafe fn hsum_i32_sse(v: core::arch::x86_64::__m128i) -> i32 { + use core::arch::x86_64::*; + let v = _mm_add_epi32(v, _mm_shuffle_epi32(v, 0x4E)); + let v = _mm_add_epi32(v, _mm_shuffle_epi32(v, 0xB1)); + _mm_cvtsi128_si32(v) +} + +// ------------------------------------------------------------------ +// score_4bit_internal — both operands signed, so we can't reuse +// `maddubs` / `VPDPBUSD`. The honest path widens the signed codebook +// bytes to i16 and uses `madd_epi16` (signed × signed → i32 pair-sum); +// AVX-512 VNNI's `VPDPWSSD` is the fused equivalent on ZMM. +// +// We load `CODEBOOK_U8` and XOR with 0x80 to recover the signed i8 form +// (= c_u − 128). The resulting `c_signed` lives in [−128, 127], so: +// per product: |c_a · c_b| ≤ 128·128 = 16 384 +// madd_epi16 pair: ≤ 2·16 384 = 32 768 < i32::MAX ✓ +// i32 acc at 64K: ≤ 16 384·16 384 = 268 M ≪ i32::MAX +// Far away from saturation in every intermediate. +// ------------------------------------------------------------------ + +/// SSE4.1 + SSSE3 implementation of [`super::score_4bit_internal`]. +/// +/// # Safety +/// CPU must support `ssse3` and `sse4.1`. +#[target_feature(enable = "sse4.1,ssse3")] +pub unsafe fn score_4bit_internal_sse(a: &[u8], b: &[u8]) -> f32 { + use core::arch::x86_64::*; + + assert_eq!( + a.len(), + b.len(), + "score_4bit_internal_sse: vector length mismatch ({} vs {})", + a.len(), + b.len(), + ); + + unsafe { + let codebook_i8 = _mm_xor_si128( + _mm_loadu_si128(CODEBOOK_U8.as_ptr().cast::<__m128i>()), + _mm_set1_epi8(-128i8), + ); + let nibble_mask = _mm_set1_epi8(0x0F); + let mut acc = _mm_setzero_si128(); + + let n_full = a.len() / 8; + for i in 0..n_full { + let va = _mm_loadl_epi64(a.as_ptr().add(i * 8).cast::<__m128i>()); + let va_lo = _mm_and_si128(va, nibble_mask); + let va_hi = _mm_and_si128(_mm_srli_epi16(va, 4), nibble_mask); + let a_idx = _mm_unpacklo_epi8(va_lo, va_hi); + let c_a_i8 = _mm_shuffle_epi8(codebook_i8, a_idx); + + let vb = _mm_loadl_epi64(b.as_ptr().add(i * 8).cast::<__m128i>()); + let vb_lo = _mm_and_si128(vb, nibble_mask); + let vb_hi = _mm_and_si128(_mm_srli_epi16(vb, 4), nibble_mask); + let b_idx = _mm_unpacklo_epi8(vb_lo, vb_hi); + let c_b_i8 = _mm_shuffle_epi8(codebook_i8, b_idx); + + let c_a_lo = _mm_cvtepi8_epi16(c_a_i8); + let c_a_hi = _mm_cvtepi8_epi16(_mm_srli_si128(c_a_i8, 8)); + let c_b_lo = _mm_cvtepi8_epi16(c_b_i8); + let c_b_hi = _mm_cvtepi8_epi16(_mm_srli_si128(c_b_i8, 8)); + + let prod_lo = _mm_madd_epi16(c_a_lo, c_b_lo); + let prod_hi = _mm_madd_epi16(c_a_hi, c_b_hi); + acc = _mm_add_epi32(acc, _mm_add_epi32(prod_lo, prod_hi)); + } + + let simd_bytes = n_full * 8; + let sum = i64::from(hsum_i32_sse(acc)) + + super::score_4bit_internal_integer(&a[simd_bytes..], &b[simd_bytes..]); + sum as f32 / (CODEBOOK_SCALE * CODEBOOK_SCALE) + } +} + +/// AVX2 implementation of [`super::score_4bit_internal`]. 32 elements per +/// iteration (16 bytes from each source) using YMM widening and +/// `_mm256_madd_epi16`. +/// +/// # Safety +/// CPU must support `avx2`. +#[target_feature(enable = "avx2")] +pub unsafe fn score_4bit_internal_avx2(a: &[u8], b: &[u8]) -> f32 { + use core::arch::x86_64::*; + + assert_eq!( + a.len(), + b.len(), + "score_4bit_internal_avx2: vector length mismatch ({} vs {})", + a.len(), + b.len(), + ); + + unsafe { + let codebook_i8_128 = _mm_xor_si128( + _mm_loadu_si128(CODEBOOK_U8.as_ptr().cast::<__m128i>()), + _mm_set1_epi8(-128i8), + ); + let codebook_i8 = _mm256_broadcastsi128_si256(codebook_i8_128); + let nibble_mask = _mm_set1_epi8(0x0F); + let mut acc = _mm256_setzero_si256(); + + let n_iters = a.len() / 16; + for i in 0..n_iters { + let va = _mm_loadu_si128(a.as_ptr().add(16 * i).cast::<__m128i>()); + let va_lo = _mm_and_si128(va, nibble_mask); + let va_hi = _mm_and_si128(_mm_srli_epi16(va, 4), nibble_mask); + let a_idx_0 = _mm_unpacklo_epi8(va_lo, va_hi); + let a_idx_1 = _mm_unpackhi_epi8(va_lo, va_hi); + let a_idx_256 = _mm256_inserti128_si256(_mm256_castsi128_si256(a_idx_0), a_idx_1, 1); + let c_a_i8 = _mm256_shuffle_epi8(codebook_i8, a_idx_256); + + let vb = _mm_loadu_si128(b.as_ptr().add(16 * i).cast::<__m128i>()); + let vb_lo = _mm_and_si128(vb, nibble_mask); + let vb_hi = _mm_and_si128(_mm_srli_epi16(vb, 4), nibble_mask); + let b_idx_0 = _mm_unpacklo_epi8(vb_lo, vb_hi); + let b_idx_1 = _mm_unpackhi_epi8(vb_lo, vb_hi); + let b_idx_256 = _mm256_inserti128_si256(_mm256_castsi128_si256(b_idx_0), b_idx_1, 1); + let c_b_i8 = _mm256_shuffle_epi8(codebook_i8, b_idx_256); + + // Widen i8×32 into two i16×16 halves. + let c_a_lo = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(c_a_i8)); + let c_a_hi = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(c_a_i8, 1)); + let c_b_lo = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(c_b_i8)); + let c_b_hi = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(c_b_i8, 1)); + + let prod_lo = _mm256_madd_epi16(c_a_lo, c_b_lo); + let prod_hi = _mm256_madd_epi16(c_a_hi, c_b_hi); + acc = _mm256_add_epi32(acc, _mm256_add_epi32(prod_lo, prod_hi)); + } + + let acc_lo = _mm256_castsi256_si128(acc); + let acc_hi = _mm256_extracti128_si256(acc, 1); + let simd_bytes = n_iters * 16; + let sum = i64::from(hsum_i32_sse(_mm_add_epi32(acc_lo, acc_hi))) + + super::score_4bit_internal_integer(&a[simd_bytes..], &b[simd_bytes..]); + sum as f32 / (CODEBOOK_SCALE * CODEBOOK_SCALE) + } +} + +/// AVX-512 VNNI implementation of [`super::score_4bit_internal`]. Uses +/// `VPDPWSSD` — the signed-signed fused counterpart of `VPDPBUSD` — +/// available as part of AVX-512 VNNI on Ice Lake Xeon+, Sapphire Rapids, +/// and Zen 4+. 32 elements per iteration: widen 32 i8 → 32 i16 in one +/// ZMM each, then a single `dpwssd` accumulates 16 pair-products into +/// 16 i32 lanes. +/// +/// # Safety +/// CPU must support `avx512f`, `avx512bw`, and `avx512vnni`. +#[target_feature(enable = "avx512f,avx512bw,avx512vnni")] +pub unsafe fn score_4bit_internal_avx512_vnni(a: &[u8], b: &[u8]) -> f32 { + use core::arch::x86_64::*; + + assert_eq!( + a.len(), + b.len(), + "score_4bit_internal_avx512_vnni: vector length mismatch ({} vs {})", + a.len(), + b.len(), + ); + + unsafe { + let codebook_i8_128 = _mm_xor_si128( + _mm_loadu_si128(CODEBOOK_U8.as_ptr().cast::<__m128i>()), + _mm_set1_epi8(-128i8), + ); + let codebook_i8_256 = _mm256_broadcastsi128_si256(codebook_i8_128); + let nibble_mask = _mm_set1_epi8(0x0F); + let mut acc = _mm512_setzero_si512(); + + // 32 elements (= 16 bytes from each source) per iter. Shuffle stays + // on 256-bit (32 codebook values fit exactly), then we widen to 512. + let n_iters = a.len() / 16; + for i in 0..n_iters { + let va = _mm_loadu_si128(a.as_ptr().add(16 * i).cast::<__m128i>()); + let va_lo = _mm_and_si128(va, nibble_mask); + let va_hi = _mm_and_si128(_mm_srli_epi16(va, 4), nibble_mask); + let a_idx_0 = _mm_unpacklo_epi8(va_lo, va_hi); + let a_idx_1 = _mm_unpackhi_epi8(va_lo, va_hi); + let a_idx_256 = _mm256_inserti128_si256(_mm256_castsi128_si256(a_idx_0), a_idx_1, 1); + let c_a_i8_256 = _mm256_shuffle_epi8(codebook_i8_256, a_idx_256); + + let vb = _mm_loadu_si128(b.as_ptr().add(16 * i).cast::<__m128i>()); + let vb_lo = _mm_and_si128(vb, nibble_mask); + let vb_hi = _mm_and_si128(_mm_srli_epi16(vb, 4), nibble_mask); + let b_idx_0 = _mm_unpacklo_epi8(vb_lo, vb_hi); + let b_idx_1 = _mm_unpackhi_epi8(vb_lo, vb_hi); + let b_idx_256 = _mm256_inserti128_si256(_mm256_castsi128_si256(b_idx_0), b_idx_1, 1); + let c_b_i8_256 = _mm256_shuffle_epi8(codebook_i8_256, b_idx_256); + + // Widen 32 i8 → 32 i16, one ZMM each. + let c_a_i16 = _mm512_cvtepi8_epi16(c_a_i8_256); + let c_b_i16 = _mm512_cvtepi8_epi16(c_b_i8_256); + + // VPDPWSSD: acc[lane] += a[2·lane]·b[2·lane] + a[2·lane+1]·b[2·lane+1] + // with every operand interpreted as signed i16. + acc = _mm512_dpwssd_epi32(acc, c_a_i16, c_b_i16); + } + + let acc_256_lo = _mm512_castsi512_si256(acc); + let acc_256_hi = _mm512_extracti64x4_epi64(acc, 1); + let acc_256 = _mm256_add_epi32(acc_256_lo, acc_256_hi); + let acc_128 = _mm_add_epi32( + _mm256_castsi256_si128(acc_256), + _mm256_extracti128_si256(acc_256, 1), + ); + let simd_bytes = n_iters * 16; + let sum = i64::from(hsum_i32_sse(acc_128)) + + super::score_4bit_internal_integer(&a[simd_bytes..], &b[simd_bytes..]); + sum as f32 / (CODEBOOK_SCALE * CODEBOOK_SCALE) + } +} + +#[cfg(test)] +mod tests { + use rand::SeedableRng as _; + use rand::prelude::StdRng; + + use super::super::super::shared::pack_codes; + use super::super::shared::{PARITY_DIMS, random_inputs}; + use super::super::{Query4bitSimd, score_4bit_internal_scalar}; + use super::{ + score_4bit_internal_avx2, score_4bit_internal_avx512_vnni, score_4bit_internal_sse, + }; + + #[test] + fn test_sse_matches_scalar() { + if !std::is_x86_feature_detected!("ssse3") || !std::is_x86_feature_detected!("sse4.1") { + return; + } + let mut rng = StdRng::seed_from_u64(7); + for &dim in PARITY_DIMS { + let (simd_query, vector) = random_inputs(&mut rng, dim); + let scalar = simd_query.dotprod_raw(&vector); + let sse = unsafe { simd_query.dotprod_raw_sse(&vector) }; + assert_eq!(scalar, sse, "scalar {scalar} != sse {sse} at dim {dim}"); + } + } + + #[test] + fn test_avx2_matches_scalar() { + if !std::is_x86_feature_detected!("avx2") { + return; + } + let mut rng = StdRng::seed_from_u64(7); + for &dim in PARITY_DIMS { + let (simd_query, vector) = random_inputs(&mut rng, dim); + let scalar = simd_query.dotprod_raw(&vector); + let avx2 = unsafe { simd_query.dotprod_raw_avx2(&vector) }; + assert_eq!(scalar, avx2, "scalar {scalar} != avx2 {avx2} at dim {dim}"); + } + } + + #[test] + fn test_avx512_vnni_matches_scalar() { + if !(std::is_x86_feature_detected!("avx512f") + && std::is_x86_feature_detected!("avx512bw") + && std::is_x86_feature_detected!("avx512vnni")) + { + return; + } + let mut rng = StdRng::seed_from_u64(7); + for &dim in PARITY_DIMS { + let (simd_query, vector) = random_inputs(&mut rng, dim); + let scalar = simd_query.dotprod_raw(&vector); + let vnni512 = unsafe { simd_query.dotprod_raw_avx512_vnni(&vector) }; + assert_eq!( + scalar, vnni512, + "scalar {scalar} != avx512_vnni {vnni512} at dim {dim}" + ); + } + } + + /// Single saturation-safety check at an extreme dim (64K) with the + /// worst-case combination: query maxed out and every lane of the vector + /// pointing at the extreme-magnitude codebook slot. Scalar is the + /// reference (i64 throughout, saturation-free by construction); each + /// SIMD path must match it exactly. A mismatch proves that some + /// intermediate integer saturated or overflowed. + #[test] + fn test_saturation_safety_64k() { + let dim = 65_536; + let query = vec![1.0_f32; dim]; + let indices: Vec = vec![15; dim]; // CODEBOOK_U8[15] = 255 (max magnitude) + let vector = pack_codes(&indices, 4); + + let q = Query4bitSimd::new(&query); + let scalar = q.dotprod_raw(&vector); + + unsafe { + if std::is_x86_feature_detected!("ssse3") && std::is_x86_feature_detected!("sse4.1") { + let sse = q.dotprod_raw_sse(&vector); + assert_eq!(scalar, sse, "sse disagrees at dim={dim}"); + } + if std::is_x86_feature_detected!("avx2") { + let avx2 = q.dotprod_raw_avx2(&vector); + assert_eq!(scalar, avx2, "avx2 disagrees at dim={dim}"); + } + if std::is_x86_feature_detected!("avx512f") + && std::is_x86_feature_detected!("avx512bw") + && std::is_x86_feature_detected!("avx512vnni") + { + let v512 = q.dotprod_raw_avx512_vnni(&vector); + assert_eq!(scalar, v512, "avx512_vnni disagrees at dim={dim}"); + } + } + } + + /// Parity: each x86 `score_4bit_internal` variant must reproduce the + /// scalar reference bit-exactly. Both sides compute `Σ c_signed_a · + /// c_signed_b / c_scale²` with deterministic ordering, so the f32 + /// outputs are identical. + #[test] + fn test_score_sse_matches_scalar() { + if !std::is_x86_feature_detected!("ssse3") || !std::is_x86_feature_detected!("sse4.1") { + return; + } + let mut rng = StdRng::seed_from_u64(7); + for &dim in PARITY_DIMS { + let (_, vec_a) = random_inputs(&mut rng, dim); + let (_, vec_b) = random_inputs(&mut rng, dim); + let scalar = score_4bit_internal_scalar(&vec_a, &vec_b); + let sse = unsafe { score_4bit_internal_sse(&vec_a, &vec_b) }; + assert_eq!( + scalar, sse, + "score: scalar {scalar} != sse {sse} at dim {dim}" + ); + } + } + + #[test] + fn test_score_avx2_matches_scalar() { + if !std::is_x86_feature_detected!("avx2") { + return; + } + let mut rng = StdRng::seed_from_u64(7); + for &dim in PARITY_DIMS { + let (_, vec_a) = random_inputs(&mut rng, dim); + let (_, vec_b) = random_inputs(&mut rng, dim); + let scalar = score_4bit_internal_scalar(&vec_a, &vec_b); + let avx2 = unsafe { score_4bit_internal_avx2(&vec_a, &vec_b) }; + assert_eq!( + scalar, avx2, + "score: scalar {scalar} != avx2 {avx2} at dim {dim}" + ); + } + } + + #[test] + fn test_score_avx512_vnni_matches_scalar() { + if !(std::is_x86_feature_detected!("avx512f") + && std::is_x86_feature_detected!("avx512bw") + && std::is_x86_feature_detected!("avx512vnni")) + { + return; + } + let mut rng = StdRng::seed_from_u64(7); + for &dim in PARITY_DIMS { + let (_, vec_a) = random_inputs(&mut rng, dim); + let (_, vec_b) = random_inputs(&mut rng, dim); + let scalar = score_4bit_internal_scalar(&vec_a, &vec_b); + let vnni512 = unsafe { score_4bit_internal_avx512_vnni(&vec_a, &vec_b) }; + assert_eq!( + scalar, vnni512, + "score: scalar {scalar} != avx512_vnni {vnni512} at dim {dim}" + ); + } + } + + /// Saturation-safety at 64K for all x86 score paths simultaneously. + /// Both vectors are every index 15 → `c_signed = 127`, every product is + /// `127² = 16 129`, every madd pair ≤ 32 258 (i32-safe, nowhere near i16 + /// because we already widen). Total = 16 384·16 129 ≈ 264 M, fits i32 + /// with ~8× headroom. + #[test] + fn test_score_saturation_safety_64k() { + let dim = 65_536; + let indices: Vec = vec![15; dim]; // CODEBOOK_U8[15] = 255 → signed 127 + let vec_a = pack_codes(&indices, 4); + let vec_b = pack_codes(&indices, 4); + + let scalar = score_4bit_internal_scalar(&vec_a, &vec_b); + + unsafe { + if std::is_x86_feature_detected!("ssse3") && std::is_x86_feature_detected!("sse4.1") { + let sse = score_4bit_internal_sse(&vec_a, &vec_b); + assert_eq!(scalar, sse, "score sse disagrees at dim={dim}"); + } + if std::is_x86_feature_detected!("avx2") { + let avx2 = score_4bit_internal_avx2(&vec_a, &vec_b); + assert_eq!(scalar, avx2, "score avx2 disagrees at dim={dim}"); + } + if std::is_x86_feature_detected!("avx512f") + && std::is_x86_feature_detected!("avx512bw") + && std::is_x86_feature_detected!("avx512vnni") + { + let vnni = score_4bit_internal_avx512_vnni(&vec_a, &vec_b); + assert_eq!(scalar, vnni, "score avx512_vnni disagrees at dim={dim}"); + } + } + } +} diff --git a/tools/codespell.toml b/tools/codespell.toml index d7635be361..80407c6511 100644 --- a/tools/codespell.toml +++ b/tools/codespell.toml @@ -17,6 +17,7 @@ ignore-words-list = [ "INDX", # bitvec crate const "generall", # The boss "hel", # Query token + "inout", # Rust inline-asm operand keyword "mmapped", # Memory mapped "ser", # Serde package ]