//! Adaptive search strategy selection and timing metrics. //! //! Automatically selects the best search strategy based on collection size //! and available hardware (SIMD via clawhdf5_accel, rayon parallelism, GPU). use std::time::Instant; use crate::vector_search; /// Search strategy selection based on collection size and hardware. /// /// ```text /// < 1K: Scalar (overhead of SIMD dispatch not worth it) /// 1K-100K: Accelerate (AMX/cblas_sgemv) > BLAS (matrixmultiply) > SIMD prenorm /// 1K-10K: SIMD brute force with pre-computed norms (fallback) /// 10K-50K: Rayon parallel SIMD (if available) OR GPU (if available) /// 50K-500K: GPU (if available) OR IVF-PQ /// > 100K: IVF-PQ always (regardless of BLAS/GPU) /// ``` #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum SearchStrategy { /// Plain scalar search for tiny collections (< 1K). Scalar, /// SIMD brute force with pre-computed norms (1K-10K). SimdBruteForce, /// BLAS batch matrix-vector multiply (1K-100K, requires `fast-math` feature). Blas, /// Apple Accelerate / OpenBLAS cblas_sgemv (1K-100K, requires `accelerate`/`openblas`). Accelerate, /// Rayon parallel SIMD search (10K-50K, requires `parallel` feature). RayonParallel, /// GPU-accelerated search (10K-500K, requires `gpu` feature). Gpu, /// IVF-PQ approximate search for large collections. IvfPq, } impl std::fmt::Display for SearchStrategy { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { SearchStrategy::Scalar => write!(f, "scalar"), SearchStrategy::SimdBruteForce => write!(f, "simd"), SearchStrategy::Blas => write!(f, "blas"), SearchStrategy::Accelerate => write!(f, "accelerate"), SearchStrategy::RayonParallel => write!(f, "rayon"), SearchStrategy::Gpu => write!(f, "gpu"), SearchStrategy::IvfPq => write!(f, "ivf-pq"), } } } /// Metrics collected during a search operation. #[derive(Debug, Clone)] pub struct SearchMetrics { /// Strategy used for this search ("scalar", "simd", "rayon", "gpu", "ivf-pq"). pub strategy: String, /// Total search time in microseconds. pub search_time_us: u64, /// Number of candidate vectors scanned. pub candidates_scanned: usize, /// Re-ranking time in microseconds (for IVF-PQ). pub rerank_time_us: Option, /// Active SIMD/GPU backend (e.g., "neon", "avx2", "avx512", "gpu-metal"). pub backend: String, } /// Configuration flags for strategy selection. #[derive(Debug, Clone, Copy)] pub struct HardwareCapabilities { /// Whether the `parallel` feature is enabled and rayon is available. pub rayon_available: bool, /// Whether the `gpu` feature is enabled and GPU hardware is detected. pub gpu_available: bool, /// Whether the `fast-math` feature is enabled (BLAS batch matmul). pub blas_available: bool, /// Whether the `accelerate` or `openblas` feature is enabled (cblas_sgemv). pub accelerate_available: bool, } impl HardwareCapabilities { /// Detect available hardware capabilities at runtime. pub fn detect() -> Self { Self { rayon_available: cfg!(feature = "parallel"), gpu_available: { #[cfg(feature = "gpu")] { clawhdf5_gpu::GpuAccelerator::is_available() } #[cfg(not(feature = "gpu"))] { false } }, blas_available: cfg!(feature = "fast-math"), accelerate_available: cfg!(any(feature = "accelerate", feature = "openblas")), } } } /// Return the name of the active SIMD/acceleration backend. pub fn active_backend_name(gpu_active: bool) -> String { if gpu_active { return "gpu".to_owned(); } let backend = clawhdf5_accel::detect_backend(); format!("{backend:?}").to_lowercase() } /// Auto-select the best search strategy based on collection size and hardware. /// /// Updated hierarchy with Accelerate/BLAS support: /// ```text /// < 1K: Scalar /// 1K-100K: Accelerate (AMX sgemv) > BLAS (matrixmultiply) > Rayon > GPU > SIMD /// > 500K: IVF-PQ always /// ``` pub fn auto_select_strategy(num_vectors: usize, hw: &HardwareCapabilities) -> SearchStrategy { if num_vectors > 500_000 { return SearchStrategy::IvfPq; } if num_vectors > 50_000 { if hw.accelerate_available { return SearchStrategy::Accelerate; } if hw.blas_available { return SearchStrategy::Blas; } if hw.gpu_available { return SearchStrategy::Gpu; } return SearchStrategy::IvfPq; } if num_vectors > 10_000 { if hw.accelerate_available { return SearchStrategy::Accelerate; } if hw.blas_available { return SearchStrategy::Blas; } if hw.rayon_available { return SearchStrategy::RayonParallel; } if hw.gpu_available { return SearchStrategy::Gpu; } return SearchStrategy::SimdBruteForce; } if num_vectors >= 1_000 { if hw.accelerate_available { return SearchStrategy::Accelerate; } if hw.blas_available { return SearchStrategy::Blas; } return SearchStrategy::SimdBruteForce; } SearchStrategy::Scalar } /// Execute a search using the given strategy and return results with metrics. /// /// This dispatches to the appropriate search implementation based on the /// selected strategy. For IVF-PQ, an index must be provided externally /// (this function uses brute-force fallback if no IVF-PQ index is available). /// /// `vectors_flat` is `vectors` flattened into one contiguous `[N × dim]` /// row-major buffer (e.g. `MemoryCache::embeddings_flat`, maintained /// incrementally alongside `vectors`). It's only consulted by the /// `Blas`/`Accelerate` strategies, which otherwise re-flatten the whole /// corpus on every call — passing the already-flat buffer skips that copy. #[allow(clippy::too_many_arguments)] pub fn search_with_metrics( query: &[f32], vectors: &[Vec], vectors_flat: &[f32], norms: &[f32], tombstones: &[u8], k: usize, strategy: SearchStrategy, #[cfg(feature = "gpu")] gpu_backend: Option<&crate::gpu_search::GpuSearchBackend>, #[cfg(not(feature = "gpu"))] _gpu_backend: Option<&()>, ) -> (Vec<(usize, f32)>, SearchMetrics) { // Only read by the Blas/Accelerate arms below, which are themselves // feature-gated — reference it unconditionally so a build with neither // feature enabled doesn't warn about an unused parameter. let _ = vectors_flat; let start = Instant::now(); let active_count = tombstones.iter().filter(|&&t| t == 0).count(); let gpu_active; let results = match strategy { SearchStrategy::Scalar => { gpu_active = false; scalar_search(query, vectors, tombstones, k) } SearchStrategy::SimdBruteForce => { gpu_active = false; let all = vector_search::cosine_similarity_batch_prenorm(query, vectors, norms, tombstones); vector_search::top_k(all, k) } SearchStrategy::Blas => { gpu_active = false; #[cfg(feature = "fast-math")] { crate::blas_search::blas_cosine_batch_flat( query, vectors_flat, norms, tombstones, query.len(), k, ) } #[cfg(not(feature = "fast-math"))] { let all = vector_search::cosine_similarity_batch_prenorm( query, vectors, norms, tombstones, ); vector_search::top_k(all, k) } } SearchStrategy::Accelerate => { gpu_active = false; #[cfg(any(feature = "accelerate", feature = "openblas"))] { crate::accelerate_search::accelerate_cosine_batch( query, vectors_flat, norms, tombstones, query.len(), k, ) } #[cfg(not(any(feature = "accelerate", feature = "openblas")))] { let all = vector_search::cosine_similarity_batch_prenorm( query, vectors, norms, tombstones, ); vector_search::top_k(all, k) } } SearchStrategy::RayonParallel => { gpu_active = false; #[cfg(feature = "parallel")] { vector_search::parallel_cosine_batch_prenorm(query, vectors, norms, tombstones, k) } #[cfg(not(feature = "parallel"))] { let all = vector_search::cosine_similarity_batch_prenorm( query, vectors, norms, tombstones, ); vector_search::top_k(all, k) } } SearchStrategy::Gpu => { #[cfg(feature = "gpu")] { if let Some(backend) = gpu_backend { gpu_active = backend.is_available(); backend.search_cosine(query, vectors, norms, tombstones, k) } else { gpu_active = false; let all = vector_search::cosine_similarity_batch_prenorm( query, vectors, norms, tombstones, ); vector_search::top_k(all, k) } } #[cfg(not(feature = "gpu"))] { gpu_active = false; let all = vector_search::cosine_similarity_batch_prenorm( query, vectors, norms, tombstones, ); vector_search::top_k(all, k) } } SearchStrategy::IvfPq => { gpu_active = false; // IVF-PQ requires an external index; fall back to prenorm brute force // when called through this generic interface. let all = vector_search::cosine_similarity_batch_prenorm(query, vectors, norms, tombstones); vector_search::top_k(all, k) } }; let elapsed = start.elapsed(); let metrics = SearchMetrics { strategy: strategy.to_string(), search_time_us: elapsed.as_micros() as u64, candidates_scanned: active_count, rerank_time_us: None, backend: active_backend_name(gpu_active), }; (results, metrics) } /// Plain scalar cosine similarity for very small collections. fn scalar_search( query: &[f32], vectors: &[Vec], tombstones: &[u8], k: usize, ) -> Vec<(usize, f32)> { let query_norm = clawhdf5_accel::vector_norm(query); if query_norm == 0.0 { return Vec::new(); } let mut results: Vec<(usize, f32)> = Vec::with_capacity(vectors.len()); for (i, vec) in vectors.iter().enumerate() { if i < tombstones.len() && tombstones[i] != 0 { continue; } // Use clawhdf5_accel even for scalar strategy — it's always available and // the "scalar" name refers to the strategy tier, not the implementation. let vec_norm = clawhdf5_accel::vector_norm(vec); let score = crate::cosine_similarity_prenorm(query, query_norm, vec, vec_norm); results.push((i, score)); } results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); results.truncate(k); results } #[cfg(test)] mod tests { use super::*; fn make_vectors(n: usize, dim: usize, seed: u32) -> Vec> { let mut s = seed; let mut next = || -> f32 { s = s.wrapping_mul(1103515245).wrapping_add(12345); ((s >> 16) as f32) / 65536.0 - 0.5 }; (0..n).map(|_| (0..dim).map(|_| next()).collect()).collect() } fn flatten(vectors: &[Vec]) -> Vec { vectors.iter().flatten().copied().collect() } // --- auto_select_strategy tests --- #[test] fn strategy_scalar_under_1k() { let hw = HardwareCapabilities { rayon_available: false, gpu_available: false, blas_available: false, accelerate_available: false, }; assert_eq!(auto_select_strategy(0, &hw), SearchStrategy::Scalar); assert_eq!(auto_select_strategy(500, &hw), SearchStrategy::Scalar); assert_eq!(auto_select_strategy(999, &hw), SearchStrategy::Scalar); } #[test] fn strategy_simd_1k_to_10k() { let hw = HardwareCapabilities { rayon_available: false, gpu_available: false, blas_available: false, accelerate_available: false, }; assert_eq!( auto_select_strategy(1_000, &hw), SearchStrategy::SimdBruteForce ); assert_eq!( auto_select_strategy(5_000, &hw), SearchStrategy::SimdBruteForce ); assert_eq!( auto_select_strategy(10_000, &hw), SearchStrategy::SimdBruteForce ); } #[test] fn strategy_rayon_10k_to_50k_when_available() { let hw = HardwareCapabilities { rayon_available: true, gpu_available: false, blas_available: false, accelerate_available: false, }; assert_eq!( auto_select_strategy(10_001, &hw), SearchStrategy::RayonParallel ); assert_eq!( auto_select_strategy(30_000, &hw), SearchStrategy::RayonParallel ); assert_eq!( auto_select_strategy(50_000, &hw), SearchStrategy::RayonParallel ); } #[test] fn strategy_gpu_10k_to_50k_when_no_rayon() { let hw = HardwareCapabilities { rayon_available: false, gpu_available: true, blas_available: false, accelerate_available: false, }; assert_eq!(auto_select_strategy(10_001, &hw), SearchStrategy::Gpu); assert_eq!(auto_select_strategy(50_000, &hw), SearchStrategy::Gpu); } #[test] fn strategy_gpu_50k_to_500k() { let hw = HardwareCapabilities { rayon_available: true, gpu_available: true, blas_available: false, accelerate_available: false, }; assert_eq!(auto_select_strategy(50_001, &hw), SearchStrategy::Gpu); assert_eq!(auto_select_strategy(200_000, &hw), SearchStrategy::Gpu); assert_eq!(auto_select_strategy(500_000, &hw), SearchStrategy::Gpu); } #[test] fn strategy_ivfpq_over_500k() { let hw = HardwareCapabilities { rayon_available: true, gpu_available: true, blas_available: true, accelerate_available: true, }; assert_eq!(auto_select_strategy(500_001, &hw), SearchStrategy::IvfPq); assert_eq!(auto_select_strategy(1_000_000, &hw), SearchStrategy::IvfPq); } #[test] fn strategy_ivfpq_fallback_50k_no_gpu() { let hw = HardwareCapabilities { rayon_available: false, gpu_available: false, blas_available: false, accelerate_available: false, }; assert_eq!(auto_select_strategy(50_001, &hw), SearchStrategy::IvfPq); } #[test] fn strategy_simd_fallback_10k_no_parallel_no_gpu() { let hw = HardwareCapabilities { rayon_available: false, gpu_available: false, blas_available: false, accelerate_available: false, }; assert_eq!( auto_select_strategy(15_000, &hw), SearchStrategy::SimdBruteForce ); } // --- SearchStrategy Display --- #[test] fn strategy_display_names() { assert_eq!(SearchStrategy::Scalar.to_string(), "scalar"); assert_eq!(SearchStrategy::SimdBruteForce.to_string(), "simd"); assert_eq!(SearchStrategy::Blas.to_string(), "blas"); assert_eq!(SearchStrategy::Accelerate.to_string(), "accelerate"); assert_eq!(SearchStrategy::RayonParallel.to_string(), "rayon"); assert_eq!(SearchStrategy::Gpu.to_string(), "gpu"); assert_eq!(SearchStrategy::IvfPq.to_string(), "ivf-pq"); } // --- SearchMetrics --- #[test] fn search_metrics_strategy_name() { let metrics = SearchMetrics { strategy: "simd".to_owned(), search_time_us: 100, candidates_scanned: 1000, rerank_time_us: None, backend: "neon".to_owned(), }; assert_eq!(metrics.strategy, "simd"); assert_eq!(metrics.candidates_scanned, 1000); assert_eq!(metrics.backend, "neon"); } // --- search_with_metrics --- #[test] fn search_with_metrics_scalar() { let vectors = make_vectors(50, 16, 42); let norms: Vec = vectors .iter() .map(|v| clawhdf5_accel::vector_norm(v)) .collect(); let tombstones = vec![0u8; 50]; let query = vectors[0].clone(); let (results, metrics) = search_with_metrics( &query, &vectors, &flatten(&vectors), &norms, &tombstones, 5, SearchStrategy::Scalar, None, ); assert_eq!(results.len(), 5); assert_eq!(metrics.strategy, "scalar"); assert!(metrics.search_time_us > 0 || metrics.candidates_scanned > 0); assert_eq!(metrics.candidates_scanned, 50); assert!(metrics.rerank_time_us.is_none()); assert!(!metrics.backend.is_empty()); // First result should be the query itself assert_eq!(results[0].0, 0); } #[test] fn search_with_metrics_simd() { let vectors = make_vectors(100, 32, 42); let norms: Vec = vectors .iter() .map(|v| clawhdf5_accel::vector_norm(v)) .collect(); let tombstones = vec![0u8; 100]; let query = vectors[0].clone(); let (results, metrics) = search_with_metrics( &query, &vectors, &flatten(&vectors), &norms, &tombstones, 10, SearchStrategy::SimdBruteForce, None, ); assert_eq!(metrics.strategy, "simd"); assert!(!results.is_empty()); assert_eq!(results[0].0, 0); } #[test] fn search_with_metrics_timing_nonzero() { let vectors = make_vectors(1000, 64, 42); let norms: Vec = vectors .iter() .map(|v| clawhdf5_accel::vector_norm(v)) .collect(); let tombstones = vec![0u8; 1000]; let query = vectors[0].clone(); let (_, metrics) = search_with_metrics( &query, &vectors, &flatten(&vectors), &norms, &tombstones, 10, SearchStrategy::SimdBruteForce, None, ); // With 1000 vectors, search should take > 0 microseconds assert!(metrics.search_time_us < 1_000_000); // under 1 second assert_eq!(metrics.candidates_scanned, 1000); } #[test] fn search_with_metrics_results_match_direct() { let vectors = make_vectors(200, 32, 42); let norms: Vec = vectors .iter() .map(|v| clawhdf5_accel::vector_norm(v)) .collect(); let tombstones = vec![0u8; 200]; let query = vectors[3].clone(); let (results, _) = search_with_metrics( &query, &vectors, &flatten(&vectors), &norms, &tombstones, 10, SearchStrategy::SimdBruteForce, None, ); let direct = vector_search::cosine_similarity_batch_prenorm(&query, &vectors, &norms, &tombstones); let direct_top = vector_search::top_k(direct, 10); assert_eq!(results.len(), direct_top.len()); for (r, d) in results.iter().zip(&direct_top) { assert_eq!(r.0, d.0); assert!((r.1 - d.1).abs() < 1e-6); } } #[test] fn search_with_metrics_respects_tombstones() { let vectors = make_vectors(100, 16, 42); let norms: Vec = vectors .iter() .map(|v| clawhdf5_accel::vector_norm(v)) .collect(); let mut tombstones = vec![0u8; 100]; tombstones[0] = 1; tombstones[1] = 1; let query = vectors[2].clone(); let (results, metrics) = search_with_metrics( &query, &vectors, &flatten(&vectors), &norms, &tombstones, 100, SearchStrategy::Scalar, None, ); assert!(results.iter().all(|r| r.0 != 0 && r.0 != 1)); assert_eq!(metrics.candidates_scanned, 98); } #[test] fn hardware_capabilities_detect() { let hw = HardwareCapabilities::detect(); // Just verify it doesn't panic and returns something let _ = hw.rayon_available; let _ = hw.gpu_available; } #[test] fn active_backend_name_returns_valid() { let name = active_backend_name(false); assert!(!name.is_empty()); // Should be one of the known backends let valid = ["neon", "avx2", "avx512", "sse4", "wasmsimd128", "scalar"]; assert!( valid.iter().any(|v| name.contains(v)), "unexpected backend: {name}" ); } #[test] fn search_metrics_has_backend_field() { let vectors = make_vectors(50, 16, 42); let norms: Vec = vectors .iter() .map(|v| clawhdf5_accel::vector_norm(v)) .collect(); let tombstones = vec![0u8; 50]; let query = vectors[0].clone(); let (_, metrics) = search_with_metrics( &query, &vectors, &flatten(&vectors), &norms, &tombstones, 5, SearchStrategy::Scalar, None, ); assert!(!metrics.backend.is_empty()); } // --- BLAS strategy selection tests --- #[test] fn strategy_blas_preferred_1k_to_100k() { let hw = HardwareCapabilities { rayon_available: true, gpu_available: true, blas_available: true, accelerate_available: false, }; // BLAS should be preferred over rayon/gpu/simd in the 1K-100K range assert_eq!(auto_select_strategy(1_000, &hw), SearchStrategy::Blas); assert_eq!(auto_select_strategy(5_000, &hw), SearchStrategy::Blas); assert_eq!(auto_select_strategy(10_001, &hw), SearchStrategy::Blas); assert_eq!(auto_select_strategy(50_000, &hw), SearchStrategy::Blas); assert_eq!(auto_select_strategy(100_000, &hw), SearchStrategy::Blas); } #[test] fn strategy_blas_not_for_small() { let hw = HardwareCapabilities { rayon_available: false, gpu_available: false, blas_available: true, accelerate_available: false, }; // Under 1K, still use scalar assert_eq!(auto_select_strategy(500, &hw), SearchStrategy::Scalar); assert_eq!(auto_select_strategy(999, &hw), SearchStrategy::Scalar); } #[test] fn strategy_fallback_without_blas() { let hw = HardwareCapabilities { rayon_available: false, gpu_available: false, blas_available: false, accelerate_available: false, }; // Without BLAS, falls back to SIMD/IVF-PQ assert_eq!( auto_select_strategy(5_000, &hw), SearchStrategy::SimdBruteForce ); assert_eq!(auto_select_strategy(50_001, &hw), SearchStrategy::IvfPq); } #[cfg(feature = "fast-math")] #[test] fn search_with_metrics_blas() { let vectors = make_vectors(200, 32, 42); let norms: Vec = vectors .iter() .map(|v| clawhdf5_accel::vector_norm(v)) .collect(); let tombstones = vec![0u8; 200]; let query = vectors[0].clone(); let (results, metrics) = search_with_metrics( &query, &vectors, &flatten(&vectors), &norms, &tombstones, 10, SearchStrategy::Blas, None, ); assert_eq!(metrics.strategy, "blas"); assert!(!results.is_empty()); assert_eq!(results[0].0, 0); } #[cfg(feature = "parallel")] #[test] fn search_with_metrics_rayon() { let vectors = make_vectors(500, 32, 42); let norms: Vec = vectors .iter() .map(|v| clawhdf5_accel::vector_norm(v)) .collect(); let tombstones = vec![0u8; 500]; let query = vectors[0].clone(); let (results, metrics) = search_with_metrics( &query, &vectors, &flatten(&vectors), &norms, &tombstones, 10, SearchStrategy::RayonParallel, None, ); assert_eq!(metrics.strategy, "rayon"); assert!(!results.is_empty()); assert_eq!(results[0].0, 0); } // --- Accelerate strategy selection tests --- #[test] fn strategy_accelerate_preferred_over_blas() { let hw = HardwareCapabilities { rayon_available: true, gpu_available: true, blas_available: true, accelerate_available: true, }; // Accelerate should be preferred over BLAS/rayon/gpu in the 1K-100K range assert_eq!(auto_select_strategy(1_000, &hw), SearchStrategy::Accelerate); assert_eq!(auto_select_strategy(5_000, &hw), SearchStrategy::Accelerate); assert_eq!( auto_select_strategy(10_001, &hw), SearchStrategy::Accelerate ); assert_eq!( auto_select_strategy(50_000, &hw), SearchStrategy::Accelerate ); assert_eq!( auto_select_strategy(100_000, &hw), SearchStrategy::Accelerate ); } #[test] fn strategy_accelerate_not_for_small() { let hw = HardwareCapabilities { rayon_available: false, gpu_available: false, blas_available: false, accelerate_available: true, }; // Under 1K, still use scalar assert_eq!(auto_select_strategy(500, &hw), SearchStrategy::Scalar); assert_eq!(auto_select_strategy(999, &hw), SearchStrategy::Scalar); } #[test] fn strategy_accelerate_not_for_huge() { let hw = HardwareCapabilities { rayon_available: true, gpu_available: true, blas_available: true, accelerate_available: true, }; // Over 500K, always IVF-PQ assert_eq!(auto_select_strategy(500_001, &hw), SearchStrategy::IvfPq); } #[cfg(any(feature = "accelerate", feature = "openblas"))] #[test] fn search_with_metrics_accelerate() { let vectors = make_vectors(200, 32, 42); let norms: Vec = vectors .iter() .map(|v| clawhdf5_accel::vector_norm(v)) .collect(); let tombstones = vec![0u8; 200]; let query = vectors[0].clone(); let (results, metrics) = search_with_metrics( &query, &vectors, &flatten(&vectors), &norms, &tombstones, 10, SearchStrategy::Accelerate, None, ); assert_eq!(metrics.strategy, "accelerate"); assert!(!results.is_empty()); assert_eq!(results[0].0, 0); } }