mirror of
https://github.com/qdrant/qdrant.git
synced 2026-07-23 11:11:00 -05:00
TQ SIMD (#8749)
* TQ 4 bit SIMD * more optimizations * unroll test * remove tries * revert avx 512 * final simd * use precomputed codebooks * close to finish * split to files * are you happy fmt * score internal * 1bit * 1bit tails * score_1bit_internal_avx2 tail * 1bit simd * 2bit case * fix 2bit * fix features * tails * fix tests * less benches * 1bit tails * 4bit tails simd * 64k overflow test * reuse packing and constants from dev after rebase * are yoy happy fmt * fix x64 build * integration * symmetric score SIMD * fix codespell * better docs * are you happy clippy * review remarks * review remarks
This commit is contained in:
Ivan Pleshkov
@@ -64,3 +64,7 @@ harness = false
|
||||
[[bench]]
|
||||
name = "turboquant"
|
||||
harness = false
|
||||
|
||||
[[bench]]
|
||||
name = "turbo_simd"
|
||||
harness = false
|
||||
|
||||
669
lib/quantization/benches/turbo_simd.rs
Normal file
669
lib/quantization/benches/turbo_simd.rs
Normal file
@@ -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<u8>,
|
||||
indices: Vec<u32>,
|
||||
/// 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<u8> = (0..count * packed_bytes)
|
||||
.map(|_| rng.random_range(0..=u8::MAX))
|
||||
.collect();
|
||||
let mut indices: Vec<u32> = (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<f32> {
|
||||
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<f32> = (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);
|
||||
<u8 as BitsStoreType>::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<u8> {
|
||||
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);
|
||||
@@ -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<Item = f64>) {
|
||||
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);
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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<f64>);
|
||||
|
||||
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<f32> = 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 `<query, v>` 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<I, J>(left: I, right: J) -> f32
|
||||
where
|
||||
I: Iterator<Item = f64>,
|
||||
J: Iterator<Item = f64>,
|
||||
{
|
||||
left.zip(right).map(|(q, u)| q * u).sum::<f64>() 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<f32> = 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}",
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
104
lib/quantization/src/turboquant/simd/mod.rs
Normal file
104
lib/quantization/src/turboquant/simd/mod.rs
Normal file
@@ -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<u8> {
|
||||
(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<f32> {
|
||||
(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<u8> {
|
||||
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<u8> {
|
||||
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
|
||||
}
|
||||
}
|
||||
228
lib/quantization/src/turboquant/simd/query1bit/arm.rs
Normal file
228
lib/quantization/src/turboquant/simd/query1bit/arm.rs
Normal file
@@ -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<const BITS: usize> super::Query1bitSimd<BITS> {
|
||||
/// 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<const BITS: usize>(dim: usize, seed: u64) {
|
||||
let mut rng = StdRng::seed_from_u64(seed);
|
||||
let query: Vec<f32> = (0..dim).map(|_| StandardNormal.sample(&mut rng)).collect();
|
||||
let data = random_bytes(&mut rng, dim / 8);
|
||||
let q = Query1bitSimd::<BITS>::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}",
|
||||
);
|
||||
}
|
||||
}
|
||||
472
lib/quantization/src/turboquant/simd/query1bit/mod.rs
Normal file
472
lib/quantization/src/turboquant/simd/query1bit/mod.rs
Normal file
@@ -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<BITS>`] — 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<const BITS: usize = 8> {
|
||||
/// 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<u8>,
|
||||
/// 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<const BITS: usize> Query1bitSimd<BITS> {
|
||||
/// 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<u8> = 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<f32> = (0..dim).map(|_| StandardNormal.sample(&mut rng)).collect();
|
||||
let data: Vec<u8> = 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}",
|
||||
);
|
||||
}
|
||||
}
|
||||
502
lib/quantization/src/turboquant/simd/query1bit/x64.rs
Normal file
502
lib/quantization/src/turboquant/simd/query1bit/x64.rs
Normal file
@@ -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<const BITS: usize> super::Query1bitSimd<BITS> {
|
||||
/// 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::<BITS>(&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::<BITS>(&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<const BITS: usize>(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<const BITS: usize>(dim: usize, seed: u64) {
|
||||
let mut rng = StdRng::seed_from_u64(seed);
|
||||
let query: Vec<f32> = (0..dim).map(|_| StandardNormal.sample(&mut rng)).collect();
|
||||
let data = random_bytes(&mut rng, dim / 8);
|
||||
let q = Query1bitSimd::<BITS>::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<const BITS: usize>(query: &[f32], data: &[u8]) {
|
||||
let q = Query1bitSimd::<BITS>::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);
|
||||
}
|
||||
}
|
||||
420
lib/quantization/src/turboquant/simd/query2bit/arm.rs
Normal file
420
lib/quantization/src/turboquant/simd/query2bit/arm.rs
Normal file
@@ -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::<u32>()));
|
||||
// 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<u8> = 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<u8> = 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");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
515
lib/quantization/src/turboquant/simd/query2bit/mod.rs
Normal file
515
lib/quantization/src/turboquant/simd/query2bit/mod.rs
Normal file
@@ -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<u8>) {
|
||||
let query = sample_normal_vec(rng, dim);
|
||||
let mut indices: Vec<u8> = (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<f32> = 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}",
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
646
lib/quantization/src/turboquant/simd/query2bit/x64.rs
Normal file
646
lib/quantization/src/turboquant/simd/query2bit/x64.rs
Normal file
@@ -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::<i32>().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<u8> = 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<u8> = 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(),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
467
lib/quantization/src/turboquant/simd/query4bit/arm.rs
Normal file
467
lib/quantization/src/turboquant/simd/query4bit/arm.rs
Normal file
@@ -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<u8> = 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<u8> = 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}");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
682
lib/quantization/src/turboquant/simd/query4bit/mod.rs
Normal file
682
lib/quantization/src/turboquant/simd/query4bit/mod.rs
Normal file
@@ -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<u8>) {
|
||||
let query = sample_normal_vec(rng, dim);
|
||||
let mut indices: Vec<u8> = (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<f32> = 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<f32> = 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}",
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
644
lib/quantization/src/turboquant/simd/query4bit/x64.rs
Normal file
644
lib/quantization/src/turboquant/simd/query4bit/x64.rs
Normal file
@@ -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<u8> = 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<u8> = 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}");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user