diff --git a/BENCHMARKS.md b/BENCHMARKS.md index 0a80afe..94a7c0a 100644 --- a/BENCHMARKS.md +++ b/BENCHMARKS.md @@ -176,6 +176,51 @@ index incrementally. | 10000 | 120 | 2916 | 33.1 | 14.0 | 15.4 | 2.15 | 3.30 | 421.2 | | 100000 | 1591 | 40515 | 747.3 | 324.7 | 158.9 | 23.07 | 30.42 | 41.3 | +### After: unit-vector dot product, reusable visited set + +Cosine distance recomputed both vector norms on every evaluation; the index now +stores unit vectors and uses a plain dot product. The per-call `HashSet` of +visited nodes became a reusable epoch-stamped array. Recall is unchanged. +Build: **2.75 -> 1.89 s** (10K), **~38 -> 21 s** (100K). QPS at `ef = 64`: +**22.7K -> 39K** (10K), **10.4K -> 14K** (100K). + +### HNSW, N = 1000, dim = 384, M = 16, ef_construction = 64 + +build: 113.4 ms (8821 vectors/s) · exact scan: 4375 QPS, p50 225 µs + +| ef | recall@10 | QPS | p50 µs | p99 µs | +|---:|---:|---:|---:|---:| +| 16 | 0.9990 | 144379 | 7 | 15 | +| 32 | 1.0000 | 110654 | 9 | 17 | +| 64 | 1.0000 | 80446 | 12 | 25 | +| 128 | 1.0000 | 38220 | 26 | 36 | +| 256 | 1.0000 | 20041 | 50 | 62 | + +### HNSW, N = 10000, dim = 384, M = 16, ef_construction = 64 + +build: 1519.4 ms (6581 vectors/s) · exact scan: 422 QPS, p50 2368 µs + +| ef | recall@10 | QPS | p50 µs | p99 µs | +|---:|---:|---:|---:|---:| +| 16 | 0.9975 | 54608 | 15 | 45 | +| 32 | 1.0000 | 66009 | 14 | 24 | +| 64 | 1.0000 | 49854 | 19 | 31 | +| 128 | 1.0000 | 22403 | 45 | 57 | +| 256 | 1.0000 | 10096 | 100 | 120 | + +### HNSW, N = 100000, dim = 384, M = 16, ef_construction = 64 + +build: 21084.6 ms (4743 vectors/s) · exact scan: 39 QPS, p50 24739 µs + +| ef | recall@10 | QPS | p50 µs | p99 µs | +|---:|---:|---:|---:|---:| +| 16 | 0.9235 | 15139 | 61 | 154 | +| 32 | 0.9675 | 18181 | 53 | 121 | +| 64 | 0.9840 | 13980 | 70 | 139 | +| 128 | 0.9990 | 10959 | 86 | 174 | +| 256 | 0.9990 | 3731 | 254 | 697 | + + ## Vector Search Latency Brute-force cosine similarity over 384-dimensional embeddings (OpenAI text-embedding-3-small size). diff --git a/CHANGELOG.md b/CHANGELOG.md index 28a74fa..5efef92 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -32,6 +32,15 @@ replayed from the WAL join the loaded index incrementally; a replayed update or delete invalidates it. `snapshot()` copies it. Batch saves no longer force a full index rebuild. +- `clawhdf5-ann`: faster HNSW build and search with identical recall. The + cosine metric stores unit vectors and compares them with a plain dot product + (it re-derived both norms on every distance evaluation), and the per-call + `HashSet` of visited nodes is a reusable epoch-stamped array. Build 2.75 -> + 1.89 s at 10K and ~38 -> 21 s at 100K; QPS at `ef = 64` 22.7K -> 39K at 10K. + Distances returned by `search` are unchanged (1 - cosine). Indexes loaded + from older HDF5 files are normalised on load. +- `clawhdf5-accel`: the SIMD backend is detected once per process instead of + on every kernel call. - `clawhdf5-ann`: `HnswIndex::graph_to_bytes` / `from_graph_bytes` — graph-only serialization (checksummed, every neighbour id and level validated on load). - `clawhdf5-agent`: BM25 results are deterministic (ties break by record id), diff --git a/crates/clawhdf5-accel/src/lib.rs b/crates/clawhdf5-accel/src/lib.rs index 02208cf..95987ef 100644 --- a/crates/clawhdf5-accel/src/lib.rs +++ b/crates/clawhdf5-accel/src/lib.rs @@ -61,8 +61,14 @@ pub enum Backend { Scalar, } -/// Detect the best available SIMD backend at runtime. +/// The best available SIMD backend, detected once per process. Every kernel +/// dispatches through this, so it sits in the innermost loop of every search. pub fn detect_backend() -> Backend { + static BACKEND: std::sync::OnceLock = std::sync::OnceLock::new(); + *BACKEND.get_or_init(detect_backend_uncached) +} + +fn detect_backend_uncached() -> Backend { #[cfg(target_arch = "aarch64")] { return Backend::Neon; // Always available on aarch64 diff --git a/crates/clawhdf5-ann/src/hnsw.rs b/crates/clawhdf5-ann/src/hnsw.rs index 2644add..b9e7ce3 100644 --- a/crates/clawhdf5-ann/src/hnsw.rs +++ b/crates/clawhdf5-ann/src/hnsw.rs @@ -1,6 +1,6 @@ //! HNSW index implementation with HDF5 serialization. -use std::collections::{BinaryHeap, HashSet}; +use std::collections::BinaryHeap; use clawhdf5_format::attribute::extract_attributes_full; use clawhdf5_format::data_layout::DataLayout; @@ -51,10 +51,30 @@ impl DistanceMetric { fn compute_distance(a: &[f32], b: &[f32], metric: DistanceMetric) -> f32 { match metric { DistanceMetric::L2 => clawhdf5_accel::l2_distance(a, b), - DistanceMetric::Cosine => 1.0 - clawhdf5_accel::cosine_similarity(a, b), + // Both sides are unit length (see `prepare`), so cosine similarity is + // the plain dot product. Computing it as dot / (|a| * |b|) re-derived + // both norms on every call — three reductions instead of one, in the + // innermost loop of both build and search. + DistanceMetric::Cosine => 1.0 - clawhdf5_accel::dot_product(a, b), } } +/// Put a vector in the form the index stores and compares: unit length for the +/// cosine metric, unchanged for L2. A zero vector stays zero, giving distance 1 +/// to everything — what the cosine kernel reports for a degenerate input. +fn prepare(mut v: Vec, metric: DistanceMetric) -> Vec { + if metric == DistanceMetric::Cosine { + let norm = clawhdf5_accel::vector_norm(&v); + if norm > f32::EPSILON { + let inv = 1.0 / norm; + v.iter_mut().for_each(|x| *x *= inv); + } else { + v.iter_mut().for_each(|x| *x = 0.0); + } + } + v +} + /// Assign a random level to a new node based on the HNSW probability distribution. /// /// Uses a deterministic approach based on the node index for reproducibility. @@ -204,6 +224,8 @@ impl HnswIndex { let m_max0 = m * 2; let n = vectors.len(); + let prepared: Vec> = vectors.iter().map(|v| prepare(v.clone(), metric)).collect(); + let vectors: &[Vec] = &prepared; // Assign levels to all nodes let mut node_levels = Vec::with_capacity(n); @@ -288,7 +310,7 @@ impl HnswIndex { } Self { - vectors: vectors.to_vec(), + vectors: prepared, graph, deleted: vec![false; n], entry_point, @@ -327,6 +349,7 @@ impl HnswIndex { /// # Panics /// Panics if `vector`'s dimension does not match the existing vectors. pub fn insert(&mut self, vector: Vec) -> usize { + let vector = prepare(vector, self.metric); let id = self.vectors.len(); // Seed an empty index. @@ -483,6 +506,8 @@ impl HnswIndex { "query dimension mismatch" ); let ef = ef.max(k); + let prepared_query = prepare(query.to_vec(), self.metric); + let query = prepared_query.as_slice(); let mut ep = self.entry_point; let top_layer = self.graph.len().saturating_sub(1); @@ -633,7 +658,9 @@ impl HnswIndex { actual: flat_vectors.len(), }); } - vectors.push(flat_vectors[start..end].to_vec()); + // Files written before vectors were stored unit-length hold the + // raw ones; preparing is idempotent, so this handles both. + vectors.push(prepare(flat_vectors[start..end].to_vec(), metric)); } // Read graph layers @@ -835,7 +862,7 @@ impl HnswIndex { } Ok(Self { - vectors, + vectors: vectors.into_iter().map(|v| prepare(v, metric)).collect(), graph, deleted, entry_point, @@ -937,9 +964,62 @@ fn search_layer( distance: ep_dist, }); - let mut visited = HashSet::new(); - visited.insert(ep); + VISITED.with_borrow_mut(|visited| { + visited.begin(vectors.len()); + visited.insert(ep); + search_layer_visit( + vectors, layer, query, ef, metric, visited, candidates, results, + ) + }) +} +/// Which nodes a layer search has already seen. A `HashSet` allocated per call +/// was the hottest non-arithmetic cost in both build and query; this is one +/// `u32` stamp per node, reused across calls: a node is visited iff its stamp +/// equals the current epoch, so "clearing" is just bumping the epoch. +#[derive(Default)] +struct Visited { + stamps: Vec, + epoch: u32, +} + +impl Visited { + fn begin(&mut self, n: usize) { + if self.stamps.len() < n { + self.stamps.resize(n, 0); + } + self.epoch = self.epoch.wrapping_add(1); + if self.epoch == 0 { + // Wrapped: stale stamps could collide with the new epoch. + self.stamps.iter_mut().for_each(|s| *s = 0); + self.epoch = 1; + } + } + + /// Mark `id` visited; `true` if it was not already. + fn insert(&mut self, id: usize) -> bool { + let seen = self.stamps[id] == self.epoch; + self.stamps[id] = self.epoch; + !seen + } +} + +thread_local! { + /// Per-thread scratch, so `search(&self)` stays shareable across threads. + static VISITED: std::cell::RefCell = std::cell::RefCell::new(Visited::default()); +} + +#[allow(clippy::too_many_arguments)] +fn search_layer_visit( + vectors: &[Vec], + layer: &[Vec], + query: &[f32], + ef: usize, + metric: DistanceMetric, + visited: &mut Visited, + mut candidates: BinaryHeap, + mut results: BinaryHeap, +) -> Vec { while let Some(closest) = candidates.pop() { let furthest_dist = results.peek().map_or(f32::MAX, |f| f.distance); if closest.distance > furthest_dist && results.len() >= ef { @@ -947,10 +1027,9 @@ fn search_layer( } for &neighbor in &layer[closest.id] { - if visited.contains(&neighbor) { + if !visited.insert(neighbor) { continue; } - visited.insert(neighbor); let d = compute_distance(query, &vectors[neighbor], metric); let furthest_dist = results.peek().map_or(f32::MAX, |f| f.distance); @@ -1236,6 +1315,7 @@ fn get_attr_string(attrs: &[(String, AttrValue)], name: &str) -> Result