upstream-archive/qdrant
1
0
Fork 0
mirror of https://github.com/qdrant/qdrant.git synced 2026-09-21 05:27:39 -05:00
Code Issues Packages Projects Releases Wiki Activity

Turbo4 query simd (#10391)

Browse Source 
* TQ SIMD: one backend ladder, resolved once per query

The 2- and 4-bit kernels share the same preference order (AVX-512 VNNI
→ AVX2 → SSE → NEON + SDOT → NEON → scalar), spelled out six times as
chains of `is_x86_feature_detected!` — and `Query{2,4}bitSimd::dotprod`
re-ran its chain for every vector scored.

Move the ladder into one `simd::SimdBackend` enum with a single `detect()`.
The query types resolve it in `new()` and dispatch on the stored value;
the symmetric `score_{2,4}bit_internal*` entry points dispatch on
`SimdBackend::detect()`.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* QuerySimd: one query layout and scalar reference for every packing width

`Query{1,2,4}bitSimd` are three copies of the same idea — quantize the
query into i8 halves, multiply them against an integer codebook — each
with its own query layout and its own set of SIMD kernels.  Introduce
`simd::query::QuerySimd<PLANES>`, generic over the number of codes per
packed byte (2, 4 or 8), which the three widths will share.

The query halves are stored as planes, one per code position within a
byte: plane `k` entry `j` is the half of query dim `PLANES · j + k`.
That is the order the codes come out of raw data bytes with a shift and
a mask, so a kernel never has to unpack them into dim order.  Planes are
zero-padded to the widest SIMD block, so a partial last block on the
data side multiplies against zeros.

The widths contribute only their integer encoding (`Encoding`: codebook
table, offset, scale and query range); the 1-bit width gets one here —
`{0, 128}` with offset 64 on x86_64, `∓127` on aarch64 — chosen so the
query keeps full i8 halves.  Only the scalar reference exists yet; the
SIMD kernels follow, and the width types switch over once they're in.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* QuerySimd: AVX-512 VNNI kernel on the query planes

One ZMM of packed codes per block; for each plane the codes are shifted
down by the code width, masked and looked up in the codebook with one
`vpshufb`, then `VPDPBUSD` folds them into the plane's low and high
accumulators.  Two accumulator pairs per vector keep the VNNI latency
off the critical path at every width.  The last partial block is a
masked load whose dead lanes multiply against the planes' zero padding.

The shift count is an immediate, so the shift-by-width helper spells
out the three widths in a `match` — the only place the kernel is not
literally generic over `PLANES`.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* QuerySimd: NEON SDOT kernel on the query planes

The AVX-512 kernel's shape on 128-bit registers: one `TBL` codebook
lookup per plane, `SDOT` (inline asm — `vdotq_s32` is still unstable)
into two accumulator pairs.  The last partial block runs on a
zero-padded copy of the remaining bytes.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* QuerySimd: AVX2 kernel on the query planes

One YMM of packed codes per block, one `vpshufb` lookup per plane and
`maddubs → madd` against ones into the same two accumulator pairs as
the VNNI kernel.  The `maddubs` pair sums stay inside i16 by the
per-width query bounds.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* QuerySimd: SSE and plain NEON kernels on the query planes

The 128-bit forms of the AVX2 and SDOT kernels: `maddubs → madd` on
XMM, `vmull_s8 → vpadalq_s16` on NEON without `dotprod`.  Every backend
of the shared query type now has its kernel.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* Query4bitSimd: score through QuerySimd<2>

`Query4bitSimd` becomes an alias of the shared query type; its own
chunk-and-tail query layout and the per-backend kernels built on it go
away, along with the accuracy tests the shared module now runs for
every width.  What stays in `query4bit` is the 4-bit encoding and the
symmetric `score_4bit_internal*` paths.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* Query2bitSimd: score through QuerySimd<4>

The 2-bit asymmetric kernels unpacked every 4 packed bytes into 16
centroid bytes through two `pshufb` / `TBL` pair-table lookups and a
zip before a single multiply-accumulate step — on AVX-512 that was four
128-bit unpacks and six lane inserts per pair of `VPDPBUSD`.  On the
query planes the same 16 codes cost one shift, one mask and one lookup
per plane, straight from a full-width load.

`Query2bitSimd` becomes an alias of the shared query type; its chunk
layout and per-backend kernels go away.  The pair-table unpack stays
for the symmetric `score_2bit_internal*` paths.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* Query1bitSimd: score through QuerySimd<8>

The 1-bit asymmetric kernels bit-plane-transposed the query and scored
`Σ_b 2^b · popcount(data AND plane_b)` per 16-byte block — eight
AND + popcount + add steps on XMM (even with AVX-512, through the VL
forms) and a `BITS`-deep accumulator array.  On the query planes a
sign bit is just a one-bit code: shift, mask, a two-entry codebook
lookup and the same multiply-accumulate as the wider widths, on full
256-/512-bit registers.

`Query1bitSimd` becomes an alias of the shared query type.  Its query
width was a const parameter (8 bits by default, 16 for TQ+ through the
`Bits1Wide` variant); the shared encoding always carries 16-bit halves,
so the variant and the TQ+ special case go away.  The popcount kernels
stay for the symmetric `score_1bit_internal`, where XOR + popcount is
the right tool.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* Bench: cold per-vector rows for every width

`query{4,2,1}bit_dotprod_cold` from one generic body — scalar reference,
public `dotprod` and each backend — so the widths can be compared on one
host.  `TURBO_SIMD_DIMS` narrows or widens the dims of a run and
`TURBO_SIMD_POOL_KB` shrinks the pool to L1 for hot-kernel numbers.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* QuerySimd: query bytes as a parameter; 8-bit queries for the 1-bit width

The shared kernels always carried two query bytes (a ~16-bit query),
which cost the 1-bit width its 8-bit-query speed: the old bit-plane
kernel scored a 16-byte vector in 4.9 ns hot (34.7 ns cold) against
9.1 ns (54.7 ns) through the planes, the difference being the second
byte's multiply-accumulates on a vector that fills a quarter of one
block.  Above 512 dims the planes win either way.

Make the number of query bytes a parameter: `QuerySimd<PLANES,
QUERY_BYTES>` with one plane per query byte and code position, and one
accumulator pair per query byte.  A one-byte query is scaled to the
range of a single byte, `RADIX / 2 − 1`.  `Query1bitSimd` is the
one-byte instance — at parity with the old kernel at small dims (cold
36.9 / 38.3 / 38.4 ns at d = 128 / 256 / 512) and 1.8× faster at 1536
(68 vs 121 ns) — and `Query1bitWideSimd` the two-byte one, which TQ+
selects through the `Bits1Wide` variant as before.  The 2- and 4-bit
widths keep two bytes.

Bench: `query1bit_wide_dotprod_cold` and a `query1bit_wide` row next
to BQ.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Ivan Pleshkov
2026-09-03 12:45:56 +02:00
committed by timvisee
co-authored by Claude Fable 5
parent d931e0c2b8
commit 82b9ed8d32
16 changed files with 1821 additions and 2570 deletions
Download Patch File Download Diff File Expand all files Collapse all files
lib/quantization/benches/turbo_simd.rs
+57 -125
View File
@@ -3,7 +3,7 @@ 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,
Query1bitSimd, Query1bitWideSimd, QuerySimd, score_1bit_internal, score_1bit_internal_scalar,
score_2bit_internal, score_2bit_internal_scalar, score_4bit_internal,
score_4bit_internal_scalar,
};
@@ -32,9 +32,22 @@ const DIMS_2BIT: &[usize] = &[128, 1532, 1536]; // 1532 = 95 chunks (odd) + 12-d
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.
/// vectors into cache — each iteration pays a real DRAM fetch. Override with
/// `TURBO_SIMD_POOL_KB` to measure hot kernels (a pool that fits L1).
const POOL_BYTES: usize = 64 * 1024 * 1024;
fn pool_bytes() -> usize {
match std::env::var("TURBO_SIMD_POOL_KB") {
Ok(kb) => {
kb.trim()
.parse::<usize>()
.expect("TURBO_SIMD_POOL_KB: not a size")
* 1024
}
Err(_) => POOL_BYTES,
}
}
struct VectorPool {
buf: Vec<u8>,
indices: Vec<u32>,
@@ -44,7 +57,7 @@ struct VectorPool {
impl VectorPool {
fn with_packed_bytes(packed_bytes: usize, seed: u64) -> Self {
let count = (POOL_BYTES / packed_bytes).max(1024);
let count = (pool_bytes() / packed_bytes).max(64);
let mut rng = SmallRng::seed_from_u64(seed);
let buf: Vec<u8> = (0..count * packed_bytes)
.map(|_| rng.random_range(0..=u8::MAX))
@@ -88,12 +101,33 @@ fn make_query(dim: usize) -> Vec<f32> {
(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 {
/// Dims for a bench group: the built-in list, or `TURBO_SIMD_DIMS` (comma
/// separated) to narrow or widen a run.
fn dims(default: &[usize]) -> Vec<usize> {
match std::env::var("TURBO_SIMD_DIMS") {
Ok(list) => list
.split(',')
.map(|dim| dim.trim().parse().expect("TURBO_SIMD_DIMS: not a dim"))
.collect(),
Err(_) => default.to_vec(),
}
}
/// Cold-cache query-vs-vector dotprod at the width packing `PLANES` codes
/// per byte: a hot encoded query against data vectors drawn from a shuffled
/// pool ≫ cache, so every call pays a real DRAM fetch — the HNSW scoring
/// pattern. `scalar` is the reference kernel, `dotprod` the public path
/// (best backend + float reconstruction), the rest the individual backends.
fn dotprod_cold<const PLANES: usize, const QUERY_BYTES: usize>(
c: &mut Criterion,
group: &str,
default_dims: &[usize],
) {
let mut group = c.benchmark_group(group);
for dim in dims(default_dims) {
let q = make_query(dim);
let query = Query4bitSimd::new(&q);
let pool = VectorPool::new_4bit(dim, 7);
let query = QuerySimd::<PLANES, QUERY_BYTES>::new(&q);
let pool = VectorPool::with_packed_bytes(dim / PLANES, 7);
group.throughput(Throughput::Elements(dim as u64));
@@ -106,10 +140,6 @@ fn bench_dotprod_cold(c: &mut Criterion) {
});
});
// 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(|| {
@@ -184,6 +214,13 @@ fn bench_dotprod_cold(c: &mut Criterion) {
group.finish();
}
fn bench_dotprod_cold(c: &mut Criterion) {
dotprod_cold::<2, 2>(c, "query4bit_dotprod_cold", DIMS_4BIT);
dotprod_cold::<4, 2>(c, "query2bit_dotprod_cold", DIMS_2BIT);
dotprod_cold::<8, 1>(c, "query1bit_dotprod_cold", DIMS_1BIT);
dotprod_cold::<8, 2>(c, "query1bit_wide_dotprod_cold", DIMS_1BIT);
}
/// 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
@@ -379,11 +416,10 @@ fn bench_score_1bit_cold(c: &mut Criterion) {
/// 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
/// Compares `Query1bitSimd` (8-bit query) and `Query1bitWideSimd` (16-bit)
/// 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 = SmallRng::seed_from_u64(42);
@@ -393,37 +429,27 @@ fn bench_query1bit_vs_bq_hot(c: &mut Criterion) {
.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 query = Query1bitSimd::new(&query_floats);
let query_wide = Query1bitWideSimd::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, _| {
group.bench_with_input(BenchmarkId::new("query1bit", 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))
query.dotprod(black_box(v))
});
});
group.bench_with_input(BenchmarkId::new("query1bit_12bit", dim), &dim, |b, _| {
group.bench_with_input(BenchmarkId::new("query1bit_wide", 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))
query_wide.dotprod(black_box(v))
});
});
@@ -467,99 +493,6 @@ fn encode_bq_scalar8bits(query: &[f32]) -> Vec<u8> {
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) {
@@ -661,7 +594,6 @@ 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,
lib/quantization/src/turboquant/mod.rs
+4 -6
View File
@@ -1,6 +1,6 @@
use serde::{Deserialize, Serialize};
use crate::turboquant::simd::{Query1bitSimd, Query2bitSimd, Query4bitSimd};
use crate::turboquant::simd::{Query1bitSimd, Query1bitWideSimd, Query2bitSimd, Query4bitSimd};
pub mod encoding;
pub mod lloyd_max;
@@ -123,15 +123,13 @@ pub struct EncodedQueryTQ {
/// on architectures without a matching SIMD instruction set the scalar
/// reference kernel inside each type takes over automatically.
///
/// `Bits1Wide` is the same kernel as `Bits1` but with 16-bit query
/// quantization instead of the default 8-bit (the kernel's max). Used in
/// `Bits1Wide` is `Bits1` with a 16-bit query instead of 8-bit. Used in
/// TQ+ for 1-bit storage: the per-coord `D'` pre-scaling pushes some query
/// coords into the bottom of the 8-bit integer range, where rounding noise
/// is large relative to the signal — 16 bits gives the most headroom the
/// existing kernel supports.
/// is large relative to the signal.
pub enum EncodedQueryTQData {
Bits1(Query1bitSimd),
Bits1Wide(Query1bitSimd<16>),
Bits1Wide(Query1bitWideSimd),
Bits2(Query2bitSimd),
Bits4(Query4bitSimd),
}
lib/quantization/src/turboquant/quantization.rs
+10 -10
View File
@@ -4,9 +4,9 @@ use crate::DistanceType;
use crate::turboquant::encoding::TqVectorExtras;
use crate::turboquant::rotation::HadamardRotation;
use crate::turboquant::simd::{
CODEBOOK_SCALE_SQ_2BIT, CODEBOOK_SCALE_SQ_4BIT, Query1bitSimd, Query2bitSimd, Query4bitSimd,
score_1bit_internal, score_2bit_internal, score_2bit_internal_weighted, score_4bit_internal,
score_4bit_internal_weighted,
CODEBOOK_SCALE_SQ_2BIT, CODEBOOK_SCALE_SQ_4BIT, Query1bitSimd, Query1bitWideSimd,
Query2bitSimd, Query4bitSimd, score_1bit_internal, score_2bit_internal,
score_2bit_internal_weighted, score_4bit_internal, score_4bit_internal_weighted,
};
use crate::turboquant::{EncodedQueryTQ, EncodedQueryTQData, TQBits, TQMode, TQRotation};
@@ -540,18 +540,18 @@ impl TurboQuantizer {
// has no downstream benefit here).
let rotated_f32: Vec<f32> = rotated.iter().map(|&x| x as f32).collect();
// For TQ+ + Bits1 storage, widen query quantization from the default
// 8 bits to the kernel's max of 16. The per-coord `D'` pre-scaling
// can push some coords toward the small end of the integer range;
// 8 bits loses too much there.
// For TQ+ + Bits1 storage, widen the query from 8 to 16 bits: the
// per-coord `D'` pre-scaling can push some coords toward the small
// end of the integer range, where 8 bits lose too much.
let use_wide_query =
self.error_correction.is_some() && matches!(self.bits, TQBits::Bits1 | TQBits::Bits1_5);
let data = match self.bits {
TQBits::Bits1 | TQBits::Bits1_5 if use_wide_query => {
EncodedQueryTQData::Bits1Wide(Query1bitSimd::<16>::new(&rotated_f32))
EncodedQueryTQData::Bits1Wide(Query1bitWideSimd::new(&rotated_f32))
}
TQBits::Bits1 | TQBits::Bits1_5 => {
EncodedQueryTQData::Bits1(Query1bitSimd::new(&rotated_f32))
}
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)),
};
lib/quantization/src/turboquant/simd/mod.rs
+65 -11
View File
@@ -3,36 +3,90 @@
//! 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.
//! precomputed for fast asymmetric scoring (original-query × PQ-vector):
//! the shared [`query::QuerySimd`] kernels instantiated at the width, which
//! dispatch 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:
//! Available SIMD backends:
//!
//! | 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 |
//! | Path | x86_64 | aarch64 |
//! |-------------------|---------------------------------------|-------------------|
//! | asymmetric, 1/2/4 | AVX-512 VNNI, AVX2, SSE4.1+SSSE3 | NEON + SDOT, NEON |
//! | symmetric, 1 | AVX-512 VPOPCNTDQ, AVX2, SSE4.1+SSSE3 | NEON |
//! | symmetric, 2/4 | AVX-512 VNNI, AVX2, SSE4.1+SSSE3 | NEON + SDOT, NEON |
//!
//! On any other target the scalar reference kernels in each module take over.
//! On any other target the scalar reference kernels take over.
pub mod hadamard;
pub mod query;
pub mod query1bit;
pub mod query2bit;
pub mod query4bit;
/// Best multiply-accumulate backend the host CPU supports, in preference
/// order AVX-512 VNNI → AVX2 → SSE → NEON + SDOT → NEON → scalar. Shared by
/// every kernel built on `u8 × i8` / `i8 × i8` products (the asymmetric
/// paths of every width and the symmetric 2- and 4-bit paths); resolve it
/// once with [`SimdBackend::detect`] and dispatch on the value, so a scoring
/// loop doesn't re-run CPU feature detection per vector.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum SimdBackend {
#[cfg(target_arch = "x86_64")]
Avx512Vnni,
#[cfg(target_arch = "x86_64")]
Avx2,
#[cfg(target_arch = "x86_64")]
Sse,
#[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
NeonSdot,
#[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
Neon,
Scalar,
}
impl SimdBackend {
pub(crate) fn detect() -> Self {
#[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 SimdBackend::Avx512Vnni;
}
if std::is_x86_feature_detected!("avx2") {
return SimdBackend::Avx2;
}
if std::is_x86_feature_detected!("sse4.1") && std::is_x86_feature_detected!("ssse3") {
return SimdBackend::Sse;
}
}
#[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
{
if std::arch::is_aarch64_feature_detected!("dotprod") {
return SimdBackend::NeonSdot;
}
return SimdBackend::Neon;
}
#[allow(unreachable_code)]
SimdBackend::Scalar
}
}
// 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.
pub use query::QuerySimd;
#[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};
pub use query1bit::{
Query1bitSimd, Query1bitWideSimd, 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,
lib/quantization/src/turboquant/simd/query/arm.rs
+362
View File
@@ -0,0 +1,362 @@
//! NEON kernels for [`QuerySimd`] on aarch64.
//!
//! The codebook is stored signed (`Encoding::codebook` as `i8`), so
//! `vmull_s8` and `SDOT` operate on true `i8 × i8` products — no bias
//! correction — and the query bytes use the full i8 range (see the module
//! docs of [`super`]).
use core::arch::aarch64::*;
use super::{Code, PLANE_BLOCK, QueryPlanes, QuerySimd, encoding};
/// Packed data bytes per NEON block: one 128-bit register of codes.
const BLOCK_128: usize = 16;
const _: () = assert!(PLANE_BLOCK.is_multiple_of(BLOCK_128));
/// The codebook as a `TBL` table for the width packing `PLANES` codes per
/// byte.
const fn codebook<const PLANES: usize>() -> [Code; 16] {
encoding(PLANES).codebook
}
/// Mask of one code in the low bits of a byte.
const fn code_mask<const PLANES: usize>() -> u8 {
((1u16 << (8 / PLANES)) - 1) as u8
}
/// One [`BLOCK_128`]-byte block of every query plane. Loaded once per
/// block and shared by all vectors scored against it.
#[derive(Clone, Copy)]
struct QueryBlock128<const PLANES: usize, const QUERY_BYTES: usize> {
bytes: [[int8x16_t; PLANES]; QUERY_BYTES],
}
impl<const PLANES: usize, const QUERY_BYTES: usize> QueryBlock128<PLANES, QUERY_BYTES> {
/// # Safety
/// `offset + BLOCK_128` must not exceed the plane length.
#[inline]
#[target_feature(enable = "neon")]
unsafe fn load(planes: &QueryPlanes<PLANES, QUERY_BYTES>, offset: usize) -> Self {
let mut block = Self {
bytes: [[vdupq_n_s8(0); PLANES]; QUERY_BYTES],
};
for (b, byte_planes) in block.bytes.iter_mut().enumerate() {
for (k, plane) in byte_planes.iter_mut().enumerate() {
*plane = unsafe { load_plane_128(&planes.bytes[b][k], offset) };
}
}
block
}
}
/// # Safety
/// `offset + BLOCK_128 <= plane.len()`.
#[inline]
#[target_feature(enable = "neon")]
unsafe fn load_plane_128(plane: &[i8], offset: usize) -> int8x16_t {
debug_assert!(offset + BLOCK_128 <= plane.len());
unsafe { vld1q_s8(plane.as_ptr().add(offset)) }
}
/// The next plane's codes moved into the low bits of every byte. The shift
/// count is an immediate, hence the match over the widths.
#[inline]
#[target_feature(enable = "neon")]
fn next_plane_128<const PLANES: usize>(codes: uint8x16_t) -> uint8x16_t {
match PLANES {
2 => vshrq_n_u8(codes, 4),
4 => vshrq_n_u8(codes, 2),
_ => vshrq_n_u8(codes, 1),
}
}
/// Codebook values addressed by the low code of every byte of `codes`.
#[inline]
#[target_feature(enable = "neon")]
fn lookup_codes<const PLANES: usize>(codes: uint8x16_t) -> int8x16_t {
let table = const { codebook::<PLANES>() };
let codebook = unsafe { vld1q_s8(table.as_ptr()) };
vqtbl1q_s8(
codebook,
vandq_u8(codes, vdupq_n_u8(const { code_mask::<PLANES>() })),
)
}
/// `acc[lane] += Σ₄ a · b` over each lane's four i8 pairs (`SDOT`). Inline
/// asm because `vdotq_s32` is still unstable (rust-lang/rust#117224).
///
/// # Safety
/// CPU must support `dotprod`.
#[inline]
#[target_feature(enable = "neon,dotprod")]
unsafe fn sdot(mut acc: int32x4_t, a: int8x16_t, b: int8x16_t) -> int32x4_t {
unsafe {
core::arch::asm!(
"sdot {acc:v}.4s, {a:v}.16b, {b:v}.16b",
acc = inout(vreg) acc,
a = in(vreg) a,
b = in(vreg) b,
options(pure, nomem, nostack, preserves_flags),
);
}
acc
}
/// `acc[lane] += Σ₄ a · b` without `dotprod`: widening i8 multiplies
/// (`vmull_s8`, exact in i16) pairwise-added into the i32 lanes
/// (`vpadalq_s16`). Same per-lane bound as [`sdot`].
#[inline]
#[target_feature(enable = "neon")]
fn mul_add(acc: int32x4_t, a: int8x16_t, b: int8x16_t) -> int32x4_t {
let acc = vpadalq_s16(acc, vmull_s8(vget_low_s8(a), vget_low_s8(b)));
vpadalq_s16(acc, vmull_high_s8(a, b))
}
/// i32 accumulators of one vector: per query byte two independent chains,
/// plane `k` feeding chain `k & 1`, which keeps the multiply-accumulate
/// latency off the critical path at every width.
///
/// i32 lane bound: each block adds at most `4 · 127 · 128 = 65 024` per
/// plane to a lane, so overflow needs ~33 K blocks ≈ 500 K packed bytes.
#[derive(Clone, Copy)]
struct Acc128<const QUERY_BYTES: usize> {
bytes: [[int32x4_t; 2]; QUERY_BYTES],
}
impl<const QUERY_BYTES: usize> Acc128<QUERY_BYTES> {
#[inline]
#[target_feature(enable = "neon")]
fn zero() -> Self {
Self {
bytes: [[vdupq_n_s32(0); 2]; QUERY_BYTES],
}
}
/// Fold one block of packed `codes` (16 bytes) into the accumulators
/// with `SDOT`: plane by plane, the codes are shifted down, masked to
/// one code per byte and looked up in the codebook.
///
/// # Safety
/// CPU must support `dotprod`.
#[inline]
#[target_feature(enable = "neon,dotprod")]
unsafe fn accumulate_sdot<const PLANES: usize>(
&mut self,
codes: uint8x16_t,
query: QueryBlock128<PLANES, QUERY_BYTES>,
) {
let mut shifted = codes;
for k in 0..PLANES {
let values = lookup_codes::<PLANES>(shifted);
for (acc, planes) in self.bytes.iter_mut().zip(&query.bytes) {
acc[k & 1] = unsafe { sdot(acc[k & 1], values, planes[k]) };
}
shifted = next_plane_128::<PLANES>(shifted);
}
}
/// [`Self::accumulate_sdot`] for CPUs without `dotprod`.
#[inline]
#[target_feature(enable = "neon")]
fn accumulate_mull<const PLANES: usize>(
&mut self,
codes: uint8x16_t,
query: QueryBlock128<PLANES, QUERY_BYTES>,
) {
let mut shifted = codes;
for k in 0..PLANES {
let values = lookup_codes::<PLANES>(shifted);
for (acc, planes) in self.bytes.iter_mut().zip(&query.bytes) {
acc[k & 1] = mul_add(acc[k & 1], values, planes[k]);
}
shifted = next_plane_128::<PLANES>(shifted);
}
}
/// Per-lane totals of every query byte.
#[inline]
#[target_feature(enable = "neon")]
fn fold(self) -> [int32x4_t; QUERY_BYTES] {
self.bytes.map(|[a, b]| vaddq_s32(a, b))
}
}
impl<const PLANES: usize, const QUERY_BYTES: usize> QuerySimd<PLANES, QUERY_BYTES> {
/// ARM NEON over the query planes for CPUs without `dotprod`: the
/// [`Self::dotprod_raw_neon_sdot`] block loop with `vmull_s8 →
/// vpadalq_s16` in place of `SDOT`.
///
/// # 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 {
assert_eq!(
vector.len(),
self.vector_bytes,
"QuerySimd<{PLANES}, {QUERY_BYTES}>::dotprod_raw_neon: vector length mismatch ({} \
vs expected {})",
vector.len(),
self.vector_bytes,
);
unsafe { Self::reduce_neon(self.accumulate_neon(vector.as_ptr())) }
}
/// [`Self::accumulate_neon_sdot`] for CPUs without `dotprod`.
///
/// # Safety
/// `data` must be readable for `self.vector_bytes` bytes.
#[inline]
#[target_feature(enable = "neon")]
unsafe fn accumulate_neon(&self, data: *const u8) -> Acc128<QUERY_BYTES> {
unsafe {
let mut acc = Acc128::zero();
let full_blocks = self.vector_bytes / BLOCK_128;
for block in 0..full_blocks {
let offset = block * BLOCK_128;
let query = QueryBlock128::load(&self.planes, offset);
acc.accumulate_mull(vld1q_u8(data.add(offset)), query);
}
let tail = self.vector_bytes % BLOCK_128;
if tail > 0 {
let offset = full_blocks * BLOCK_128;
let query = QueryBlock128::load(&self.planes, offset);
let mut block = [0u8; BLOCK_128];
std::ptr::copy_nonoverlapping(data.add(offset), block.as_mut_ptr(), tail);
acc.accumulate_mull(vld1q_u8(block.as_ptr()), query);
}
acc
}
}
/// ARMv8.2-A Dot Product variant over the query planes: one 128-bit
/// register of packed codes per block, one `TBL` codebook lookup per
/// plane and `SDOT` into independent accumulators (see [`Acc128`]).
/// The last partial block runs on a zero-padded copy of the remaining
/// bytes.
///
/// # Safety
/// CPU must support `neon` and `dotprod`.
#[target_feature(enable = "neon,dotprod")]
pub unsafe fn dotprod_raw_neon_sdot(&self, vector: &[u8]) -> i64 {
assert_eq!(
vector.len(),
self.vector_bytes,
"QuerySimd<{PLANES}, {QUERY_BYTES}>::dotprod_raw_neon_sdot: vector length mismatch \
({} vs expected {})",
vector.len(),
self.vector_bytes,
);
unsafe { Self::reduce_neon(self.accumulate_neon_sdot(vector.as_ptr())) }
}
/// Block loop of the SDOT kernels over the vector at `data`.
///
/// # Safety
/// CPU must support `dotprod`; `data` must be readable for
/// `self.vector_bytes` bytes.
#[inline]
#[target_feature(enable = "neon,dotprod")]
unsafe fn accumulate_neon_sdot(&self, data: *const u8) -> Acc128<QUERY_BYTES> {
unsafe {
let mut acc = Acc128::zero();
let full_blocks = self.vector_bytes / BLOCK_128;
for block in 0..full_blocks {
let offset = block * BLOCK_128;
let query = QueryBlock128::load(&self.planes, offset);
acc.accumulate_sdot(vld1q_u8(data.add(offset)), query);
}
let tail = self.vector_bytes % BLOCK_128;
if tail > 0 {
let offset = full_blocks * BLOCK_128;
let query = QueryBlock128::load(&self.planes, offset);
let mut block = [0u8; BLOCK_128];
std::ptr::copy_nonoverlapping(data.add(offset), block.as_mut_ptr(), tail);
acc.accumulate_sdot(vld1q_u8(block.as_ptr()), query);
}
acc
}
}
/// Raw dot product from one vector's accumulators. Horizontal adds are
/// single instructions here, so nothing is gained by fusing them.
#[inline]
#[target_feature(enable = "neon")]
fn reduce_neon(acc: Acc128<QUERY_BYTES>) -> i64 {
Self::combine_bytes(acc.fold().map(|total| i64::from(vaddvq_s32(total))))
}
}
#[cfg(test)]
mod tests {
use rand::SeedableRng as _;
use rand::prelude::StdRng;
use super::super::QuerySimd;
use super::super::shared::{parity_dims, random_inputs};
/// Every kernel the host supports must reproduce the scalar reference
/// bit-exactly at every parity dim.
fn kernels_match_scalar<const PLANES: usize, const QUERY_BYTES: usize>() {
let has_dotprod = std::arch::is_aarch64_feature_detected!("dotprod");
let mut rng = StdRng::seed_from_u64(7);
for dim in parity_dims::<PLANES>() {
let (query, vector) = random_inputs::<PLANES, QUERY_BYTES>(&mut rng, dim);
let scalar = query.dotprod_raw(&vector);
let tag = format!("PLANES={PLANES} QUERY_BYTES={QUERY_BYTES} dim={dim}");
unsafe {
let neon = query.dotprod_raw_neon(&vector);
assert_eq!(scalar, neon, "{tag}: scalar {scalar} != neon {neon}");
if has_dotprod {
let sdot = query.dotprod_raw_neon_sdot(&vector);
assert_eq!(scalar, sdot, "{tag}: scalar {scalar} != sdot {sdot}");
}
}
}
}
#[test]
fn test_kernels_match_scalar() {
kernels_match_scalar::<2, 2>();
kernels_match_scalar::<4, 2>();
kernels_match_scalar::<8, 1>();
kernels_match_scalar::<8, 2>();
}
/// Saturation safety at an extreme dim (64K) under the worst-case load:
/// the query maxed out and every code at the max-magnitude codebook
/// slot (all-ones bytes at every width). The scalar reference is i64
/// throughout; a SIMD mismatch proves some intermediate saturated or
/// overflowed.
fn saturation_safety_64k<const PLANES: usize, const QUERY_BYTES: usize>() {
let dim = 65_536;
let query = QuerySimd::<PLANES, QUERY_BYTES>::new(&vec![1.0_f32; dim]);
let vector = vec![0xFF_u8; dim / PLANES];
let scalar = query.dotprod_raw(&vector);
let tag = format!("PLANES={PLANES} QUERY_BYTES={QUERY_BYTES}");
unsafe {
let neon = query.dotprod_raw_neon(&vector);
assert_eq!(scalar, neon, "{tag}: neon disagrees");
if std::arch::is_aarch64_feature_detected!("dotprod") {
let sdot = query.dotprod_raw_neon_sdot(&vector);
assert_eq!(scalar, sdot, "{tag}: sdot disagrees");
}
}
}
#[test]
fn test_saturation_safety_64k() {
saturation_safety_64k::<2, 2>();
saturation_safety_64k::<4, 2>();
saturation_safety_64k::<8, 1>();
saturation_safety_64k::<8, 2>();
}
}
lib/quantization/src/turboquant/simd/query/mod.rs
+488
View File
@@ -0,0 +1,488 @@
//! Asymmetric scoring — an f32 query against vectors of packed codebook
//! indices — generic over the packing width and the query precision.
//!
//! [`QuerySimd<PLANES, QUERY_BYTES>`] scores against vectors whose codes are
//! packed `PLANES` per byte: two 4-bit codes, four 2-bit codes or eight
//! 1-bit codes (`PLANES = 8 / bits`). The widths differ only in their
//! integer codebook (see [`Encoding`]); the query layout, the scalar
//! reference and every SIMD kernel are shared.
//!
//! # Query layout
//!
//! The query is quantized to signed integers of `QUERY_BYTES` i8 bytes,
//! `q_signed = Σ_b K^b · byte_b` (one byte for an 8-bit query, two for a
//! ~16-bit one), which the kernels multiply against the codebook values
//! with `u8 × i8` (x86_64) or `i8 × i8` (aarch64) instructions. The bytes
//! are stored as [`QueryPlanes`]: one plane per query byte and code
//! position within a packed data byte, so a wide load of raw data bytes
//! lines up with the query without unpacking into dim order.
//!
//! # Integer encoding
//!
//! Both architectures share the reconstruction
//! `postprocess_scale · (dot_raw − bias_correction)`, but encode the
//! codebook and the query differently to get the most precision out of
//! their instruction sets:
//!
//! * **aarch64** — `vmull_s8` and `SDOT` are true `i8 × i8 → i16/i32`
//! signed multiplies, so the codebook is stored as signed `i8 ∈ [−127,
//! 127]` with no offset. Query bytes are full `i8` combined with
//! `K = 256` (~15.9-bit precision for a two-byte query), and there is no
//! bias.
//!
//! * **x86_64** — `maddubs` / `VPDPBUSD` consume one `u8` and one `i8`
//! operand. The codebook is stored unsigned (`c_u = c_signed + offset`),
//! and the query bytes are kept narrow enough that the `maddubs` pair
//! sum stays inside i16: with the full `u8` codebooks of the 2- and 4-bit
//! widths (`c_u ≤ 255`) the bytes are 7-bit (`K = 128`, `|pair| ≤
//! 2·255·64 = 32 640`); the 1-bit width only needs `c_u ∈ {0, 128}` and
//! keeps full i8 bytes. The offset contributes a per-query bias
//! `offset · Σ q_signed`, subtracted once per vector.
use super::{SimdBackend, query1bit, query2bit, query4bit};
/// Codebook entry in the storage form the kernels multiply: signed on
/// aarch64 (`vmull_s8` / `SDOT`), unsigned everywhere else (`maddubs` /
/// `VPDPBUSD` take a `u8` operand; the scalar fallback mirrors x86_64).
#[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
pub(crate) type Code = i8;
#[cfg(not(all(target_arch = "aarch64", target_feature = "neon")))]
pub(crate) type Code = u8;
/// Integer encoding of one packing width: the codebook the kernels multiply
/// and the query quantization that keeps every intermediate inside its lane.
#[derive(Clone, Copy)]
pub(crate) struct Encoding {
/// Codebook value per code in the arch-native storage form, padded to
/// a full 16-entry shuffle table; entries past the width's code range
/// are never addressed.
pub codebook: [Code; 16],
/// `c_signed = codebook[k] − offset`; `0` where the codebook is stored
/// signed.
pub offset: i64,
/// Integer units per centroid unit: `c_signed ≈ c_float · scale`.
pub scale: f32,
/// Radix of the query bytes: `q_signed = query_high_coef · high + low`
/// for a two-byte query, each byte in `[−K/2, K/2)`.
pub query_high_coef: i64,
/// Largest `|q_signed|` a two-byte query is scaled to — the symmetric
/// range both bytes cover.
pub query_abs_max: f32,
}
/// Pad a width's codebook to the 16-entry shuffle table [`Encoding`] holds.
pub(crate) const fn pad_codebook<const N: usize>(codebook: [Code; N]) -> [Code; 16] {
let mut padded = [0; 16];
let mut k = 0;
while k < N {
padded[k] = codebook[k];
k += 1;
}
padded
}
/// The encoding of the width packing `planes` codes per byte.
const fn encoding(planes: usize) -> Encoding {
match planes {
2 => query4bit::ENCODING,
4 => query2bit::ENCODING,
8 => query1bit::ENCODING,
_ => panic!("QuerySimd: PLANES must be 2, 4 or 8"),
}
}
/// Widest block of packed data bytes any kernel consumes at once (AVX-512:
/// 64 bytes). Every query plane is padded to a multiple of it.
const PLANE_BLOCK: usize = 64;
/// Query bytes regrouped by code position — the layout the SIMD kernels
/// consume.
///
/// A packed data byte `j` holds the codes of dims `PLANES · j + k` for
/// `k ∈ 0..PLANES`, code `k` in bits `[k · bits, (k + 1) · bits)`. Plane
/// `bytes[b][k]` lines query byte `b` up with exactly that order: its entry
/// `j` is byte `b` of query dim `PLANES · j + k`. So a kernel shifts the
/// raw data bytes by `k · bits`, masks off one code per byte and multiplies
/// against the planes of `k` — no unpacking into dim order.
///
/// Each plane is zero-padded to a multiple of [`PLANE_BLOCK`] bytes, so a
/// kernel can always load a whole block on the query side; whatever the
/// data lanes past the vector's end hold, they multiply against zeros.
pub(crate) struct QueryPlanes<const PLANES: usize, const QUERY_BYTES: usize> {
bytes: [[Vec<i8>; PLANES]; QUERY_BYTES],
}
impl<const PLANES: usize, const QUERY_BYTES: usize> QueryPlanes<PLANES, QUERY_BYTES> {
/// Build the planes from the encoded query, one `[i8; QUERY_BYTES]` per
/// dim; the number of dims must be a multiple of `PLANES`.
fn new(encoded: impl ExactSizeIterator<Item = [i8; QUERY_BYTES]>) -> Self {
let dim = encoded.len();
debug_assert!(dim.is_multiple_of(PLANES));
let len = (dim / PLANES).next_multiple_of(PLANE_BLOCK);
let mut bytes: [[Vec<i8>; PLANES]; QUERY_BYTES] =
std::array::from_fn(|_| std::array::from_fn(|_| vec![0; len]));
for (i, dim_bytes) in encoded.enumerate() {
let (j, k) = (i / PLANES, i % PLANES);
for (b, &byte) in dim_bytes.iter().enumerate() {
bytes[b][k][j] = byte;
}
}
Self { bytes }
}
}
/// Encoded query for asymmetric scoring against vectors packing `PLANES`
/// codes per byte, quantized to `QUERY_BYTES` (1 or 2) i8 bytes per dim.
///
/// The f32 query is scaled to the widest range its bytes can hold and split
/// into i8 bytes stored as [`QueryPlanes`] (see the module docs for the
/// per-arch encoding). Any dim that is a multiple of `PLANES` works, so a
/// matryoshka-trimmed model fits without re-encoding.
///
/// `dotprod` computes `dot_raw = Σ_j q_signed[j] · codebook[v[j]]` and
/// returns `postprocess_scale · (dot_raw − bias_correction)`.
pub struct QuerySimd<const PLANES: usize, const QUERY_BYTES: usize> {
/// Query bytes, regrouped by code position.
planes: QueryPlanes<PLANES, QUERY_BYTES>,
/// Packed bytes per encoded vector: `dim / PLANES`.
vector_bytes: usize,
/// `1 / (q_scale · c_scale)` — prefactor from integer to float dot
/// product.
postprocess_scale: f32,
/// `offset · Σ q_signed[j]` over all dims, subtracted from `dot_raw` to
/// recover the signed dot product; `0` where the codebook is signed.
bias_correction: i64,
/// SIMD backend resolved once at construction, so scoring doesn't re-run
/// CPU feature detection for every vector.
backend: SimdBackend,
}
impl<const PLANES: usize, const QUERY_BYTES: usize> QuerySimd<PLANES, QUERY_BYTES> {
const ENCODING: Encoding = encoding(PLANES);
/// Bits per code.
const BITS: usize = 8 / PLANES;
/// Radix of the query bytes.
const RADIX: i64 = Self::ENCODING.query_high_coef;
/// Largest `|q_signed|` the query is scaled to: the encoding's two-byte
/// range, or for a one-byte query the range of a single byte.
const QUERY_ABS_MAX: f32 = match QUERY_BYTES {
1 => (Self::RADIX / 2 - 1) as f32,
2 => Self::ENCODING.query_abs_max,
_ => panic!("QuerySimd: QUERY_BYTES must be 1 or 2"),
};
/// Encode `data`; its length must be a multiple of `PLANES`.
pub fn new(data: &[f32]) -> Self {
assert!(
data.len().is_multiple_of(PLANES),
"QuerySimd<{PLANES}, {QUERY_BYTES}> requires query dim to be a multiple of {PLANES} \
(got {})",
data.len(),
);
let encoding = Self::ENCODING;
let q_abs_max = data
.iter()
.copied()
.map(f32::abs)
.fold(0.0_f32, f32::max)
.max(f32::EPSILON);
let q_scale = Self::QUERY_ABS_MAX / q_abs_max;
let k = Self::RADIX as i32;
let half_k = k / 2;
let clamp_hi = Self::QUERY_ABS_MAX;
let clamp_lo = -Self::QUERY_ABS_MAX;
// Balanced signed split: every byte but the last takes the remainder
// in `[−k/2, k/2)`, the last takes what is left.
let mut sum_q_signed: i64 = 0;
let planes = QueryPlanes::new(data.iter().map(|&value| {
let q_signed = (value * q_scale).round().clamp(clamp_lo, clamp_hi) as i32;
sum_q_signed += i64::from(q_signed);
let mut rest = q_signed;
std::array::from_fn(|b| {
if b + 1 == QUERY_BYTES {
return rest as i8;
}
let r = rest.rem_euclid(k);
let byte = if r >= half_k { r - k } else { r };
rest = (rest - byte) / k;
byte as i8
})
}));
Self {
planes,
vector_bytes: data.len() / PLANES,
postprocess_scale: 1.0 / (q_scale * encoding.scale),
bias_correction: encoding.offset * sum_q_signed,
backend: SimdBackend::detect(),
}
}
/// Packed bytes per encoded vector: `dim / PLANES`.
#[inline]
pub fn vector_bytes(&self) -> usize {
self.vector_bytes
}
/// Score the encoded query against a `vector` of packed codes (`PLANES`
/// codes per byte, code `k` of byte `j` in bits `[k · bits, (k + 1) ·
/// bits)`). `vector.len()` must equal [`Self::vector_bytes`].
///
/// Dispatches to the SIMD backend resolved at construction (see
/// [`SimdBackend`]).
pub fn dotprod(&self, vector: &[u8]) -> f32 {
self.postprocess(self.dotprod_raw_best(vector))
}
/// Float reconstruction of a raw integer dot product.
#[inline]
fn postprocess(&self, dot_raw: i64) -> f32 {
self.postprocess_scale * (dot_raw - self.bias_correction) as f32
}
#[inline]
fn dotprod_raw_best(&self, vector: &[u8]) -> i64 {
match self.backend {
#[cfg(target_arch = "x86_64")]
SimdBackend::Avx512Vnni => unsafe { self.dotprod_raw_avx512_vnni(vector) },
#[cfg(target_arch = "x86_64")]
SimdBackend::Avx2 => unsafe { self.dotprod_raw_avx2(vector) },
#[cfg(target_arch = "x86_64")]
SimdBackend::Sse => unsafe { self.dotprod_raw_sse(vector) },
#[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
SimdBackend::NeonSdot => unsafe { self.dotprod_raw_neon_sdot(vector) },
#[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
SimdBackend::Neon => unsafe { self.dotprod_raw_neon(vector) },
SimdBackend::Scalar => self.dotprod_raw(vector),
}
}
/// `Σ_b K^b · sums[b]` — the raw dot product from per-byte totals.
#[inline]
fn combine_bytes(sums: [i64; QUERY_BYTES]) -> i64 {
sums.iter()
.rev()
.fold(0, |total, &sum| total * Self::RADIX + sum)
}
/// Scalar reference: `Σ_j q_signed[j] · codebook[v[j]]` over all dims,
/// computed per query byte with i64 accumulators and combined at the
/// end — saturation-free by construction, which is what the SIMD parity
/// tests check against.
pub fn dotprod_raw(&self, vector: &[u8]) -> i64 {
assert_eq!(
vector.len(),
self.vector_bytes,
"QuerySimd<{PLANES}, {QUERY_BYTES}>::dotprod_raw: vector length mismatch ({} vs \
expected {})",
vector.len(),
self.vector_bytes,
);
let encoding = Self::ENCODING;
let mask = (1u8 << Self::BITS) - 1;
let mut sums = [0i64; QUERY_BYTES];
for (j, &byte) in vector.iter().enumerate() {
for k in 0..PLANES {
let code = (byte >> (k * Self::BITS)) & mask;
let c = i64::from(encoding.codebook[code as usize]);
for (sum, planes) in sums.iter_mut().zip(&self.planes.bytes) {
*sum += i64::from(planes[k][j]) * c;
}
}
}
Self::combine_bytes(sums)
}
}
#[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
mod arm;
#[cfg(target_arch = "x86_64")]
mod x64;
/// Test helpers shared by the kernel parity tests of every backend.
#[cfg(test)]
pub(crate) mod shared {
use rand::prelude::StdRng;
use super::super::shared::{random_bytes, sample_normal_vec};
use super::QuerySimd;
/// Corner-case vector lengths (packed bytes) for the block loops of
/// every backend — blocks of 16 (SSE, NEON), 32 (AVX2) and 64 (AVX-512)
/// bytes, with the last partial block padded or masked:
/// • `16, 64, 128, 512, 1024` — whole blocks only, at every width.
/// • `1` — a single data byte.
/// • `15, 31, 63` — one byte short of a 16-, 32- and 64-byte block;
/// `65` — one byte past a 64-byte block.
/// • `8, 9, 23, 24, 135, 513, 1023` — assorted partial blocks,
/// including a realistic matryoshka slice.
pub const PARITY_BYTES: &[usize] = &[
1, 8, 9, 15, 16, 23, 24, 31, 63, 64, 65, 128, 135, 512, 513, 1023, 1024,
];
/// The dims of [`PARITY_BYTES`] at a width packing `PLANES` codes per
/// byte.
pub fn parity_dims<const PLANES: usize>() -> impl Iterator<Item = usize> {
PARITY_BYTES.iter().map(move |&bytes| bytes * PLANES)
}
/// Parity-test inputs: a query ~ N(0, 1) and a vector of uniformly
/// random codes (random bytes are random packed codes at every width).
pub fn random_inputs<const PLANES: usize, const QUERY_BYTES: usize>(
rng: &mut StdRng,
dim: usize,
) -> (QuerySimd<PLANES, QUERY_BYTES>, Vec<u8>) {
let query = sample_normal_vec(rng, dim);
(QuerySimd::new(&query), random_bytes(rng, dim / PLANES))
}
}
/// Accuracy tests of the public `QuerySimd` API against float ground truth,
/// at every width and query precision. Per-backend parity tests (SIMD
/// kernel vs the scalar reference `dotprod_raw`) 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::QuerySimd;
use crate::turboquant::TQBits;
/// The width packing `PLANES` codes per byte.
fn bits<const PLANES: usize>() -> TQBits {
match PLANES {
2 => TQBits::Bits4,
4 => TQBits::Bits2,
8 => TQBits::Bits1,
_ => unreachable!(),
}
}
/// Reconstruction accuracy on realistic PQ inputs: query ∼ N(0,1), vector
/// drawn from N(0,1) then mapped to its nearest centroid. `dotprod` is
/// compared against the ideal PQ dot (`Σ q[j] · c[v[j]]` with
/// float-precision centroid lookup) — the error the integer encoding
/// adds over a hypothetical perfect-precision PQ should be tiny.
///
/// Parameterized over matryoshka-style corner-case dims to exercise the
/// tail handling end-to-end (not just bit-exact parity).
fn dotprod_matches_float<const PLANES: usize, const QUERY_BYTES: usize>(dim: usize) {
let mut rng = StdRng::seed_from_u64(42);
let n_trials = 64;
let bits = bits::<PLANES>();
let centroids = bits.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 = QuerySimd::<PLANES, QUERY_BYTES>::new(&query)
.dotprod(&pack_codes(&indices, bits.bit_size()));
// Error scales roughly like √dim · σ_q · ε_c. Allow a tolerance
// that is comfortably above the 3σ tail for dim up to ~2K; an
// 8-bit query quantizes ~64× coarser than a two-byte one.
let base = if QUERY_BYTES == 1 { 0.06 } else { 0.03 };
let tol = (0.5_f32).max(base * (dim as f32).sqrt());
assert!(
(pq_dot - simd_dot).abs() < tol,
"PLANES={PLANES} QUERY_BYTES={QUERY_BYTES} dim={dim}: simd_dot {simd_dot} too \
far from ideal PQ dot {pq_dot} (tol={tol})",
);
}
}
/// Corner-case packed lengths: whole blocks, small and maximal tails,
/// odd block counts and a realistic matryoshka slice.
#[rstest::rstest]
#[case::full_blocks(128)]
#[case::small_tail(9)]
#[case::max_tail(15)]
#[case::odd_blocks_only(24)]
#[case::odd_blocks_plus_tail(31)]
#[case::matryoshka(135)]
#[case::large_with_tail(1023)]
fn test_dotprod_matches_float(#[case] bytes: usize) {
dotprod_matches_float::<2, 2>(bytes * 2);
dotprod_matches_float::<4, 2>(bytes * 4);
dotprod_matches_float::<8, 1>(bytes * 8);
dotprod_matches_float::<8, 2>(bytes * 8);
}
/// Quantitative proof that the integer encoding is negligible next to
/// PQ centroid snapping: RMS error added by the 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.
fn simd_noise_below_pq_noise<const PLANES: usize, const QUERY_BYTES: usize>() {
let mut rng = StdRng::seed_from_u64(123);
let dim = 256;
let n_trials = 256;
let bits = bits::<PLANES>();
let centroids = bits.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(
QuerySimd::<PLANES, QUERY_BYTES>::new(&query)
.dotprod(&pack_codes(&indices, bits.bit_size())),
);
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 PLANES={PLANES} QUERY_BYTES={QUERY_BYTES} 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,
"PLANES={PLANES} QUERY_BYTES={QUERY_BYTES}: SIMD noise RMS {rms_simd_noise:.4} \
should be << PQ noise RMS {rms_pq_noise:.4} (ratio {:.2}×)",
rms_pq_noise / rms_simd_noise,
);
}
#[test]
fn test_simd_noise_below_pq_noise() {
simd_noise_below_pq_noise::<2, 2>();
simd_noise_below_pq_noise::<4, 2>();
simd_noise_below_pq_noise::<8, 1>();
simd_noise_below_pq_noise::<8, 2>();
}
}
lib/quantization/src/turboquant/simd/query/x64.rs
+655
View File
@@ -0,0 +1,655 @@
//! x86_64 kernels for [`QuerySimd`].
//!
//! The codebook is stored unsigned (`Encoding::codebook` as `u8`), which
//! is what `maddubs` / `VPDPBUSD` expect in their `u8` operand slot; the
//! signed shift `c_signed = c_u − offset` is undone by the per-query
//! `bias_correction`. See the module docs of [`super`] for the query
//! ranges that keep the `maddubs` pair sums inside i16.
use core::arch::x86_64::*;
use super::{Code, PLANE_BLOCK, QueryPlanes, QuerySimd, encoding};
/// Packed data bytes per SSE block: one XMM of codes.
const BLOCK_128: usize = 16;
const _: () = assert!(PLANE_BLOCK.is_multiple_of(BLOCK_128));
/// Packed data bytes per AVX2 block: one YMM of codes.
const BLOCK_256: usize = 32;
const _: () = assert!(PLANE_BLOCK.is_multiple_of(BLOCK_256));
/// Packed data bytes per AVX-512 block: one ZMM of codes.
const BLOCK_512: usize = 64;
const _: () = assert!(PLANE_BLOCK.is_multiple_of(BLOCK_512));
/// The codebook as a shuffle table for the width packing `PLANES` codes per
/// byte.
const fn codebook<const PLANES: usize>() -> [Code; 16] {
encoding(PLANES).codebook
}
/// Mask of one code in the low bits of a byte.
const fn code_mask<const PLANES: usize>() -> i8 {
((1u16 << (8 / PLANES)) - 1) as i8
}
/// One [`BLOCK_128`]-byte block of every query plane.
#[derive(Clone, Copy)]
struct QueryBlock128<const PLANES: usize, const QUERY_BYTES: usize> {
bytes: [[__m128i; PLANES]; QUERY_BYTES],
}
impl<const PLANES: usize, const QUERY_BYTES: usize> QueryBlock128<PLANES, QUERY_BYTES> {
/// # Safety
/// CPU must support `sse2`; `offset + BLOCK_128` must not exceed the
/// plane length.
#[inline]
#[target_feature(enable = "sse2")]
unsafe fn load(planes: &QueryPlanes<PLANES, QUERY_BYTES>, offset: usize) -> Self {
let mut block = Self {
bytes: [[_mm_setzero_si128(); PLANES]; QUERY_BYTES],
};
for (b, byte_planes) in block.bytes.iter_mut().enumerate() {
for (k, plane) in byte_planes.iter_mut().enumerate() {
*plane = unsafe { load_plane_128(&planes.bytes[b][k], offset) };
}
}
block
}
}
/// # Safety
/// CPU must support `sse2`; `offset + BLOCK_128 <= plane.len()`.
#[inline]
#[target_feature(enable = "sse2")]
unsafe fn load_plane_128(plane: &[i8], offset: usize) -> __m128i {
debug_assert!(offset + BLOCK_128 <= plane.len());
unsafe { _mm_loadu_si128(plane.as_ptr().add(offset).cast::<__m128i>()) }
}
/// [`next_plane_512`] on XMM.
#[inline]
#[target_feature(enable = "sse2")]
unsafe fn next_plane_128<const PLANES: usize>(codes: __m128i) -> __m128i {
match PLANES {
2 => _mm_srli_epi16(codes, 4),
4 => _mm_srli_epi16(codes, 2),
_ => _mm_srli_epi16(codes, 1),
}
}
/// `maddubs → madd` accumulators of one vector; the 128-bit form of
/// [`Acc256`], with the same integer bounds.
#[derive(Clone, Copy)]
struct Acc128<const QUERY_BYTES: usize> {
bytes: [[__m128i; 2]; QUERY_BYTES],
}
impl<const QUERY_BYTES: usize> Acc128<QUERY_BYTES> {
/// # Safety
/// CPU must support `sse2`.
#[inline]
#[target_feature(enable = "sse2")]
unsafe fn zero() -> Self {
Self {
bytes: [[_mm_setzero_si128(); 2]; QUERY_BYTES],
}
}
/// Fold one block of packed `codes` (16 bytes) into the accumulators.
///
/// # Safety
/// CPU must support `ssse3` and `sse4.1`.
#[inline]
#[target_feature(enable = "sse4.1,ssse3")]
unsafe fn accumulate<const PLANES: usize>(
&mut self,
codes: __m128i,
query: QueryBlock128<PLANES, QUERY_BYTES>,
) {
let table = const { codebook::<PLANES>() };
let codebook = unsafe { _mm_loadu_si128(table.as_ptr().cast::<__m128i>()) };
let mask = _mm_set1_epi8(const { code_mask::<PLANES>() });
let ones = _mm_set1_epi16(1);
let mut shifted = codes;
for k in 0..PLANES {
let values = _mm_shuffle_epi8(codebook, _mm_and_si128(shifted, mask));
for (acc, planes) in self.bytes.iter_mut().zip(&query.bytes) {
acc[k & 1] = _mm_add_epi32(
acc[k & 1],
_mm_madd_epi16(_mm_maddubs_epi16(values, planes[k]), ones),
);
}
shifted = unsafe { next_plane_128::<PLANES>(shifted) };
}
}
/// Per-lane totals of every query byte.
///
/// # Safety
/// CPU must support `sse2`.
#[inline]
#[target_feature(enable = "sse2")]
unsafe fn fold(self) -> [__m128i; QUERY_BYTES] {
self.bytes.map(|[a, b]| _mm_add_epi32(a, b))
}
}
/// One [`BLOCK_256`]-byte block of every query plane.
#[derive(Clone, Copy)]
struct QueryBlock256<const PLANES: usize, const QUERY_BYTES: usize> {
bytes: [[__m256i; PLANES]; QUERY_BYTES],
}
impl<const PLANES: usize, const QUERY_BYTES: usize> QueryBlock256<PLANES, QUERY_BYTES> {
/// # Safety
/// CPU must support `avx2`; `offset + BLOCK_256` must not exceed the
/// plane length.
#[inline]
#[target_feature(enable = "avx2")]
unsafe fn load(planes: &QueryPlanes<PLANES, QUERY_BYTES>, offset: usize) -> Self {
let mut block = Self {
bytes: [[_mm256_setzero_si256(); PLANES]; QUERY_BYTES],
};
for (b, byte_planes) in block.bytes.iter_mut().enumerate() {
for (k, plane) in byte_planes.iter_mut().enumerate() {
*plane = unsafe { load_plane_256(&planes.bytes[b][k], offset) };
}
}
block
}
}
/// # Safety
/// CPU must support `avx2`; `offset + BLOCK_256 <= plane.len()`.
#[inline]
#[target_feature(enable = "avx2")]
unsafe fn load_plane_256(plane: &[i8], offset: usize) -> __m256i {
debug_assert!(offset + BLOCK_256 <= plane.len());
unsafe { _mm256_loadu_si256(plane.as_ptr().add(offset).cast::<__m256i>()) }
}
/// [`next_plane_512`] on YMM.
#[inline]
#[target_feature(enable = "avx2")]
unsafe fn next_plane_256<const PLANES: usize>(codes: __m256i) -> __m256i {
match PLANES {
2 => _mm256_srli_epi16(codes, 4),
4 => _mm256_srli_epi16(codes, 2),
_ => _mm256_srli_epi16(codes, 1),
}
}
/// `maddubs → madd` accumulators of one vector; the same two-pair shape as
/// [`Acc512`].
///
/// The `maddubs` pair sums stay inside i16 by the module-level query bound
/// (`|pair| ≤ 2 · 255 · 64 = 32 640`, or exactly `2 · 128 · 128 = 32 768`
/// at the negative end for the 1-bit encoding); `madd` against ones then
/// adds at most 65 280 (65 536) per i32 lane per plane, the same bound as
/// VNNI.
#[derive(Clone, Copy)]
struct Acc256<const QUERY_BYTES: usize> {
bytes: [[__m256i; 2]; QUERY_BYTES],
}
impl<const QUERY_BYTES: usize> Acc256<QUERY_BYTES> {
/// # Safety
/// CPU must support `avx2`.
#[inline]
#[target_feature(enable = "avx2")]
unsafe fn zero() -> Self {
Self {
bytes: [[_mm256_setzero_si256(); 2]; QUERY_BYTES],
}
}
/// Fold one block of packed `codes` (32 bytes) into the accumulators.
///
/// # Safety
/// CPU must support `avx2`.
#[inline]
#[target_feature(enable = "avx2")]
unsafe fn accumulate<const PLANES: usize>(
&mut self,
codes: __m256i,
query: QueryBlock256<PLANES, QUERY_BYTES>,
) {
let table = const { codebook::<PLANES>() };
let codebook = _mm256_broadcastsi128_si256(unsafe {
_mm_loadu_si128(table.as_ptr().cast::<__m128i>())
});
let mask = _mm256_set1_epi8(const { code_mask::<PLANES>() });
let ones = _mm256_set1_epi16(1);
let mut shifted = codes;
for k in 0..PLANES {
let values = _mm256_shuffle_epi8(codebook, _mm256_and_si256(shifted, mask));
for (acc, planes) in self.bytes.iter_mut().zip(&query.bytes) {
acc[k & 1] = _mm256_add_epi32(
acc[k & 1],
_mm256_madd_epi16(_mm256_maddubs_epi16(values, planes[k]), ones),
);
}
shifted = unsafe { next_plane_256::<PLANES>(shifted) };
}
}
/// Per-lane totals of every query byte.
///
/// # Safety
/// CPU must support `avx2`.
#[inline]
#[target_feature(enable = "avx2")]
unsafe fn fold(self) -> [__m256i; QUERY_BYTES] {
self.bytes.map(|[a, b]| _mm256_add_epi32(a, b))
}
}
/// One [`BLOCK_512`]-byte block of every query plane.
#[derive(Clone, Copy)]
struct QueryBlock512<const PLANES: usize, const QUERY_BYTES: usize> {
bytes: [[__m512i; PLANES]; QUERY_BYTES],
}
impl<const PLANES: usize, const QUERY_BYTES: usize> QueryBlock512<PLANES, QUERY_BYTES> {
/// # Safety
/// CPU must support `avx512f`; `offset + BLOCK_512` must not exceed the
/// plane length.
#[inline]
#[target_feature(enable = "avx512f")]
unsafe fn load(planes: &QueryPlanes<PLANES, QUERY_BYTES>, offset: usize) -> Self {
let mut block = Self {
bytes: [[_mm512_setzero_si512(); PLANES]; QUERY_BYTES],
};
for (b, byte_planes) in block.bytes.iter_mut().enumerate() {
for (k, plane) in byte_planes.iter_mut().enumerate() {
*plane = unsafe { load_plane_512(&planes.bytes[b][k], offset) };
}
}
block
}
}
/// # Safety
/// CPU must support `avx512f`; `offset + BLOCK_512 <= plane.len()`.
#[inline]
#[target_feature(enable = "avx512f")]
unsafe fn load_plane_512(plane: &[i8], offset: usize) -> __m512i {
debug_assert!(offset + BLOCK_512 <= plane.len());
unsafe { _mm512_loadu_si512(plane.as_ptr().add(offset).cast::<__m512i>()) }
}
/// The next plane's codes moved into the low bits of every byte: `codes >>
/// bits` within 16-bit lanes, so the caller masks off what crosses over
/// from the neighboring byte. The shift count is an immediate, hence the
/// match over the widths.
#[inline]
#[target_feature(enable = "avx512bw")]
unsafe fn next_plane_512<const PLANES: usize>(codes: __m512i) -> __m512i {
match PLANES {
2 => _mm512_srli_epi16(codes, 4),
4 => _mm512_srli_epi16(codes, 2),
_ => _mm512_srli_epi16(codes, 1),
}
}
/// `VPDPBUSD` accumulators of one vector: per query byte two independent
/// chains, plane `k` feeding chain `k & 1`. Two chains per byte are
/// enough to keep the multiply-accumulate latency off the critical path at
/// every width (`PLANES / 2` dependent steps per block against `PLANES`
/// issued per byte).
///
/// i32 lane bound: each `VPDPBUSD` adds at most `4 · 255 · 64 = 65 280`
/// (`4 · 128 · 128 = 65 536` for the 1-bit encoding) to a lane, so overflow
/// needs ~32 K blocks ≈ 2 M packed bytes — far beyond any real input.
#[derive(Clone, Copy)]
struct Acc512<const QUERY_BYTES: usize> {
bytes: [[__m512i; 2]; QUERY_BYTES],
}
impl<const QUERY_BYTES: usize> Acc512<QUERY_BYTES> {
/// # Safety
/// CPU must support `avx512f`.
#[inline]
#[target_feature(enable = "avx512f")]
unsafe fn zero() -> Self {
Self {
bytes: [[_mm512_setzero_si512(); 2]; QUERY_BYTES],
}
}
/// Fold one block of packed `codes` (64 bytes) into the accumulators:
/// plane by plane, the codes are shifted down, masked to one code per
/// byte and looked up in the codebook.
///
/// # Safety
/// CPU must support `avx512f`, `avx512bw`, and `avx512vnni`.
#[inline]
#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
unsafe fn accumulate<const PLANES: usize>(
&mut self,
codes: __m512i,
query: QueryBlock512<PLANES, QUERY_BYTES>,
) {
let table = const { codebook::<PLANES>() };
let codebook =
_mm512_broadcast_i32x4(unsafe { _mm_loadu_si128(table.as_ptr().cast::<__m128i>()) });
let mask = _mm512_set1_epi8(const { code_mask::<PLANES>() });
let mut shifted = codes;
for k in 0..PLANES {
let values = _mm512_shuffle_epi8(codebook, _mm512_and_si512(shifted, mask));
for (acc, planes) in self.bytes.iter_mut().zip(&query.bytes) {
acc[k & 1] = _mm512_dpbusd_epi32(acc[k & 1], values, planes[k]);
}
shifted = unsafe { next_plane_512::<PLANES>(shifted) };
}
}
/// Per-lane totals of every query byte.
///
/// # Safety
/// CPU must support `avx512f`.
#[inline]
#[target_feature(enable = "avx512f")]
unsafe fn fold(self) -> [__m512i; QUERY_BYTES] {
self.bytes.map(|[a, b]| _mm512_add_epi32(a, b))
}
}
impl<const PLANES: usize, const QUERY_BYTES: usize> QuerySimd<PLANES, QUERY_BYTES> {
/// x86_64 SSE4.1 + SSSE3 over the query planes: one XMM of packed codes
/// per block, the 128-bit form of the AVX2 kernel.
///
/// # Safety
/// CPU must support `ssse3` and `sse4.1`.
#[target_feature(enable = "sse4.1,ssse3")]
pub unsafe fn dotprod_raw_sse(&self, vector: &[u8]) -> i64 {
assert_eq!(
vector.len(),
self.vector_bytes,
"QuerySimd<{PLANES}, {QUERY_BYTES}>::dotprod_raw_sse: vector length mismatch ({} vs \
expected {})",
vector.len(),
self.vector_bytes,
);
unsafe {
let totals = self.accumulate_sse(vector.as_ptr()).fold();
Self::combine_bytes(totals.map(|total| i64::from(hsum_i32_sse(total))))
}
}
/// Block loop of the SSE kernels over the vector at `data`.
///
/// # Safety
/// CPU must support `ssse3` and `sse4.1`; `data` must be readable for
/// `self.vector_bytes` bytes.
#[inline]
#[target_feature(enable = "sse4.1,ssse3")]
unsafe fn accumulate_sse(&self, data: *const u8) -> Acc128<QUERY_BYTES> {
unsafe {
let mut acc = Acc128::zero();
let full_blocks = self.vector_bytes / BLOCK_128;
for block in 0..full_blocks {
let offset = block * BLOCK_128;
let query = QueryBlock128::load(&self.planes, offset);
let codes = _mm_loadu_si128(data.add(offset).cast::<__m128i>());
acc.accumulate(codes, query);
}
let tail = self.vector_bytes % BLOCK_128;
if tail > 0 {
let offset = full_blocks * BLOCK_128;
let query = QueryBlock128::load(&self.planes, offset);
let mut block = [0u8; BLOCK_128];
std::ptr::copy_nonoverlapping(data.add(offset), block.as_mut_ptr(), tail);
let codes = _mm_loadu_si128(block.as_ptr().cast::<__m128i>());
acc.accumulate(codes, query);
}
acc
}
}
/// x86_64 AVX2 over the query planes: one YMM of packed codes per
/// block, one `vpshufb` codebook lookup per plane and `maddubs → madd`
/// products into independent accumulators (see [`Acc256`]). The last
/// partial block runs on a zero-padded copy of the remaining bytes.
///
/// # Safety
/// CPU must support `avx2`.
#[target_feature(enable = "avx2")]
pub unsafe fn dotprod_raw_avx2(&self, vector: &[u8]) -> i64 {
assert_eq!(
vector.len(),
self.vector_bytes,
"QuerySimd<{PLANES}, {QUERY_BYTES}>::dotprod_raw_avx2: vector length mismatch ({} \
vs expected {})",
vector.len(),
self.vector_bytes,
);
unsafe { Self::reduce_avx2(self.accumulate_avx2(vector.as_ptr())) }
}
/// Block loop of the AVX2 kernels over the vector at `data`.
///
/// # Safety
/// CPU must support `avx2`; `data` must be readable for
/// `self.vector_bytes` bytes.
#[inline]
#[target_feature(enable = "avx2")]
unsafe fn accumulate_avx2(&self, data: *const u8) -> Acc256<QUERY_BYTES> {
unsafe {
let mut acc = Acc256::zero();
let full_blocks = self.vector_bytes / BLOCK_256;
for block in 0..full_blocks {
let offset = block * BLOCK_256;
let query = QueryBlock256::load(&self.planes, offset);
let codes = _mm256_loadu_si256(data.add(offset).cast::<__m256i>());
acc.accumulate(codes, query);
}
let tail = self.vector_bytes % BLOCK_256;
if tail > 0 {
let offset = full_blocks * BLOCK_256;
let query = QueryBlock256::load(&self.planes, offset);
let mut block = [0u8; BLOCK_256];
std::ptr::copy_nonoverlapping(data.add(offset), block.as_mut_ptr(), tail);
let codes = _mm256_loadu_si256(block.as_ptr().cast::<__m256i>());
acc.accumulate(codes, query);
}
acc
}
}
/// Raw dot product from one vector's accumulators.
///
/// # Safety
/// CPU must support `avx2`.
#[inline]
#[target_feature(enable = "avx2")]
unsafe fn reduce_avx2(acc: Acc256<QUERY_BYTES>) -> i64 {
unsafe { Self::combine_bytes(acc.fold().map(|total| i64::from(hsum_i32_avx2(total)))) }
}
/// AVX-512 VNNI (Ice Lake Xeon+, Zen 4+) over the query planes: one ZMM
/// of packed codes per block, one `vpshufb` codebook lookup per plane
/// and `VPDPBUSD` into independent accumulators (see [`Acc512`]). The
/// last partial block is a masked load whose dead lanes multiply
/// against the planes' zero padding.
///
/// # Safety
/// CPU must support `avx512f`, `avx512bw`, and `avx512vnni`.
#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
pub unsafe fn dotprod_raw_avx512_vnni(&self, vector: &[u8]) -> i64 {
assert_eq!(
vector.len(),
self.vector_bytes,
"QuerySimd<{PLANES}, {QUERY_BYTES}>::dotprod_raw_avx512_vnni: vector length \
mismatch ({} vs expected {})",
vector.len(),
self.vector_bytes,
);
unsafe { Self::reduce_avx512(self.accumulate_avx512(vector.as_ptr())) }
}
/// Block loop of the AVX-512 kernels over the vector at `data`.
///
/// # Safety
/// CPU must support `avx512f`, `avx512bw`, and `avx512vnni`; `data` must
/// be readable for `self.vector_bytes` bytes.
#[inline]
#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
unsafe fn accumulate_avx512(&self, data: *const u8) -> Acc512<QUERY_BYTES> {
unsafe {
let mut acc = Acc512::zero();
let full_blocks = self.vector_bytes / BLOCK_512;
for block in 0..full_blocks {
let offset = block * BLOCK_512;
let query = QueryBlock512::load(&self.planes, offset);
let codes = _mm512_loadu_si512(data.add(offset).cast::<__m512i>());
acc.accumulate(codes, query);
}
let tail = self.vector_bytes % BLOCK_512;
if tail > 0 {
let offset = full_blocks * BLOCK_512;
let query = QueryBlock512::load(&self.planes, offset);
let mask: __mmask64 = (1 << tail) - 1;
let codes = _mm512_maskz_loadu_epi8(mask, data.add(offset).cast::<i8>());
acc.accumulate(codes, query);
}
acc
}
}
/// Raw dot product from one vector's accumulators.
///
/// # Safety
/// CPU must support `avx512f`.
#[inline]
#[target_feature(enable = "avx512f")]
unsafe fn reduce_avx512(acc: Acc512<QUERY_BYTES>) -> i64 {
unsafe {
Self::combine_bytes(
acc.fold()
.map(|total| i64::from(_mm512_reduce_add_epi32(total))),
)
}
}
}
#[target_feature(enable = "sse2")]
unsafe fn hsum_i32_sse(v: __m128i) -> i32 {
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)
}
#[inline]
#[target_feature(enable = "avx2")]
unsafe fn hsum_i32_avx2(v: __m256i) -> i32 {
unsafe {
hsum_i32_sse(_mm_add_epi32(
_mm256_castsi256_si128(v),
_mm256_extracti128_si256(v, 1),
))
}
}
#[cfg(test)]
mod tests {
use rand::SeedableRng as _;
use rand::prelude::StdRng;
use super::super::QuerySimd;
use super::super::shared::{parity_dims, random_inputs};
fn has_sse() -> bool {
std::is_x86_feature_detected!("ssse3") && std::is_x86_feature_detected!("sse4.1")
}
fn has_avx512_vnni() -> bool {
std::is_x86_feature_detected!("avx512f")
&& std::is_x86_feature_detected!("avx512bw")
&& std::is_x86_feature_detected!("avx512vnni")
}
/// Every kernel the host supports must reproduce the scalar reference
/// bit-exactly at every parity dim.
fn kernels_match_scalar<const PLANES: usize, const QUERY_BYTES: usize>() {
let mut rng = StdRng::seed_from_u64(7);
for dim in parity_dims::<PLANES>() {
let (query, vector) = random_inputs::<PLANES, QUERY_BYTES>(&mut rng, dim);
let scalar = query.dotprod_raw(&vector);
let tag = format!("PLANES={PLANES} QUERY_BYTES={QUERY_BYTES} dim={dim}");
unsafe {
if has_sse() {
let sse = query.dotprod_raw_sse(&vector);
assert_eq!(scalar, sse, "{tag}: scalar {scalar} != sse {sse}");
}
if std::is_x86_feature_detected!("avx2") {
let avx2 = query.dotprod_raw_avx2(&vector);
assert_eq!(scalar, avx2, "{tag}: scalar {scalar} != avx2 {avx2}");
}
if has_avx512_vnni() {
let vnni512 = query.dotprod_raw_avx512_vnni(&vector);
assert_eq!(
scalar, vnni512,
"{tag}: scalar {scalar} != avx512_vnni {vnni512}"
);
}
}
}
}
#[test]
fn test_kernels_match_scalar() {
kernels_match_scalar::<2, 2>();
kernels_match_scalar::<4, 2>();
kernels_match_scalar::<8, 1>();
kernels_match_scalar::<8, 2>();
}
/// Saturation safety at an extreme dim (64K) under the worst-case load:
/// the query maxed out and every code at the max-magnitude codebook
/// slot (all-ones bytes at every width). The scalar reference is i64
/// throughout; a SIMD mismatch proves some intermediate saturated or
/// overflowed.
fn saturation_safety_64k<const PLANES: usize, const QUERY_BYTES: usize>() {
let dim = 65_536;
let query = QuerySimd::<PLANES, QUERY_BYTES>::new(&vec![1.0_f32; dim]);
let vector = vec![0xFF_u8; dim / PLANES];
let scalar = query.dotprod_raw(&vector);
let tag = format!("PLANES={PLANES} QUERY_BYTES={QUERY_BYTES}");
unsafe {
if has_sse() {
let sse = query.dotprod_raw_sse(&vector);
assert_eq!(scalar, sse, "{tag}: sse disagrees");
}
if std::is_x86_feature_detected!("avx2") {
let avx2 = query.dotprod_raw_avx2(&vector);
assert_eq!(scalar, avx2, "{tag}: avx2 disagrees");
}
if has_avx512_vnni() {
let vnni512 = query.dotprod_raw_avx512_vnni(&vector);
assert_eq!(scalar, vnni512, "{tag}: avx512_vnni disagrees");
}
}
}
#[test]
fn test_saturation_safety_64k() {
saturation_safety_64k::<2, 2>();
saturation_safety_64k::<4, 2>();
saturation_safety_64k::<8, 1>();
saturation_safety_64k::<8, 2>();
}
}
lib/quantization/src/turboquant/simd/query1bit/arm.rs
-131
View File
@@ -45,73 +45,6 @@ pub unsafe fn score_1bit_internal_neon(a: &[u8], b: &[u8]) -> f32 {
}
}
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 _;
@@ -161,68 +94,4 @@ mod tests {
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}",
);
}
}
lib/quantization/src/turboquant/simd/query1bit/mod.rs
+46 -285
View File
@@ -10,11 +10,12 @@
//! 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.
//! * [`Query1bitSimd`] / [`Query1bitWideSimd`] — asymmetric scoring of an
//! original query (8- or 16-bit) against packed data, through the shared
//! [`QuerySimd`] kernels with the 1-bit [`Encoding`] defined here (a sign
//! bit indexes a two-entry codebook).
use super::query::{Code, Encoding, QuerySimd, pad_codebook};
/// `|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.
@@ -24,224 +25,47 @@ const CENTROID_ABS: f32 = 0.797_884_6;
/// 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;
/// Integer encoding of the 1-bit width for the shared asymmetric kernels
/// ([`super::query::QuerySimd`]). Both codebook entries are `∓c`, so any
/// symmetric integer pair represents them exactly; the pair is chosen per
/// arch for the widest query the multiply keeps inside its lane.
pub(super) const ENCODING: Encoding = Encoding {
codebook: pad_codebook(CODEBOOK),
offset: CODEBOOK_OFFSET,
scale: CODEBOOK_SCALE,
query_high_coef: 256,
query_abs_max: 32639.0,
};
/// 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,
}
/// aarch64: signed `∓127`, no offset.
#[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
const CODEBOOK: [Code; 2] = [-127, 127];
#[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
const CODEBOOK_OFFSET: i64 = 0;
#[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
const CODEBOOK_SCALE: f32 = 127.0 / CENTROID_ABS;
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(),
);
/// x86_64: `{0, 128}` with offset 64 — a codebook magnitude of 128 lets
/// the query halves span the full i8 range without the `maddubs` pair sum
/// leaving i16 (`|pair| ≤ 2 · 128 · 128 = 32 768`, reached only at the
/// negative end, which i16 holds).
#[cfg(not(all(target_arch = "aarch64", target_feature = "neon")))]
const CODEBOOK: [Code; 2] = [0, 128];
#[cfg(not(all(target_arch = "aarch64", target_feature = "neon")))]
const CODEBOOK_OFFSET: i64 = 64;
#[cfg(not(all(target_arch = "aarch64", target_feature = "neon")))]
const CODEBOOK_SCALE: f32 = 64.0 / CENTROID_ABS;
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;
/// Encoded query for asymmetric 1-bit scoring: [`QuerySimd`] over eight
/// codes per byte (the sign bit of dim `8j + k` is bit `k` of byte `j`),
/// with an 8-bit query — plenty next to the ±c codes, and half the
/// multiply-accumulate work of [`Query1bitWideSimd`].
pub type Query1bitSimd = QuerySimd<8, 1>;
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(crate) 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(crate) fn num_full_blocks(&self) -> usize {
self.num_full_blocks
}
}
/// [`Query1bitSimd`] with a 16-bit query, for TQ+: the per-coord `D'`
/// pre-scaling pushes some query coords toward the small end of the
/// integer range, where 8 bits lose too much.
pub type Query1bitWideSimd = QuerySimd<8, 2>;
/// Dot product between two already-encoded 1-bit PQ vectors.
///
@@ -326,8 +150,9 @@ pub mod shared {
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.
/// Codebook and symmetric-scoring accuracy tests. Per-arch parity tests of
/// the popcount kernels live in the `arm` / `x64` submodules; the
/// asymmetric path is covered by `super::query`.
#[cfg(test)]
mod tests {
use rand::SeedableRng as _;
@@ -405,68 +230,4 @@ mod tests {
assert!((score_1bit_internal(&a, &a) - max).abs() < 1e-3);
assert!((score_1bit_internal(&a, &not_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}",
);
}
}
lib/quantization/src/turboquant/simd/query1bit/x64.rs
-220
View File
@@ -182,160 +182,6 @@ pub unsafe fn score_1bit_internal_avx512_vpopcntdq(a: &[u8], b: &[u8]) -> f32 {
}
}
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 _;
@@ -433,70 +279,4 @@ mod tests {
}
}
}
/// 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);
}
}
lib/quantization/src/turboquant/simd/query2bit/arm.rs
+5 -204
View File
@@ -1,7 +1,5 @@
//! 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.
//! NEON kernels for the symmetric 2-bit paths (`score_2bit_internal*`) on
//! aarch64. The codebook is the signed `CODEBOOK_I8`.
//!
//! # Unpack trick
//! Each packed data byte holds 4 × 2-bit codes; consecutive pairs (`c0c1`,
@@ -21,7 +19,7 @@
//!
//! Same pipeline as 4-bit from that point on (i8 × i8 → i16 + `vpadalq_s16`).
use super::{CODEBOOK_I8, CODEBOOK_SCALE, QUERY_HIGH_COEF, Query2bitSimd};
use super::{CODEBOOK_I8, CODEBOOK_SCALE};
/// `PAIR_TABLE_EVEN[nibble]` = `CODEBOOK_I8[nibble & 0b11]`.
const PAIR_TABLE_EVEN: [i8; 16] = {
@@ -76,152 +74,6 @@ unsafe fn unpack_16_codes(bytes4: *const u8) -> core::arch::aarch64::int8x16_t {
}
}
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.
///
@@ -260,8 +112,7 @@ pub unsafe fn score_2bit_internal_neon(a: &[u8], b: &[u8]) -> f32 {
}
}
/// SDOT variant of [`score_2bit_internal_neon`]. 2× chunk unroll mirrors
/// [`Query2bitSimd::dotprod_raw_neon_sdot`].
/// SDOT variant of [`score_2bit_internal_neon`], 2× chunk-unrolled.
///
/// # Safety
/// CPU must support `neon` and `dotprod`.
@@ -398,9 +249,7 @@ mod tests {
use super::super::super::shared::pack_codes;
use super::super::shared::{PARITY_DIMS, random_inputs};
use super::super::{
Query2bitSimd, score_2bit_internal_scalar, score_2bit_internal_weighted_scalar,
};
use super::super::{score_2bit_internal_scalar, score_2bit_internal_weighted_scalar};
use super::{
score_2bit_internal_neon, score_2bit_internal_neon_sdot, score_2bit_internal_weighted_neon,
};
@@ -414,54 +263,6 @@ mod tests {
.collect()
}
#[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);
lib/quantization/src/turboquant/simd/query2bit/mod.rs
+60 -289
View File
@@ -1,22 +1,38 @@
//! 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:
//! Asymmetric scoring (query against packed codes) goes through the shared
//! [`QuerySimd`] kernels with the 2-bit [`Encoding`] defined here. The
//! symmetric paths (`score_2bit_internal*`, two packed vectors) have their
//! own kernels in the `arm` / `x64` submodules, built on a **pair-table**
//! unpack: a nibble of a packed byte holds a pair of 2-bit codes (16
//! combinations), which maps one-to-one to a 16-entry `vqtbl1q_s8` /
//! `pshufb` table. Two such lookups (even / odd code of each pair),
//! zipped, give 16 centroid bytes in natural dim order per 4 packed bytes.
//!
//! 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.
//! Both codebooks derive from `CENTROIDS_2BIT` (Lloyd-Max on N(0,1)); see
//! `test_codebook_matches_lloyd_max` for the consistency check.
use super::SimdBackend;
use super::query::{Code, Encoding, QuerySimd, pad_codebook};
/// `max|c|` over `CENTROIDS_2BIT` — the extreme centroid magnitude.
const CODEBOOK_ABS_MAX: f32 = 1.510;
/// Integer encoding of the 2-bit width for the shared asymmetric kernels
/// ([`super::query::QuerySimd`]).
pub(super) const ENCODING: Encoding = Encoding {
codebook: pad_codebook(CODEBOOK),
offset: CODEBOOK_OFFSET,
scale: CODEBOOK_SCALE,
query_high_coef: QUERY_HIGH_COEF,
query_abs_max: QUERY_ABS_MAX,
};
#[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
const CODEBOOK: [Code; 4] = CODEBOOK_I8;
#[cfg(not(all(target_arch = "aarch64", target_feature = "neon")))]
const CODEBOOK: [Code; 4] = CODEBOOK_U8;
/// 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];
@@ -100,207 +116,10 @@ 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(crate) 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(crate) 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
}
}
/// Encoded query for asymmetric 2-bit scoring: [`QuerySimd`] over four
/// codes per byte (code `k` of dim `4j + k` in bits `[2k, 2k + 2)` of byte
/// `j`).
pub type Query2bitSimd = QuerySimd<4, 2>;
/// 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.
@@ -316,30 +135,19 @@ pub fn score_2bit_internal(a: &[u8], b: &[u8]) -> f32 {
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) };
}
match SimdBackend::detect() {
#[cfg(target_arch = "x86_64")]
SimdBackend::Avx512Vnni => unsafe { x64::score_2bit_internal_avx512_vnni(a, b) },
#[cfg(target_arch = "x86_64")]
SimdBackend::Avx2 => unsafe { x64::score_2bit_internal_avx2(a, b) },
#[cfg(target_arch = "x86_64")]
SimdBackend::Sse => unsafe { x64::score_2bit_internal_sse(a, b) },
#[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
SimdBackend::NeonSdot => unsafe { arm::score_2bit_internal_neon_sdot(a, b) },
#[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
SimdBackend::Neon => unsafe { arm::score_2bit_internal_neon(a, b) },
SimdBackend::Scalar => score_2bit_internal_scalar(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
@@ -395,24 +203,19 @@ pub fn score_2bit_internal_weighted(a: &[u8], b: &[u8], weights: &[i16]) -> i64
4 * a.len(),
);
#[cfg(target_arch = "x86_64")]
{
if std::is_x86_feature_detected!("avx2")
&& std::is_x86_feature_detected!("ssse3")
&& std::is_x86_feature_detected!("sse4.1")
{
return unsafe { x64::score_2bit_internal_weighted_avx2(a, b, weights) };
}
if std::is_x86_feature_detected!("ssse3") && std::is_x86_feature_detected!("sse4.1") {
return unsafe { x64::score_2bit_internal_weighted_sse(a, b, weights) };
}
match SimdBackend::detect() {
#[cfg(target_arch = "x86_64")]
SimdBackend::Avx512Vnni | SimdBackend::Avx2 => unsafe {
x64::score_2bit_internal_weighted_avx2(a, b, weights)
},
#[cfg(target_arch = "x86_64")]
SimdBackend::Sse => unsafe { x64::score_2bit_internal_weighted_sse(a, b, weights) },
#[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
SimdBackend::NeonSdot | SimdBackend::Neon => unsafe {
arm::score_2bit_internal_weighted_neon(a, b, weights)
},
SimdBackend::Scalar => score_2bit_internal_weighted_scalar(a, b, weights),
}
#[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
{
return unsafe { arm::score_2bit_internal_weighted_neon(a, b, weights) };
}
#[allow(unreachable_code)]
score_2bit_internal_weighted_scalar(a, b, weights)
}
/// Scalar reference for [`score_2bit_internal_weighted`].
@@ -482,15 +285,16 @@ pub mod shared {
}
}
/// Accuracy / precision tests for `Query2bitSimd` and `score_2bit_internal`.
/// Per-arch SIMD parity tests live in the `arm` / `x64` submodules.
/// Codebook and symmetric-scoring accuracy tests. Per-arch parity tests of
/// the symmetric kernels live in the `arm` / `x64` submodules; the
/// asymmetric path is covered by `super::query`.
#[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 super::{CODEBOOK_ABS_MAX, score_2bit_internal_scalar};
use crate::turboquant::TQBits;
#[test]
@@ -528,39 +332,6 @@ mod tests {
}
}
#[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() {
lib/quantization/src/turboquant/simd/query2bit/x64.rs
+6 -291
View File
@@ -1,19 +1,16 @@
//! x86_64 SIMD paths for [`Query2bitSimd`] and [`super::score_2bit_internal`].
//! x86_64 kernels for the symmetric 2-bit paths (`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`]:
//! The codebook is the unsigned `CODEBOOK_U8`; the 2-bit → centroid unpack
//! 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.
//! in natural dim order.
use super::{CODEBOOK_SCALE, CODEBOOK_U8, QUERY_HIGH_COEF, Query2bitSimd};
use super::{CODEBOOK_SCALE, CODEBOOK_U8};
/// `PAIR_TABLE_EVEN_U8[nibble]` = `CODEBOOK_U8[nibble & 0b11]`.
const PAIR_TABLE_EVEN_U8: [u8; 16] = {
@@ -73,215 +70,6 @@ unsafe fn hsum_i32_sse(v: core::arch::x86_64::__m128i) -> i32 {
_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
@@ -661,9 +449,7 @@ mod tests {
use super::super::super::shared::pack_codes;
use super::super::shared::{PARITY_DIMS, random_inputs};
use super::super::{
Query2bitSimd, score_2bit_internal_scalar, score_2bit_internal_weighted_scalar,
};
use super::super::{score_2bit_internal_scalar, score_2bit_internal_weighted_scalar};
use super::*;
/// Build deterministic non-negative i16 weights of length `4 · vec_bytes`
@@ -675,77 +461,6 @@ mod tests {
.collect()
}
#[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") {
lib/quantization/src/turboquant/simd/query4bit/arm.rs
+5 -286
View File
@@ -1,233 +1,7 @@
//! 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.
//! NEON kernels for the symmetric 4-bit paths (`score_4bit_internal*`) on
//! aarch64. The codebook is stored as signed i8 (`CODEBOOK_I8`).
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))
}
}
}
use super::{CODEBOOK_I8, CODEBOOK_SCALE};
/// NEON implementation of [`super::score_4bit_internal`]. Both PQ-encoded
/// vectors are unpacked into i8 codebook values via `vqtbl1q_s8`, multiplied
@@ -283,8 +57,7 @@ pub unsafe fn score_4bit_internal_neon(a: &[u8], b: &[u8]) -> f32 {
}
/// 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`].
/// accumulators consume two chunks (32 elements) per iteration.
///
/// # Safety
/// CPU must support `neon` and `dotprod`.
@@ -447,9 +220,7 @@ mod tests {
use super::super::super::shared::pack_codes;
use super::super::shared::{PARITY_DIMS, random_inputs};
use super::super::{
Query4bitSimd, score_4bit_internal_scalar, score_4bit_internal_weighted_scalar,
};
use super::super::{score_4bit_internal_scalar, score_4bit_internal_weighted_scalar};
use super::{
score_4bit_internal_neon, score_4bit_internal_neon_sdot, score_4bit_internal_weighted_neon,
};
@@ -463,58 +234,6 @@ mod tests {
.collect()
}
#[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
lib/quantization/src/turboquant/simd/query4bit/mod.rs
+53 -389
View File
@@ -1,34 +1,34 @@
//! 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`.
//! Asymmetric scoring (query against packed codes) goes through the shared
//! [`QuerySimd`] kernels with the 4-bit [`Encoding`] defined here; the
//! symmetric paths (`score_4bit_internal*`, two packed vectors) have their
//! own kernels in the `arm` / `x64` submodules.
//!
//! Both codebooks derive from `CENTROIDS_4BIT` (Lloyd-Max on N(0,1)); see
//! `test_codebook_matches_lloyd_max` for the consistency check.
use super::SimdBackend;
use super::query::{Code, Encoding, QuerySimd};
/// `max|c|` over `CENTROIDS_4BIT` — the extreme centroid. Shared by both archs.
const CODEBOOK_ABS_MAX: f32 = 2.733;
/// Integer encoding of the 4-bit width for the shared asymmetric kernels
/// ([`super::query::QuerySimd`]).
pub(super) const ENCODING: Encoding = Encoding {
codebook: CODEBOOK,
offset: CODEBOOK_OFFSET,
scale: CODEBOOK_SCALE,
query_high_coef: QUERY_HIGH_COEF,
query_abs_max: QUERY_ABS_MAX,
};
#[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
const CODEBOOK: [Code; 16] = CODEBOOK_I8;
#[cfg(not(all(target_arch = "aarch64", target_feature = "neon")))]
const CODEBOOK: [Code; 16] = CODEBOOK_U8;
/// 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"))]
@@ -136,234 +136,9 @@ 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(crate) 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(crate) 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(crate) 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)
}
}
/// Encoded query for asymmetric 4-bit scoring: [`QuerySimd`] over two codes
/// per byte (low nibble = even dim, high nibble = odd dim).
pub type Query4bitSimd = QuerySimd<2, 2>;
/// Dot product between two already-encoded 4-bit PQ vectors. Both `a` and
/// `b` are the packed-nibble format that [`Query4bitSimd::dotprod`] takes as
@@ -387,30 +162,19 @@ pub fn score_4bit_internal(a: &[u8], b: &[u8]) -> f32 {
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) };
}
match SimdBackend::detect() {
#[cfg(target_arch = "x86_64")]
SimdBackend::Avx512Vnni => unsafe { x64::score_4bit_internal_avx512_vnni(a, b) },
#[cfg(target_arch = "x86_64")]
SimdBackend::Avx2 => unsafe { x64::score_4bit_internal_avx2(a, b) },
#[cfg(target_arch = "x86_64")]
SimdBackend::Sse => unsafe { x64::score_4bit_internal_sse(a, b) },
#[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
SimdBackend::NeonSdot => unsafe { arm::score_4bit_internal_neon_sdot(a, b) },
#[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
SimdBackend::Neon => unsafe { arm::score_4bit_internal_neon(a, b) },
SimdBackend::Scalar => score_4bit_internal_scalar(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
@@ -471,21 +235,19 @@ pub fn score_4bit_internal_weighted(a: &[u8], b: &[u8], weights: &[i16]) -> i64
2 * a.len(),
);
#[cfg(target_arch = "x86_64")]
{
if std::is_x86_feature_detected!("avx2") {
return unsafe { x64::score_4bit_internal_weighted_avx2(a, b, weights) };
}
if std::is_x86_feature_detected!("sse4.1") && std::is_x86_feature_detected!("ssse3") {
return unsafe { x64::score_4bit_internal_weighted_sse(a, b, weights) };
}
match SimdBackend::detect() {
#[cfg(target_arch = "x86_64")]
SimdBackend::Avx512Vnni | SimdBackend::Avx2 => unsafe {
x64::score_4bit_internal_weighted_avx2(a, b, weights)
},
#[cfg(target_arch = "x86_64")]
SimdBackend::Sse => unsafe { x64::score_4bit_internal_weighted_sse(a, b, weights) },
#[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
SimdBackend::NeonSdot | SimdBackend::Neon => unsafe {
arm::score_4bit_internal_weighted_neon(a, b, weights)
},
SimdBackend::Scalar => score_4bit_internal_weighted_scalar(a, b, weights),
}
#[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
{
return unsafe { arm::score_4bit_internal_weighted_neon(a, b, weights) };
}
#[allow(unreachable_code)]
score_4bit_internal_weighted_scalar(a, b, weights)
}
/// Scalar reference for [`score_4bit_internal_weighted`].
@@ -560,11 +322,9 @@ pub mod shared {
}
}
/// 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.
/// Codebook and symmetric-scoring accuracy tests. Per-arch parity tests of
/// the symmetric kernels live in the `arm` / `x64` submodules; the
/// asymmetric path is covered by `super::query`.
#[cfg(test)]
mod tests {
// Anonymous `use _` brings the trait into scope for `StdRng::seed_from_u64`
@@ -573,7 +333,7 @@ mod tests {
use rand::prelude::StdRng;
use super::super::shared::{encode_to_nearest_centroid, pack_codes, sample_normal_vec};
use super::{CODEBOOK_ABS_MAX, Query4bitSimd};
use super::CODEBOOK_ABS_MAX;
use crate::turboquant::TQBits;
/// Whichever codebook representation the current arch uses (signed i8 on
@@ -622,102 +382,6 @@ mod tests {
}
}
/// 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
lib/quantization/src/turboquant/simd/query4bit/x64.rs
+5 -323
View File
@@ -1,242 +1,9 @@
//! x86_64 SIMD paths for [`Query4bitSimd`].
//! x86_64 kernels for the symmetric 4-bit paths (`score_4bit_internal*`).
//!
//! 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.
//! The codebook is stored as unsigned u8 ∈ [0, 255] (`CODEBOOK_U8`); the
//! kernels XOR it with 0x80 to recover the signed form in-register.
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))
}
}
}
use super::{CODEBOOK_SCALE, CODEBOOK_U8};
#[target_feature(enable = "sse2")]
unsafe fn hsum_i32_sse(v: core::arch::x86_64::__m128i) -> i32 {
@@ -665,9 +432,7 @@ mod tests {
use super::super::super::shared::pack_codes;
use super::super::shared::{PARITY_DIMS, random_inputs};
use super::super::{
Query4bitSimd, score_4bit_internal_scalar, score_4bit_internal_weighted_scalar,
};
use super::super::{score_4bit_internal_scalar, score_4bit_internal_weighted_scalar};
use super::{
score_4bit_internal_avx2, score_4bit_internal_avx512_vnni, score_4bit_internal_sse,
score_4bit_internal_weighted_avx2, score_4bit_internal_weighted_sse,
@@ -682,89 +447,6 @@ mod tests {
.collect()
}
#[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