src/vector/tiled.rs
//! **The tiled exact self-join** — every row pair at or above a cosine bar,
//! computed in register blocks (4 rows × 3 rows, twelve AVX2 accumulators)
//! over cache tiles, with the SAME arithmetic per pair as the reference dot
//! it stands in for. Not one hit is gained or lost: a pair's cosine is the
//! very float the reference computes, so the `>=` decides exactly as before.
//!
//! Two references, two replicas:
//!
//! | reference | used by | replica |
//! |---|---|---|
//! | znippy-zoomies `dot_avx2` (`VectorIndex::pairs_above` on an AVX2+FMA CPU without AVX-512) | the pair index's cold build | one `fmadd` chain of 8 lanes per pair, lanes summed in order, a scalar tail |
//! | [`super::pair_index::dot`] (8 independent lanes, multiply then add) | the pair index's incremental join | one `mul` + `add` chain of 8 lanes per pair, lanes summed in order, a scalar tail |
//!
//! The order of the k-steps of ONE pair is the reference's (ascending, eight
//! at a time); tiling only changes which pairs are computed together, never
//! how one is summed. A CPU without AVX2+FMA, or with AVX-512 (whose
//! reference reduces in a tree), is not served here: the caller keeps its
//! reference path.
/// Rows per register block (the `a` side) and per tile.
const RA: usize = 4;
const RB: usize = 3;
/// Rows per cache tile: 64 rows × 768 dims × 4 B = 192 KiB, two tiles in L2.
const TILE: usize = 64;
/// Which reference the replica reproduces.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Arith {
/// znippy-zoomies `dot_avx2`: fused multiply-add.
Fma,
/// [`super::pair_index::dot`]: multiply, then add.
MulAdd,
}
/// Can the replica of `arith` run here, bit-identical to its reference?
pub fn available(arith: Arith) -> bool {
#[cfg(target_arch = "x86_64")]
{
let avx2 = std::is_x86_feature_detected!("avx2");
match arith {
// the cold reference is dot_avx512 where AVX-512 exists: not this replica
Arith::Fma => avx2 && std::is_x86_feature_detected!("fma") && !std::is_x86_feature_detected!("avx512f"),
Arith::MulAdd => avx2,
}
}
#[cfg(not(target_arch = "x86_64"))]
{
let _ = arith;
false
}
}
/// The reference dot of one pair, scalar (tails, edge rows, tests).
pub fn reference(arith: Arith, a: &[f32], b: &[f32]) -> f32 {
let mut acc = [0f32; 8];
let chunks = a.len() / 8;
for c in 0..chunks {
for l in 0..8 {
let (x, y) = (a[c * 8 + l], b[c * 8 + l]);
acc[l] = match arith {
Arith::Fma => x.mul_add(y, acc[l]),
Arith::MulAdd => acc[l] + x * y,
};
}
}
let mut s: f32 = acc.iter().sum();
for i in chunks * 8..a.len() {
s += a[i] * b[i];
}
s
}
/// Every pair `(i, j)`, `i < j`, of the `rows × dim` matrix `data` with
/// reference cosine ≥ `floor`, in ascending `(i, j)` order; `rows_from`: only
/// pairs whose `j` is at or after it (the incremental join: `j ≥ from`, any
/// `i < j`). `None` when [`available`] says no.
pub fn pairs_above(data: &[f32], dim: usize, floor: f32, from: usize, arith: Arith) -> Option<Vec<(u32, u32, f32)>> {
if !available(arith) || dim == 0 {
return None;
}
let rows = data.len() / dim;
if from >= rows {
return Some(Vec::new());
}
// tiles of j (the later row); each pairs with every i-tile up to it
let jtiles: Vec<usize> = (from / TILE * TILE..rows).step_by(TILE).collect();
let threads = std::thread::available_parallelism().map_or(1, |n| n.get()).min(jtiles.len()).max(1);
let mut out: Vec<(u32, u32, f32)> = std::thread::scope(|s| {
let handles: Vec<_> = (0..threads)
.map(|t| {
let jtiles = &jtiles;
s.spawn(move || {
let mut out = Vec::new();
// interleaved from the far end: a late j-tile pairs with many i-tiles
let mut k = t;
while k < jtiles.len() {
let j0 = jtiles[jtiles.len() - 1 - k];
let j1 = (j0 + TILE).min(rows);
let mut i0 = 0;
while i0 < j1 {
let i1 = (i0 + TILE).min(j1);
tile(data, dim, floor, from, (i0, i1), (j0, j1), arith, &mut out);
i0 = i1;
}
k += threads;
}
out
})
})
.collect();
handles.into_iter().flat_map(|h| h.join().unwrap_or_else(|p| std::panic::resume_unwind(p))).collect()
});
out.sort_unstable_by(|x, y| x.0.cmp(&y.0).then(x.1.cmp(&y.1)));
Some(out)
}
/// One tile: rows `ia..ib` against rows `ja..jb`, emitting `i < j`, `j ≥ from`.
#[allow(clippy::too_many_arguments)]
fn tile(data: &[f32], dim: usize, floor: f32, from: usize, (ia, ib): (usize, usize), (ja, jb): (usize, usize), arith: Arith, out: &mut Vec<(u32, u32, f32)>) {
let row = |r: usize| &data[r * dim..(r + 1) * dim];
let mut emit = |i: usize, j: usize, c: f32| {
if i < j && j >= from && c >= floor {
out.push((i as u32, j as u32, c));
}
};
let mut i = ia;
while i < ib {
let ni = (ib - i).min(RA);
let mut j = ja.max(from);
while j < jb {
let nj = (jb - j).min(RB);
// the block is wholly below the diagonal: nothing to emit
if j + nj <= i + 1 {
j += nj;
continue;
}
if ni == RA && nj == RB {
let mut c = [[0f32; RB]; RA];
// SAFETY: `available(arith)` confirmed AVX2 (+FMA for Fma); every row is `dim` long.
unsafe { block(data, dim, i, j, arith, &mut c) };
for (x, cr) in c.iter().enumerate() {
for (y, cv) in cr.iter().enumerate() {
emit(i + x, j + y, *cv);
}
}
} else {
for x in 0..ni {
for y in 0..nj {
if i + x < j + y {
emit(i + x, j + y, reference(arith, row(i + x), row(j + y)));
}
}
}
}
j += nj;
}
i += ni;
}
}
/// A 4 × 3 register block: twelve accumulators, one per pair, each summed in
/// the reference's order.
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2,fma")]
unsafe fn block(data: &[f32], dim: usize, i: usize, j: usize, arith: Arith, c: &mut [[f32; RB]; RA]) {
use std::arch::x86_64::*;
unsafe {
let p = data.as_ptr();
let a: [*const f32; RA] = std::array::from_fn(|x| p.add((i + x) * dim));
let b: [*const f32; RB] = std::array::from_fn(|y| p.add((j + y) * dim));
let mut acc = [[_mm256_setzero_ps(); RB]; RA];
let mut k = 0;
while k + 8 <= dim {
let vb: [__m256; RB] = std::array::from_fn(|y| _mm256_loadu_ps(b[y].add(k)));
for x in 0..RA {
let va = _mm256_loadu_ps(a[x].add(k));
for y in 0..RB {
acc[x][y] = match arith {
Arith::Fma => _mm256_fmadd_ps(va, vb[y], acc[x][y]),
Arith::MulAdd => _mm256_add_ps(acc[x][y], _mm256_mul_ps(va, vb[y])),
};
}
}
k += 8;
}
for x in 0..RA {
for y in 0..RB {
let mut tmp = [0f32; 8];
_mm256_storeu_ps(tmp.as_mut_ptr(), acc[x][y]);
let mut s = tmp.iter().sum::<f32>();
let mut t = k;
while t < dim {
s += *a[x].add(t) * *b[y].add(t);
t += 1;
}
c[x][y] = s;
}
}
}
}
/// One pair, the reference's arithmetic, on AVX2 (the int8 prefilter's recheck).
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2,fma")]
unsafe fn pair(a: &[f32], b: &[f32], arith: Arith) -> f32 {
use std::arch::x86_64::*;
unsafe {
let dim = a.len();
let mut acc = _mm256_setzero_ps();
let mut k = 0;
while k + 8 <= dim {
let (va, vb) = (_mm256_loadu_ps(a.as_ptr().add(k)), _mm256_loadu_ps(b.as_ptr().add(k)));
acc = match arith {
Arith::Fma => _mm256_fmadd_ps(va, vb, acc),
Arith::MulAdd => _mm256_add_ps(acc, _mm256_mul_ps(va, vb)),
};
k += 8;
}
let mut tmp = [0f32; 8];
_mm256_storeu_ps(tmp.as_mut_ptr(), acc);
let mut s = tmp.iter().sum::<f32>();
while k < dim {
s += a[k] * b[k];
k += 1;
}
s
}
}
#[cfg(not(target_arch = "x86_64"))]
unsafe fn pair(a: &[f32], b: &[f32], arith: Arith) -> f32 {
reference(arith, a, b)
}
#[cfg(not(target_arch = "x86_64"))]
unsafe fn block(_: &[f32], _: usize, _: usize, _: usize, _: Arith, _: &mut [[f32; RB]; RA]) {
unreachable!("not available off x86_64")
}
// ── the int8 PREFILTER (an exact f32 recheck of every survivor) ─────────────
//
// Row `x` is quantized twice: the "a" side to 8 bits (`u = round(x/s8) + 128`,
// `s8 = max|x|/127`, unsigned for `vpmaddubsw`), the "b" side to 7 bits
// (`q = round(x/s7)`, `s7 = max|x|/63`: `2·255·63 < 2¹⁵`, so the pairwise i16
// sums never saturate). `Σ u·q − 128·Σ q` is the exact integer dot of the two
// quantized rows; times `s8·s7` it is `x̂·ŷ`. With `e = x − x̂`,
//
// |x·y − x̂·ŷ| ≤ ‖x̂‖‖e_y‖ + ‖e_x‖‖ŷ‖ + ‖e_x‖‖e_y‖ (Cauchy–Schwarz)
//
// — a bound computed per pair from four per-row norms. A pair survives when
// `x̂·ŷ + bound + SLACK ≥ floor`; every survivor is decided by the reference
// f32 dot. If the reference says `≥ floor`, the real dot is within `1e-4` of
// it (768 products of unit vectors in f32), so the prefilter kept it: not one
// hit is lost, and no float differs.
/// Rounding headroom of the bound (the reference's own f32 error is < 1e-4).
const SLACK: f64 = 1e-3;
/// The quantized matrix.
pub struct Quantized {
/// Padded dims (a multiple of 32).
pub pdim: usize,
/// `rows × pdim`: the 8-bit "a" side, biased by 128 (pad = 128, i.e. 0).
pub a8: Vec<u8>,
/// `rows × pdim`: the 7-bit "b" side (pad = 0).
pub b7: Vec<i8>,
/// Per row: `s8·s7`-free factors and norms: (s8, s7, Σq7, ‖x̂8‖, ‖e8‖, ‖x̂7‖, ‖e7‖).
pub meta: Vec<(f64, f64, i64, f64, f64, f64, f64)>,
}
pub fn quantize(data: &[f32], dim: usize) -> Quantized {
let rows = data.len() / dim;
let pdim = dim.div_ceil(32) * 32;
let mut a8 = vec![128u8; rows * pdim];
let mut b7 = vec![0i8; rows * pdim];
let mut meta = Vec::with_capacity(rows);
for r in 0..rows {
let x = &data[r * dim..(r + 1) * dim];
let max = x.iter().fold(0f64, |m, v| m.max((*v as f64).abs()));
let (s8, s7) = if max > 0.0 { (max / 127.0, max / 63.0) } else { (1.0, 1.0) };
let (mut h8, mut e8, mut h7, mut e7, mut sq) = (0f64, 0f64, 0f64, 0f64, 0i64);
for (k, v) in x.iter().enumerate() {
let v = *v as f64;
let q8 = (v / s8).round().clamp(-127.0, 127.0);
let q7 = (v / s7).round().clamp(-63.0, 63.0);
a8[r * pdim + k] = (q8 as i32 + 128) as u8;
b7[r * pdim + k] = q7 as i8;
sq += q7 as i64;
h8 += (q8 * s8).powi(2);
e8 += (v - q8 * s8).powi(2);
h7 += (q7 * s7).powi(2);
e7 += (v - q7 * s7).powi(2);
}
meta.push((s8, s7, sq, h8.sqrt(), e8.sqrt(), h7.sqrt(), e7.sqrt()));
}
Quantized { pdim, a8, b7, meta }
}
impl Quantized {
/// The approximate dot of rows `i` (8-bit side) and `j` (7-bit side) from
/// the exact integer dot, and its error bound.
fn approx(&self, i: usize, j: usize, int_dot: i64) -> (f64, f64) {
let (s8, _, _, h8, e8, _, _) = self.meta[i];
let (_, s7, sq, _, _, h7, e7) = self.meta[j];
((int_dot - 128 * sq) as f64 * s8 * s7, h8 * e7 + e8 * h7 + e8 * e7)
}
}
/// [`pairs_above`] through the int8 prefilter: the same pairs, the same floats.
pub fn pairs_above_prefiltered(data: &[f32], dim: usize, floor: f32, from: usize, arith: Arith) -> Option<Vec<(u32, u32, f32)>> {
if !available(arith) || dim == 0 {
return None;
}
let rows = data.len() / dim;
if from >= rows {
return Some(Vec::new());
}
let q = quantize(data, dim);
let jtiles: Vec<usize> = (from / TILE * TILE..rows).step_by(TILE).collect();
let threads = std::thread::available_parallelism().map_or(1, |n| n.get()).min(jtiles.len()).max(1);
let q = &q;
let mut out: Vec<(u32, u32, f32)> = std::thread::scope(|s| {
let handles: Vec<_> = (0..threads)
.map(|t| {
let jtiles = &jtiles;
s.spawn(move || {
let mut out = Vec::new();
let row = |r: usize| &data[r * dim..(r + 1) * dim];
let keep = |i: usize, j: usize, d: i64, out: &mut Vec<(u32, u32, f32)>| {
if i >= j || j < from {
return;
}
let (ap, bound) = q.approx(i, j, d);
if ap + bound + SLACK >= floor as f64 {
// SAFETY: available(arith) confirmed AVX2+FMA.
let c = unsafe { pair(row(i), row(j), arith) };
if c >= floor {
out.push((i as u32, j as u32, c));
}
}
};
let mut k = t;
while k < jtiles.len() {
let j0 = jtiles[jtiles.len() - 1 - k];
let j1 = (j0 + TILE).min(rows);
let mut i0 = 0;
while i0 < j1 {
let i1 = (i0 + TILE).min(j1);
let mut i = i0;
while i < i1 {
let ni = (i1 - i).min(RA);
let mut j = j0.max(from);
while j < j1 {
let nj = (j1 - j).min(RB);
if j + nj > i + 1 {
if ni == RA && nj == RB {
let mut d = [[0i64; RB]; RA];
// SAFETY: AVX2 confirmed; rows are `pdim` long.
unsafe { block_i8(q, i, j, &mut d) };
for x in 0..RA {
for y in 0..RB {
keep(i + x, j + y, d[x][y], &mut out);
}
}
} else {
for x in 0..ni {
for y in 0..nj {
let (a, b) = (&q.a8[(i + x) * q.pdim..(i + x + 1) * q.pdim], &q.b7[(j + y) * q.pdim..(j + y + 1) * q.pdim]);
let d: i64 = a.iter().zip(b).map(|(u, v)| *u as i64 * *v as i64).sum();
keep(i + x, j + y, d, &mut out);
}
}
}
}
j += nj;
}
i += ni;
}
i0 = i1;
}
k += threads;
}
out
})
})
.collect();
handles.into_iter().flat_map(|h| h.join().unwrap_or_else(|p| std::panic::resume_unwind(p))).collect()
});
out.sort_unstable_by(|x, y| x.0.cmp(&y.0).then(x.1.cmp(&y.1)));
Some(out)
}
/// A 4 × 3 block of exact integer dots `Σ u8·i8` (`vpmaddubsw` → `vpmaddwd`).
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
unsafe fn block_i8(q: &Quantized, i: usize, j: usize, d: &mut [[i64; RB]; RA]) {
use std::arch::x86_64::*;
unsafe {
let pd = q.pdim;
let a: [*const u8; RA] = std::array::from_fn(|x| q.a8.as_ptr().add((i + x) * pd));
let b: [*const i8; RB] = std::array::from_fn(|y| q.b7.as_ptr().add((j + y) * pd));
let ones = _mm256_set1_epi16(1);
let mut acc = [[_mm256_setzero_si256(); RB]; RA];
let mut k = 0;
while k < pd {
let vb: [__m256i; RB] = std::array::from_fn(|y| _mm256_loadu_si256(b[y].add(k) as *const __m256i));
for x in 0..RA {
let va = _mm256_loadu_si256(a[x].add(k) as *const __m256i);
for y in 0..RB {
acc[x][y] = _mm256_add_epi32(acc[x][y], _mm256_madd_epi16(_mm256_maddubs_epi16(va, vb[y]), ones));
}
}
k += 32;
}
for x in 0..RA {
for y in 0..RB {
let mut tmp = [0i32; 8];
_mm256_storeu_si256(tmp.as_mut_ptr() as *mut __m256i, acc[x][y]);
d[x][y] = tmp.iter().map(|v| *v as i64).sum();
}
}
}
}
#[cfg(not(target_arch = "x86_64"))]
unsafe fn block_i8(_: &Quantized, _: usize, _: usize, _: &mut [[i64; RB]; RA]) {
unreachable!("not available off x86_64")
}
#[cfg(test)]
mod tests {
use super::*;
fn corpus(n: usize, dim: usize, seed: u64) -> Vec<f32> {
let mut s = seed;
let mut rnd = move || {
s ^= s << 13;
s ^= s >> 7;
s ^= s << 17;
(s >> 11) as f32 / (1u64 << 53) as f32 - 0.5
};
let centers: Vec<Vec<f32>> = (0..23).map(|_| (0..dim).map(|_| rnd()).collect()).collect();
let mut out = Vec::with_capacity(n * dim);
for i in 0..n {
let c = ¢ers[i % 23];
let v: Vec<f32> = c.iter().map(|x| x + 0.15 * rnd()).collect();
out.extend(crate::vector::normalized(&v));
}
out
}
/// The replica IS the reference: the scalar `reference` equals
/// znippy-zoomies' kernel (Fma, through `VectorIndex::pairs_above`) and
/// `pair_index::dot` (MulAdd) bit for bit, and the tiled join returns
/// exactly the pairs (and floats) a per-pair reference scan returns — at
/// odd sizes (tails of rows and of dims) and from an incremental `from`.
/// The bound holds: for every pair of a corpus, the exact integer dot of
/// the quantized rows plus the bound brackets the reference f32 dot, and
/// pairs a hair above or below a bar are all decided by the reference.
#[test]
fn the_int8_bound_brackets_every_reference_dot() {
let (n, dim) = (200usize, 768usize);
let data = corpus(n, dim, 99);
let q = quantize(&data, dim);
let row = |r: usize| &data[r * dim..(r + 1) * dim];
let mut worst = 0f64;
for i in 0..n {
for j in 0..n {
let d: i64 = q.a8[i * q.pdim..(i + 1) * q.pdim].iter().zip(&q.b7[j * q.pdim..(j + 1) * q.pdim]).map(|(u, v)| *u as i64 * *v as i64).sum();
let (ap, bound) = q.approx(i, j, d);
let r = reference(Arith::Fma, row(i), row(j)) as f64;
assert!((r - ap).abs() <= bound + 1e-4, "{i} {j}: ref {r} approx {ap} bound {bound}");
worst = worst.max(bound);
}
}
assert!(worst < 0.1, "the bound is tight enough to prune: {worst}");
if available(Arith::Fma) {
// bars set exactly ON some pairs' reference floats: those pairs are kept
let mut bars: Vec<f32> = (1..6).map(|k| reference(Arith::Fma, row(k), row(k + 23))).collect();
bars.push(0.9);
for bar in bars {
let want = pairs_above(&data, dim, bar, 0, Arith::Fma).unwrap();
assert_eq!(pairs_above_prefiltered(&data, dim, bar, 0, Arith::Fma).unwrap(), want, "bar {bar}");
}
}
}
#[test]
fn the_tiled_join_is_the_reference_scan_bit_for_bit() {
for (n, dim) in [(301usize, 768usize), (130, 128), (77, 13)] {
let data = corpus(n, dim, n as u64);
let row = |r: usize| &data[r * dim..(r + 1) * dim];
for r in 0..n.min(40) {
assert_eq!(reference(Arith::MulAdd, row(r), row(n - 1 - r)).to_bits(), crate::vector::pair_index::dot(row(r), row(n - 1 - r)).to_bits());
}
if available(Arith::Fma) {
let mut idx = crate::vector::VectorIndex::new(dim).unwrap();
// the index renormalizes on add: feed it what it would store, compare on that
idx.add(&data, &(0..n as u64).collect::<Vec<_>>()).unwrap();
let stored: Vec<f32> = (0..n).flat_map(|r| crate::vector::normalized(row(r))).collect();
let want: Vec<(u32, u32, u32)> = idx.pairs_above(0.5).into_iter().map(|(a, b, c)| (a as u32, b as u32, c.to_bits())).collect();
let got: Vec<(u32, u32, u32)> = pairs_above(&stored, dim, 0.5, 0, Arith::Fma).unwrap().into_iter().map(|(a, b, c)| (a, b, c.to_bits())).collect();
assert!(!want.is_empty());
assert_eq!(got, want, "n {n} dim {dim}: the cold build's pairs and floats");
}
for arith in [Arith::Fma, Arith::MulAdd] {
if !available(arith) {
continue;
}
for from in [0usize, n / 2, n - 1] {
let mut want = Vec::new();
for j in from..n {
for i in 0..j {
let c = reference(arith, row(i), row(j));
if c >= 0.6 {
want.push((i as u32, j as u32, c.to_bits()));
}
}
}
want.sort_unstable();
let got: Vec<(u32, u32, u32)> = pairs_above(&data, dim, 0.6, from, arith).unwrap().into_iter().map(|(a, b, c)| (a, b, c.to_bits())).collect();
assert_eq!(got, want, "{arith:?} n {n} dim {dim} from {from}");
let pre: Vec<(u32, u32, u32)> = pairs_above_prefiltered(&data, dim, 0.6, from, arith).unwrap().into_iter().map(|(a, b, c)| (a, b, c.to_bits())).collect();
assert_eq!(pre, want, "prefiltered {arith:?} n {n} dim {dim} from {from}");
}
}
}
}
}