From 2e7e0456c1483c1d5ec31235982cd629a59a1e46 Mon Sep 17 00:00:00 2001 From: osobh Date: Sat, 19 Sep 2026 20:17:55 -0700 Subject: [PATCH] perf(agent): store embeddings once, not twice MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MemoryCache held every embedding in two places: a `Vec>` and a flattened copy for the batched kernels, kept in lock-step on every push, update and compaction. A store loaded from disk therefore carried the corpus twice, plus one heap allocation per entry. A new `cache::Embeddings` owns just the flat `[N x dim]` buffer and indexes into it, so `embeddings[i]` still reads as a `&[f32]` row. The batch kernels take a `VectorSet` (implemented for both `Embeddings` and `Vec>`) instead of `&[Vec]`, so their callers and tests are unchanged. Loading no longer unflattens what it just read. 100k 384-dim entries, reopened from disk: 505 -> 357 MiB, 3.44x -> 2.43x the raw vectors. Recall (1.0000 at ef=64) and query latency are unchanged. Rows are now always exactly `dim` long, shorter ones zero-padded. The old representation allowed ragged rows, which silently misaligned the flattened copy — every row after a wrong-length embedding — and `update` carried a comment about falling back to a rebuild to avoid exactly that. It is now unrepresentable. A record saved without an embedding holds a zero row and is told apart by its norm, which is what `total_embeddings` now counts. Measured with a counting allocator rather than RSS: freeing a structure returns its pages to the allocator's pool, not the OS, so an RSS reading from inside the process showed the two representations as identical. Breaking: MemoryCache::embeddings changes type, embeddings_flat is replaced by flat_embeddings(), rebuild_flat() is a deprecated no-op. Co-Authored-By: Claude Opus 5 (1M context) --- BENCHMARKS.md | 29 +++ CHANGELOG.md | 18 ++ crates/clawhdf5-agent/src/cache.rs | 210 ++++++++++++++---- crates/clawhdf5-agent/src/hybrid.rs | 16 +- crates/clawhdf5-agent/src/lib.rs | 15 +- crates/clawhdf5-agent/src/openclaw.rs | 8 +- crates/clawhdf5-agent/src/schema.rs | 12 +- crates/clawhdf5-agent/src/vector_search.rs | 105 ++++++--- .../clawhdf5-bench/src/bin/search_harness.rs | 119 ++++++++++ 9 files changed, 429 insertions(+), 103 deletions(-) diff --git a/BENCHMARKS.md b/BENCHMARKS.md index 7f8944e..3f4f814 100644 --- a/BENCHMARKS.md +++ b/BENCHMARKS.md @@ -28,6 +28,35 @@ --- +## Memory footprint + +`cargo run --release -p clawhdf5-bench --bin search_harness -- --footprint --full`, +384-dim `f32`. The figure that matters is **reopened**: a store loaded from +disk, which is what a long-lived process holds. + +Measured with a counting global allocator, not RSS. RSS cannot see this from +inside one process — freeing a large structure returns its pages to the +allocator's pool rather than to the OS, so allocating the next one shows no +change at all. Measured that way a store holding the corpus twice and one +holding it once came out *identical* (1.00x both), which is how the first +attempt at this measurement went. + +| N | vectors (raw) | reopened, before | reopened, after | +|---:|---:|---:|---:| +| 1 000 | 1 MiB | 5 MiB (3.41x) | 4 MiB (2.39x) | +| 10 000 | 15 MiB | 50 MiB (3.43x) | 35 MiB (2.42x) | +| 100 000 | 146 MiB | 505 MiB (3.44x) | **357 MiB (2.43x)** | + +The cache stored every embedding twice — once as a `Vec>` and once +flattened for the batched kernels, kept in lock-step on every push, update and +compaction. Storing only the flat buffer and indexing into it gives back +almost exactly one copy of the corpus (148 MiB at 100k) and one heap +allocation per entry. Recall and query latency are unchanged. + +What remains at 2.43x: the flat vectors (1.0x), the HNSW index's own copy of +them (1.0x), and text, ids and graph (~0.4x). The index copy is the next +target — it is what a quantised or borrowed representation would address. + ## Read harness Produced by `cargo run --release -p clawhdf5-bench --bin read_harness`: a 4096 x diff --git a/CHANGELOG.md b/CHANGELOG.md index f3e103d..8b28c88 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,24 @@ ## Unreleased +### Memory +- `clawhdf5-agent`: **a loaded store holds ~30% less memory** (100k 384-dim + entries: 505 -> 357 MiB, 3.44x -> 2.43x the raw vectors). The cache kept + every embedding twice — a `Vec>` and a flattened copy for the + batched kernels, maintained in lock-step — so it now stores only the flat + buffer and indexes into it. Recall and query latency are unchanged. + **Breaking:** `MemoryCache::embeddings` is a `cache::Embeddings` rather than + a `Vec>` (indexing still yields a `&[f32]` row); `embeddings_flat` + is gone, replaced by `flat_embeddings()`; `rebuild_flat()` is a deprecated + no-op. Rows are now always exactly `dim` long — shorter ones are + zero-padded — which makes the ragged-row case that used to silently + misalign the flattened copy unrepresentable. +- `clawhdf5-bench`: `search_harness --footprint` reports live heap use per + stage, measured with a counting allocator (RSS cannot see a structure freed + into the allocator's own pool). + +## Unreleased + ### Retrieval quality - `clawhdf5-agent`: **re-ranking discarded the retrieval score.** `reranker::rerank` built its combined score from temporal decay, source diff --git a/crates/clawhdf5-agent/src/cache.rs b/crates/clawhdf5-agent/src/cache.rs index 3b86e5d..e69da1b 100644 --- a/crates/clawhdf5-agent/src/cache.rs +++ b/crates/clawhdf5-agent/src/cache.rs @@ -2,16 +2,143 @@ use crate::vector_search; +/// Every entry's embedding, in one contiguous `[N x dim]` buffer. +/// +/// Rows are always exactly `dim` long: a shorter one is zero-padded, a longer +/// one truncated. The previous `Vec>` allowed ragged rows, which +/// silently misaligned the flattened copy that the batched kernels read — a +/// single wrong-length embedding shifted every row after it. Padding makes +/// that unrepresentable. A record stored without an embedding therefore holds +/// a zero row, and is told apart by its norm being zero rather than by length. +/// +/// This used to be two fields — a `Vec>` and a flattened copy kept in +/// lock-step — which stored the whole corpus twice and cost one heap +/// allocation per entry on top. At 100k 384-dim entries that duplicate was +/// ~150 MiB. Indexing yields a `&[f32]` row, so `embeddings[i]` still reads +/// the same way. +#[derive(Debug, Clone, Default)] +pub struct Embeddings { + flat: Vec, + dim: usize, +} + +impl Embeddings { + pub fn new(dim: usize) -> Self { + Self { + flat: Vec::new(), + dim, + } + } + + /// Number of embeddings. + pub fn len(&self) -> usize { + self.flat.len().checked_div(self.dim).unwrap_or(0) + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// The whole buffer, `[N x dim]` row-major — what batched kernels read. + pub fn as_flat(&self) -> &[f32] { + &self.flat + } + + pub fn dim(&self) -> usize { + self.dim + } + + /// Row `i`, or `None` if out of range. + pub fn get(&self, i: usize) -> Option<&[f32]> { + let start = i.checked_mul(self.dim)?; + self.flat.get(start..start.checked_add(self.dim)?) + } + + pub fn iter(&self) -> impl ExactSizeIterator { + self.flat.chunks_exact(self.dim.max(1)) + } + + /// Append one embedding. A row whose length doesn't match `dim` is padded + /// or truncated, so the buffer stays rectangular whatever a caller passes. + pub fn push(&mut self, embedding: &[f32]) { + if self.dim == 0 { + return; + } + let take = embedding.len().min(self.dim); + self.flat.extend_from_slice(&embedding[..take]); + self.flat.resize(self.flat.len() + (self.dim - take), 0.0); + } + + /// Replace row `i`. Out-of-range indices are ignored. + pub fn set(&mut self, i: usize, embedding: &[f32]) { + let Some(start) = i.checked_mul(self.dim) else { + return; + }; + if start + self.dim > self.flat.len() { + return; + } + let take = embedding.len().min(self.dim); + self.flat[start..start + take].copy_from_slice(&embedding[..take]); + self.flat[start + take..start + self.dim].fill(0.0); + } + + /// Keep only the rows `keep` returns true for, preserving order. + pub fn retain(&mut self, mut keep: impl FnMut(usize) -> bool) { + if self.dim == 0 { + return; + } + let mut write = 0usize; + for read in 0..self.len() { + if keep(read) { + if write != read { + let (dst, src) = (write * self.dim, read * self.dim); + self.flat.copy_within(src..src + self.dim, dst); + } + write += 1; + } + } + self.flat.truncate(write * self.dim); + } + + /// Replace the contents with `rows`. + pub fn reset_from(&mut self, dim: usize, rows: impl IntoIterator>) { + self.dim = dim; + self.flat.clear(); + for row in rows { + self.push(&row); + } + } + + /// Adopt an already-flat buffer, trimming any partial trailing row. + pub fn set_flat(&mut self, dim: usize, mut flat: Vec) { + self.dim = dim; + match flat.len().checked_div(dim) { + Some(rows) => flat.truncate(rows * dim), + None => flat.clear(), + } + self.flat = flat; + } +} + +impl PartialEq for Embeddings { + fn eq(&self, other: &Self) -> bool { + self.dim == other.dim && self.flat == other.flat + } +} + +impl std::ops::Index for Embeddings { + type Output = [f32]; + + fn index(&self, i: usize) -> &[f32] { + self.get(i).expect("embedding index out of range") + } +} + /// In-memory cache for the /memory group data. #[derive(Debug, Clone)] pub struct MemoryCache { pub chunks: Vec, - pub embeddings: Vec>, - /// `embeddings` flattened into one contiguous `[N × embedding_dim]` - /// buffer, maintained incrementally alongside `embeddings` (push/update/ - /// compact) so BLAS/Accelerate batch search can read it directly instead - /// of re-flattening the whole corpus on every query. - pub embeddings_flat: Vec, + pub embeddings: Embeddings, pub source_channels: Vec, pub timestamps: Vec, pub session_ids: Vec, @@ -28,8 +155,7 @@ impl MemoryCache { pub fn new(embedding_dim: usize) -> Self { Self { chunks: Vec::new(), - embeddings: Vec::new(), - embeddings_flat: Vec::new(), + embeddings: Embeddings::new(embedding_dim), source_channels: Vec::new(), timestamps: Vec::new(), session_ids: Vec::new(), @@ -41,15 +167,14 @@ impl MemoryCache { } } - /// Rebuild `embeddings_flat` from `embeddings` from scratch. Callers that - /// populate `embeddings` directly (bulk loads) must call this afterward. - pub fn rebuild_flat(&mut self) { - self.embeddings_flat.clear(); - self.embeddings_flat - .reserve(self.embeddings.len() * self.embedding_dim); - for emb in &self.embeddings { - self.embeddings_flat.extend_from_slice(emb); - } + /// Kept for callers that used to have to re-flatten after a bulk load. + /// The buffer is always flat now, so there is nothing to rebuild. + #[deprecated(note = "embeddings are stored flat; this is a no-op")] + pub fn rebuild_flat(&mut self) {} + + /// The embeddings as one contiguous `[N x dim]` buffer. + pub fn flat_embeddings(&self) -> &[f32] { + self.embeddings.as_flat() } /// Total number of entries (including tombstoned). @@ -79,8 +204,7 @@ impl MemoryCache { let idx = self.chunks.len(); let norm = vector_search::compute_norm(&embedding); self.chunks.push(chunk); - self.embeddings_flat.extend_from_slice(&embedding); - self.embeddings.push(embedding); + self.embeddings.push(&embedding); self.source_channels.push(source_channel); self.timestamps.push(timestamp); self.session_ids.push(session_id); @@ -118,20 +242,7 @@ impl MemoryCache { if idx < self.chunks.len() { let norm = vector_search::compute_norm(&embedding); self.chunks[idx] = chunk; - let dim = self.embedding_dim; - let flat_start = idx * dim; - let matches_dim = - embedding.len() == dim && flat_start + dim <= self.embeddings_flat.len(); - self.embeddings[idx] = embedding; - if matches_dim { - self.embeddings_flat[flat_start..flat_start + dim] - .copy_from_slice(&self.embeddings[idx]); - } else { - // Embedding length doesn't match embedding_dim (shouldn't - // happen in practice) — fall back to a full rebuild rather - // than leave embeddings_flat misaligned with embeddings. - self.rebuild_flat(); - } + self.embeddings.set(idx, &embedding); self.source_channels[idx] = source_channel; self.timestamps[idx] = timestamp; self.session_ids[idx] = session_id; @@ -183,7 +294,7 @@ impl MemoryCache { new_idx += 1; let norm = vector_search::compute_norm(&self.embeddings[i]); new_chunks.push(self.chunks[i].clone()); - new_embeddings.push(self.embeddings[i].clone()); + new_embeddings.push(self.embeddings[i].to_vec()); new_source_channels.push(self.source_channels[i].clone()); new_timestamps.push(self.timestamps[i]); new_session_ids.push(self.session_ids[i].clone()); @@ -196,7 +307,8 @@ impl MemoryCache { let removed = old_len - new_chunks.len(); self.chunks = new_chunks; - self.embeddings = new_embeddings; + self.embeddings + .reset_from(self.embedding_dim, new_embeddings); self.source_channels = new_source_channels; self.timestamps = new_timestamps; self.session_ids = new_session_ids; @@ -204,16 +316,14 @@ impl MemoryCache { self.tombstones = new_tombstones; self.norms = new_norms; self.activation_weights = new_activation_weights; - self.rebuild_flat(); (removed, index_map) } - /// Flatten all embeddings into a single Vec for HDF5 storage. - /// `embeddings_flat` is already maintained incrementally, so this just - /// clones it — kept as a method for callers that want an owned copy. - pub fn flat_embeddings(&self) -> Vec { - self.embeddings_flat.clone() + /// All embeddings as one owned `[N x dim]` buffer, for HDF5 storage. + /// Prefer [`MemoryCache::flat_embeddings`] where a borrow will do. + pub fn flat_embeddings_owned(&self) -> Vec { + self.embeddings.as_flat().to_vec() } } @@ -224,7 +334,7 @@ mod tests { /// `embeddings_flat` must always equal a from-scratch flatten of `embeddings`. fn assert_flat_in_sync(cache: &MemoryCache) { let expected: Vec = cache.embeddings.iter().flatten().copied().collect(); - assert_eq!(cache.embeddings_flat, expected); + assert_eq!(cache.embeddings.as_flat(), expected); } #[test] @@ -247,7 +357,10 @@ mod tests { String::new(), ); assert_flat_in_sync(&cache); - assert_eq!(cache.embeddings_flat, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]); + assert_eq!( + cache.embeddings.as_flat(), + vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + ); } #[test] @@ -279,7 +392,7 @@ mod tests { ); assert_flat_in_sync(&cache); assert_eq!( - cache.embeddings_flat, + cache.embeddings.as_flat(), vec![7.0, 8.0, 9.0, 4.0, 5.0, 6.0], "update must overwrite the correct flat slice, not just append" ); @@ -315,14 +428,15 @@ mod tests { cache.mark_deleted(1); cache.compact(); assert_flat_in_sync(&cache); - assert_eq!(cache.embeddings_flat, vec![1.0, 1.0, 3.0, 3.0]); + assert_eq!(cache.embeddings.as_flat(), vec![1.0, 1.0, 3.0, 3.0]); } #[test] fn rebuild_flat_matches_manual_flatten() { let mut cache = MemoryCache::new(2); - cache.embeddings = vec![vec![1.0, 2.0], vec![3.0, 4.0]]; - cache.rebuild_flat(); - assert_eq!(cache.embeddings_flat, vec![1.0, 2.0, 3.0, 4.0]); + cache + .embeddings + .reset_from(2, vec![vec![1.0, 2.0], vec![3.0, 4.0]]); + assert_eq!(cache.embeddings.as_flat(), vec![1.0, 2.0, 3.0, 4.0]); } } diff --git a/crates/clawhdf5-agent/src/hybrid.rs b/crates/clawhdf5-agent/src/hybrid.rs index 80ea344..21d94d5 100644 --- a/crates/clawhdf5-agent/src/hybrid.rs +++ b/crates/clawhdf5-agent/src/hybrid.rs @@ -28,7 +28,7 @@ use crate::vector_search; pub fn hybrid_search( query_embedding: &[f32], query_text: &str, - vectors: &[Vec], + vectors: &(impl crate::vector_search::VectorSet + Sync + ?Sized), chunks: &[String], tombstones: &[u8], bm25_index: &BM25Index, @@ -56,7 +56,7 @@ pub fn hybrid_search( pub fn hybrid_search_fused( query_embedding: &[f32], query_text: &str, - vectors: &[Vec], + vectors: &(impl crate::vector_search::VectorSet + Sync + ?Sized), _chunks: &[String], tombstones: &[u8], bm25_index: &BM25Index, @@ -69,12 +69,12 @@ pub fn hybrid_search_fused( let vec_scores = { #[cfg(feature = "parallel")] { - if vectors.len() > 10_000 { + if vectors.count() > 10_000 { vector_search::parallel_cosine_batch( query_embedding, vectors, tombstones, - vectors.len(), + vectors.count(), ) } else { vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones) @@ -270,7 +270,7 @@ fn normalize_scores(scores: &[(usize, f32)]) -> Vec<(usize, f32)> { pub fn rrf_hybrid_search( query_embedding: &[f32], query_text: &str, - vectors: &[Vec], + vectors: &(impl crate::vector_search::VectorSet + Sync + ?Sized), _chunks: &[String], tombstones: &[u8], bm25_index: &BM25Index, @@ -282,12 +282,12 @@ pub fn rrf_hybrid_search( let mut vec_scores = { #[cfg(feature = "parallel")] { - if vectors.len() > 10_000 { + if vectors.count() > 10_000 { vector_search::parallel_cosine_batch( query_embedding, vectors, tombstones, - vectors.len(), + vectors.count(), ) } else { vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones) @@ -298,7 +298,7 @@ pub fn rrf_hybrid_search( vector_search::cosine_similarity_batch(query_embedding, vectors, tombstones) } }; - let mut kw_scores = bm25_index.search(query_text, vectors.len()); + let mut kw_scores = bm25_index.search(query_text, vectors.count()); // Sort both lists descending so rank 1 = best. vec_scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); diff --git a/crates/clawhdf5-agent/src/lib.rs b/crates/clawhdf5-agent/src/lib.rs index 0c69042..93ef424 100644 --- a/crates/clawhdf5-agent/src/lib.rs +++ b/crates/clawhdf5-agent/src/lib.rs @@ -565,7 +565,9 @@ impl HDF5Memory { if u64::from_le_bytes(stamp.try_into().ok()?) != generation { return None; } - let vectors = cache.embeddings.get(..n_checkpoint)?.to_vec(); + let vectors: Vec> = (0..n_checkpoint) + .map(|i| cache.embeddings.get(i).map(<[f32]>::to_vec)) + .collect::>()?; let mut index = HnswIndex::from_graph_bytes(graph, vectors).ok()?; if index.dimension() != cache.embedding_dim { return None; @@ -573,7 +575,7 @@ impl HDF5Memory { // Records appended since (replayed from the WAL) join incrementally. for id in n_checkpoint..cache.embeddings.len() { if cache.embeddings[id].len() != index.dimension() - || index.insert(cache.embeddings[id].clone()) != id + || index.insert(cache.embeddings[id].to_vec()) != id { return None; } @@ -816,8 +818,11 @@ impl HDF5Memory { if self.cache.embeddings.iter().any(|e| e.len() != dim) { return None; } + // The index owns its vectors, so it needs rows rather than the cache's + // flat buffer. This copy is the index's own; the cache keeps one. + let rows: Vec> = self.cache.embeddings.iter().map(<[f32]>::to_vec).collect(); let mut index = HnswIndex::build_with_metric( - &self.cache.embeddings, + &rows, HNSW_M, HNSW_EF_CONSTRUCTION, DistanceMetric::Cosine, @@ -846,7 +851,7 @@ impl HDF5Memory { let dim = index.dimension(); let appended = (self.hnsw_synced_len..n).all(|id| { self.cache.embeddings[id].len() == dim - && index.insert(self.cache.embeddings[id].clone()) == id + && index.insert(self.cache.embeddings[id].to_vec()) == id }); if appended { for id in self.hnsw_synced_len..n { @@ -877,7 +882,7 @@ impl HDF5Memory { let emb_len = self.cache.embeddings[idx].len(); match self.hnsw.as_mut() { Some(index) if emb_len == index.dimension() => { - let id = index.insert(self.cache.embeddings[idx].clone()); + let id = index.insert(self.cache.embeddings[idx].to_vec()); if id == idx { self.hnsw_synced_len = self.cache.embeddings.len(); } else { diff --git a/crates/clawhdf5-agent/src/openclaw.rs b/crates/clawhdf5-agent/src/openclaw.rs index 59f4382..40e58de 100644 --- a/crates/clawhdf5-agent/src/openclaw.rs +++ b/crates/clawhdf5-agent/src/openclaw.rs @@ -466,7 +466,7 @@ impl ClawhdfBackend { let record = MemoryRecord { id: i as u64, chunk: cache.chunks[i].clone(), - embedding: cache.embeddings[i].clone(), + embedding: cache.embeddings[i].to_vec(), tier: MemoryTier::Working, importance: cache.activation_weights[i], access_count: 0, @@ -717,11 +717,13 @@ impl MemoryBackend for ClawhdfBackend { let total_records = cache.count_active(); + // A record saved without an embedding occupies a zero row, so "has an + // embedding" is "has a non-zero norm" rather than "row is non-empty". let total_embeddings = cache - .embeddings + .norms .iter() .enumerate() - .filter(|(i, emb)| cache.tombstones[*i] == 0 && !emb.is_empty()) + .filter(|(i, norm)| cache.tombstones[*i] == 0 && **norm > 0.0) .count(); let file_size_bytes = std::fs::metadata(&self.hdf5_path) diff --git a/crates/clawhdf5-agent/src/schema.rs b/crates/clawhdf5-agent/src/schema.rs index c7cb92a..3128364 100644 --- a/crates/clawhdf5-agent/src/schema.rs +++ b/crates/clawhdf5-agent/src/schema.rs @@ -153,7 +153,7 @@ fn build_memory_group( { let ds = group .create_dataset("embeddings") - .with_f32_data(&flat) + .with_f32_data(flat) .with_shape(&[n, d]); // Chunk size tuning: target ~256KB per chunk for optimal I/O @@ -563,12 +563,7 @@ fn load_memory_group( .collect(), }; - // Unflatten embeddings - let embeddings: Vec> = flat_embeddings - .chunks(embedding_dim) - .map(|c| c.to_vec()) - .collect(); - + // No unflattening: the cache stores the buffer as it is on disk. // Read activation_weights if present, default to vec![1.0; N] for backward compat let activation_weights = match read_f32_dataset(&group, "activation_weights") { Ok(w) if w.len() == n => w, @@ -576,7 +571,7 @@ fn load_memory_group( }; cache.chunks = chunks; - cache.embeddings = embeddings; + cache.embeddings.set_flat(embedding_dim, flat_embeddings); cache.source_channels = source_channels; cache.timestamps = timestamps; cache.session_ids = session_ids; @@ -584,7 +579,6 @@ fn load_memory_group( cache.tombstones = tombstones; cache.norms = norms; cache.activation_weights = activation_weights; - cache.rebuild_flat(); Ok(cache) } diff --git a/crates/clawhdf5-agent/src/vector_search.rs b/crates/clawhdf5-agent/src/vector_search.rs index 1693e1b..015d91a 100644 --- a/crates/clawhdf5-agent/src/vector_search.rs +++ b/crates/clawhdf5-agent/src/vector_search.rs @@ -4,6 +4,44 @@ //! `clawhdf5_accel`, with optional float16 support via the `half` crate. //! Supports pre-computed norms for eliminating redundant norm computations. +/// A corpus of equal-length embeddings addressable by index. +/// +/// Lets the batch kernels read either the cache's flat `[N x dim]` buffer or a +/// plain `Vec>` without either side owning a second copy. +pub trait VectorSet { + /// Number of embeddings. + fn count(&self) -> usize; + /// Embedding `i`; callers only index below [`VectorSet::count`]. + fn row(&self, i: usize) -> &[f32]; +} + +impl VectorSet for [Vec] { + fn count(&self) -> usize { + self.len() + } + fn row(&self, i: usize) -> &[f32] { + &self[i] + } +} + +impl VectorSet for Vec> { + fn count(&self) -> usize { + self.len() + } + fn row(&self, i: usize) -> &[f32] { + &self[i] + } +} + +impl VectorSet for crate::cache::Embeddings { + fn count(&self) -> usize { + self.len() + } + fn row(&self, i: usize) -> &[f32] { + &self[i] + } +} + /// Compute cosine similarity between two f32 slices. /// /// Returns 0.0 if either vector has zero magnitude. @@ -22,7 +60,7 @@ pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 { /// Returns `(index, score)` pairs sorted by score descending. pub fn cosine_similarity_batch( query: &[f32], - vectors: &[Vec], + vectors: &(impl VectorSet + ?Sized), tombstones: &[u8], ) -> Vec<(usize, f32)> { let query_norm = clawhdf5_accel::vector_norm(query); @@ -30,7 +68,7 @@ pub fn cosine_similarity_batch( return Vec::new(); } - let n = vectors.len(); + let n = vectors.count(); let mut results: Vec<(usize, f32)> = Vec::with_capacity(n); // Process 4 vectors at a time where possible @@ -42,8 +80,9 @@ pub fn cosine_similarity_batch( if i < tombstones.len() && tombstones[i] != 0 { continue; } - let vec_norm = clawhdf5_accel::vector_norm(&vectors[i]); - let score = crate::cosine_similarity_prenorm(query, query_norm, &vectors[i], vec_norm); + let vec_norm = clawhdf5_accel::vector_norm(vectors.row(i)); + let score = + crate::cosine_similarity_prenorm(query, query_norm, vectors.row(i), vec_norm); results.push((i, score)); } } @@ -53,8 +92,8 @@ pub fn cosine_similarity_batch( if i < tombstones.len() && tombstones[i] != 0 { continue; } - let vec_norm = clawhdf5_accel::vector_norm(&vectors[i]); - let score = crate::cosine_similarity_prenorm(query, query_norm, &vectors[i], vec_norm); + let vec_norm = clawhdf5_accel::vector_norm(vectors.row(i)); + let score = crate::cosine_similarity_prenorm(query, query_norm, vectors.row(i), vec_norm); results.push((i, score)); } @@ -68,7 +107,7 @@ pub fn cosine_similarity_batch( /// collections. Uses `score = dot(query, vec) / (query_norm * stored_norm)`. pub fn cosine_similarity_batch_prenorm( query: &[f32], - vectors: &[Vec], + vectors: &(impl VectorSet + ?Sized), norms: &[f32], tombstones: &[u8], ) -> Vec<(usize, f32)> { @@ -77,7 +116,7 @@ pub fn cosine_similarity_batch_prenorm( return Vec::new(); } - let n = vectors.len(); + let n = vectors.count(); let mut results: Vec<(usize, f32)> = Vec::with_capacity(n); for i in 0..n { @@ -85,7 +124,7 @@ pub fn cosine_similarity_batch_prenorm( continue; } let vec_norm = norms[i]; - let score = crate::cosine_similarity_prenorm(query, query_norm, &vectors[i], vec_norm); + let score = crate::cosine_similarity_prenorm(query, query_norm, vectors.row(i), vec_norm); results.push((i, score)); } @@ -162,7 +201,7 @@ pub fn cosine_similarity_f16( #[cfg(feature = "parallel")] pub fn parallel_cosine_batch( query: &[f32], - vectors: &[Vec], + vectors: &(impl VectorSet + Sync + ?Sized), tombstones: &[u8], k: usize, ) -> Vec<(usize, f32)> { @@ -174,24 +213,27 @@ pub fn parallel_cosine_batch( } let num_cores = rayon::current_num_threads().max(1); - let chunk_size = vectors.len().div_ceil(num_cores); + let chunk_size = vectors.count().div_ceil(num_cores); if chunk_size == 0 { return Vec::new(); } - let mut all_results: Vec<(usize, f32)> = vectors - .par_chunks(chunk_size) - .enumerate() - .flat_map(|(chunk_idx, chunk)| { + // Chunk over index ranges: the corpus may be one flat buffer rather than + // a slice of rows, so there is nothing to `par_chunks` over. + let n = vectors.count(); + let mut all_results: Vec<(usize, f32)> = (0..n.div_ceil(chunk_size)) + .into_par_iter() + .flat_map(|chunk_idx| { let base = chunk_idx * chunk_size; - let mut local: Vec<(usize, f32)> = Vec::with_capacity(chunk.len()); - for (j, vec) in chunk.iter().enumerate() { - let i = base + j; + let end = (base + chunk_size).min(n); + let mut local: Vec<(usize, f32)> = Vec::with_capacity(end - base); + for i in base..end { if i < tombstones.len() && tombstones[i] != 0 { continue; } - let vec_norm = clawhdf5_accel::vector_norm(vec); - let score = crate::cosine_similarity_prenorm(query, query_norm, vec, vec_norm); + let vec_norm = clawhdf5_accel::vector_norm(vectors.row(i)); + let score = + crate::cosine_similarity_prenorm(query, query_norm, vectors.row(i), vec_norm); local.push((i, score)); } local.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); @@ -209,7 +251,7 @@ pub fn parallel_cosine_batch( #[cfg(feature = "parallel")] pub fn parallel_cosine_batch_prenorm( query: &[f32], - vectors: &[Vec], + vectors: &(impl VectorSet + Sync + ?Sized), norms: &[f32], tombstones: &[u8], k: usize, @@ -222,23 +264,26 @@ pub fn parallel_cosine_batch_prenorm( } let num_cores = rayon::current_num_threads().max(1); - let chunk_size = vectors.len().div_ceil(num_cores); + let chunk_size = vectors.count().div_ceil(num_cores); if chunk_size == 0 { return Vec::new(); } - let mut all_results: Vec<(usize, f32)> = vectors - .par_chunks(chunk_size) - .enumerate() - .flat_map(|(chunk_idx, chunk)| { + // Chunk over index ranges: the corpus may be one flat buffer rather than + // a slice of rows, so there is nothing to `par_chunks` over. + let n = vectors.count(); + let mut all_results: Vec<(usize, f32)> = (0..n.div_ceil(chunk_size)) + .into_par_iter() + .flat_map(|chunk_idx| { let base = chunk_idx * chunk_size; - let mut local: Vec<(usize, f32)> = Vec::with_capacity(chunk.len()); - for (j, vec) in chunk.iter().enumerate() { - let i = base + j; + let end = (base + chunk_size).min(n); + let mut local: Vec<(usize, f32)> = Vec::with_capacity(end - base); + for i in base..end { if i < tombstones.len() && tombstones[i] != 0 { continue; } - let score = crate::cosine_similarity_prenorm(query, query_norm, vec, norms[i]); + let score = + crate::cosine_similarity_prenorm(query, query_norm, vectors.row(i), norms[i]); local.push((i, score)); } local.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); diff --git a/crates/clawhdf5-bench/src/bin/search_harness.rs b/crates/clawhdf5-bench/src/bin/search_harness.rs index d058a60..b8902a8 100644 --- a/crates/clawhdf5-bench/src/bin/search_harness.rs +++ b/crates/clawhdf5-bench/src/bin/search_harness.rs @@ -198,6 +198,57 @@ fn summarize(mut samples: Vec) -> Latency { } } +/// Counts live heap bytes, so a structure's cost can be measured by +/// difference. +/// +/// RSS cannot do this from inside one process: freeing a large structure +/// returns its pages to the allocator's pool rather than to the OS, so +/// allocating the next one shows no change. Measured that way, a store that +/// holds the corpus twice and one that holds it once look identical. +struct CountingAllocator; + +static LIVE_BYTES: std::sync::atomic::AtomicI64 = std::sync::atomic::AtomicI64::new(0); + +// SAFETY: every method forwards to the system allocator with the same layout +// it was given, and only adds bookkeeping around it. +unsafe impl std::alloc::GlobalAlloc for CountingAllocator { + unsafe fn alloc(&self, layout: std::alloc::Layout) -> *mut u8 { + let ptr = unsafe { std::alloc::System.alloc(layout) }; + if !ptr.is_null() { + LIVE_BYTES.fetch_add(layout.size() as i64, std::sync::atomic::Ordering::Relaxed); + } + ptr + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: std::alloc::Layout) { + LIVE_BYTES.fetch_sub(layout.size() as i64, std::sync::atomic::Ordering::Relaxed); + unsafe { std::alloc::System.dealloc(ptr, layout) } + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: std::alloc::Layout, new_size: usize) -> *mut u8 { + let new_ptr = unsafe { std::alloc::System.realloc(ptr, layout, new_size) }; + if !new_ptr.is_null() { + LIVE_BYTES.fetch_add( + new_size as i64 - layout.size() as i64, + std::sync::atomic::Ordering::Relaxed, + ); + } + new_ptr + } +} + +#[global_allocator] +static ALLOCATOR: CountingAllocator = CountingAllocator; + +/// Live heap bytes right now. +fn heap_bytes() -> u64 { + LIVE_BYTES.load(std::sync::atomic::Ordering::Relaxed).max(0) as u64 +} + +fn mib(bytes: u64) -> f64 { + bytes as f64 / (1 << 20) as f64 +} + fn micros(d: Duration) -> f64 { d.as_secs_f64() * 1e6 } @@ -450,6 +501,62 @@ fn fusion_study(n: usize) { } } +/// What an in-memory store costs, stage by stage. The vectors are the floor: +/// everything above it is bookkeeping that could in principle be shared. +fn bench_footprint(n: usize) { + let data = make_dataset(n, 0xF007 ^ n as u64); + let mut rng = Rng(11); + let dir = tempfile::TempDir::new().unwrap(); + let path = dir.path().join("footprint.h5"); + + let base = heap_bytes(); + let entries: Vec = data + .vectors + .iter() + .enumerate() + .map(|(i, v)| MemoryEntry { + chunk: text_for(data.cluster_of[i], i, &mut rng), + embedding: v.clone(), + source_channel: "bench".into(), + timestamp: i as f64, + session_id: format!("s{}", i % 50), + tags: format!("t{i}"), + }) + .collect(); + let after_entries = heap_bytes(); + + let mut mem = HDF5Memory::create(MemoryConfig::new(path, "bench", DIM)).unwrap(); + mem.save_batch(entries).unwrap(); + let after_store = heap_bytes(); + + // First query builds the vector and keyword indexes. + std::hint::black_box(mem.hybrid_search(&data.queries[0], "record", 0.7, 0.3, K)); + let after_indexes = heap_bytes(); + + // Reopening is the figure that matters for a long-lived process, and the + // only one RSS reports honestly: memory freed when the ingest buffers went + // away stays in the allocator's pool, so the stage deltas above understate + // what was given back. + let path = mem.config().path.clone(); + drop(mem); + let before_open = heap_bytes(); + let reopened = HDF5Memory::open(&path).unwrap(); + let after_open = heap_bytes(); + let loaded = after_open.saturating_sub(before_open); + drop(reopened); + + let raw = (n * DIM * 4) as u64; + println!( + "| {n} | {:.0} | {:.0} | {:.0} | {:.0} | {:.0} | {:.2}x |", + mib(raw), + mib(after_entries.saturating_sub(base)), + mib(after_store.saturating_sub(after_entries)), + mib(after_indexes.saturating_sub(after_store)), + mib(loaded), + loaded as f64 / raw as f64, + ); +} + fn main() { let args: Vec = std::env::args().skip(1).collect(); let full = args.iter().any(|a| a == "--full"); @@ -485,6 +592,18 @@ fn main() { let mut json = Vec::new(); println!("## Search harness"); + + if args.iter().any(|a| a == "--footprint") { + println!("\n### Resident memory, {DIM}-dim f32\n"); + println!( + "| N | vectors (raw) | entries MiB | store MiB | indexes MiB | reopened MiB | reopened / raw |" + ); + println!("|---:|---:|---:|---:|---:|---:|---:|"); + for &n in sizes { + bench_footprint(n); + } + return; + } // `--e2e-only` skips the index benchmarks, so the end-to-end section runs // in a process that has not already spun up a thread pool. if !args.iter().any(|a| a == "--e2e-only") {