dtourolle
6b1aac477d
docs/ had 26 developer documents flat beside the manual, and the two audiences are very differently sized: most readers want the manual and the gesture reference, a few want the register, the designs and the measurements. The manual and gestures.md stay at the top; everything for someone changing the code moves to docs/dev/, and the two documents that name their own successors — the v0.1 milestone and the UI-refinement plan — go to docs/dev/archive/ rather than being deleted, since both are still cited. docs/README.md is the index, users first. Every reference follows: code comments, Cargo manifests, the workflows, the pre-commit hook, the bench and traceability tools (which locate the repo root by docs/dev/requirements.md now), packaging, the Docker READMEs, CLAUDE.md, CONTRIBUTING.md and the README. The matrix links one level deeper and is regenerated. Links out of the moved documents into the tree gain a level; a link checker over every Markdown file finds none broken.
848 lines
31 KiB
Rust
848 lines
31 KiB
Rust
//! TRACES: FR-CULL-9 | FR-CULL-10
|
||
//! Finding the face pairs that could possibly be the same person.
|
||
//!
|
||
//! [`crate::cluster`] used to begin by computing every pairwise cosine and
|
||
//! holding the lot as an `n²` matrix of `f32`. At 1,800 faces that is 13 MB,
|
||
//! which is why it survived; at 25,000 it is 2.5 GB, which is why it could not
|
||
//! keep surviving.
|
||
//!
|
||
//! Almost all of that matrix is thrown away unread. Clustering only ever asks
|
||
//! whether a pair is *above* the merge threshold, and in a real library the
|
||
//! answer is no for well over 99% of pairs — the reference library's 1,813
|
||
//! faces produced 7,875 qualifying pairs out of 1.6 million. So this module
|
||
//! answers the only question that is actually asked — **which pairs clear the
|
||
//! bar** — and returns that sparse list. Memory goes from `O(n²)` to `O(edges)`
|
||
//! and the caller never has to hold a matrix at all.
|
||
//!
|
||
//! # Exact, not approximate
|
||
//!
|
||
//! The usual way to make this fast is an approximate nearest-neighbour index,
|
||
//! which trades recall for speed: it *misses* some true neighbours, and a
|
||
//! missed neighbour here is a face that silently never joins its person.
|
||
//! Nothing surfaces that — the screen just quietly shows one person as two —
|
||
//! so it is a poor trade for a feature whose whole job is to be trusted. This
|
||
//! module is exact, and the clusters it produces are identical to those from a
|
||
//! full scan.
|
||
//!
|
||
//! # An exact index was tried, measured, and removed
|
||
//!
|
||
//! Worth recording so it is not rediscovered as a good idea. The obvious exact
|
||
//! index is IVF with a triangle-inequality bound: group the embeddings into
|
||
//! cells, and skip a whole cell **pair** when the geometry proves no member of
|
||
//! one can reach any member of the other. On the sphere,
|
||
//!
|
||
//! ```text
|
||
//! angle(x, y) >= angle(c_P, c_Q) - radius(P) - radius(Q)
|
||
//! ```
|
||
//!
|
||
//! so a cell pair is impossible when `cos` of that lower bound falls below the
|
||
//! threshold. Exact, no recall loss, and it prunes beautifully on synthetic
|
||
//! clusters.
|
||
//!
|
||
//! It prunes **nothing at all** on real face embeddings. Measured over the
|
||
//! 1,813-face reference library, at √n = 43 cells:
|
||
//!
|
||
//! | quantity | measured |
|
||
//! |---|---|
|
||
//! | median pair angle | 88.5° (cosine 0.026) |
|
||
//! | merge threshold | 66.2° (cosine 0.403) |
|
||
//! | median cell radius | 80.4° |
|
||
//! | median centroid separation | 85.0° |
|
||
//! | cell pairs surviving the bound | **946 of 946 — 100%** |
|
||
//!
|
||
//! The arithmetic is not close. For the bound to exclude a typical cell pair it
|
||
//! needs `radius(P) + radius(Q) < 85° - 66° = 19°`, so cells of radius under
|
||
//! ~10°. But two photographs of the *same person* sit 36–60° apart, so even a
|
||
//! perfect single-identity cell has a radius three times too large. No
|
||
//! ball-based partition of this space can have cells tight enough for the
|
||
//! inequality to bite — 512-d embeddings are near-orthogonal, and that is the
|
||
//! curse of dimensionality doing exactly what it says.
|
||
//!
|
||
//! So the scan stayed exhaustive, and the effort went where it actually pays:
|
||
//! not materialising the matrix, an unrolled dot product, and spreading the
|
||
//! blocks across cores. That is `O(n²)` time and `O(edges)` memory, which for
|
||
//! this problem is the honest answer.
|
||
|
||
use crate::calibrate::Calibration;
|
||
|
||
/// Rows of the similarity triangle handed to one thread at a time.
|
||
///
|
||
/// Small enough that the tail of the triangle divides evenly across cores —
|
||
/// row `i` does `n - i` comparisons, so equal *row counts* are very unequal
|
||
/// work — and large enough that the per-block overhead disappears.
|
||
const BLOCK: usize = 64;
|
||
|
||
/// Columns compared against one row block before moving on.
|
||
///
|
||
/// The other half of the tiling: 64 embeddings of 512 floats is 128 KB, which
|
||
/// sits in L2 beside the row block instead of being re-read from memory for
|
||
/// every row. See [`scan_rows`] for what it was worth.
|
||
const COLUMN_TILE: usize = 64;
|
||
|
||
/// Below this many faces, do the whole thing on the calling thread.
|
||
///
|
||
/// Spawning threads for a set this small costs more than the scan.
|
||
const THREADS_ABOVE: usize = 2048;
|
||
|
||
/// One face pair that clears the merge threshold.
|
||
///
|
||
/// `i < j` always, and the probability is carried because the caller would
|
||
/// otherwise recompute the sigmoid it took a dot product to reach.
|
||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||
pub struct Pair {
|
||
pub i: usize,
|
||
pub j: usize,
|
||
pub probability: f32,
|
||
}
|
||
|
||
/// What the caller has to tell us about each face.
|
||
///
|
||
/// Deliberately not [`crate::cluster::Candidate`]: this module has no business
|
||
/// knowing what a person or a photograph is, and taking the three arrays it
|
||
/// actually reads keeps it testable on bare vectors.
|
||
pub struct Faces<'a> {
|
||
/// Every embedding end to end, L2-normalised, [`Faces::dim`] floats each.
|
||
///
|
||
/// **Flat, not a slice of vectors**, and the difference is measurable: a
|
||
/// `&[Vec<f32>]` is one heap allocation per face and the scan chases a
|
||
/// pointer per row, which defeats both the prefetcher and the tiling this
|
||
/// module does to stay in cache. One buffer is also the layout a GPU would
|
||
/// want, which is where this is eventually going.
|
||
pub embeddings: &'a [f32],
|
||
/// Floats per embedding — [`crate::EMBEDDING_DIM`] in practice, a parameter
|
||
/// so the tests can work in 64 dimensions.
|
||
pub dim: usize,
|
||
/// Source pixels across the aligned crop, for the calibration's size term.
|
||
pub crop_px: &'a [f32],
|
||
/// Which photograph each face came from. Two faces in one frame are not
|
||
/// the same person, so those pairs are never returned (docs/dev/faces.md §9).
|
||
pub images: &'a [u64],
|
||
/// Which faces may be compared *against* — the gallery
|
||
/// ([`crate::embedding::MIN_GALLERY_QUALITY`]).
|
||
///
|
||
/// A pair needs at least one gallery side: a probe measured against a
|
||
/// reference is a comparison, two short vectors measured against each
|
||
/// other is noise agreeing with noise, and those pairs are never returned.
|
||
/// Filtered here rather than by the caller for the same reason
|
||
/// co-occurrence is: what this module leaves out of the list stays out of
|
||
/// the graph, the components and the merge order, so nothing downstream
|
||
/// has to remember the rule.
|
||
pub gallery: &'a [bool],
|
||
}
|
||
|
||
impl Faces<'_> {
|
||
/// How many faces there are.
|
||
pub fn len(&self) -> usize {
|
||
if self.dim == 0 {
|
||
0
|
||
} else {
|
||
self.embeddings.len() / self.dim
|
||
}
|
||
}
|
||
|
||
pub fn is_empty(&self) -> bool {
|
||
self.len() == 0
|
||
}
|
||
|
||
#[inline(always)]
|
||
fn row(&self, i: usize) -> &[f32] {
|
||
&self.embeddings[i * self.dim..(i + 1) * self.dim]
|
||
}
|
||
}
|
||
|
||
/// Every pair whose calibrated probability reaches `min_probability`.
|
||
///
|
||
/// Excludes pairs from the same photograph, which the clusterer would refuse
|
||
/// anyway — dropping them here keeps them out of the graph the caller builds
|
||
/// and out of the connected components it derives from it.
|
||
///
|
||
/// Ordered by `(i, j)`, which is what the caller's determinism rests on.
|
||
pub fn above_threshold(faces: &Faces, cal: &Calibration, min_probability: f32) -> Vec<Pair> {
|
||
let n = faces.len();
|
||
if n < 2 {
|
||
return Vec::new();
|
||
}
|
||
|
||
// The loosest cosine that could clear the bar for *any* pair in the set.
|
||
// Cheaper than the sigmoid by far, and it rejects almost everything.
|
||
let scan = Scan {
|
||
faces,
|
||
cal,
|
||
min_probability,
|
||
tau: loosest_cosine(faces.crop_px, cal, min_probability),
|
||
dot: fastest_dot(),
|
||
};
|
||
|
||
let blocks: Vec<(usize, usize)> = (0..n)
|
||
.step_by(BLOCK)
|
||
.map(|start| (start, (start + BLOCK).min(n)))
|
||
.collect();
|
||
|
||
if n < THREADS_ABOVE {
|
||
let mut out = Vec::new();
|
||
for &(from, to) in &blocks {
|
||
scan_rows(&scan, from, to, &mut out);
|
||
}
|
||
return out;
|
||
}
|
||
|
||
// One worker per core bar one. This runs on a background thread behind a
|
||
// button the user pressed, and NFR-ARCH-2 puts it behind the UI: taking
|
||
// every core would stall the window it is reporting progress to.
|
||
let workers = std::thread::available_parallelism()
|
||
.map(|p| p.get().saturating_sub(1).max(1))
|
||
.unwrap_or(1)
|
||
.min(blocks.len());
|
||
|
||
let next = std::sync::atomic::AtomicUsize::new(0);
|
||
let mut parts: Vec<Vec<Vec<Pair>>> = std::thread::scope(|scope| {
|
||
let handles: Vec<_> = (0..workers)
|
||
.map(|_| {
|
||
let next = &next;
|
||
let blocks = &blocks;
|
||
scope.spawn(move || {
|
||
// Results stay tagged with their block index, so the order
|
||
// of the output does not depend on which thread got there
|
||
// first. Determinism is a promise this module keeps.
|
||
let mut mine: Vec<(usize, Vec<Pair>)> = Vec::new();
|
||
loop {
|
||
let b = next.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
|
||
let Some(&(from, to)) = blocks.get(b) else {
|
||
break;
|
||
};
|
||
let mut out = Vec::new();
|
||
scan_rows(&scan, from, to, &mut out);
|
||
mine.push((b, out));
|
||
}
|
||
mine
|
||
})
|
||
})
|
||
.collect();
|
||
|
||
let mut slots: Vec<Vec<Vec<Pair>>> = vec![Vec::new(); blocks.len()];
|
||
for h in handles {
|
||
for (b, pairs) in h.join().unwrap_or_default() {
|
||
slots[b].push(pairs);
|
||
}
|
||
}
|
||
slots
|
||
});
|
||
|
||
let mut out = Vec::new();
|
||
for slot in &mut parts {
|
||
for pairs in slot.drain(..) {
|
||
out.extend(pairs);
|
||
}
|
||
}
|
||
out
|
||
}
|
||
|
||
/// Everything [`scan_rows`] needs that does not change between blocks.
|
||
///
|
||
/// A struct rather than eight arguments, and the grouping is real: these five
|
||
/// are fixed for a whole scan and only the row range moves.
|
||
#[derive(Clone, Copy)]
|
||
struct Scan<'a> {
|
||
faces: &'a Faces<'a>,
|
||
cal: &'a Calibration,
|
||
min_probability: f32,
|
||
tau: f32,
|
||
dot: DotFn,
|
||
}
|
||
|
||
/// Compare rows `from..to` against everything after them.
|
||
///
|
||
/// The upper triangle, split by rows, and **tiled on the column side too**.
|
||
/// Walking `j` from `i + 1` to `n` for one row at a time streams the whole
|
||
/// embedding array past the core once per row — 336 GB of traffic for an
|
||
/// 18,000-face library — where a column tile small enough to sit in L2 is read
|
||
/// once per *tile* of rows. Measured on that library, the tiling alone took the
|
||
/// scan from 4.64 s to 2.81 s before any change to the kernel.
|
||
///
|
||
/// Pairs come out ordered by `(i, j)`: the tiles are walked in ascending order
|
||
/// but the row loop is inside them, so the block's own output is sorted before
|
||
/// it is returned. Blocks are concatenated in row order, so the whole list is
|
||
/// ordered — which is the promise the caller's determinism rests on.
|
||
fn scan_rows(scan: &Scan, from: usize, to: usize, out: &mut Vec<Pair>) {
|
||
let Scan {
|
||
faces,
|
||
cal,
|
||
min_probability,
|
||
tau,
|
||
dot,
|
||
} = *scan;
|
||
let n = faces.len();
|
||
for tile in (from..n).step_by(COLUMN_TILE) {
|
||
let tile_end = (tile + COLUMN_TILE).min(n);
|
||
for i in from..to {
|
||
// The diagonal: nothing at or before `i` is this row's business.
|
||
let start = tile.max(i + 1);
|
||
if start >= tile_end {
|
||
continue;
|
||
}
|
||
let a = faces.row(i);
|
||
let crop_a = faces.crop_px[i];
|
||
let image_a = faces.images[i];
|
||
let gallery_a = faces.gallery[i];
|
||
for j in start..tile_end {
|
||
if image_a == faces.images[j] || !(gallery_a || faces.gallery[j]) {
|
||
continue;
|
||
}
|
||
let cos = dot(a, faces.row(j));
|
||
// The cheap rejection, and it takes well over 99% of pairs.
|
||
if cos < tau {
|
||
continue;
|
||
}
|
||
let probability = cal.probability(cos, crop_a.min(faces.crop_px[j]), 0.0);
|
||
if probability >= min_probability {
|
||
out.push(Pair { i, j, probability });
|
||
}
|
||
}
|
||
}
|
||
}
|
||
out.sort_unstable_by_key(|p| (p.i, p.j));
|
||
}
|
||
|
||
/// The lowest cosine that could yield `min_probability` for any pair in the set.
|
||
///
|
||
/// The calibration is `sigmoid(a·cos + b + w_size·log2(crop))`, so for a fixed
|
||
/// size term the cosine boundary is exact. The size term is *not* fixed — it
|
||
/// varies per pair with the smaller of the two faces — so the safe bound uses
|
||
/// whichever face size pushes the boundary lowest: the largest face when
|
||
/// `w_size` is positive, the smallest when it is negative.
|
||
///
|
||
/// Returns [`f32::NEG_INFINITY`] — a filter that rejects nothing — where no
|
||
/// boundary exists: a non-positive steepness, for which probability does not
|
||
/// increase with cosine, or a threshold at the ends of the sigmoid. Those are
|
||
/// degenerate calibrations rather than impossible ones, and the right response
|
||
/// is to stop pruning, not to guess.
|
||
fn loosest_cosine(crop_px: &[f32], cal: &Calibration, min_probability: f32) -> f32 {
|
||
if cal.a <= 0.0 || !(min_probability > 0.0 && min_probability < 1.0) {
|
||
return f32::NEG_INFINITY;
|
||
}
|
||
|
||
let (mut lo, mut hi) = (f32::INFINITY, 0.0_f32);
|
||
for &c in crop_px {
|
||
let c = c.max(1.0);
|
||
lo = lo.min(c);
|
||
hi = hi.max(c);
|
||
}
|
||
if !lo.is_finite() {
|
||
return f32::NEG_INFINITY;
|
||
}
|
||
|
||
let extreme = if cal.w_size >= 0.0 { hi } else { lo };
|
||
let tau = cal.boundary_at(min_probability, extreme, 0.0);
|
||
if tau.is_nan() {
|
||
return f32::NEG_INFINITY;
|
||
}
|
||
// Cosines never exceed 1, so a boundary above it legitimately rejects
|
||
// everything. Clamped rather than left free so the comparison stays cheap.
|
||
tau.min(1.0)
|
||
}
|
||
|
||
/// Cosine of two L2-normalised embeddings.
|
||
///
|
||
/// Eight accumulators rather than one. Floating-point addition is not
|
||
/// associative, so the compiler may not re-associate a single running total and
|
||
/// the loop serialises on the adder's latency; eight independent chains give it
|
||
/// something to pipeline and vectorise. The order is fixed and identical on
|
||
/// every run, which is what the caller's determinism needs — it is a different
|
||
/// order from the naive sum, not a variable one.
|
||
/// A dot product over two equal-length, L2-normalised rows.
|
||
///
|
||
/// Chosen once per scan rather than per pair — see [`fastest_dot`].
|
||
pub(crate) type DotFn = fn(&[f32], &[f32]) -> f32;
|
||
|
||
/// The widest dot product this machine can actually run.
|
||
///
|
||
/// # Why this is worth unsafe code
|
||
///
|
||
/// The scan is the arithmetic floor of the whole subsystem and it was running
|
||
/// at **0.7 flops per cycle**. The workspace builds for baseline `x86-64`,
|
||
/// which is SSE2 and no FMA, and the portable loop below was not being
|
||
/// vectorised into even that. Measured over a real 18,143-face library, on
|
||
/// twenty cores:
|
||
///
|
||
/// | kernel | scan | GFLOP/s |
|
||
/// |---|---|---|
|
||
/// | portable, untiled (what this replaced) | 4.64 s | 36 |
|
||
/// | portable, tiled | 2.81 s | 60 |
|
||
/// | **AVX2 + FMA, tiled** | **0.86 s** | **195 |
|
||
///
|
||
/// Identical pair lists — 1,531,969 — all three ways.
|
||
///
|
||
/// **Runtime detection on x86-64, unconditional on aarch64.** Advanced SIMD is
|
||
/// in the aarch64 baseline, so every Android device that runs this has NEON and
|
||
/// there is nothing to detect; on x86-64 AVX2 is not baseline and a binary that
|
||
/// assumed it would not start on older hardware.
|
||
///
|
||
/// The three kernels sum in different orders, so a cosine may differ in its
|
||
/// last bit between them. That only matters for a pair sitting exactly on the
|
||
/// threshold, and the portable kernel already sums in eight accumulators rather
|
||
/// than one, so the module was never bit-comparable with a naive sum.
|
||
pub(crate) fn fastest_dot() -> DotFn {
|
||
#[cfg(target_arch = "x86_64")]
|
||
{
|
||
if std::arch::is_x86_feature_detected!("avx2") && std::arch::is_x86_feature_detected!("fma")
|
||
{
|
||
return dot_avx2;
|
||
}
|
||
}
|
||
#[cfg(target_arch = "aarch64")]
|
||
{
|
||
return dot_neon;
|
||
}
|
||
#[allow(unreachable_code)]
|
||
dot
|
||
}
|
||
|
||
/// Eight-wide fused multiply-add, on the half of desktops that have it.
|
||
#[cfg(target_arch = "x86_64")]
|
||
fn dot_avx2(a: &[f32], b: &[f32]) -> f32 {
|
||
// SAFETY: `fastest_dot` is the only thing that hands this out, and only
|
||
// after `is_x86_feature_detected!` has said both features are present.
|
||
unsafe { dot_avx2_inner(a, b) }
|
||
}
|
||
|
||
#[cfg(target_arch = "x86_64")]
|
||
#[target_feature(enable = "avx2", enable = "fma")]
|
||
unsafe fn dot_avx2_inner(a: &[f32], b: &[f32]) -> f32 {
|
||
use std::arch::x86_64::*;
|
||
|
||
let n = a.len().min(b.len());
|
||
let (mut acc0, mut acc1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
|
||
let mut k = 0;
|
||
// Two accumulators, because one FMA cannot start until the previous one
|
||
// retires and the unit is pipelined several deep.
|
||
while k + 16 <= n {
|
||
// SAFETY: `k + 16 <= n`, and `n` is within both slices.
|
||
acc0 = _mm256_fmadd_ps(
|
||
_mm256_loadu_ps(a.as_ptr().add(k)),
|
||
_mm256_loadu_ps(b.as_ptr().add(k)),
|
||
acc0,
|
||
);
|
||
acc1 = _mm256_fmadd_ps(
|
||
_mm256_loadu_ps(a.as_ptr().add(k + 8)),
|
||
_mm256_loadu_ps(b.as_ptr().add(k + 8)),
|
||
acc1,
|
||
);
|
||
k += 16;
|
||
}
|
||
|
||
let mut lanes = [0.0_f32; 8];
|
||
_mm256_storeu_ps(lanes.as_mut_ptr(), _mm256_add_ps(acc0, acc1));
|
||
let mut total = ((lanes[0] + lanes[1]) + (lanes[2] + lanes[3]))
|
||
+ ((lanes[4] + lanes[5]) + (lanes[6] + lanes[7]));
|
||
for k in k..n {
|
||
total += a[k] * b[k];
|
||
}
|
||
total
|
||
}
|
||
|
||
/// The same, four-wide, for the phone and the tablet.
|
||
///
|
||
/// No feature detection and no `target_feature`: Advanced SIMD is mandatory in
|
||
/// the aarch64 baseline, so this compiles for every Android target the app
|
||
/// builds for. The explicit `vfmaq` matters — LLVM will not fuse a multiply and
|
||
/// an add on its own without fast-math, which is most of the win.
|
||
#[cfg(target_arch = "aarch64")]
|
||
fn dot_neon(a: &[f32], b: &[f32]) -> f32 {
|
||
use std::arch::aarch64::*;
|
||
|
||
let n = a.len().min(b.len());
|
||
// SAFETY: every load below is bounded by `k + 16 <= n`, and `n` is within
|
||
// both slices. NEON needs no feature detection on aarch64.
|
||
unsafe {
|
||
let mut acc = [vdupq_n_f32(0.0); 4];
|
||
let mut k = 0;
|
||
while k + 16 <= n {
|
||
for (l, slot) in acc.iter_mut().enumerate() {
|
||
*slot = vfmaq_f32(
|
||
*slot,
|
||
vld1q_f32(a.as_ptr().add(k + l * 4)),
|
||
vld1q_f32(b.as_ptr().add(k + l * 4)),
|
||
);
|
||
}
|
||
k += 16;
|
||
}
|
||
|
||
let mut total = vaddvq_f32(vaddq_f32(
|
||
vaddq_f32(acc[0], acc[1]),
|
||
vaddq_f32(acc[2], acc[3]),
|
||
));
|
||
for k in k..n {
|
||
total += a[k] * b[k];
|
||
}
|
||
total
|
||
}
|
||
}
|
||
|
||
/// The portable fallback, and the definition the others have to agree with.
|
||
///
|
||
/// Eight accumulators so the adds are independent; that is as much as can be
|
||
/// asked of a loop that has to compile for anything.
|
||
pub(crate) fn dot(a: &[f32], b: &[f32]) -> f32 {
|
||
const LANES: usize = 8;
|
||
let mut acc = [0.0_f32; LANES];
|
||
let n = a.len().min(b.len());
|
||
let chunks = n / LANES;
|
||
|
||
for c in 0..chunks {
|
||
let base = c * LANES;
|
||
for (l, slot) in acc.iter_mut().enumerate() {
|
||
*slot += a[base + l] * b[base + l];
|
||
}
|
||
}
|
||
let mut total =
|
||
((acc[0] + acc[1]) + (acc[2] + acc[3])) + ((acc[4] + acc[5]) + (acc[6] + acc[7]));
|
||
for k in chunks * LANES..n {
|
||
total += a[k] * b[k];
|
||
}
|
||
total
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
const DIM: usize = 64;
|
||
|
||
fn cal() -> Calibration {
|
||
Calibration {
|
||
a: 30.0,
|
||
b: -30.0 * 0.35,
|
||
w_size: 0.0,
|
||
valid: true,
|
||
positive_pairs: 1000,
|
||
negative_pairs: 10_000,
|
||
}
|
||
}
|
||
|
||
/// A unit vector in a reproducible pseudo-random direction.
|
||
///
|
||
/// Hashed from the seed rather than drawn from an RNG, so a failure is
|
||
/// reproducible from the test alone.
|
||
fn vector(seed: u64) -> Vec<f32> {
|
||
let mut s = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1;
|
||
let mut v = Vec::with_capacity(DIM);
|
||
for _ in 0..DIM {
|
||
s ^= s << 13;
|
||
s ^= s >> 7;
|
||
s ^= s << 17;
|
||
v.push(((s >> 11) as f64 / (1u64 << 53) as f64) as f32 - 0.5);
|
||
}
|
||
normalise(v)
|
||
}
|
||
|
||
fn normalise(mut v: Vec<f32>) -> Vec<f32> {
|
||
let n = v.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||
for x in &mut v {
|
||
*x /= n;
|
||
}
|
||
v
|
||
}
|
||
|
||
/// A vector a known cosine away from `base`.
|
||
fn near(base: &[f32], other: &[f32], cosine: f32) -> Vec<f32> {
|
||
let d = dot(base, other);
|
||
let mut perp: Vec<f32> = other.iter().zip(base).map(|(o, b)| o - d * b).collect();
|
||
let n = perp.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||
for x in &mut perp {
|
||
*x /= n;
|
||
}
|
||
let s = (1.0 - cosine * cosine).max(0.0).sqrt();
|
||
normalise(
|
||
base.iter()
|
||
.zip(&perp)
|
||
.map(|(b, p)| cosine * b + s * p)
|
||
.collect(),
|
||
)
|
||
}
|
||
|
||
struct Set {
|
||
embeddings: Vec<f32>,
|
||
crop_px: Vec<f32>,
|
||
images: Vec<u64>,
|
||
gallery: Vec<bool>,
|
||
}
|
||
|
||
impl Set {
|
||
fn faces(&self) -> Faces<'_> {
|
||
Faces {
|
||
embeddings: &self.embeddings,
|
||
dim: DIM,
|
||
crop_px: &self.crop_px,
|
||
images: &self.images,
|
||
gallery: &self.gallery,
|
||
}
|
||
}
|
||
|
||
fn len(&self) -> usize {
|
||
self.embeddings.len() / DIM
|
||
}
|
||
}
|
||
|
||
/// `groups` identities, `per` faces each, every face in its own photograph.
|
||
fn population(groups: usize, per: usize, tightness: f32) -> Set {
|
||
let mut embeddings = Vec::new();
|
||
let mut images = Vec::new();
|
||
let mut image = 0u64;
|
||
for g in 0..groups {
|
||
let base = vector(g as u64 + 1);
|
||
let off = vector(g as u64 + 9_999);
|
||
for m in 0..per {
|
||
embeddings.push(if m == 0 {
|
||
base.clone()
|
||
} else {
|
||
near(&base, &off, tightness)
|
||
});
|
||
images.push(image);
|
||
image += 1;
|
||
}
|
||
}
|
||
let crop_px = vec![150.0; embeddings.len()];
|
||
let gallery = vec![true; embeddings.len()];
|
||
Set {
|
||
embeddings: embeddings.concat(),
|
||
crop_px,
|
||
images,
|
||
gallery,
|
||
}
|
||
}
|
||
|
||
/// The unpruned, unthreaded, unblocked definition of the answer.
|
||
fn reference(faces: &Faces, cal: &Calibration, min_probability: f32) -> Vec<Pair> {
|
||
let n = faces.len();
|
||
let mut out = Vec::new();
|
||
for i in 0..n {
|
||
for j in i + 1..n {
|
||
if faces.images[i] == faces.images[j] || !(faces.gallery[i] || faces.gallery[j]) {
|
||
continue;
|
||
}
|
||
let cos: f32 = faces
|
||
.row(i)
|
||
.iter()
|
||
.zip(faces.row(j))
|
||
.map(|(x, y)| x * y)
|
||
.sum();
|
||
let p = cal.probability(cos, faces.crop_px[i].min(faces.crop_px[j]), 0.0);
|
||
if p >= min_probability {
|
||
out.push(Pair {
|
||
i,
|
||
j,
|
||
probability: p,
|
||
});
|
||
}
|
||
}
|
||
}
|
||
out
|
||
}
|
||
|
||
fn same_pairs(a: &[Pair], b: &[Pair]) -> bool {
|
||
a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.i == y.i && x.j == y.j)
|
||
}
|
||
|
||
/// The guard on the SIMD kernels, and the only check the aarch64 one gets
|
||
/// on a machine that is not aarch64: whatever [`fastest_dot`] picked has to
|
||
/// agree with the portable definition. A wrong lane index or a mishandled
|
||
/// tail would show up here as a wildly different number, not a rounding
|
||
/// difference.
|
||
#[test]
|
||
fn the_fastest_kernel_agrees_with_the_portable_one() {
|
||
let fast = fastest_dot();
|
||
for seed in 0..64u64 {
|
||
let a = vector(seed);
|
||
let b = vector(seed + 1_000);
|
||
let (want, got) = (dot(&a, &b), fast(&a, &b));
|
||
assert!(
|
||
(want - got).abs() < 1e-5,
|
||
"kernel disagreed on seed {seed}: {want} vs {got}"
|
||
);
|
||
}
|
||
|
||
// A length that is not a multiple of the widest step, so the tail is
|
||
// exercised rather than assumed away.
|
||
let a: Vec<f32> = (0..37).map(|k| k as f32 * 0.01).collect();
|
||
let b: Vec<f32> = (0..37).map(|k| 1.0 - k as f32 * 0.02).collect();
|
||
assert!((dot(&a, &b) - fast(&a, &b)).abs() < 1e-5, "tail mishandled");
|
||
}
|
||
|
||
#[test]
|
||
fn nothing_to_pair_is_no_pairs() {
|
||
let s = population(1, 1, 0.9);
|
||
assert!(above_threshold(&s.faces(), &cal(), 0.9).is_empty());
|
||
}
|
||
|
||
#[test]
|
||
fn the_blocked_scan_finds_exactly_what_the_reference_does() {
|
||
let s = population(60, 8, 0.97);
|
||
let f = s.faces();
|
||
let got = above_threshold(&f, &cal(), 0.9);
|
||
let want = reference(&f, &cal(), 0.9);
|
||
assert!(!want.is_empty(), "the reference found nothing to check");
|
||
assert!(same_pairs(&got, &want), "{} vs {}", got.len(), want.len());
|
||
}
|
||
|
||
/// Past `THREADS_ABOVE` the work is split across cores and stitched back
|
||
/// together, and the stitching is where an order bug would live.
|
||
#[test]
|
||
fn the_threaded_scan_finds_exactly_what_the_reference_does() {
|
||
let s = population(300, 8, 0.97);
|
||
assert!(
|
||
s.len() > THREADS_ABOVE,
|
||
"population is below the threading cutoff"
|
||
);
|
||
let f = s.faces();
|
||
let got = above_threshold(&f, &cal(), 0.9);
|
||
let want = reference(&f, &cal(), 0.9);
|
||
assert!(!want.is_empty());
|
||
assert!(same_pairs(&got, &want), "{} vs {}", got.len(), want.len());
|
||
}
|
||
|
||
/// The size term moves the cosine boundary per pair, so the pre-filter has
|
||
/// to be built from the most permissive size in the set or it will drop a
|
||
/// pair that would have qualified.
|
||
#[test]
|
||
fn a_size_weighted_calibration_still_matches_the_reference() {
|
||
let mut s = population(60, 8, 0.97);
|
||
for (i, c) in s.crop_px.iter_mut().enumerate() {
|
||
*c = 40.0 + (i % 17) as f32 * 30.0;
|
||
}
|
||
let sized = Calibration {
|
||
w_size: 0.5,
|
||
b: -30.0 * 0.35 - 0.5 * 7.0,
|
||
..cal()
|
||
};
|
||
let f = s.faces();
|
||
assert!(same_pairs(
|
||
&above_threshold(&f, &sized, 0.9),
|
||
&reference(&f, &sized, 0.9)
|
||
));
|
||
}
|
||
|
||
/// A negative size weight flips which extreme is permissive. Cheap to get
|
||
/// wrong and silent when it is, so it gets its own case.
|
||
#[test]
|
||
fn a_negative_size_weight_prunes_from_the_other_end() {
|
||
let mut s = population(60, 8, 0.97);
|
||
for (i, c) in s.crop_px.iter_mut().enumerate() {
|
||
*c = 40.0 + (i % 17) as f32 * 30.0;
|
||
}
|
||
let sized = Calibration {
|
||
w_size: -0.5,
|
||
b: -30.0 * 0.35 + 0.5 * 7.0,
|
||
..cal()
|
||
};
|
||
let f = s.faces();
|
||
assert!(same_pairs(
|
||
&above_threshold(&f, &sized, 0.9),
|
||
&reference(&f, &sized, 0.9)
|
||
));
|
||
}
|
||
|
||
/// A degenerate calibration has no cosine boundary to prune against, and
|
||
/// must stop pruning rather than prune on a bound that does not hold.
|
||
#[test]
|
||
fn a_flat_calibration_prunes_nothing_and_still_agrees() {
|
||
let flat = Calibration {
|
||
a: 0.0,
|
||
b: 4.0,
|
||
..cal()
|
||
};
|
||
assert_eq!(loosest_cosine(&[150.0], &flat, 0.9), f32::NEG_INFINITY);
|
||
let s = population(20, 6, 0.97);
|
||
let f = s.faces();
|
||
assert!(same_pairs(
|
||
&above_threshold(&f, &flat, 0.9),
|
||
&reference(&f, &flat, 0.9)
|
||
));
|
||
}
|
||
|
||
#[test]
|
||
fn two_faces_in_one_photograph_are_never_paired() {
|
||
let mut s = population(1, 2, 1.0);
|
||
s.images = vec![7, 7];
|
||
assert!(above_threshold(&s.faces(), &cal(), 0.9).is_empty());
|
||
}
|
||
|
||
/// A probe against a reference is a comparison; two probes against each
|
||
/// other is not. The rule lives here so that nothing downstream sees the
|
||
/// pair at all.
|
||
#[test]
|
||
fn two_faces_outside_the_gallery_are_never_paired() {
|
||
let mut s = population(1, 3, 1.0);
|
||
s.gallery = vec![false, false, true];
|
||
let pairs = above_threshold(&s.faces(), &cal(), 0.9);
|
||
assert!(
|
||
!pairs.iter().any(|p| p.i == 0 && p.j == 1),
|
||
"two probes were paired with each other"
|
||
);
|
||
// Each probe is still measured against the one reference.
|
||
assert!(pairs.iter().any(|p| p.i == 0 && p.j == 2));
|
||
assert!(pairs.iter().any(|p| p.i == 1 && p.j == 2));
|
||
}
|
||
|
||
#[test]
|
||
fn the_gallery_rule_matches_the_reference_at_scale() {
|
||
let mut s = population(60, 8, 0.97);
|
||
for (i, g) in s.gallery.iter_mut().enumerate() {
|
||
*g = i % 3 != 0;
|
||
}
|
||
let f = s.faces();
|
||
let got = above_threshold(&f, &cal(), 0.9);
|
||
let want = reference(&f, &cal(), 0.9);
|
||
assert!(same_pairs(&got, &want), "{} vs {}", got.len(), want.len());
|
||
}
|
||
|
||
#[test]
|
||
fn pairs_come_back_in_index_order() {
|
||
let s = population(300, 8, 0.97);
|
||
let pairs = above_threshold(&s.faces(), &cal(), 0.9);
|
||
assert!(pairs
|
||
.windows(2)
|
||
.all(|w| (w[0].i, w[0].j) < (w[1].i, w[1].j)));
|
||
assert!(pairs.iter().all(|p| p.i < p.j));
|
||
}
|
||
|
||
/// Threads must not make the answer depend on which one finished first.
|
||
#[test]
|
||
fn the_same_input_yields_the_same_pairs() {
|
||
let s = population(300, 8, 0.97);
|
||
assert_eq!(
|
||
above_threshold(&s.faces(), &cal(), 0.9),
|
||
above_threshold(&s.faces(), &cal(), 0.9)
|
||
);
|
||
}
|
||
|
||
/// The unrolled dot has to agree with the obvious one, tail included — the
|
||
/// lengths here are deliberately not multiples of the lane count.
|
||
#[test]
|
||
fn the_unrolled_dot_matches_the_naive_one() {
|
||
for len in [1usize, 7, 8, 9, 63, 64, 65, 512] {
|
||
let a: Vec<f32> = (0..len).map(|i| (i as f32 * 0.37).sin()).collect();
|
||
let b: Vec<f32> = (0..len).map(|i| (i as f32 * 0.11).cos()).collect();
|
||
let naive: f32 = a.iter().zip(&b).map(|(x, y)| x * y).sum();
|
||
assert!(
|
||
(dot(&a, &b) - naive).abs() < 1e-4,
|
||
"len {len}: {} vs {naive}",
|
||
dot(&a, &b)
|
||
);
|
||
}
|
||
}
|
||
|
||
/// The pre-filter is the whole speed story, so it is worth asserting it
|
||
/// actually rejects the bulk of the population rather than trusting it to.
|
||
#[test]
|
||
fn the_threshold_filter_rejects_almost_everything() {
|
||
let s = population(60, 8, 0.97);
|
||
let n = s.embeddings.len();
|
||
let total = n * (n - 1) / 2;
|
||
let kept = above_threshold(&s.faces(), &cal(), 0.9).len();
|
||
assert!(
|
||
kept * 20 < total,
|
||
"kept {kept} of {total} pairs, which is not sparse"
|
||
);
|
||
}
|
||
}
|