1207 lines
36 KiB
Rust
1207 lines
36 KiB
Rust
use clawhdf5_agent::bm25::BM25Index;
|
|
use clawhdf5_agent::cache::MemoryCache;
|
|
use clawhdf5_agent::hybrid::hybrid_search;
|
|
use clawhdf5_agent::ivf::{IVFIndex, IVFPQIndex};
|
|
use clawhdf5_agent::pq::ProductQuantizer;
|
|
use clawhdf5_agent::strategy::{self, HardwareCapabilities, SearchStrategy};
|
|
use clawhdf5_agent::vector_search;
|
|
use clawhdf5_agent::{AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry};
|
|
use criterion::{BatchSize, Criterion, criterion_group, criterion_main};
|
|
use tempfile::TempDir;
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Simple deterministic PRNG (LCG)
|
|
// ---------------------------------------------------------------------------
|
|
|
|
struct Rng(u32);
|
|
|
|
impl Rng {
|
|
fn new(seed: u32) -> Self {
|
|
Self(seed)
|
|
}
|
|
fn next_u32(&mut self) -> u32 {
|
|
self.0 = self.0.wrapping_mul(1103515245).wrapping_add(12345);
|
|
self.0 >> 16
|
|
}
|
|
fn next_f32(&mut self) -> f32 {
|
|
self.next_u32() as f32 / 65536.0 - 0.5
|
|
}
|
|
fn next_usize(&mut self, max: usize) -> usize {
|
|
self.next_u32() as usize % max
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Data generation helpers
|
|
// ---------------------------------------------------------------------------
|
|
|
|
const WORDS: &[&str] = &[
|
|
"the",
|
|
"quick",
|
|
"brown",
|
|
"fox",
|
|
"jumps",
|
|
"over",
|
|
"lazy",
|
|
"dog",
|
|
"rust",
|
|
"programming",
|
|
"memory",
|
|
"vector",
|
|
"search",
|
|
"index",
|
|
"data",
|
|
"system",
|
|
"agent",
|
|
"knowledge",
|
|
"graph",
|
|
"neural",
|
|
"network",
|
|
"machine",
|
|
"learning",
|
|
"deep",
|
|
"embedding",
|
|
"cosine",
|
|
"similarity",
|
|
"token",
|
|
"chunk",
|
|
"session",
|
|
"channel",
|
|
"hybrid",
|
|
];
|
|
|
|
fn make_vec(rng: &mut Rng, dim: usize) -> Vec<f32> {
|
|
(0..dim).map(|_| rng.next_f32()).collect()
|
|
}
|
|
|
|
fn make_vecs(n: usize, dim: usize, seed: u32) -> Vec<Vec<f32>> {
|
|
let mut rng = Rng::new(seed);
|
|
(0..n).map(|_| make_vec(&mut rng, dim)).collect()
|
|
}
|
|
|
|
fn make_text(rng: &mut Rng, word_count: usize) -> String {
|
|
(0..word_count)
|
|
.map(|_| WORDS[rng.next_usize(WORDS.len())])
|
|
.collect::<Vec<_>>()
|
|
.join(" ")
|
|
}
|
|
|
|
fn make_texts(n: usize, word_count: usize, seed: u32) -> Vec<String> {
|
|
let mut rng = Rng::new(seed);
|
|
(0..n).map(|_| make_text(&mut rng, word_count)).collect()
|
|
}
|
|
|
|
fn make_entry(rng: &mut Rng, dim: usize) -> MemoryEntry {
|
|
MemoryEntry {
|
|
chunk: make_text(rng, 20),
|
|
embedding: make_vec(rng, dim),
|
|
source_channel: "bench".into(),
|
|
timestamp: 1_000_000.0,
|
|
session_id: "bench-session".into(),
|
|
tags: "bench".into(),
|
|
}
|
|
}
|
|
|
|
fn make_entries(n: usize, dim: usize, seed: u32) -> Vec<MemoryEntry> {
|
|
let mut rng = Rng::new(seed);
|
|
(0..n).map(|_| make_entry(&mut rng, dim)).collect()
|
|
}
|
|
|
|
fn make_config(dir: &TempDir, dim: usize) -> MemoryConfig {
|
|
MemoryConfig::new(dir.path().join("bench.h5"), "bench-agent", dim)
|
|
}
|
|
|
|
fn nearest_centroid_bench(
|
|
vector: &[f32],
|
|
centroids: &[f32],
|
|
num_clusters: usize,
|
|
dim: usize,
|
|
) -> usize {
|
|
let mut best = 0;
|
|
let mut best_sim = f32::NEG_INFINITY;
|
|
let vnorm = vector_search::compute_norm(vector);
|
|
for c in 0..num_clusters {
|
|
let centroid = ¢roids[c * dim..(c + 1) * dim];
|
|
let cnorm = vector_search::compute_norm(centroid);
|
|
let sim = clawhdf5_agent::cosine_similarity_prenorm(vector, vnorm, centroid, cnorm);
|
|
if sim > best_sim {
|
|
best_sim = sim;
|
|
best = c;
|
|
}
|
|
}
|
|
best
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Save benchmarks
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn save_benches(c: &mut Criterion) {
|
|
let dim = 384;
|
|
|
|
c.bench_function("save_single", |b| {
|
|
b.iter_batched(
|
|
|| {
|
|
let dir = TempDir::new().unwrap();
|
|
let config = make_config(&dir, dim);
|
|
let mem = HDF5Memory::create(config).unwrap();
|
|
let mut rng = Rng::new(42);
|
|
let entry = make_entry(&mut rng, dim);
|
|
(dir, mem, entry)
|
|
},
|
|
|(_dir, mut mem, entry)| {
|
|
mem.save(entry).unwrap();
|
|
},
|
|
BatchSize::SmallInput,
|
|
);
|
|
});
|
|
|
|
c.bench_function("save_batch_100", |b| {
|
|
b.iter_batched(
|
|
|| {
|
|
let dir = TempDir::new().unwrap();
|
|
let config = make_config(&dir, dim);
|
|
let mem = HDF5Memory::create(config).unwrap();
|
|
let entries = make_entries(100, dim, 42);
|
|
(dir, mem, entries)
|
|
},
|
|
|(_dir, mut mem, entries)| {
|
|
mem.save_batch(entries).unwrap();
|
|
},
|
|
BatchSize::SmallInput,
|
|
);
|
|
});
|
|
|
|
c.bench_function("save_batch_1000", |b| {
|
|
b.iter_batched(
|
|
|| {
|
|
let dir = TempDir::new().unwrap();
|
|
let config = make_config(&dir, dim);
|
|
let mem = HDF5Memory::create(config).unwrap();
|
|
let entries = make_entries(1000, dim, 42);
|
|
(dir, mem, entries)
|
|
},
|
|
|(_dir, mut mem, entries)| {
|
|
mem.save_batch(entries).unwrap();
|
|
},
|
|
BatchSize::LargeInput,
|
|
);
|
|
});
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Vector search benchmarks (SIMD vs baseline)
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn vector_search_benches(c: &mut Criterion) {
|
|
let dim = 384;
|
|
let query = make_vec(&mut Rng::new(99), dim);
|
|
|
|
{
|
|
let vectors = make_vecs(1_000, dim, 42);
|
|
let tombstones = vec![0u8; 1_000];
|
|
c.bench_function("vector_search_1k", |b| {
|
|
b.iter(|| vector_search::cosine_similarity_batch(&query, &vectors, &tombstones));
|
|
});
|
|
}
|
|
|
|
{
|
|
let vectors = make_vecs(10_000, dim, 42);
|
|
let tombstones = vec![0u8; 10_000];
|
|
c.bench_function("simd_cosine_10k", |b| {
|
|
b.iter(|| vector_search::cosine_similarity_batch(&query, &vectors, &tombstones));
|
|
});
|
|
}
|
|
|
|
{
|
|
let vectors = make_vecs(100_000, dim, 42);
|
|
let tombstones = vec![0u8; 100_000];
|
|
c.bench_function("simd_cosine_100k", |b| {
|
|
b.iter(|| vector_search::cosine_similarity_batch(&query, &vectors, &tombstones));
|
|
});
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Pre-computed norms benchmark
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn prenorm_benches(c: &mut Criterion) {
|
|
let dim = 384;
|
|
let n = 10_000;
|
|
let query = make_vec(&mut Rng::new(99), dim);
|
|
let vectors = make_vecs(n, dim, 42);
|
|
let norms: Vec<f32> = vectors
|
|
.iter()
|
|
.map(|v| vector_search::compute_norm(v))
|
|
.collect();
|
|
let tombstones = vec![0u8; n];
|
|
|
|
c.bench_function("prenorm_search_10k", |b| {
|
|
b.iter(|| {
|
|
vector_search::cosine_similarity_batch_prenorm(&query, &vectors, &norms, &tombstones)
|
|
});
|
|
});
|
|
|
|
c.bench_function("full_norm_search_10k", |b| {
|
|
b.iter(|| vector_search::cosine_similarity_batch(&query, &vectors, &tombstones));
|
|
});
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// PQ search benchmark
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn pq_benches(c: &mut Criterion) {
|
|
let dim = 384;
|
|
let n = 100_000;
|
|
let vectors = make_vecs(n, dim, 42);
|
|
let query = make_vec(&mut Rng::new(99), dim);
|
|
let tombstones = vec![0u8; n];
|
|
|
|
// Train PQ on subset for speed
|
|
let train_set: Vec<Vec<f32>> = vectors[..1000].to_vec();
|
|
let pq = ProductQuantizer::train(&train_set, dim, 48, 256);
|
|
let all_codes = pq.encode_all(&vectors);
|
|
|
|
c.bench_function("pq_search_100k", |b| {
|
|
b.iter(|| pq.search(&query, &all_codes, &tombstones, 10));
|
|
});
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// IVF-PQ search benchmark
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn ivf_pq_benches(c: &mut Criterion) {
|
|
let dim = 384;
|
|
let n = 100_000;
|
|
let vectors = make_vecs(n, dim, 42);
|
|
let norms: Vec<f32> = vectors
|
|
.iter()
|
|
.map(|v| vector_search::compute_norm(v))
|
|
.collect();
|
|
let query = make_vec(&mut Rng::new(99), dim);
|
|
let tombstones = vec![0u8; n];
|
|
|
|
// Train on subset
|
|
let train_set: Vec<Vec<f32>> = vectors[..2000].to_vec();
|
|
let ivf = IVFIndex::train(&train_set, dim, 100);
|
|
// Re-assign all vectors
|
|
let nc = ivf.num_clusters;
|
|
let mut inverted_lists = vec![Vec::new(); nc];
|
|
for (i, v) in vectors.iter().enumerate() {
|
|
let c = nearest_centroid_bench(v, &ivf.centroids, nc, dim);
|
|
inverted_lists[c].push(i);
|
|
}
|
|
let seil_lists = inverted_lists
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(c, list)| {
|
|
list.iter()
|
|
.map(|&idx| clawhdf5_agent::ivf::SharedEntry {
|
|
vector_idx: idx,
|
|
primary_list: c,
|
|
})
|
|
.collect()
|
|
})
|
|
.collect();
|
|
let ivf_full = IVFIndex {
|
|
centroids: ivf.centroids.clone(),
|
|
num_clusters: nc,
|
|
dim,
|
|
inverted_lists,
|
|
redundancy_factor: 1,
|
|
seil_lists,
|
|
};
|
|
|
|
c.bench_function("ivf_search_100k_nprobe10", |b| {
|
|
b.iter(|| ivf_full.search(&query, &vectors, &norms, &tombstones, 10, 10));
|
|
});
|
|
|
|
// IVF-PQ combined
|
|
let pq = ProductQuantizer::train(&train_set, dim, 48, 256);
|
|
let codes = pq.encode_all(&vectors);
|
|
let ivfpq = IVFPQIndex {
|
|
ivf: ivf_full,
|
|
pq,
|
|
codes,
|
|
};
|
|
|
|
c.bench_function("ivf_pq_search_100k", |b| {
|
|
b.iter(|| ivfpq.search(&query, &vectors, &norms, &tombstones, 10, 100, 10));
|
|
});
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// RAIRS benchmark
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn rairs_benches(c: &mut Criterion) {
|
|
let dim = 128;
|
|
let n = 10_000;
|
|
let vectors = make_vecs(n, dim, 42);
|
|
let norms: Vec<f32> = vectors
|
|
.iter()
|
|
.map(|v| vector_search::compute_norm(v))
|
|
.collect();
|
|
let query = make_vec(&mut Rng::new(99), dim);
|
|
let tombstones = vec![0u8; n];
|
|
|
|
// Train RAIRS index with rf=2
|
|
let ivf_rairs = IVFIndex::train_rairs(&vectors, dim, 100, 2);
|
|
|
|
c.bench_function("rairs_search_10k_nprobe10", |b| {
|
|
b.iter(|| ivf_rairs.search_rairs(&query, &vectors, &norms, &tombstones, 10, 10));
|
|
});
|
|
|
|
// Standard IVF for comparison
|
|
let ivf_std = IVFIndex::train(&vectors, dim, 100);
|
|
|
|
c.bench_function("ivf_search_10k_nprobe10", |b| {
|
|
b.iter(|| ivf_std.search(&query, &vectors, &norms, &tombstones, 10, 10));
|
|
});
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// BM25 benchmarks
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn bm25_benches(c: &mut Criterion) {
|
|
{
|
|
let docs = make_texts(1_000, 50, 42);
|
|
let tombstones = vec![0u8; 1_000];
|
|
let index = BM25Index::build(&docs, &tombstones);
|
|
c.bench_function("bm25_search_1k", |b| {
|
|
b.iter(|| index.search("rust programming memory", 10));
|
|
});
|
|
}
|
|
|
|
{
|
|
let docs = make_texts(10_000, 50, 42);
|
|
let tombstones = vec![0u8; 10_000];
|
|
let index = BM25Index::build(&docs, &tombstones);
|
|
c.bench_function("bm25_search_10k", |b| {
|
|
b.iter(|| index.search("rust programming memory", 10));
|
|
});
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Hybrid search benchmark
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn hybrid_benches(c: &mut Criterion) {
|
|
let dim = 384;
|
|
let n = 10_000;
|
|
let vectors = make_vecs(n, dim, 42);
|
|
let docs = make_texts(n, 50, 77);
|
|
let tombstones = vec![0u8; n];
|
|
let bm25 = BM25Index::build(&docs, &tombstones);
|
|
let query_vec = make_vec(&mut Rng::new(99), dim);
|
|
|
|
c.bench_function("hybrid_search_10k", |b| {
|
|
b.iter(|| {
|
|
hybrid_search(
|
|
&query_vec,
|
|
"rust memory search",
|
|
&vectors,
|
|
&docs,
|
|
&tombstones,
|
|
&bm25,
|
|
0.7,
|
|
0.3,
|
|
10,
|
|
)
|
|
});
|
|
});
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Compact benchmark
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn compact_benches(c: &mut Criterion) {
|
|
let dim = 384;
|
|
let n = 10_000;
|
|
let mut rng = Rng::new(42);
|
|
|
|
let mut cache = MemoryCache::new(dim);
|
|
for _ in 0..n {
|
|
cache.push(
|
|
make_text(&mut rng, 20),
|
|
make_vec(&mut rng, dim),
|
|
"bench".into(),
|
|
1_000_000.0,
|
|
"session".into(),
|
|
"tag".into(),
|
|
);
|
|
}
|
|
// Tombstone 10%
|
|
for i in (0..n).step_by(10) {
|
|
cache.mark_deleted(i);
|
|
}
|
|
|
|
c.bench_function("compact_10pct_tombstoned", |b| {
|
|
b.iter_batched(
|
|
|| cache.clone(),
|
|
|mut c| {
|
|
c.compact();
|
|
},
|
|
BatchSize::LargeInput,
|
|
);
|
|
});
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// I/O benchmarks (open, snapshot, mmap)
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn io_benches(_c: &mut Criterion) {
|
|
// Skipped: pre-existing file reopen issue prevents HDF5Memory::open() in bench
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Rayon parallel search benchmarks
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn rayon_benches(c: &mut Criterion) {
|
|
let dim = 384;
|
|
let query = make_vec(&mut Rng::new(99), dim);
|
|
|
|
{
|
|
let n = 10_000;
|
|
let vectors = make_vecs(n, dim, 42);
|
|
let norms: Vec<f32> = vectors
|
|
.iter()
|
|
.map(|v| vector_search::compute_norm(v))
|
|
.collect();
|
|
let tombstones = vec![0u8; n];
|
|
|
|
c.bench_function("rayon_cosine_10k", |b| {
|
|
b.iter(|| {
|
|
use rayon::prelude::*;
|
|
let query_norm = vector_search::compute_norm(&query);
|
|
let num_cores = rayon::current_num_threads().max(1);
|
|
let chunk_size = (n + num_cores - 1) / num_cores;
|
|
let mut results: Vec<(usize, f32)> = vectors
|
|
.par_chunks(chunk_size)
|
|
.enumerate()
|
|
.flat_map(|(ci, chunk)| {
|
|
let base = ci * chunk_size;
|
|
chunk
|
|
.iter()
|
|
.enumerate()
|
|
.filter_map(|(j, v)| {
|
|
let i = base + j;
|
|
if tombstones[i] != 0 {
|
|
return None;
|
|
}
|
|
let score = clawhdf5_agent::cosine_similarity_prenorm(
|
|
&query, query_norm, v, norms[i],
|
|
);
|
|
Some((i, score))
|
|
})
|
|
.collect::<Vec<_>>()
|
|
})
|
|
.collect();
|
|
results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
|
results.truncate(10);
|
|
results
|
|
});
|
|
});
|
|
|
|
c.bench_function("sequential_cosine_10k", |b| {
|
|
b.iter(|| {
|
|
vector_search::cosine_similarity_batch_prenorm(
|
|
&query,
|
|
&vectors,
|
|
&norms,
|
|
&tombstones,
|
|
)
|
|
});
|
|
});
|
|
}
|
|
|
|
{
|
|
let n = 100_000;
|
|
let vectors = make_vecs(n, dim, 42);
|
|
let norms: Vec<f32> = vectors
|
|
.iter()
|
|
.map(|v| vector_search::compute_norm(v))
|
|
.collect();
|
|
let tombstones = vec![0u8; n];
|
|
|
|
c.bench_function("rayon_cosine_100k", |b| {
|
|
b.iter(|| {
|
|
use rayon::prelude::*;
|
|
let query_norm = vector_search::compute_norm(&query);
|
|
let num_cores = rayon::current_num_threads().max(1);
|
|
let chunk_size = (n + num_cores - 1) / num_cores;
|
|
let mut results: Vec<(usize, f32)> = vectors
|
|
.par_chunks(chunk_size)
|
|
.enumerate()
|
|
.flat_map(|(ci, chunk)| {
|
|
let base = ci * chunk_size;
|
|
chunk
|
|
.iter()
|
|
.enumerate()
|
|
.filter_map(|(j, v)| {
|
|
let i = base + j;
|
|
if tombstones[i] != 0 {
|
|
return None;
|
|
}
|
|
let score = clawhdf5_agent::cosine_similarity_prenorm(
|
|
&query, query_norm, v, norms[i],
|
|
);
|
|
Some((i, score))
|
|
})
|
|
.collect::<Vec<_>>()
|
|
})
|
|
.collect();
|
|
results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
|
results.truncate(10);
|
|
results
|
|
});
|
|
});
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// BLAS search benchmarks
|
|
// ---------------------------------------------------------------------------
|
|
|
|
#[cfg(feature = "fast-math")]
|
|
fn blas_benches(c: &mut Criterion) {
|
|
let dim = 384;
|
|
let query = make_vec(&mut Rng::new(99), dim);
|
|
|
|
{
|
|
let n = 10_000;
|
|
let vectors = make_vecs(n, dim, 42);
|
|
let norms: Vec<f32> = vectors
|
|
.iter()
|
|
.map(|v| vector_search::compute_norm(v))
|
|
.collect();
|
|
let tombstones = vec![0u8; n];
|
|
|
|
c.bench_function("blas_cosine_10k", |b| {
|
|
b.iter(|| {
|
|
vector_search::cosine_similarity_batch_blas(
|
|
&query,
|
|
&vectors,
|
|
&norms,
|
|
&tombstones,
|
|
10,
|
|
)
|
|
});
|
|
});
|
|
}
|
|
|
|
{
|
|
let n = 100_000;
|
|
let vectors = make_vecs(n, dim, 42);
|
|
let norms: Vec<f32> = vectors
|
|
.iter()
|
|
.map(|v| vector_search::compute_norm(v))
|
|
.collect();
|
|
let tombstones = vec![0u8; n];
|
|
|
|
c.bench_function("blas_cosine_100k", |b| {
|
|
b.iter(|| {
|
|
vector_search::cosine_similarity_batch_blas(
|
|
&query,
|
|
&vectors,
|
|
&norms,
|
|
&tombstones,
|
|
10,
|
|
)
|
|
});
|
|
});
|
|
}
|
|
|
|
// Batch norms benchmark
|
|
{
|
|
let n = 10_000;
|
|
let vectors = make_vecs(n, dim, 42);
|
|
let flat: Vec<f32> = vectors.iter().flat_map(|v| v.iter().copied()).collect();
|
|
|
|
c.bench_function("blas_batch_norms_10k", |b| {
|
|
b.iter(|| clawhdf5_agent::blas_search::blas_batch_norms(&flat, dim));
|
|
});
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Accelerate (cblas_sgemv / AMX) search benchmarks
|
|
// ---------------------------------------------------------------------------
|
|
|
|
#[cfg(any(feature = "accelerate", feature = "openblas"))]
|
|
fn accelerate_benches(c: &mut Criterion) {
|
|
let dim = 384;
|
|
let query = make_vec(&mut Rng::new(99), dim);
|
|
|
|
{
|
|
let n = 10_000;
|
|
let vectors = make_vecs(n, dim, 42);
|
|
let flat: Vec<f32> = vectors.iter().flat_map(|v| v.iter().copied()).collect();
|
|
let norms: Vec<f32> = vectors
|
|
.iter()
|
|
.map(|v| vector_search::compute_norm(v))
|
|
.collect();
|
|
let tombstones = vec![0u8; n];
|
|
|
|
c.bench_function("accelerate_cosine_10k", |b| {
|
|
b.iter(|| {
|
|
clawhdf5_agent::accelerate_search::accelerate_cosine_batch(
|
|
&query,
|
|
&flat,
|
|
&norms,
|
|
&tombstones,
|
|
dim,
|
|
10,
|
|
)
|
|
});
|
|
});
|
|
|
|
// Compare: SIMD prenorm on same data
|
|
c.bench_function("simd_prenorm_cosine_10k", |b| {
|
|
b.iter(|| {
|
|
vector_search::cosine_similarity_batch_prenorm(
|
|
&query,
|
|
&vectors,
|
|
&norms,
|
|
&tombstones,
|
|
)
|
|
});
|
|
});
|
|
|
|
// Compare: Accelerate via Vec<Vec> wrapper
|
|
c.bench_function("accelerate_vecs_cosine_10k", |b| {
|
|
b.iter(|| {
|
|
clawhdf5_agent::accelerate_search::accelerate_cosine_batch_vecs(
|
|
&query,
|
|
&vectors,
|
|
&norms,
|
|
&tombstones,
|
|
10,
|
|
)
|
|
});
|
|
});
|
|
}
|
|
|
|
{
|
|
let n = 100_000;
|
|
let vectors = make_vecs(n, dim, 42);
|
|
let flat: Vec<f32> = vectors.iter().flat_map(|v| v.iter().copied()).collect();
|
|
let norms: Vec<f32> = vectors
|
|
.iter()
|
|
.map(|v| vector_search::compute_norm(v))
|
|
.collect();
|
|
let tombstones = vec![0u8; n];
|
|
|
|
c.bench_function("accelerate_cosine_100k", |b| {
|
|
b.iter(|| {
|
|
clawhdf5_agent::accelerate_search::accelerate_cosine_batch(
|
|
&query,
|
|
&flat,
|
|
&norms,
|
|
&tombstones,
|
|
dim,
|
|
10,
|
|
)
|
|
});
|
|
});
|
|
}
|
|
|
|
// Batch norms via cblas_snrm2
|
|
{
|
|
let n = 10_000;
|
|
let vectors = make_vecs(n, dim, 42);
|
|
let flat: Vec<f32> = vectors.iter().flat_map(|v| v.iter().copied()).collect();
|
|
|
|
c.bench_function("accelerate_batch_norms_10k", |b| {
|
|
b.iter(|| clawhdf5_agent::accelerate_search::accelerate_batch_norms(&flat, dim));
|
|
});
|
|
}
|
|
|
|
// vDSP comparison (macOS only)
|
|
#[cfg(target_os = "macos")]
|
|
{
|
|
let n = 10_000;
|
|
let vectors = make_vecs(n, dim, 42);
|
|
let flat: Vec<f32> = vectors.iter().flat_map(|v| v.iter().copied()).collect();
|
|
let norms: Vec<f32> = vectors
|
|
.iter()
|
|
.map(|v| vector_search::compute_norm(v))
|
|
.collect();
|
|
let tombstones = vec![0u8; n];
|
|
|
|
c.bench_function("vdsp_cosine_10k", |b| {
|
|
b.iter(|| {
|
|
clawhdf5_agent::accelerate_search::vdsp_cosine_batch(
|
|
&query,
|
|
&flat,
|
|
&norms,
|
|
&tombstones,
|
|
dim,
|
|
10,
|
|
)
|
|
});
|
|
});
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Adaptive strategy benchmarks
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn adaptive_benches(c: &mut Criterion) {
|
|
let dim = 384;
|
|
let n = 10_000;
|
|
let query = make_vec(&mut Rng::new(99), dim);
|
|
let vectors = make_vecs(n, dim, 42);
|
|
let norms: Vec<f32> = vectors
|
|
.iter()
|
|
.map(|v| vector_search::compute_norm(v))
|
|
.collect();
|
|
let tombstones = vec![0u8; n];
|
|
|
|
c.bench_function("adaptive_search_10k", |b| {
|
|
let hw = HardwareCapabilities::detect();
|
|
let strat = strategy::auto_select_strategy(n, &hw);
|
|
b.iter(|| {
|
|
strategy::search_with_metrics(&query, &vectors, &norms, &tombstones, 10, strat, None)
|
|
});
|
|
});
|
|
|
|
// Strategy comparison: run all CPU strategies on same dataset
|
|
c.bench_function("strategy_scalar_10k", |b| {
|
|
b.iter(|| {
|
|
strategy::search_with_metrics(
|
|
&query,
|
|
&vectors,
|
|
&norms,
|
|
&tombstones,
|
|
10,
|
|
SearchStrategy::Scalar,
|
|
None,
|
|
)
|
|
});
|
|
});
|
|
|
|
c.bench_function("strategy_simd_10k", |b| {
|
|
b.iter(|| {
|
|
strategy::search_with_metrics(
|
|
&query,
|
|
&vectors,
|
|
&norms,
|
|
&tombstones,
|
|
10,
|
|
SearchStrategy::SimdBruteForce,
|
|
None,
|
|
)
|
|
});
|
|
});
|
|
|
|
c.bench_function("strategy_rayon_10k", |b| {
|
|
b.iter(|| {
|
|
strategy::search_with_metrics(
|
|
&query,
|
|
&vectors,
|
|
&norms,
|
|
&tombstones,
|
|
10,
|
|
SearchStrategy::RayonParallel,
|
|
None,
|
|
)
|
|
});
|
|
});
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Criterion group & main
|
|
// ---------------------------------------------------------------------------
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// WAL benchmarks — save latency with/without WAL
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn wal_benches(c: &mut Criterion) {
|
|
let dim = 384;
|
|
|
|
c.bench_function("save_with_wal_single", |b| {
|
|
b.iter_batched(
|
|
|| {
|
|
let dir = TempDir::new().unwrap();
|
|
let mut config = make_config(&dir, dim);
|
|
config.wal_enabled = true;
|
|
let mem = HDF5Memory::create(config).unwrap();
|
|
let mut rng = Rng::new(42);
|
|
let entry = make_entry(&mut rng, dim);
|
|
(dir, mem, entry)
|
|
},
|
|
|(_dir, mut mem, entry)| {
|
|
mem.save(entry).unwrap();
|
|
},
|
|
BatchSize::SmallInput,
|
|
);
|
|
});
|
|
|
|
c.bench_function("save_without_wal_single", |b| {
|
|
b.iter_batched(
|
|
|| {
|
|
let dir = TempDir::new().unwrap();
|
|
let mut config = make_config(&dir, dim);
|
|
config.wal_enabled = false;
|
|
let mem = HDF5Memory::create(config).unwrap();
|
|
let mut rng = Rng::new(42);
|
|
let entry = make_entry(&mut rng, dim);
|
|
(dir, mem, entry)
|
|
},
|
|
|(_dir, mut mem, entry)| {
|
|
mem.save(entry).unwrap();
|
|
},
|
|
BatchSize::SmallInput,
|
|
);
|
|
});
|
|
|
|
// WAL save into a store with 1K existing entries
|
|
c.bench_function("save_wal_1k_existing", |b| {
|
|
b.iter_batched(
|
|
|| {
|
|
let dir = TempDir::new().unwrap();
|
|
let mut config = make_config(&dir, dim);
|
|
config.wal_enabled = true;
|
|
config.wal_max_entries = 5000;
|
|
let mut mem = HDF5Memory::create(config).unwrap();
|
|
let entries = make_entries(1000, dim, 99);
|
|
mem.save_batch(entries).unwrap();
|
|
let mut rng = Rng::new(42);
|
|
let entry = make_entry(&mut rng, dim);
|
|
(dir, mem, entry)
|
|
},
|
|
|(_dir, mut mem, entry)| {
|
|
mem.save(entry).unwrap();
|
|
},
|
|
BatchSize::LargeInput,
|
|
);
|
|
});
|
|
|
|
c.bench_function("wal_flush_100_entries", |b| {
|
|
b.iter_batched(
|
|
|| {
|
|
let dir = TempDir::new().unwrap();
|
|
let mut config = make_config(&dir, dim);
|
|
config.wal_enabled = true;
|
|
config.wal_max_entries = 5000;
|
|
let mut mem = HDF5Memory::create(config).unwrap();
|
|
let entries = make_entries(100, dim, 77);
|
|
for e in entries {
|
|
mem.save(e).unwrap();
|
|
}
|
|
(dir, mem)
|
|
},
|
|
|(_dir, mut mem)| {
|
|
mem.flush_wal().unwrap();
|
|
},
|
|
BatchSize::LargeInput,
|
|
);
|
|
});
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Decision Gate benchmarks
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn gate_benches(c: &mut Criterion) {
|
|
use clawhdf5_agent::decision_gate::{DecisionGate, GateConfig};
|
|
|
|
let gate = DecisionGate::new(GateConfig::default());
|
|
|
|
c.bench_function("gate_trivial_skip", |b| {
|
|
b.iter(|| gate.should_save("ok"));
|
|
});
|
|
|
|
c.bench_function("gate_nontrivial_pass", |b| {
|
|
b.iter(|| gate.should_save("Tell me about the deployment architecture for the new system"));
|
|
});
|
|
|
|
c.bench_function("gate_short_phrase_skip", |b| {
|
|
b.iter(|| gate.should_save("got it"));
|
|
});
|
|
|
|
c.bench_function("gate_ratio_check", |b| {
|
|
b.iter(|| gate.should_save("yes yes definitely sure absolutely right"));
|
|
});
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Hebbian activation benchmarks
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn hebbian_benches(c: &mut Criterion) {
|
|
let dim = 384;
|
|
|
|
c.bench_function("tick_session_1k", |b| {
|
|
b.iter_batched(
|
|
|| {
|
|
let dir = TempDir::new().unwrap();
|
|
let config = make_config(&dir, dim);
|
|
let mut mem = HDF5Memory::create(config).unwrap();
|
|
let entries = make_entries(1000, dim, 55);
|
|
mem.save_batch(entries).unwrap();
|
|
(dir, mem)
|
|
},
|
|
|(_dir, mut mem)| {
|
|
mem.tick_session().unwrap();
|
|
},
|
|
BatchSize::LargeInput,
|
|
);
|
|
});
|
|
|
|
c.bench_function("tick_session_10k", |b| {
|
|
b.iter_batched(
|
|
|| {
|
|
let dir = TempDir::new().unwrap();
|
|
let config = make_config(&dir, dim);
|
|
let mut mem = HDF5Memory::create(config).unwrap();
|
|
let entries = make_entries(10000, dim, 55);
|
|
mem.save_batch(entries).unwrap();
|
|
(dir, mem)
|
|
},
|
|
|(_dir, mut mem)| {
|
|
mem.tick_session().unwrap();
|
|
},
|
|
BatchSize::LargeInput,
|
|
);
|
|
});
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Entity alias resolution benchmarks
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn alias_benches(c: &mut Criterion) {
|
|
use clawhdf5_agent::knowledge::KnowledgeCache;
|
|
|
|
// Build a knowledge cache with 50 entities and 100 aliases
|
|
let mut kg = KnowledgeCache::new();
|
|
let names = [
|
|
"Alice",
|
|
"Bob",
|
|
"Charlie",
|
|
"David",
|
|
"Eve",
|
|
"Frank",
|
|
"Grace",
|
|
"Henry",
|
|
"Irene",
|
|
"Jack",
|
|
"Karen",
|
|
"Leo",
|
|
"Mona",
|
|
"Nick",
|
|
"Olivia",
|
|
"Peter",
|
|
"Quinn",
|
|
"Rachel",
|
|
"Sam",
|
|
"Tina",
|
|
"Uma",
|
|
"Victor",
|
|
"Wendy",
|
|
"Xander",
|
|
"Yara",
|
|
"Zack",
|
|
"Acme Corp",
|
|
"BigCo",
|
|
"CloudNet",
|
|
"DataSys",
|
|
"EdgeAI",
|
|
"FinTech",
|
|
"GlobalX",
|
|
"HyperNet",
|
|
"InfoSec",
|
|
"JetOps",
|
|
"KernelDev",
|
|
"LogicAI",
|
|
"MeshNet",
|
|
"NanoTech",
|
|
"OpenStack",
|
|
"PlatformX",
|
|
"QuantumAI",
|
|
"RoboSys",
|
|
"SkyNet",
|
|
"TechFlow",
|
|
"UniCloud",
|
|
"VirtNet",
|
|
"WaveAI",
|
|
"XenoData",
|
|
];
|
|
for (i, name) in names.iter().enumerate() {
|
|
let id = kg.add_entity(name, "entity", -1);
|
|
kg.add_alias(&format!("alias_{i}"), id as i64);
|
|
kg.add_alias(&name.to_lowercase(), id as i64);
|
|
}
|
|
|
|
c.bench_function("alias_resolve_short_query", |b| {
|
|
b.iter(|| kg.resolve_aliases("what does alias_7 do at alias_30?"));
|
|
});
|
|
|
|
c.bench_function("alias_resolve_no_match", |b| {
|
|
b.iter(|| kg.resolve_aliases("tell me about the weather forecast for tomorrow"));
|
|
});
|
|
|
|
c.bench_function("alias_resolve_long_query", |b| {
|
|
b.iter(|| kg.resolve_aliases("I need alias_0 to talk to alias_15 about the alias_27 project and also check with alias_42"));
|
|
});
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// MemoryStrategy benchmarks
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn strategy_benches(c: &mut Criterion) {
|
|
use clawhdf5_agent::SearchResult;
|
|
use clawhdf5_agent::memory_strategy::*;
|
|
|
|
struct EmptyStore;
|
|
impl MemoryStoreView for EmptyStore {
|
|
fn search(&self, _: &[f32], _: usize) -> Vec<SearchResult> {
|
|
vec![]
|
|
}
|
|
fn memory_count(&self) -> usize {
|
|
0
|
|
}
|
|
fn entity_count(&self) -> usize {
|
|
0
|
|
}
|
|
}
|
|
|
|
let save_every = SaveEveryExchange::default();
|
|
let shift = SaveOnSemanticShift {
|
|
gate: clawhdf5_agent::decision_gate::DecisionGate::new(
|
|
clawhdf5_agent::decision_gate::GateConfig::default(),
|
|
),
|
|
shift_threshold: 0.25,
|
|
lookback_k: 5,
|
|
};
|
|
|
|
let exchange = Exchange {
|
|
user_turn: "What is the deployment architecture for the new microservices?".into(),
|
|
agent_turn:
|
|
"The deployment uses Kubernetes with a service mesh for inter-service communication."
|
|
.into(),
|
|
session_id: "bench-session".into(),
|
|
turn_number: 1,
|
|
timestamp: 1_000_000.0,
|
|
user_embedding: Some(vec![0.1; 384]),
|
|
agent_embedding: Some(vec![0.2; 384]),
|
|
};
|
|
|
|
let trivial_exchange = Exchange {
|
|
user_turn: "ok".into(),
|
|
agent_turn: "Got it!".into(),
|
|
session_id: "bench-session".into(),
|
|
turn_number: 2,
|
|
timestamp: 1_000_001.0,
|
|
user_embedding: None,
|
|
agent_embedding: None,
|
|
};
|
|
|
|
let store = EmptyStore;
|
|
|
|
c.bench_function("strategy_save_every_substantive", |b| {
|
|
b.iter(|| save_every.evaluate(&exchange, &store));
|
|
});
|
|
|
|
c.bench_function("strategy_save_every_trivial", |b| {
|
|
b.iter(|| save_every.evaluate(&trivial_exchange, &store));
|
|
});
|
|
|
|
c.bench_function("strategy_semantic_shift_empty_store", |b| {
|
|
b.iter(|| shift.evaluate(&exchange, &store));
|
|
});
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// AGENTS.md generation benchmark
|
|
// ---------------------------------------------------------------------------
|
|
|
|
fn agents_md_benches(c: &mut Criterion) {
|
|
let dim = 384;
|
|
|
|
c.bench_function("agents_md_generate_1k", |b| {
|
|
b.iter_batched(
|
|
|| {
|
|
let dir = TempDir::new().unwrap();
|
|
let config = make_config(&dir, dim);
|
|
let mut mem = HDF5Memory::create(config).unwrap();
|
|
let entries = make_entries(1000, dim, 42);
|
|
mem.save_batch(entries).unwrap();
|
|
mem.add_entity("Alice", "person", -1).unwrap();
|
|
mem.add_entity("ProjectX", "project", -1).unwrap();
|
|
(dir, mem)
|
|
},
|
|
|(_dir, mem)| {
|
|
mem.generate_agents_md();
|
|
},
|
|
BatchSize::LargeInput,
|
|
);
|
|
});
|
|
}
|
|
|
|
criterion_group!(
|
|
benches,
|
|
save_benches,
|
|
vector_search_benches,
|
|
prenorm_benches,
|
|
pq_benches,
|
|
ivf_pq_benches,
|
|
rairs_benches,
|
|
bm25_benches,
|
|
hybrid_benches,
|
|
compact_benches,
|
|
io_benches,
|
|
rayon_benches,
|
|
adaptive_benches,
|
|
wal_benches,
|
|
gate_benches,
|
|
hebbian_benches,
|
|
alias_benches,
|
|
strategy_benches,
|
|
agents_md_benches,
|
|
);
|
|
|
|
#[cfg(feature = "fast-math")]
|
|
criterion_group!(blas_bench_group, blas_benches);
|
|
|
|
#[cfg(any(feature = "accelerate", feature = "openblas"))]
|
|
criterion_group!(accelerate_bench_group, accelerate_benches);
|
|
|
|
// Criterion main: include all available bench groups
|
|
#[cfg(all(
|
|
feature = "fast-math",
|
|
any(feature = "accelerate", feature = "openblas")
|
|
))]
|
|
criterion_main!(benches, blas_bench_group, accelerate_bench_group);
|
|
|
|
#[cfg(all(
|
|
feature = "fast-math",
|
|
not(any(feature = "accelerate", feature = "openblas"))
|
|
))]
|
|
criterion_main!(benches, blas_bench_group);
|
|
|
|
#[cfg(all(
|
|
not(feature = "fast-math"),
|
|
any(feature = "accelerate", feature = "openblas")
|
|
))]
|
|
criterion_main!(benches, accelerate_bench_group);
|
|
|
|
#[cfg(all(
|
|
not(feature = "fast-math"),
|
|
not(any(feature = "accelerate", feature = "openblas"))
|
|
))]
|
|
criterion_main!(benches);
|