451 lines
13 KiB
Rust
451 lines
13 KiB
Rust
use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion};
|
|
use std::collections::HashMap;
|
|
use std::time::Duration;
|
|
|
|
/// Benchmark embedding generation
|
|
fn bench_embedding_generation(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("embedding_generation");
|
|
group.measurement_time(Duration::from_secs(8));
|
|
|
|
let image_sizes = [(224, 224), (384, 384), (512, 512)];
|
|
|
|
for (w, h) in image_sizes {
|
|
let image_data = generate_test_image(w, h);
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("generate", format!("{}x{}", w, h)),
|
|
&image_data,
|
|
|b, img| {
|
|
b.iter(|| black_box(generate_embedding(black_box(img))));
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark similarity search (vector search)
|
|
fn bench_similarity_search(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("similarity_search");
|
|
group.measurement_time(Duration::from_secs(10));
|
|
|
|
// Create cache with varying sizes
|
|
let cache_sizes = [100, 1000, 10000];
|
|
|
|
for cache_size in cache_sizes {
|
|
let cache = create_embedding_cache(cache_size);
|
|
let query_embedding = generate_random_embedding(512);
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("linear_search", cache_size),
|
|
&(&cache, &query_embedding),
|
|
|b, (cache, query)| {
|
|
b.iter(|| {
|
|
black_box(linear_similarity_search(
|
|
black_box(cache),
|
|
black_box(query),
|
|
10,
|
|
))
|
|
});
|
|
},
|
|
);
|
|
|
|
// Approximate nearest neighbor search
|
|
group.bench_with_input(
|
|
BenchmarkId::new("ann_search", cache_size),
|
|
&(&cache, &query_embedding),
|
|
|b, (cache, query)| {
|
|
b.iter(|| {
|
|
black_box(ann_similarity_search(
|
|
black_box(cache),
|
|
black_box(query),
|
|
10,
|
|
))
|
|
});
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark cache hit latency
|
|
fn bench_cache_hit_latency(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("cache_hit_latency");
|
|
group.measurement_time(Duration::from_secs(5));
|
|
|
|
let cache = create_embedding_cache(1000);
|
|
let query = generate_random_embedding(512);
|
|
|
|
group.bench_function("exact_match", |b| {
|
|
let cached_embedding = cache.values().next().unwrap();
|
|
b.iter(|| {
|
|
black_box(find_exact_match(
|
|
black_box(&cache),
|
|
black_box(cached_embedding),
|
|
))
|
|
});
|
|
});
|
|
|
|
group.bench_function("similarity_threshold", |b| {
|
|
b.iter(|| {
|
|
black_box(find_by_similarity_threshold(
|
|
black_box(&cache),
|
|
black_box(&query),
|
|
0.95,
|
|
))
|
|
});
|
|
});
|
|
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark cache miss latency
|
|
fn bench_cache_miss_latency(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("cache_miss_latency");
|
|
group.measurement_time(Duration::from_secs(8));
|
|
|
|
let cache = create_embedding_cache(1000);
|
|
let new_image = generate_test_image(384, 384);
|
|
|
|
group.bench_function("miss_with_generation", |b| {
|
|
b.iter(|| {
|
|
let query_embedding = generate_embedding(black_box(&new_image));
|
|
let result = linear_similarity_search(black_box(&cache), &query_embedding, 1);
|
|
if result.is_empty() || result[0].1 < 0.95 {
|
|
// Cache miss - would need to process
|
|
black_box(process_new_image(black_box(&new_image)))
|
|
} else {
|
|
black_box(result[0].2.clone())
|
|
}
|
|
});
|
|
});
|
|
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark cache insertion
|
|
fn bench_cache_insertion(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("cache_insertion");
|
|
group.measurement_time(Duration::from_secs(8));
|
|
|
|
group.bench_function("insert_new_entry", |b| {
|
|
let mut cache = create_embedding_cache(1000);
|
|
let mut counter = 0;
|
|
|
|
b.iter(|| {
|
|
let embedding = generate_random_embedding(512);
|
|
let key = format!("key_{}", counter);
|
|
cache.insert(key.clone(), embedding);
|
|
counter += 1;
|
|
black_box(&cache)
|
|
});
|
|
});
|
|
|
|
group.bench_function("insert_with_eviction", |b| {
|
|
let mut cache = LRUCache::new(1000);
|
|
let mut counter = 0;
|
|
|
|
b.iter(|| {
|
|
let embedding = generate_random_embedding(512);
|
|
let key = format!("key_{}", counter);
|
|
cache.insert(key, embedding);
|
|
counter += 1;
|
|
black_box(&cache)
|
|
});
|
|
});
|
|
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark cache update operations
|
|
fn bench_cache_updates(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("cache_updates");
|
|
group.measurement_time(Duration::from_secs(5));
|
|
|
|
let mut cache = create_embedding_cache(1000);
|
|
let keys: Vec<_> = cache.keys().cloned().collect();
|
|
|
|
group.bench_function("update_existing", |b| {
|
|
let mut idx = 0;
|
|
b.iter(|| {
|
|
let key = &keys[idx % keys.len()];
|
|
let new_embedding = generate_random_embedding(512);
|
|
cache.insert(key.clone(), new_embedding);
|
|
idx += 1;
|
|
black_box(&cache)
|
|
});
|
|
});
|
|
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark batch cache operations
|
|
fn bench_batch_cache_ops(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("batch_cache_operations");
|
|
group.measurement_time(Duration::from_secs(10));
|
|
|
|
let batch_sizes = [10, 50, 100];
|
|
|
|
for batch_size in batch_sizes {
|
|
let cache = create_embedding_cache(1000);
|
|
let queries: Vec<_> = (0..batch_size)
|
|
.map(|_| generate_random_embedding(512))
|
|
.collect();
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("batch_search", batch_size),
|
|
&(&cache, &queries),
|
|
|b, (cache, queries)| {
|
|
b.iter(|| {
|
|
let results: Vec<_> = queries
|
|
.iter()
|
|
.map(|q| linear_similarity_search(black_box(cache), q, 10))
|
|
.collect();
|
|
black_box(results)
|
|
});
|
|
},
|
|
);
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("batch_insert", batch_size),
|
|
&queries,
|
|
|b, queries| {
|
|
b.iter_with_setup(
|
|
|| create_embedding_cache(1000),
|
|
|mut cache| {
|
|
for (i, embedding) in queries.iter().enumerate() {
|
|
cache.insert(format!("batch_{}", i), embedding.clone());
|
|
}
|
|
black_box(cache)
|
|
},
|
|
);
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark cache statistics and monitoring
|
|
fn bench_cache_statistics(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("cache_statistics");
|
|
group.measurement_time(Duration::from_secs(5));
|
|
|
|
let cache = create_embedding_cache(10000);
|
|
|
|
group.bench_function("compute_stats", |b| {
|
|
b.iter(|| black_box(compute_cache_statistics(black_box(&cache))));
|
|
});
|
|
|
|
group.bench_function("memory_usage", |b| {
|
|
b.iter(|| black_box(estimate_cache_memory(black_box(&cache))));
|
|
});
|
|
|
|
group.finish();
|
|
}
|
|
|
|
// Mock implementations
|
|
|
|
type Embedding = Vec<f32>;
|
|
|
|
struct LRUCache {
|
|
capacity: usize,
|
|
cache: HashMap<String, Embedding>,
|
|
access_order: Vec<String>,
|
|
}
|
|
|
|
impl LRUCache {
|
|
fn new(capacity: usize) -> Self {
|
|
Self {
|
|
capacity,
|
|
cache: HashMap::new(),
|
|
access_order: Vec::new(),
|
|
}
|
|
}
|
|
|
|
fn insert(&mut self, key: String, value: Embedding) {
|
|
if self.cache.len() >= self.capacity && !self.cache.contains_key(&key) {
|
|
if let Some(lru_key) = self.access_order.first().cloned() {
|
|
self.cache.remove(&lru_key);
|
|
self.access_order.remove(0);
|
|
}
|
|
}
|
|
|
|
self.cache.insert(key.clone(), value);
|
|
self.access_order.retain(|k| k != &key);
|
|
self.access_order.push(key);
|
|
}
|
|
}
|
|
|
|
fn generate_test_image(width: u32, height: u32) -> Vec<u8> {
|
|
vec![128u8; (width * height * 3) as usize]
|
|
}
|
|
|
|
fn generate_random_embedding(dim: usize) -> Embedding {
|
|
(0..dim).map(|i| (i as f32 * 0.001) % 1.0).collect()
|
|
}
|
|
|
|
fn generate_embedding(image_data: &[u8]) -> Embedding {
|
|
// Simulate embedding generation from image
|
|
let dim = 512;
|
|
let mut embedding = Vec::with_capacity(dim);
|
|
|
|
for i in 0..dim {
|
|
let idx = (i * image_data.len() / dim) % image_data.len();
|
|
embedding.push(image_data[idx] as f32 / 255.0);
|
|
}
|
|
|
|
// Normalize
|
|
let norm: f32 = embedding.iter().map(|x| x * x).sum::<f32>().sqrt();
|
|
embedding.iter_mut().for_each(|x| *x /= norm);
|
|
|
|
embedding
|
|
}
|
|
|
|
fn create_embedding_cache(size: usize) -> HashMap<String, Embedding> {
|
|
let mut cache = HashMap::new();
|
|
for i in 0..size {
|
|
let embedding = generate_random_embedding(512);
|
|
cache.insert(format!("image_{}", i), embedding);
|
|
}
|
|
cache
|
|
}
|
|
|
|
fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
|
|
let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
|
|
let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
|
|
let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
|
|
|
|
if norm_a > 0.0 && norm_b > 0.0 {
|
|
dot / (norm_a * norm_b)
|
|
} else {
|
|
0.0
|
|
}
|
|
}
|
|
|
|
fn linear_similarity_search(
|
|
cache: &HashMap<String, Embedding>,
|
|
query: &Embedding,
|
|
top_k: usize,
|
|
) -> Vec<(String, f32, Embedding)> {
|
|
let mut results: Vec<_> = cache
|
|
.iter()
|
|
.map(|(key, embedding)| {
|
|
let similarity = cosine_similarity(query, embedding);
|
|
(key.clone(), similarity, embedding.clone())
|
|
})
|
|
.collect();
|
|
|
|
results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
|
|
results.truncate(top_k);
|
|
results
|
|
}
|
|
|
|
fn ann_similarity_search(
|
|
cache: &HashMap<String, Embedding>,
|
|
query: &Embedding,
|
|
top_k: usize,
|
|
) -> Vec<(String, f32, Embedding)> {
|
|
// Simplified ANN using random sampling
|
|
let sample_size = (cache.len() / 10).max(100).min(cache.len());
|
|
let mut results: Vec<_> = cache
|
|
.iter()
|
|
.enumerate()
|
|
.filter(|(i, _)| i % (cache.len() / sample_size.max(1)) == 0)
|
|
.map(|(_, (key, embedding))| {
|
|
let similarity = cosine_similarity(query, embedding);
|
|
(key.clone(), similarity, embedding.clone())
|
|
})
|
|
.collect();
|
|
|
|
results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
|
|
results.truncate(top_k);
|
|
results
|
|
}
|
|
|
|
fn find_exact_match(cache: &HashMap<String, Embedding>, query: &Embedding) -> Option<String> {
|
|
cache.iter().find_map(|(key, embedding)| {
|
|
if embedding.len() == query.len()
|
|
&& embedding
|
|
.iter()
|
|
.zip(query.iter())
|
|
.all(|(a, b)| (a - b).abs() < 1e-6)
|
|
{
|
|
Some(key.clone())
|
|
} else {
|
|
None
|
|
}
|
|
})
|
|
}
|
|
|
|
fn find_by_similarity_threshold(
|
|
cache: &HashMap<String, Embedding>,
|
|
query: &Embedding,
|
|
threshold: f32,
|
|
) -> Option<(String, f32)> {
|
|
cache
|
|
.iter()
|
|
.filter_map(|(key, embedding)| {
|
|
let similarity = cosine_similarity(query, embedding);
|
|
if similarity >= threshold {
|
|
Some((key.clone(), similarity))
|
|
} else {
|
|
None
|
|
}
|
|
})
|
|
.max_by(|a, b| a.1.partial_cmp(&b.1).unwrap())
|
|
}
|
|
|
|
fn process_new_image(_image_data: &[u8]) -> String {
|
|
// Simulate OCR processing
|
|
std::thread::sleep(Duration::from_millis(50));
|
|
"processed_result".to_string()
|
|
}
|
|
|
|
struct CacheStatistics {
|
|
size: usize,
|
|
avg_embedding_norm: f32,
|
|
memory_bytes: usize,
|
|
}
|
|
|
|
fn compute_cache_statistics(cache: &HashMap<String, Embedding>) -> CacheStatistics {
|
|
let size = cache.len();
|
|
let avg_norm = if size > 0 {
|
|
let total_norm: f32 = cache
|
|
.values()
|
|
.map(|emb| emb.iter().map(|x| x * x).sum::<f32>().sqrt())
|
|
.sum();
|
|
total_norm / size as f32
|
|
} else {
|
|
0.0
|
|
};
|
|
|
|
let memory_bytes = estimate_cache_memory(cache);
|
|
|
|
CacheStatistics {
|
|
size,
|
|
avg_embedding_norm: avg_norm,
|
|
memory_bytes,
|
|
}
|
|
}
|
|
|
|
fn estimate_cache_memory(cache: &HashMap<String, Embedding>) -> usize {
|
|
let key_bytes: usize = cache.keys().map(|k| k.len()).sum();
|
|
let embedding_bytes: usize = cache.values().map(|e| e.len() * 4).sum();
|
|
key_bytes + embedding_bytes + cache.len() * 64 // HashMap overhead
|
|
}
|
|
|
|
criterion_group!(
|
|
benches,
|
|
bench_embedding_generation,
|
|
bench_similarity_search,
|
|
bench_cache_hit_latency,
|
|
bench_cache_miss_latency,
|
|
bench_cache_insertion,
|
|
bench_cache_updates,
|
|
bench_batch_cache_ops,
|
|
bench_cache_statistics
|
|
);
|
|
criterion_main!(benches);
|