src/vector/pair_index.rs
//! **The pair index** — the exact cosine self-join of the function vectors,
//! kept in the process and brought forward by what changed.
//!
//! `similar_code` and `semantic_clones` both need EVERY pair of function
//! vectors above a cosine bar (0.92 and 0.80). The scan that answers it is
//! `n²/2` dot products over every function node — on the live workspace
//! ~10⁵ nodes, so ~5·10⁹ dots, every time a lens is asked, and a lens is asked
//! again after every commit (its answer is kept per room version).
//!
//! Two facts make almost all of that work repeat itself:
//!
//! * **Identical bodies share a vector.** A body is embedded once (content
//! hash), and every node cut from the same text carries the same floats.
//! The join runs over the DISTINCT vectors; pairs of nodes are expanded
//! from pairs of rows, and nodes of one row pair with each other at the
//! row's own cosine (≈ 1).
//! * **A commit changes few bodies.** The rows of this process are kept with
//! every pair of them at or above [`FLOOR`]. A new vector is joined against
//! the rows already here (`k × n` dots, not `n²/2`); a vector that left is
//! simply not asked about. When the rows nobody asks about outweigh the
//! ones asked about, the index is built again.
//!
//! The join stays **exact** — no ANN. A pair the report misses is code
//! nobody looks at, so every pair at or above the bar is in the answer. A bar
//! below [`FLOOR`] is not served here (the caller scans in full).
//!
//! CPU only: the forge that reads it has no GPU. The cold build is the
//! `VectorIndex::pairs_above` scan (runtime-selected SIMD kernel, gatling
//! fork-join); the incremental join is a lane-split dot over scoped threads.
use std::collections::HashMap;
use std::hash::{Hash, Hasher};
use std::sync::Mutex;
/// The lowest cosine the index keeps pairs at: `semantic_clones`' default
/// floor. `similar_code` (0.92) reads a subset.
pub const FLOOR: f32 = 0.80;
/// Below this many vectors the plain scan is cheap and the index is not used.
pub const MIN_ROWS: usize = 2_048;
/// What one call did, for the answer's `knn` block and the log.
#[derive(Debug, Clone, Default, PartialEq)]
pub struct Served {
/// Vectors asked about / distinct among them.
pub vectors: usize,
pub distinct: usize,
/// Rows the index held after the call (distinct vectors ever seen since the last build).
pub rows: usize,
/// Rows this call added (embedded since the last ask).
pub added: usize,
/// `cold` (built from nothing), `incremental` (joined `added` rows), `warm` (nothing new).
pub mode: &'static str,
/// Dot products this call computed.
pub dots: u64,
/// Row pairs at or above [`FLOOR`] held.
pub row_pairs: usize,
pub ms: u64,
}
struct Index {
dim: usize,
/// `rows × dim`, L2-normalized.
data: Vec<f32>,
/// hash of the raw vector's bits → rows with that hash (a collision is checked).
by_key: HashMap<u64, Vec<u32>>,
/// `(a, b, cosine)`, `a < b`, every row pair at or above [`FLOOR`].
pairs: Vec<(u32, u32, f32)>,
}
impl Index {
fn new(dim: usize) -> Self {
Self { dim, data: Vec::new(), by_key: HashMap::new(), pairs: Vec::new() }
}
fn rows(&self) -> usize {
self.data.len() / self.dim.max(1)
}
fn row(&self, r: u32) -> &[f32] {
&self.data[r as usize * self.dim..(r as usize + 1) * self.dim]
}
/// The row of `v`, or `None`; `norm` is `v` normalized (the stored form).
fn find(&self, key: u64, norm: &[f32]) -> Option<u32> {
self.by_key.get(&key)?.iter().copied().find(|r| self.row(*r) == norm)
}
}
static INDEX: Mutex<Option<Index>> = Mutex::new(None);
/// The indexes of the OTHER vector spaces (the clone signals' role vectors),
/// one per space name: a space never evicts another's rows.
static SPACES: Mutex<Option<HashMap<&'static str, Index>>> = Mutex::new(None);
fn key_of(v: &[f32]) -> u64 {
let mut h = std::collections::hash_map::DefaultHasher::new();
for x in v {
x.to_bits().hash(&mut h);
}
h.finish()
}
/// The dot of two equal-length rows: eight independent lanes, so the
/// compiler vectorizes it without reassociating one sum.
pub fn dot(a: &[f32], b: &[f32]) -> f32 {
let mut acc = [0f32; 8];
let (ca, cb) = (a.chunks_exact(8), b.chunks_exact(8));
let (ra, rb) = (ca.remainder(), cb.remainder());
for (x, y) in ca.zip(cb) {
for l in 0..8 {
acc[l] += x[l] * y[l];
}
}
let mut s: f32 = acc.iter().sum();
for (x, y) in ra.iter().zip(rb) {
s += x * y;
}
s
}
/// Pairs `(j, i, cos)`, `j < i`, for every `i` in `from..rows`, `cos ≥ floor`:
/// every pair that involves a row from `from` on, exactly once.
fn join_from(data: &[f32], dim: usize, from: usize, floor: f32) -> Vec<(u32, u32, f32)> {
let rows = data.len() / dim;
if from >= rows {
return Vec::new();
}
// the tiled replica of `dot` (same floats, register blocks over cache tiles)
if let Some(p) = tiled_join(data, dim, floor, from, super::tiled::Arith::MulAdd) {
return p;
}
let threads = std::thread::available_parallelism().map_or(1, |n| n.get()).min(rows - from).max(1);
let row = |r: usize| &data[r * dim..(r + 1) * dim];
std::thread::scope(|s| {
let handles: Vec<_> = (0..threads)
.map(|t| {
s.spawn(move || {
let mut out = Vec::new();
// interleaved: row i costs i dots, so bands of consecutive rows would be unequal
let mut i = from + t;
while i < rows {
let a = row(i);
for j in 0..i {
let c = dot(a, row(j));
if c >= floor {
out.push((j as u32, i as u32, c));
}
}
i += threads;
}
out
})
})
.collect();
handles.into_iter().flat_map(|h| h.join().unwrap_or_else(|p| std::panic::resume_unwind(p))).collect()
})
}
/// The join kernel: `tiled` (register blocks over cache tiles), `prefilter`
/// (the same, behind the int8 prefilter) — both the reference's floats — or
/// `reference` (the per-pair scan, for measuring). `NORNIR_PAIR_KERNEL`
/// picks one for a measurement; the default is the measured winner.
fn tiled_join(data: &[f32], dim: usize, floor: f32, from: usize, arith: super::tiled::Arith) -> Option<Vec<(u32, u32, f32)>> {
match std::env::var("NORNIR_PAIR_KERNEL").as_deref() {
Ok("reference") => None,
Ok("tiled") => super::tiled::pairs_above(data, dim, floor, from, arith),
_ => super::tiled::pairs_above_prefiltered(data, dim, floor, from, arith),
}
}
/// The cold build: every row pair at or above `floor`, through the exact
/// SIMD scan.
fn join_all(data: &[f32], dim: usize, floor: f32) -> Vec<(u32, u32, f32)> {
let rows = data.len() / dim;
// the tiled replica of the index's kernel, over the rows exactly as the
// index would store them (it normalizes on add): the same floats
if super::tiled::available(super::tiled::Arith::Fma) {
let stored: Vec<f32> = data.chunks_exact(dim).flat_map(super::normalized).collect();
if let Some(p) = tiled_join(&stored, dim, floor, 0, super::tiled::Arith::Fma) {
return p;
}
}
let Ok(mut idx) = super::VectorIndex::new(dim) else { return join_from(data, dim, 0, floor) };
let ids: Vec<u64> = (0..rows as u64).collect();
if idx.add(data, &ids).is_err() {
return join_from(data, dim, 0, floor);
}
idx.pairs_above(floor).into_iter().map(|(a, b, c)| (a as u32, b as u32, c)).collect()
}
/// **Every pair `(i, j)`, `i < j`, of `vectors` whose cosine is ≥ `threshold`**,
/// in ascending `(i, j)` order — through the index. `None` when the index does
/// not serve the ask (a bar below [`FLOOR`], fewer than [`MIN_ROWS`] vectors,
/// vectors of unequal or zero length): the caller scans.
pub fn pairs_above<'a>(vectors: impl ExactSizeIterator<Item = &'a [f32]> + Clone, threshold: f32) -> Option<(Vec<(usize, usize, f32)>, Served)> {
let n = vectors.len();
if threshold < FLOOR || n < MIN_ROWS {
return None;
}
let dim = vectors.clone().next().map_or(0, <[f32]>::len);
if dim == 0 || vectors.clone().any(|v| v.len() != dim) {
return None;
}
let mut guard = INDEX.lock().unwrap_or_else(|e| e.into_inner());
if guard.as_ref().is_some_and(|i| i.dim != dim) {
*guard = None;
}
let served = serve(&mut guard, vectors, threshold, dim);
*LAST.lock().unwrap_or_else(|e| e.into_inner()) = Some(served.1.clone());
Some(served)
}
/// [`pairs_above`] in a vector space of its own (`space`: e.g. `"role"`, the
/// clone signals' graph vectors): the same exact join, kept and brought
/// forward the same way, in an index that never evicts the text vectors'.
/// `None` in the same cases as [`pairs_above`].
pub fn pairs_above_in<'a>(space: &'static str, vectors: impl ExactSizeIterator<Item = &'a [f32]> + Clone, threshold: f32) -> Option<(Vec<(usize, usize, f32)>, Served)> {
let n = vectors.len();
if threshold < FLOOR || n < MIN_ROWS {
return None;
}
let dim = vectors.clone().next().map_or(0, <[f32]>::len);
if dim == 0 || vectors.clone().any(|v| v.len() != dim) {
return None;
}
let mut spaces = SPACES.lock().unwrap_or_else(|e| e.into_inner());
let map = spaces.get_or_insert_with(HashMap::new);
let mut slot = map.remove(space).filter(|i| i.dim == dim);
let served = serve(&mut slot, vectors, threshold, dim);
if let Some(idx) = slot {
map.insert(space, idx);
}
Some(served)
}
fn serve<'a>(guard: &mut Option<Index>, vectors: impl ExactSizeIterator<Item = &'a [f32]> + Clone, threshold: f32, dim: usize) -> (Vec<(usize, usize, f32)>, Served) {
let n = vectors.len();
let t0 = std::time::Instant::now();
let mut served = Served { vectors: n, ..Default::default() };
// map every vector to its row; the new ones are appended
let mut place = |idx: &mut Index, added: &mut usize| -> Vec<u32> {
vectors
.clone()
.map(|v| {
let key = key_of(v);
let norm = super::normalized(v);
match idx.find(key, &norm) {
Some(r) => r,
None => {
let r = idx.rows() as u32;
idx.data.extend_from_slice(&norm);
idx.by_key.entry(key).or_default().push(r);
*added += 1;
r
}
}
})
.collect()
};
let idx = guard.get_or_insert_with(|| Index::new(dim));
let before = idx.rows();
let mut added = 0usize;
let mut row_of = place(idx, &mut added);
let mut distinct: Vec<u32> = row_of.clone();
distinct.sort_unstable();
distinct.dedup();
// rows nobody asks about any more outweigh the live ones: build again
if before > 0 && idx.rows() > 2 * distinct.len() {
*idx = Index::new(dim);
added = 0;
row_of = place(idx, &mut added);
distinct = row_of.clone();
distinct.sort_unstable();
distinct.dedup();
}
let rows = idx.rows();
// most rows new (the first backfill after a few early members): the SIMD
// scan over everything beats the lane join over the new ones (MEASURED on
// dwarf: 58,228 of 66,259 rows joined incrementally took 280 s)
if added * 4 > rows && rows > added {
idx.pairs.clear();
added = rows;
}
if idx.pairs.is_empty() && rows == added {
idx.pairs = join_all(&idx.data, dim, FLOOR);
served.mode = "cold";
served.dots = (rows as u64) * (rows.saturating_sub(1) as u64) / 2;
} else if added > 0 {
let from = rows - added;
idx.pairs.extend(join_from(&idx.data, dim, from, FLOOR));
served.mode = "incremental";
served.dots = (from..rows).map(|i| i as u64).sum();
} else {
served.mode = "warm";
}
// ---- expand row pairs into vector pairs
let mut members: Vec<Vec<usize>> = vec![Vec::new(); rows];
for (i, r) in row_of.iter().enumerate() {
members[*r as usize].push(i);
}
let mut out: Vec<(usize, usize, f32)> = Vec::new();
for &(a, b, c) in &idx.pairs {
if c < threshold {
continue;
}
let (ma, mb) = (&members[a as usize], &members[b as usize]);
for &i in ma {
for &j in mb {
out.push(if i < j { (i, j, c) } else { (j, i, c) });
}
}
}
for r in &distinct {
let m = &members[*r as usize];
if m.len() < 2 {
continue;
}
let row = idx.row(*r);
let c = dot(row, row);
if c >= threshold {
for x in 0..m.len() {
for y in x + 1..m.len() {
out.push((m[x], m[y], c));
}
}
}
}
out.sort_unstable_by(|x, y| x.0.cmp(&y.0).then(x.1.cmp(&y.1)));
served.distinct = distinct.len();
served.rows = rows;
served.added = added;
served.row_pairs = idx.pairs.len();
served.ms = t0.elapsed().as_millis() as u64;
(out, served)
}
static LAST: Mutex<Option<Served>> = Mutex::new(None);
/// What the newest call did (`None`: the index has not served yet).
pub fn last() -> Option<Served> {
LAST.lock().unwrap_or_else(|e| e.into_inner()).clone()
}
/// Forget everything (tests).
#[cfg(test)]
pub(crate) fn reset() {
*INDEX.lock().unwrap_or_else(|e| e.into_inner()) = None;
*SPACES.lock().unwrap_or_else(|e| e.into_inner()) = None;
}
/// The index is one per process: tests that use it take this turn.
#[cfg(test)]
pub(crate) static TEST_TURN: Mutex<()> = Mutex::new(());
#[cfg(test)]
mod tests {
use super::*;
/// Deterministic pseudo-random rows in 48 clusters, some bodies repeated
/// verbatim (the same text in several crates).
fn corpus(n: usize, dim: usize, seed: u64) -> Vec<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..48).map(|_| (0..dim).map(|_| rnd()).collect()).collect();
let mut out: Vec<Vec<f32>> = Vec::with_capacity(n);
for i in 0..n {
if i % 17 == 5 && i > 0 {
out.push(out[i / 3].clone()); // a verbatim copy
continue;
}
let c = ¢ers[i % 48];
let v: Vec<f32> = c.iter().map(|x| x + 0.18 * rnd()).collect();
out.push(super::super::normalized(&v));
}
out
}
fn naive(v: &[Vec<f32>], t: f32) -> Vec<(usize, usize)> {
let norm: Vec<Vec<f32>> = v.iter().map(|x| super::super::normalized(x)).collect();
let mut out = Vec::new();
for i in 0..v.len() {
for j in i + 1..v.len() {
if dot(&norm[i], &norm[j]) >= t {
out.push((i, j));
}
}
}
out
}
/// Pairs within this distance of the bar may land on either side
/// depending on the kernel's summation order; they are not compared.
fn away_from_the_bar(v: &[Vec<f32>], t: f32) -> impl Fn(&(usize, usize)) -> bool + '_ {
move |&(i, j)| (dot(&super::super::normalized(&v[i]), &super::super::normalized(&v[j])) - t).abs() > 1e-5
}
/// Cold, warm, incremental (new bodies arrive, old ones leave) and the
/// rebuild all answer exactly what the full scan answers — the same pairs
/// in the same order — at both lens bars.
#[test]
fn the_index_answers_what_the_full_scan_answers_through_every_change() {
let _turn = TEST_TURN.lock().unwrap_or_else(|e| e.into_inner());
reset();
let dim = 64;
let all = corpus(MIN_ROWS + 900, dim, 7);
let check = |v: &[Vec<f32>], mode: &str| {
for t in [0.92f32, 0.80] {
let (got, served) = pairs_above(v.iter().map(Vec::as_slice), t).expect("served");
let got: Vec<(usize, usize)> = got.iter().map(|p| (p.0, p.1)).filter(away_from_the_bar(v, t)).collect();
let want: Vec<(usize, usize)> = naive(v, t).into_iter().filter(away_from_the_bar(v, t)).collect();
assert!(!want.is_empty(), "the corpus must have pairs above {t}");
assert_eq!(got, want, "{mode} at {t}");
if t == 0.92 {
assert_eq!(served.mode, mode, "{served:?}");
}
}
};
// cold
let first: Vec<Vec<f32>> = all[..MIN_ROWS + 300].to_vec();
check(&first, "cold");
// warm: the same vectors again, other order
let mut again = first.clone();
again.reverse();
check(&again, "warm");
// a commit: 40 bodies leave, 600 arrive
let next: Vec<Vec<f32>> = all[40..].to_vec();
check(&next, "incremental");
let s = last().unwrap();
assert_eq!(s.added, 0, "the second bar asked nothing new: {s:?}");
// most rows gone: built again
let small: Vec<Vec<f32>> = all[..MIN_ROWS].iter().map(|v| v.iter().map(|x| -x).collect()).collect();
check(&small, "cold");
}
#[test]
fn a_bar_below_the_floor_or_a_small_set_is_not_served() {
let v = corpus(MIN_ROWS, 16, 3);
assert!(pairs_above(v.iter().map(Vec::as_slice), FLOOR - 0.01).is_none());
assert!(pairs_above(v[..MIN_ROWS - 1].iter().map(Vec::as_slice), 0.9).is_none());
}
#[test]
fn the_lane_dot_is_the_dot() {
let a: Vec<f32> = (0..771).map(|i| (i as f32 * 0.37).sin()).collect();
let b: Vec<f32> = (0..771).map(|i| (i as f32 * 0.11).cos()).collect();
let plain: f64 = a.iter().zip(&b).map(|(x, y)| *x as f64 * *y as f64).sum();
assert!((dot(&a, &b) as f64 - plain).abs() < 1e-3);
}
}