From ffcbd452e64550e236d9f08c5af6d0d8ced8ce29 Mon Sep 17 00:00:00 2001 From: pasha Date: Tue, 11 Aug 2026 18:34:23 +0300 Subject: [PATCH] perf(vector): parallelize large IVF scans --- hermes-core/src/segment/ann_disk.rs | 991 ++++++++++++++++++++++++++-- 1 file changed, 928 insertions(+), 63 deletions(-) diff --git a/hermes-core/src/segment/ann_disk.rs b/hermes-core/src/segment/ann_disk.rs index cb28dc3c..41adb7b4 100644 --- a/hermes-core/src/segment/ann_disk.rs +++ b/hermes-core/src/segment/ann_disk.rs @@ -55,6 +55,11 @@ const TQ_PRUNE_ESTIMATE_BOUND: f32 = 1.3; const TQ_PARALLEL_SCAN_MIN_VECTORS: usize = 65_536; #[cfg(feature = "native")] const TQ_PARALLEL_SCAN_CHUNK_BLOCKS: usize = 512; +/// IVF-TQ scans are already parallel across segments. Fan out inside one +/// segment only when the selected leaves contain enough postings to amortize +/// Rayon scheduling and worker-local top-k state. +const IVF_PARALLEL_SCAN_MIN_POSTINGS: usize = 65_536; +const IVF_TQ_PARALLEL_SCAN_CHUNK_BLOCKS: usize = 512; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum AnnKind { @@ -125,6 +130,21 @@ struct AnnRun { codes: Range, } +#[derive(Clone, Copy)] +struct IvfTqScanTask<'a> { + run: &'a AnnRun, + cluster_dot: f32, + first_block: usize, + last_block: usize, +} + +#[derive(Clone, Copy)] +struct BinaryScanTask<'a> { + run: &'a AnnRun, + first_index: usize, + count: usize, +} + /// Mmap-backed searchable ANN payload. Only the fixed-size run directory is /// heap-resident; all corpus-sized columns remain zero-copy file slices. pub(crate) struct AnnDiskIndex { @@ -508,6 +528,48 @@ impl AnnDiskIndex { &self.runs[start..end] } + #[cfg(feature = "native")] + fn ivf_tq_scan_tasks<'a>( + &'a self, + plan: &'a crate::structures::TqIvfQueryPlan, + block_bytes: usize, + chunk_blocks: usize, + ) -> Vec> { + debug_assert!(chunk_blocks > 0); + let mut tasks = Vec::new(); + for (cluster_id, cluster_dot) in plan.cluster_dots() { + for run in self.cluster_runs(cluster_id) { + let block_count = run.codes.len() / block_bytes; + for first_block in (0..block_count).step_by(chunk_blocks) { + tasks.push(IvfTqScanTask { + run, + cluster_dot, + first_block, + last_block: (first_block + chunk_blocks).min(block_count), + }); + } + } + } + tasks + } + + #[cfg(feature = "native")] + fn binary_scan_tasks<'a>(&'a self, cluster_ids: &[u32]) -> Vec> { + let mut tasks = Vec::new(); + for &cluster_id in cluster_ids { + for run in self.cluster_runs(cluster_id) { + for first_index in (0..run.count).step_by(BINARY_SCORE_BATCH) { + tasks.push(BinaryScanTask { + run, + first_index, + count: BINARY_SCORE_BATCH.min(run.count - first_index), + }); + } + } + } + tasks + } + /// Prefetch exactly the mmap ranges needed by the selected IVF leaves. /// /// A pure-copy merge preserves each source payload as one physical extent, @@ -620,6 +682,23 @@ impl AnnDiskIndex { k: usize, plan: &crate::structures::TqIvfQueryPlan, combiner: crate::query::MultiValueCombiner, + ) -> io::Result> { + self.search_ivf_tq_combined_documents_with_tuning( + k, + plan, + combiner, + IVF_PARALLEL_SCAN_MIN_POSTINGS, + IVF_TQ_PARALLEL_SCAN_CHUNK_BLOCKS, + ) + } + + fn search_ivf_tq_combined_documents_with_tuning( + &self, + k: usize, + plan: &crate::structures::TqIvfQueryPlan, + combiner: crate::query::MultiValueCombiner, + parallel_min_postings: usize, + parallel_chunk_blocks: usize, ) -> io::Result> { use crate::structures::vector::quantization::{ TQ_BLOCK_LANES, tq_ivf_block_bytes, tq_score_ivf_block, @@ -628,6 +707,8 @@ impl AnnDiskIndex { validate_combined_search(combiner)?; let tq_plan = plan.tq_plan(); self.validate_ivf_tq_query_plan(plan)?; + #[cfg(not(feature = "native"))] + let _ = (parallel_min_postings, parallel_chunk_blocks); if k == 0 { return Ok(Vec::new()); } @@ -636,9 +717,37 @@ impl AnnDiskIndex { let block_bytes = tq_ivf_block_bytes(self.header.code_size); let bytes = self.raw.as_slice(); + let posting_count = probed_posting_count(self, &plan.cluster_ids)?; + + #[cfg(feature = "native")] + if rayon::current_num_threads() > 1 && posting_count >= parallel_min_postings { + use rayon::prelude::*; + let tasks = self.ivf_tq_scan_tasks(plan, block_bytes, parallel_chunk_blocks); + let ordinal_scores = tasks + .par_iter() + .try_fold( + || Vec::with_capacity(parallel_chunk_blocks * TQ_BLOCK_LANES), + |mut ordinal_scores, task| { + score_ivf_tq_combined_blocks( + bytes, + tq_plan, + block_bytes, + *task, + &mut ordinal_scores, + )?; + Ok::<_, io::Error>(ordinal_scores) + }, + ) + .try_reduce(Vec::new, |mut left, mut right| { + left.append(&mut right); + Ok(left) + })?; + return Ok(combine_scored_ordinals(ordinal_scores, k, combiner)); + } + let mut ordinal_scores = Vec::new(); ordinal_scores - .try_reserve_exact(probed_posting_count(self, &plan.cluster_ids)?) + .try_reserve_exact(posting_count) .map_err(|_| invalid_data("IVF-TQ combined score buffer allocation failed"))?; let mut scores = [0.0f32; TQ_BLOCK_LANES]; for (cluster_id, cluster_dot) in plan.cluster_dots() { @@ -679,8 +788,27 @@ impl AnnDiskIndex { query: &[u8], cluster_ids: &[u32], combiner: crate::query::MultiValueCombiner, + ) -> io::Result { + self.search_binary_combined_documents_with_tuning( + k, + query, + cluster_ids, + combiner, + IVF_PARALLEL_SCAN_MIN_POSTINGS, + ) + } + + fn search_binary_combined_documents_with_tuning( + &self, + k: usize, + query: &[u8], + cluster_ids: &[u32], + combiner: crate::query::MultiValueCombiner, + parallel_min_postings: usize, ) -> io::Result { validate_combined_search(combiner)?; + #[cfg(not(feature = "native"))] + let _ = parallel_min_postings; if query.len() != self.header.code_size { return Err(io::Error::new( io::ErrorKind::InvalidInput, @@ -694,10 +822,51 @@ impl AnnDiskIndex { self.prefetch_cluster_runs(cluster_ids); let bytes = self.raw.as_slice(); + let posting_count = probed_posting_count(self, cluster_ids)?; + #[cfg(feature = "native")] + if rayon::current_num_threads() > 1 && posting_count >= parallel_min_postings { + use rayon::prelude::*; + let tasks = self.binary_scan_tasks(cluster_ids); + let (ordinal_scores, _) = tasks + .par_iter() + .try_fold( + || { + ( + Vec::with_capacity(BINARY_SCORE_BATCH), + vec![0.0f32; BINARY_SCORE_BATCH], + ) + }, + |(mut ordinal_scores, mut scores), task| { + score_binary_task( + bytes, + query, + self.header.dim, + self.header.code_size, + *task, + &mut scores, + &mut ordinal_scores, + )?; + Ok::<_, io::Error>((ordinal_scores, scores)) + }, + ) + .try_reduce( + || (Vec::new(), Vec::new()), + |(mut left, scores), (mut right, _)| { + left.append(&mut right); + Ok((left, scores)) + }, + )?; + return Ok(combine_scored_ordinals_retaining( + ordinal_scores, + k, + combiner, + )); + } + let mut score_batch = vec![0.0f32; BINARY_SCORE_BATCH.min(self.header.vector_count)]; let mut ordinal_scores = Vec::new(); ordinal_scores - .try_reserve_exact(probed_posting_count(self, cluster_ids)?) + .try_reserve_exact(posting_count) .map_err(|_| invalid_data("binary combined score buffer allocation failed"))?; score_binary_cluster_runs( self, @@ -820,53 +989,121 @@ impl AnnDiskIndex { k: usize, plan: &crate::structures::TqIvfQueryPlan, ) -> io::Result> { - use crate::structures::vector::quantization::{ - TQ_BLOCK_LANES, tq_ivf_block_bytes, tq_score_ivf_block, - }; + self.search_ivf_tq_distinct_with_tuning( + k, + plan, + IVF_PARALLEL_SCAN_MIN_POSTINGS, + IVF_TQ_PARALLEL_SCAN_CHUNK_BLOCKS, + ) + } + + fn search_ivf_tq_distinct_with_tuning( + &self, + k: usize, + plan: &crate::structures::TqIvfQueryPlan, + parallel_min_postings: usize, + parallel_chunk_blocks: usize, + ) -> io::Result> { + use crate::structures::vector::quantization::tq_ivf_block_bytes; let tq_plan = plan.tq_plan(); self.validate_ivf_tq_query_plan(plan)?; + #[cfg(not(feature = "native"))] + let _ = (parallel_min_postings, parallel_chunk_blocks); + if k == 0 { + return Ok(Vec::new()); + } #[cfg(feature = "native")] self.prefetch_cluster_runs(&plan.cluster_ids); let block_bytes = tq_ivf_block_bytes(self.header.code_size); let bytes = self.raw.as_slice(); + + // Score one chunk first so every worker starts with a top-k backed by + // real documents. Its k-th score is therefore a safe pruning floor. + // Chunking also exposes parallelism when a skewed leaf is one large + // physical run. Nested Rayon stays on the caller's bounded search pool. + #[cfg(feature = "native")] + if rayon::current_num_threads() > 1 + && probed_posting_count(self, &plan.cluster_ids)? >= parallel_min_postings + { + use rayon::prelude::*; + let tasks = self.ivf_tq_scan_tasks(plan, block_bytes, parallel_chunk_blocks); + + if let Some((pilot, remaining)) = tasks.split_first() { + let mut pilot_collector = BoundedAnnCollector::::new(k); + let (pilot_pruned, pilot_scored) = + score_ivf_tq_blocks(bytes, tq_plan, block_bytes, *pilot, &mut pilot_collector)?; + let seed = pilot_collector.into_sorted_results(); + let seeded_collector = || { + let mut collector = BoundedAnnCollector::::new(k); + for &(doc_id, ordinal, score) in &seed { + collector.insert(doc_id, ordinal, score); + } + collector + }; + let (collector, pruned_blocks, scored_blocks) = remaining + .par_iter() + .try_fold( + || (seeded_collector(), 0usize, 0usize), + |(mut collector, pruned, scored), task| { + let (task_pruned, task_scored) = score_ivf_tq_blocks( + bytes, + tq_plan, + block_bytes, + *task, + &mut collector, + )?; + Ok::<_, io::Error>(( + collector, + pruned + task_pruned, + scored + task_scored, + )) + }, + ) + .try_reduce( + || (seeded_collector(), 0usize, 0usize), + |(mut collector, left_pruned, left_scored), + (partial, right_pruned, right_scored)| { + collector.merge_from(partial); + Ok(( + collector, + left_pruned + right_pruned, + left_scored + right_scored, + )) + }, + )?; + log_ivf_tq_pruning( + pilot_pruned + pruned_blocks, + pilot_scored + scored_blocks, + true, + ); + return Ok(collector.into_sorted_results()); + } + } + let mut collector = BoundedAnnCollector::::new(k); - let mut scores = [0.0f32; TQ_BLOCK_LANES]; let mut pruned_blocks = 0usize; let mut scored_blocks = 0usize; for (cluster_id, cluster_dot) in plan.cluster_dots() { for run in self.cluster_runs(cluster_id) { - let codes = &bytes[run.codes.clone()]; - let block_count = codes.len() / block_bytes; - for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() { - if let Some(threshold) = collector.pruning_threshold() { - let max_scale = tq_ivf_block_max_scale(block); - if cluster_dot + max_scale * TQ_PRUNE_ESTIMATE_BOUND <= threshold { - pruned_blocks += block_count - block_index; - break; - } - } - scored_blocks += 1; - tq_score_ivf_block(tq_plan, block, cluster_dot, &mut scores); - let lane_base = block_index * TQ_BLOCK_LANES; - let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base)); - for (lane, &score) in scores.iter().enumerate().take(lanes) { - let index = lane_base + lane; - collector.insert( - run_doc_id(bytes, run, index)?, - read_u16(bytes, run.ordinals.start + index * 2), - score, - ); - } - } + let block_count = run.codes.len() / block_bytes; + let (run_pruned, run_scored) = score_ivf_tq_blocks( + bytes, + tq_plan, + block_bytes, + IvfTqScanTask { + run, + cluster_dot, + first_block: 0, + last_block: block_count, + }, + &mut collector, + )?; + pruned_blocks += run_pruned; + scored_blocks += run_scored; } } - if pruned_blocks > 0 { - log::debug!( - "[search_ivf_tq] pruned {pruned_blocks} of {} blocks via scale bounds", - pruned_blocks + scored_blocks, - ); - } + log_ivf_tq_pruning(pruned_blocks, scored_blocks, false); Ok(collector.into_sorted_results()) } @@ -899,15 +1136,87 @@ impl AnnDiskIndex { k: usize, cluster_ids: &[u32], ) -> io::Result> { + self.search_binary_clusters_with_tuning::( + query, + k, + cluster_ids, + IVF_PARALLEL_SCAN_MIN_POSTINGS, + ) + } + + fn search_binary_clusters_with_tuning( + &self, + query: &[u8], + k: usize, + cluster_ids: &[u32], + parallel_min_postings: usize, + ) -> io::Result> { + #[cfg(not(feature = "native"))] + let _ = parallel_min_postings; if query.len() != self.header.code_size { return Err(io::Error::new( io::ErrorKind::InvalidInput, "binary ANN query has the wrong byte length", )); } + if k == 0 { + return Ok(Vec::new()); + } #[cfg(feature = "native")] self.prefetch_cluster_runs(cluster_ids); let bytes = self.raw.as_slice(); + + // A probe plan returns each leaf once and every indexed vector belongs + // to exactly one leaf, so `(doc_id, ordinal)` keys are unique in the + // common single-value path. + debug_assert!( + BY_DOCUMENT || { + let mut seen = rustc_hash::FxHashSet::default(); + cluster_ids + .iter() + .all(|cluster_id| seen.insert(*cluster_id)) + }, + "an IVF probe plan must not repeat a cluster", + ); + + #[cfg(feature = "native")] + if rayon::current_num_threads() > 1 + && probed_posting_count(self, cluster_ids)? >= parallel_min_postings + { + use rayon::prelude::*; + let tasks = self.binary_scan_tasks(cluster_ids); + let (collector, _) = tasks + .par_iter() + .try_fold( + || { + ( + BoundedAnnCollector::::new(k), + vec![0.0f32; BINARY_SCORE_BATCH], + ) + }, + |(mut collector, mut scores), task| { + score_binary_task( + bytes, + query, + self.header.dim, + self.header.code_size, + *task, + &mut scores, + &mut collector, + )?; + Ok::<_, io::Error>((collector, scores)) + }, + ) + .try_reduce( + || (BoundedAnnCollector::::new(k), Vec::new()), + |(mut collector, scores), (partial, _)| { + collector.merge_from(partial); + Ok((collector, scores)) + }, + )?; + return Ok(collector.into_sorted_results()); + } + let mut scores = vec![0.0f32; BINARY_SCORE_BATCH.min(self.header.vector_count)]; if BY_DOCUMENT { let mut collector = BoundedAnnCollector::::new(k); @@ -922,24 +1231,95 @@ impl AnnDiskIndex { return Ok(collector.into_sorted_results()); } - // A probe plan returns each leaf once and every indexed vector belongs - // to exactly one leaf, so `(doc_id, ordinal)` keys are unique here. - // Avoid the deduplication hash map in the common single-value path. - debug_assert!( - { - let mut seen = rustc_hash::FxHashSet::default(); - cluster_ids - .iter() - .all(|cluster_id| seen.insert(*cluster_id)) - }, - "an IVF probe plan must not repeat a cluster", - ); + // Avoid the deduplication hash map in the serial single-value path. let mut collector = BoundedUniqueAnnCollector::::new(k); score_binary_cluster_runs(self, bytes, query, cluster_ids, &mut scores, &mut collector)?; Ok(collector.into_sorted_results()) } } +/// Score one contiguous block range from an IVF-TQ run. A task may start in +/// the middle of a run because residual scales descend across the complete +/// run: the first block is still an upper bound for the remainder of the task. +fn score_ivf_tq_blocks( + bytes: &[u8], + plan: &crate::structures::TqQueryPlan, + block_bytes: usize, + task: IvfTqScanTask<'_>, + collector: &mut BoundedAnnCollector, +) -> io::Result<(usize, usize)> { + use crate::structures::vector::quantization::{TQ_BLOCK_LANES, tq_score_ivf_block}; + + let run = task.run; + let codes = &bytes[run.codes.clone()]; + let mut scores = [0.0f32; TQ_BLOCK_LANES]; + let mut scored_blocks = 0usize; + for block_index in task.first_block..task.last_block { + let block = &codes[block_index * block_bytes..][..block_bytes]; + if let Some(threshold) = collector.pruning_threshold() + && task.cluster_dot + tq_ivf_block_max_scale(block) * TQ_PRUNE_ESTIMATE_BOUND + <= threshold + { + return Ok((task.last_block - block_index, scored_blocks)); + } + scored_blocks += 1; + tq_score_ivf_block(plan, block, task.cluster_dot, &mut scores); + let lane_base = block_index * TQ_BLOCK_LANES; + let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base)); + for (lane, &score) in scores.iter().enumerate().take(lanes) { + let index = lane_base + lane; + collector.insert( + run_doc_id(bytes, run, index)?, + read_u16(bytes, run.ordinals.start + index * 2), + score, + ); + } + } + Ok((0, scored_blocks)) +} + +#[cfg(feature = "native")] +fn score_ivf_tq_combined_blocks( + bytes: &[u8], + plan: &crate::structures::TqQueryPlan, + block_bytes: usize, + task: IvfTqScanTask<'_>, + ordinal_scores: &mut Vec<(u32, u16, f32)>, +) -> io::Result<()> { + use crate::structures::vector::quantization::{TQ_BLOCK_LANES, tq_score_ivf_block}; + + let run = task.run; + let codes = &bytes[run.codes.clone()]; + let mut scores = [0.0f32; TQ_BLOCK_LANES]; + for block_index in task.first_block..task.last_block { + let block = &codes[block_index * block_bytes..][..block_bytes]; + tq_score_ivf_block(plan, block, task.cluster_dot, &mut scores); + let lane_base = block_index * TQ_BLOCK_LANES; + let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base)); + for (lane, &score) in scores.iter().enumerate().take(lanes) { + if score.is_finite() { + let index = lane_base + lane; + ordinal_scores.push(( + run_doc_id(bytes, run, index)?, + read_u16(bytes, run.ordinals.start + index * 2), + score, + )); + } + } + } + Ok(()) +} + +fn log_ivf_tq_pruning(pruned_blocks: usize, scored_blocks: usize, parallel: bool) { + if pruned_blocks > 0 { + log::debug!( + "[search_ivf_tq] pruned {pruned_blocks} of {} blocks via scale bounds ({})", + pruned_blocks + scored_blocks, + if parallel { "parallel" } else { "serial" }, + ); + } +} + #[cfg(feature = "native")] fn coalesce_prefetch_ranges(ranges: &mut Vec>) { if ranges.len() < 2 { @@ -1035,24 +1415,48 @@ fn score_binary_run( collector: &mut impl AnnScoreSink, ) -> io::Result<()> { for batch_start in (0..run.count).step_by(BINARY_SCORE_BATCH) { - let batch_count = BINARY_SCORE_BATCH.min(run.count - batch_start); - let code_start = run.codes.start + batch_start * code_size; - let code_end = code_start + batch_count * code_size; - crate::structures::simd::batch_hamming_scores( + score_binary_task( + bytes, query, - &bytes[code_start..code_end], - code_size, dim_bits, - &mut scores[..batch_count], + code_size, + BinaryScanTask { + run, + first_index: batch_start, + count: BINARY_SCORE_BATCH.min(run.count - batch_start), + }, + scores, + collector, + )?; + } + Ok(()) +} + +fn score_binary_task( + bytes: &[u8], + query: &[u8], + dim_bits: usize, + code_size: usize, + task: BinaryScanTask<'_>, + scores: &mut [f32], + collector: &mut impl AnnScoreSink, +) -> io::Result<()> { + let code_start = task.run.codes.start + task.first_index * code_size; + let code_end = code_start + task.count * code_size; + crate::structures::simd::batch_hamming_scores( + query, + &bytes[code_start..code_end], + code_size, + dim_bits, + &mut scores[..task.count], + ); + for (batch_index, &score) in scores.iter().enumerate().take(task.count) { + let index = task.first_index + batch_index; + collector.insert_score( + run_doc_id(bytes, task.run, index)?, + read_u16(bytes, task.run.ordinals.start + index * 2), + score, ); - for (batch_index, &score) in scores.iter().enumerate().take(batch_count) { - let index = batch_start + batch_index; - collector.insert_score( - run_doc_id(bytes, run, index)?, - read_u16(bytes, run.ordinals.start + index * 2), - score, - ); - } } Ok(()) } @@ -1102,7 +1506,18 @@ fn combine_scored_ordinals_retaining( return (Vec::new(), Vec::new()); } scores.retain(|entry| entry.2.is_finite()); - scores.sort_unstable_by(|left, right| left.0.cmp(&right.0).then_with(|| left.1.cmp(&right.1))); + let by_document_ordinal = |left: &(u32, u16, f32), right: &(u32, u16, f32)| { + left.0.cmp(&right.0).then_with(|| left.1.cmp(&right.1)) + }; + #[cfg(feature = "native")] + if rayon::current_num_threads() > 1 && scores.len() >= IVF_PARALLEL_SCAN_MIN_POSTINGS { + use rayon::prelude::*; + scores.par_sort_unstable_by(by_document_ordinal); + } else { + scores.sort_unstable_by(by_document_ordinal); + } + #[cfg(not(feature = "native"))] + scores.sort_unstable_by(by_document_ordinal); // Compact duplicates in place. IVF-TQ SOAR copies can have slightly // different residual estimates; preserving the highest one matches the @@ -2622,6 +3037,16 @@ mod tests { .collect(); docs.sort_unstable(); assert_eq!(docs, [0, 1, 2, 3]); + let serial = merged + .search_binary_clusters_with_tuning::(&[0], 4, &[0, 1], usize::MAX) + .unwrap(); + let parallel = rayon::ThreadPoolBuilder::new() + .num_threads(4) + .build() + .unwrap() + .install(|| merged.search_binary_clusters_with_tuning::(&[0], 4, &[0, 1], 1)) + .unwrap(); + assert_eq!(parallel, serial, "parallel binary top-k changed results"); // A merged source's directory is cluster-sorted while its payload is // source-order. A later merge must follow physical offsets and still @@ -2710,6 +3135,17 @@ mod tests { let mut bytes = Vec::new(); write_built_runs(binary_header(6), &runs, &mut bytes).unwrap(); let disk = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::BinaryIvf, 3).unwrap(); + let parallel_pool = rayon::ThreadPoolBuilder::new() + .num_threads(4) + .build() + .unwrap(); + let serial_documents = disk + .search_binary_clusters_with_tuning::(&[0], 3, &[0, 1], usize::MAX) + .unwrap(); + let parallel_documents = parallel_pool + .install(|| disk.search_binary_clusters_with_tuning::(&[0], 3, &[0, 1], 1)) + .unwrap(); + assert_eq!(parallel_documents, serial_documents); // Under Sum the SOAR duplicate decides the winner outright: doc 0 // counted twice (2.0) would beat doc 1's two distinct values (1.5); @@ -2718,6 +3154,10 @@ mod tests { let (result, probed) = disk .search_binary_combined_documents(1, &[0], &[0, 1], sum) .unwrap(); + let parallel_sum = parallel_pool + .install(|| disk.search_binary_combined_documents_with_tuning(1, &[0], &[0, 1], sum, 1)) + .unwrap(); + assert_eq!(parallel_sum, (result.clone(), probed.clone())); assert_eq!(result.len(), 1, "combined search must honor k"); assert_eq!( result[0].doc_id, 1, @@ -2743,6 +3183,12 @@ mod tests { let (result, probed) = disk .search_binary_combined_documents(1, &[0], &[0, 1], smooth_max) .unwrap(); + let parallel_smooth_max = parallel_pool + .install(|| { + disk.search_binary_combined_documents_with_tuning(1, &[0], &[0, 1], smooth_max, 1) + }) + .unwrap(); + assert_eq!(parallel_smooth_max, (result.clone(), probed.clone())); assert_eq!(result.len(), 1, "combined search must honor k"); assert_eq!(result[0].doc_id, 0, "{result:?}"); let expected = smooth_max.combine(&[(0, 1.0), (1, 0.0)]); @@ -3079,6 +3525,10 @@ mod tests { } } let k = 10; + let parallel_pool = rayon::ThreadPoolBuilder::new() + .num_threads(4) + .build() + .unwrap(); for query_seed in [1u64, 9, 42] { let mut qstate = query_seed; let mut qnext = move || { @@ -3125,11 +3575,21 @@ mod tests { } let pruned = disk.search_ivf_tq_distinct(k, &plan).unwrap(); + let forced_parallel = parallel_pool.install(|| { + // Small chunks force several tasks from this compact test + // payload instead of allocating a production-sized one. + disk.search_ivf_tq_distinct_with_tuning(k, &plan, 1, 4) + .unwrap() + }); let reference = reference.into_sorted_results(); assert_eq!( pruned, reference, "scale-bound pruning must not change the estimated top-k (seed {query_seed})" ); + assert_eq!( + forced_parallel, reference, + "parallel IVF-TQ pruning must match the unpruned top-k (seed {query_seed})" + ); for combiner in [ crate::query::MultiValueCombiner::Max, crate::query::MultiValueCombiner::Sum, @@ -3141,16 +3601,421 @@ mod tests { let combined = disk .search_ivf_tq_combined_documents(k, &plan, combiner) .unwrap(); + let parallel_combined = parallel_pool + .install(|| { + disk.search_ivf_tq_combined_documents_with_tuning(k, &plan, combiner, 1, 4) + }) + .unwrap(); assert_eq!( combined, expected, "combined IVF-TQ scan diverged from the unpruned reference \ for {combiner:?} (seed {query_seed})", ); + assert_eq!( + parallel_combined, expected, + "parallel combined IVF-TQ scan diverged from the unpruned reference \ + for {combiner:?} (seed {query_seed})", + ); assert!(combined.len() <= k); } } } + /// Focused single-segment scan benchmark for the adaptive IVF-TQ fan-out. + /// Run with: + /// + /// `cargo test --release -p hermes-core ivf_tq_parallel_scan_benchmark -- --ignored --nocapture` + /// + /// `IVF_SCAN_BENCH_DOCS`, `IVF_SCAN_BENCH_DIM`, `IVF_SCAN_BENCH_ITERS`, + /// and `IVF_SCAN_BENCH_THREADS` override the defaults. + #[test] + #[ignore] + fn ivf_tq_parallel_scan_benchmark() { + use crate::structures::vector::ivf::CoarseCentroids; + use crate::structures::{IvfTqIndex, TqCodec, TqIvfEncodeScratch, TqIvfQueryPlan}; + + fn env_usize(name: &str, default: usize) -> usize { + std::env::var(name) + .ok() + .map(|value| { + value + .parse::() + .unwrap_or_else(|_| panic!("{name} must be a positive integer")) + }) + .unwrap_or(default) + .max(1) + } + + fn median_ms(samples: &mut [f64]) -> f64 { + samples.sort_unstable_by(f64::total_cmp); + samples[samples.len() / 2] + } + + let dim = env_usize("IVF_SCAN_BENCH_DIM", 128); + let count = env_usize("IVF_SCAN_BENCH_DOCS", 262_144); + let iterations = env_usize("IVF_SCAN_BENCH_ITERS", 30); + let threads = env_usize("IVF_SCAN_BENCH_THREADS", num_cpus::get().min(8)); + let codec = std::sync::Arc::new(TqCodec::new(dim)); + let version = crate::structures::mark_ivf_tq_cosine_generation(0x51ca_0001); + let mut centroid = vec![0.0f32; dim]; + centroid[0] = 1.0; + let centroids = CoarseCentroids { + num_clusters: 1, + dim, + centroids: centroid, + version, + soar_config: None, + routing_index: None, + }; + let mut index = IvfTqIndex::new( + dim, + crate::dsl::IvfRoutingMode::Flat, + version, + std::sync::Arc::clone(&codec), + ); + let mut scratch = TqIvfEncodeScratch::default(); + let mut state = 0x5eed_u64; + let mut vector = vec![0.0f32; dim]; + for posting_index in 0..count { + for value in &mut vector { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + *value = ((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5; + } + let norm = vector.iter().map(|value| value * value).sum::().sqrt(); + vector.iter_mut().for_each(|value| *value /= norm); + index.add_vector( + ¢roids, + u32::try_from(posting_index / 2).unwrap(), + u16::try_from(posting_index % 2).unwrap(), + &vector, + &mut scratch, + ); + } + let mut bytes = Vec::new(); + write_built_ivf_tq(&index, 1, &mut bytes).unwrap(); + let disk = AnnDiskIndex::open( + OwnedBytes::new(bytes), + AnnKind::IvfTq, + u32::try_from(count.div_ceil(2)).unwrap(), + ) + .unwrap(); + let query = vec![1.0f32; dim]; + let plan = TqIvfQueryPlan::build( + ¢roids, + &codec, + &query, + 1, + crate::dsl::IvfRoutingMode::Flat, + ); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap(); + let serial_pool = rayon::ThreadPoolBuilder::new() + .num_threads(1) + .build() + .unwrap(); + + let serial = pool + .install(|| disk.search_ivf_tq_distinct_with_tuning(20, &plan, usize::MAX, 512)) + .unwrap(); + let parallel = pool + .install(|| disk.search_ivf_tq_distinct_with_tuning(20, &plan, 1, 512)) + .unwrap(); + assert_eq!(parallel, serial); + + let mut serial_ms = Vec::with_capacity(iterations); + let mut parallel_ms = Vec::with_capacity(iterations); + for _ in 0..iterations { + let start = std::time::Instant::now(); + std::hint::black_box( + pool.install(|| { + disk.search_ivf_tq_distinct_with_tuning(20, &plan, usize::MAX, 512) + }) + .unwrap(), + ); + serial_ms.push(start.elapsed().as_secs_f64() * 1_000.0); + + let start = std::time::Instant::now(); + std::hint::black_box( + pool.install(|| disk.search_ivf_tq_distinct_with_tuning(20, &plan, 1, 512)) + .unwrap(), + ); + parallel_ms.push(start.elapsed().as_secs_f64() * 1_000.0); + } + let serial_p50 = median_ms(&mut serial_ms); + let parallel_p50 = median_ms(&mut parallel_ms); + println!( + "IVF-TQ distinct scan: postings={count} dim={dim} threads={threads} \ + serial_p50={serial_p50:.3}ms parallel_p50={parallel_p50:.3}ms speedup={:.2}x", + serial_p50 / parallel_p50, + ); + + let combiner = crate::query::MultiValueCombiner::Sum; + let serial_combined = serial_pool + .install(|| { + disk.search_ivf_tq_combined_documents_with_tuning(20, &plan, combiner, 1, 512) + }) + .unwrap(); + let parallel_combined = pool + .install(|| { + disk.search_ivf_tq_combined_documents_with_tuning(20, &plan, combiner, 1, 512) + }) + .unwrap(); + assert_eq!(parallel_combined, serial_combined); + let mut serial_combined_ms = Vec::with_capacity(iterations); + let mut parallel_combined_ms = Vec::with_capacity(iterations); + for _ in 0..iterations { + let start = std::time::Instant::now(); + std::hint::black_box( + serial_pool + .install(|| { + disk.search_ivf_tq_combined_documents_with_tuning( + 20, &plan, combiner, 1, 512, + ) + }) + .unwrap(), + ); + serial_combined_ms.push(start.elapsed().as_secs_f64() * 1_000.0); + + let start = std::time::Instant::now(); + std::hint::black_box( + pool.install(|| { + disk.search_ivf_tq_combined_documents_with_tuning(20, &plan, combiner, 1, 512) + }) + .unwrap(), + ); + parallel_combined_ms.push(start.elapsed().as_secs_f64() * 1_000.0); + } + let serial_combined_p50 = median_ms(&mut serial_combined_ms); + let parallel_combined_p50 = median_ms(&mut parallel_combined_ms); + println!( + "IVF-TQ combined scan: postings={count} dim={dim} threads={threads} \ + serial_p50={serial_combined_p50:.3}ms parallel_p50={parallel_combined_p50:.3}ms \ + speedup={:.2}x", + serial_combined_p50 / parallel_combined_p50, + ); + } + + /// Focused binary-IVF benchmark for single-value top-k and multi-value + /// combined scoring. Environment overrides use the `BINARY_SCAN_BENCH_*` + /// prefix with `POSTINGS`, `DIM_BITS`, `ITERS`, and `THREADS` suffixes. + #[test] + #[ignore] + fn binary_ivf_parallel_scan_benchmark() { + fn env_usize(name: &str, default: usize) -> usize { + std::env::var(name) + .ok() + .map(|value| value.parse::().unwrap()) + .unwrap_or(default) + .max(1) + } + + fn median_ms(samples: &mut [f64]) -> f64 { + samples.sort_unstable_by(f64::total_cmp); + samples[samples.len() / 2] + } + + let count = env_usize("BINARY_SCAN_BENCH_POSTINGS", 131_072); + let dim_bits = env_usize("BINARY_SCAN_BENCH_DIM_BITS", 2_048); + assert!(dim_bits.is_multiple_of(8)); + let code_size = dim_bits / 8; + let iterations = env_usize("BINARY_SCAN_BENCH_ITERS", 30); + let threads = env_usize("BINARY_SCAN_BENCH_THREADS", num_cpus::get().min(8)); + let mut state = 0xb1a4_5eed_u64; + let mut codes = vec![0u8; count * code_size]; + for code in &mut codes { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + *code = (state >> 56) as u8; + } + let query = vec![0xa5u8; code_size]; + let unique_docs: Vec = (0..count) + .map(|index| u32::try_from(index).unwrap()) + .collect(); + let unique_ordinals = vec![0u16; count]; + let multi_docs: Vec = (0..count) + .map(|index| u32::try_from(index / 2).unwrap()) + .collect(); + let multi_ordinals: Vec = (0..count) + .map(|index| u16::try_from(index % 2).unwrap()) + .collect(); + let header = AnnDiskHeader { + kind: AnnKind::BinaryIvf, + routing: IvfRoutingMode::Flat, + dim: dim_bits, + code_size, + num_clusters: 1, + quantizer_version: 0xb1a4, + codebook_version: 0, + vector_count: count, + }; + let mut unique_bytes = Vec::new(); + write_built_runs( + header.clone(), + &[BuildRun { + cluster_id: 0, + doc_ids: &unique_docs, + ordinals: &unique_ordinals, + codes: &codes, + }], + &mut unique_bytes, + ) + .unwrap(); + let unique_disk = AnnDiskIndex::open( + OwnedBytes::new(unique_bytes), + AnnKind::BinaryIvf, + u32::try_from(count).unwrap(), + ) + .unwrap(); + let mut multi_bytes = Vec::new(); + write_built_runs( + header, + &[BuildRun { + cluster_id: 0, + doc_ids: &multi_docs, + ordinals: &multi_ordinals, + codes: &codes, + }], + &mut multi_bytes, + ) + .unwrap(); + let multi_disk = AnnDiskIndex::open( + OwnedBytes::new(multi_bytes), + AnnKind::BinaryIvf, + u32::try_from(count.div_ceil(2)).unwrap(), + ) + .unwrap(); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .build() + .unwrap(); + let serial_pool = rayon::ThreadPoolBuilder::new() + .num_threads(1) + .build() + .unwrap(); + + let serial = pool + .install(|| { + unique_disk.search_binary_clusters_with_tuning::( + &query, + 20, + &[0], + usize::MAX, + ) + }) + .unwrap(); + let parallel = pool + .install(|| { + unique_disk.search_binary_clusters_with_tuning::(&query, 20, &[0], 1) + }) + .unwrap(); + assert_eq!(parallel, serial); + let mut serial_ms = Vec::with_capacity(iterations); + let mut parallel_ms = Vec::with_capacity(iterations); + for _ in 0..iterations { + let start = std::time::Instant::now(); + std::hint::black_box( + pool.install(|| { + unique_disk.search_binary_clusters_with_tuning::( + &query, + 20, + &[0], + usize::MAX, + ) + }) + .unwrap(), + ); + serial_ms.push(start.elapsed().as_secs_f64() * 1_000.0); + + let start = std::time::Instant::now(); + std::hint::black_box( + pool.install(|| { + unique_disk.search_binary_clusters_with_tuning::(&query, 20, &[0], 1) + }) + .unwrap(), + ); + parallel_ms.push(start.elapsed().as_secs_f64() * 1_000.0); + } + let serial_p50 = median_ms(&mut serial_ms); + let parallel_p50 = median_ms(&mut parallel_ms); + println!( + "binary-IVF distinct scan: postings={count} dim_bits={dim_bits} threads={threads} \ + serial_p50={serial_p50:.3}ms parallel_p50={parallel_p50:.3}ms speedup={:.2}x", + serial_p50 / parallel_p50, + ); + + let combiner = crate::query::MultiValueCombiner::Sum; + let serial_combined = serial_pool + .install(|| { + multi_disk.search_binary_combined_documents_with_tuning( + 20, + &query, + &[0], + combiner, + 1, + ) + }) + .unwrap(); + let parallel_combined = pool + .install(|| { + multi_disk.search_binary_combined_documents_with_tuning( + 20, + &query, + &[0], + combiner, + 1, + ) + }) + .unwrap(); + assert_eq!(parallel_combined, serial_combined); + let mut serial_combined_ms = Vec::with_capacity(iterations); + let mut parallel_combined_ms = Vec::with_capacity(iterations); + for _ in 0..iterations { + let start = std::time::Instant::now(); + std::hint::black_box( + serial_pool + .install(|| { + multi_disk.search_binary_combined_documents_with_tuning( + 20, + &query, + &[0], + combiner, + 1, + ) + }) + .unwrap(), + ); + serial_combined_ms.push(start.elapsed().as_secs_f64() * 1_000.0); + + let start = std::time::Instant::now(); + std::hint::black_box( + pool.install(|| { + multi_disk.search_binary_combined_documents_with_tuning( + 20, + &query, + &[0], + combiner, + 1, + ) + }) + .unwrap(), + ); + parallel_combined_ms.push(start.elapsed().as_secs_f64() * 1_000.0); + } + let serial_combined_p50 = median_ms(&mut serial_combined_ms); + let parallel_combined_p50 = median_ms(&mut parallel_combined_ms); + println!( + "binary-IVF combined scan: postings={count} dim_bits={dim_bits} threads={threads} \ + serial_p50={serial_combined_p50:.3}ms parallel_p50={parallel_combined_p50:.3}ms \ + speedup={:.2}x", + serial_combined_p50 / parallel_combined_p50, + ); + } + #[test] fn open_rejects_tq_payload_with_inconsistent_geometry() { let (bytes, _) = build_tq_payload(20, 4, 9);