diff --git a/src/Cargo.lock b/src/Cargo.lock index 961a2555..8b205a0d 100644 --- a/src/Cargo.lock +++ b/src/Cargo.lock @@ -1567,6 +1567,7 @@ dependencies = [ "obicompactvec", "obikindex", "obikseq", + "obipipeline", "rayon", "tempfile", ] diff --git a/src/obikidxcache/Cargo.toml b/src/obikidxcache/Cargo.toml index 29a7f586..661ffc3a 100644 --- a/src/obikidxcache/Cargo.toml +++ b/src/obikidxcache/Cargo.toml @@ -7,6 +7,7 @@ edition = "2024" obikindex = { path = "../obikindex"} obikseq = { path = "../obikseq"} obicompactvec = { path = "../obicompactvec" } +obipipeline = { path = "../obipipeline" } ndarray = "0.17" rayon = "1" diff --git a/src/obikidxcache/src/index_cache.rs b/src/obikidxcache/src/index_cache.rs index c5ef6661..2f72e141 100644 --- a/src/obikidxcache/src/index_cache.rs +++ b/src/obikidxcache/src/index_cache.rs @@ -3,6 +3,8 @@ use std::sync::Arc; use obikindex::{KmerIndex, layer::KmerLayer}; use obikseq::CanonicalKmer; +use obipipeline::Throttle; +use rayon::prelude::*; use crate::meta_cache::MetaCache; @@ -188,4 +190,86 @@ impl IndexCache { let partition = kmer.partition(self.raw_index.n_partitions()); self.find_in_partition(partition, kmer) } + + /// Parallel map-reduce over every cached `KmerLayer` (across every + /// cached partition), bounded to `max_concurrent` layers scanned at + /// once — a cap on how many layers are actively being mapped + /// simultaneously, independent of `rayon`'s own thread count (a large + /// per-layer working set, e.g. a whole matrix scan, can make "one + /// layer per thread" too much resident memory at once even when the + /// thread pool itself is much bigger). + /// + /// Safe only when `map` is independent per layer (no cross-layer or + /// cross-partition state) — when a caller's own per-layer work already + /// parallelizes internally at a finer grain (e.g. across partitions, + /// like `obikphylo`'s sibling-annex traversal), stacking this on top + /// fights that inner parallelism for cache/mmap locality instead of + /// helping; use the plain sequential [`iter`](Self::iter) there + /// instead. + pub fn par_map_reduce( + &self, + max_concurrent: usize, + identity: impl Fn() -> T + Sync, + map: M, + reduce: R, + ) -> T + where + T: Send, + M: Fn(&KmerLayer) -> T + Sync, + R: Fn(T, T) -> T + Sync, + { + let throttle = Arc::new(Throttle::new(max_concurrent)); + self.layer_cache + .par_iter() + .flat_map(|(_, layers)| layers.par_iter()) + .map(|layer| { + let _guard = throttle.acquire_guard(); + map(layer) + }) + .reduce(&identity, &reduce) + } + + /// Parallel, side-effecting foreach over every cached `KmerLayer`, + /// bounded to `max_concurrent` layers at once — same bounding + /// reasoning and same "independent per layer" safety requirement as + /// [`par_map_reduce`](Self::par_map_reduce), just without a value to + /// combine back. + pub fn par_for_each(&self, max_concurrent: usize, f: F) + where + F: Fn(&KmerLayer) + Sync, + { + let throttle = Arc::new(Throttle::new(max_concurrent)); + self.layer_cache + .par_iter() + .flat_map(|(_, layers)| layers.par_iter()) + .for_each(|layer| { + let _guard = throttle.acquire_guard(); + f(layer); + }); + } + + /// Bounded parallel iterator over every cached `KmerLayer` — for + /// callers that want to chain their own `rayon` combinators rather + /// than fit their computation into [`par_map_reduce`](Self::par_map_reduce)/ + /// [`par_for_each`](Self::par_for_each)'s shape. + /// + /// Each item is paired with a [`obipipeline::ThrottleGuard`] (same + /// `Throttled` carrier `obipipeline::throttle` uses for a plain + /// sequential source) that releases its slot on drop — keep the whole + /// `Throttled` alive for as long as you're "using" that layer (e.g. + /// through your own `.map()`/`.for_each()`), the same discipline + /// `obipipeline`'s own consumers already follow. + pub fn par_iter_bounded( + &self, + max_concurrent: usize, + ) -> impl ParallelIterator> { + let throttle = Arc::new(Throttle::new(max_concurrent)); + self.layer_cache + .par_iter() + .flat_map(|(_, layers)| layers.par_iter()) + .map(move |layer| obipipeline::Throttled { + item: layer, + guard: throttle.acquire_guard(), + }) + } } diff --git a/src/obikphylo/src/distance.rs b/src/obikphylo/src/distance.rs index 511765d2..39731ee6 100644 --- a/src/obikphylo/src/distance.rs +++ b/src/obikphylo/src/distance.rs @@ -2,16 +2,17 @@ //! (summing per-layer partials into one global partial) and the final //! distance calculation are both phylogenetic domain logic, so both live //! here, not in `obikidxcache` (which only ever hands out already-opened -//! `KmerLayer`s via `IndexCache::iter()` — plain exposition, no -//! computation of its own). +//! `KmerLayer`s, via plain exposition — `IndexCache::iter()` — or bounded +//! parallel map-reduce over them — `IndexCache::par_map_reduce()` — neither +//! of which computes anything phylo-specific itself). //! //! [`AggregatedCount`]/[`AggregatedBit`] exist purely to "hang" -//! `obicompactvec`'s `CountPartials`/`BitPartials` traits on a borrowed -//! collection of `KmerLayer`s, so their already-written finalisation -//! methods (`bray_dist_matrix`, `jaccard_dist_matrix`, ...) can be reused -//! as-is instead of re-derived here — see `KmerLayer`'s own -//! `ColumnWeights`/`CountPartials`/`BitPartials` impls (`obikindex`) for -//! the per-layer primitives these sum. +//! `obicompactvec`'s `CountPartials`/`BitPartials` traits on a bounded +//! parallel reduction over `IndexCache`'s layers, so their already-written +//! finalisation methods (`bray_dist_matrix`, `jaccard_dist_matrix`, ...) can +//! be reused as-is instead of re-derived here — see `KmerLayer`'s own +//! `ColumnWeights`/`CountPartials`/`BitPartials` impls (`obikindex`) for the +//! per-layer primitives these sum. use ndarray::{Array1, Array2}; use obicompactvec::{BitPartials, ColumnWeights, CountPartials}; @@ -64,77 +65,118 @@ impl DistanceMetric { } } -// ── Aggregators — sum every cached layer's own partials into one global -// partial, then hand the result off to `CountPartials`/`BitPartials`'s own -// provided finalisation methods. +// ── Aggregators — bounded parallel reduction of every cached layer's own +// partial (of the wanted content kind) into one global partial, then handed +// off to `CountPartials`/`BitPartials`'s own provided finalisation methods. +// Layers of the *other* content kind contribute their zero element, same +// as a layer that's simply absent — cheaper than pre-collecting a filtered +// `Vec` first, since `IndexCache::par_map_reduce` never allocates one. -fn sum_array1(n: usize, items: impl Iterator>) -> Array1 { - items.fold(Array1::zeros(n), |acc, x| acc + x) -} -fn sum_array2_u64(n: usize, items: impl Iterator>) -> Array2 { - items.fold(Array2::zeros((n, n)), |acc, x| acc + x) -} -fn sum_array2_f64(n: usize, items: impl Iterator>) -> Array2 { - items.fold(Array2::zeros((n, n)), |acc, x| acc + x) -} +/// Layers scanned at once per `par_map_reduce` call — bounds resident +/// memory (each layer's own partial-computation working set), independent +/// of how many `rayon` threads exist. See `IndexCache::par_map_reduce`'s +/// own docs for why this is a separate knob from thread count. +const MAX_CONCURRENT_LAYERS: usize = 8; struct AggregatedCount<'a> { + cache: &'a IndexCache, n_genomes: usize, - layers: Vec<&'a KmerLayer>, +} + +impl AggregatedCount<'_> { + fn reduce( + &self, + identity: impl Fn() -> T + Sync, + per_layer: impl Fn(&KmerLayer) -> T + Sync, + combine: impl Fn(T, T) -> T + Sync, + ) -> T { + self.cache.par_map_reduce( + MAX_CONCURRENT_LAYERS, + &identity, + |layer| if layer.content() == LayerContent::Count { per_layer(layer) } else { identity() }, + &combine, + ) + } } impl ColumnWeights for AggregatedCount<'_> { fn col_weights(&self) -> Array1 { - sum_array1(self.n_genomes, self.layers.iter().map(|l| l.col_weights())) + let n = self.n_genomes; + self.reduce(|| Array1::zeros(n), |l| l.col_weights(), |a, b| a + b) } } impl CountPartials for AggregatedCount<'_> { fn partial_bray(&self) -> Array2 { - sum_array2_u64(self.n_genomes, self.layers.iter().map(|l| l.partial_bray())) + let n = self.n_genomes; + self.reduce(|| Array2::zeros((n, n)), |l| l.partial_bray(), |a, b| a + b) } fn partial_euclidean(&self) -> Array2 { - sum_array2_f64(self.n_genomes, self.layers.iter().map(|l| l.partial_euclidean())) + let n = self.n_genomes; + self.reduce(|| Array2::zeros((n, n)), |l| l.partial_euclidean(), |a, b| a + b) } fn partial_threshold_jaccard(&self, threshold: u32) -> (Array2, Array2) { let n = self.n_genomes; - self.layers.iter().map(|l| l.partial_threshold_jaccard(threshold)).fold( - (Array2::zeros((n, n)), Array2::zeros((n, n))), - |(inter, union), (i, u)| (inter + i, union + u), + self.reduce( + || (Array2::zeros((n, n)), Array2::zeros((n, n))), + |l| l.partial_threshold_jaccard(threshold), + |(ia, ua), (ib, ub)| (ia + ib, ua + ub), ) } fn partial_relfreq_bray(&self, global: &Array1) -> Array2 { - sum_array2_f64(self.n_genomes, self.layers.iter().map(|l| l.partial_relfreq_bray(global))) + let n = self.n_genomes; + self.reduce(|| Array2::zeros((n, n)), |l| l.partial_relfreq_bray(global), |a, b| a + b) } fn partial_relfreq_euclidean(&self, global: &Array1) -> Array2 { - sum_array2_f64(self.n_genomes, self.layers.iter().map(|l| l.partial_relfreq_euclidean(global))) + let n = self.n_genomes; + self.reduce(|| Array2::zeros((n, n)), |l| l.partial_relfreq_euclidean(global), |a, b| a + b) } fn partial_hellinger(&self, global: &Array1) -> Array2 { - sum_array2_f64(self.n_genomes, self.layers.iter().map(|l| l.partial_hellinger(global))) + let n = self.n_genomes; + self.reduce(|| Array2::zeros((n, n)), |l| l.partial_hellinger(global), |a, b| a + b) } } struct AggregatedBit<'a> { + cache: &'a IndexCache, n_genomes: usize, - layers: Vec<&'a KmerLayer>, +} + +impl AggregatedBit<'_> { + fn reduce( + &self, + identity: impl Fn() -> T + Sync, + per_layer: impl Fn(&KmerLayer) -> T + Sync, + combine: impl Fn(T, T) -> T + Sync, + ) -> T { + self.cache.par_map_reduce( + MAX_CONCURRENT_LAYERS, + &identity, + |layer| if layer.content() == LayerContent::Presence { per_layer(layer) } else { identity() }, + &combine, + ) + } } impl ColumnWeights for AggregatedBit<'_> { fn col_weights(&self) -> Array1 { - sum_array1(self.n_genomes, self.layers.iter().map(|l| l.col_weights())) + let n = self.n_genomes; + self.reduce(|| Array1::zeros(n), |l| l.col_weights(), |a, b| a + b) } } impl BitPartials for AggregatedBit<'_> { fn partial_jaccard(&self) -> (Array2, Array2) { let n = self.n_genomes; - self.layers.iter().map(|l| l.partial_jaccard()).fold( - (Array2::zeros((n, n)), Array2::zeros((n, n))), - |(inter, union), (i, u)| (inter + i, union + u), + self.reduce( + || (Array2::zeros((n, n)), Array2::zeros((n, n))), + |l| l.partial_jaccard(), + |(ia, ua), (ib, ub)| (ia + ib, ua + ub), ) } fn partial_hamming(&self) -> Array2 { - sum_array2_u64(self.n_genomes, self.layers.iter().map(|l| l.partial_hamming())) + let n = self.n_genomes; + self.reduce(|| Array2::zeros((n, n)), |l| l.partial_hamming(), |a, b| a + b) } } @@ -167,8 +209,7 @@ impl Metrics for IndexCache { let kmer_size = self.meta().config().kmer_size; if use_counts { - let layers: Vec<&KmerLayer> = self.iter().filter(|l| l.content() == LayerContent::Count).collect(); - let agg = AggregatedCount { n_genomes, layers }; + let agg = AggregatedCount { cache: self, n_genomes }; let matrix = match metric { DistanceMetric::BrayCurtis => agg.bray_dist_matrix(), @@ -195,8 +236,7 @@ impl Metrics for IndexCache { Ok(DistanceOutput { matrix, shared_kmers: shared }) } else { - let layers: Vec<&KmerLayer> = self.iter().filter(|l| l.content() == LayerContent::Presence).collect(); - let agg = AggregatedBit { n_genomes, layers }; + let agg = AggregatedBit { cache: self, n_genomes }; let matrix = match metric { DistanceMetric::Jaccard => agg.jaccard_dist_matrix(), diff --git a/src/obikstats/src/stats.rs b/src/obikstats/src/stats.rs index 7324c166..28fa82f0 100644 --- a/src/obikstats/src/stats.rs +++ b/src/obikstats/src/stats.rs @@ -2,13 +2,16 @@ use std::fs; use std::path::Path; use std::sync::Arc; -use rayon::prelude::*; - use obikalgorithm::Algorithm; use obikidxcache::index_cache::IndexCache; use obikindex::KmerIndex; use obikindex::layer::KmerLayer; +/// Layers scanned at once per `par_map_reduce` call — see +/// `obikidxcache::IndexCache::par_map_reduce`'s own docs for why this is a +/// separate knob from thread count. +const MAX_CONCURRENT_LAYERS: usize = 8; + /// Bits per kmer broken down by index component. pub struct IndexBitsPerKmer { /// Total distinct k-mers across all partitions and layers. @@ -105,16 +108,16 @@ impl Algorithm for BitsPerKmer { fn run(&mut self) -> obikalgorithm::Result { let n_genomes = self.index.meta().genomes()?.len().max(1); let cache = IndexCache::new(Arc::clone(&self.index), None); - let layers: Vec<&KmerLayer> = cache.iter().collect(); - let (n_kmers, mphf_b, evidence_b, matrix_b) = layers - .par_iter() - .map(|layer| layer_bytes(layer)) - .map(|lb| (lb.n_kmers, lb.mphf, lb.evidence, lb.matrix)) - .reduce( - || (0usize, 0u64, 0u64, 0u64), - |a, b| (a.0 + b.0, a.1 + b.1, a.2 + b.2, a.3 + b.3), - ); + let (n_kmers, mphf_b, evidence_b, matrix_b) = cache.par_map_reduce( + MAX_CONCURRENT_LAYERS, + || (0usize, 0u64, 0u64, 0u64), + |layer| { + let lb = layer_bytes(layer); + (lb.n_kmers, lb.mphf, lb.evidence, lb.matrix) + }, + |a, b| (a.0 + b.0, a.1 + b.1, a.2 + b.2, a.3 + b.3), + ); if n_kmers == 0 { return Ok(IndexBitsPerKmer { @@ -166,19 +169,26 @@ impl Algorithm for GenomeKmerCounts { fn run(&mut self) -> obikalgorithm::Result<(usize, Vec)> { let n_genomes = self.index.meta().genomes()?.len(); let cache = IndexCache::new(Arc::clone(&self.index), None); - let layers: Vec<&KmerLayer> = cache.iter().collect(); - let total_kmers: usize = layers.iter().map(|l| l.n()).sum(); - - let mut total_counts = vec![0u64; n_genomes]; - let per_layer_weights: Vec<_> = layers.par_iter().map(|l| l.col_weights()).collect(); - for weights in per_layer_weights { - for (g, &v) in weights.iter().enumerate() { - if g < n_genomes { - total_counts[g] += v; + let (total_kmers, total_counts) = cache.par_map_reduce( + MAX_CONCURRENT_LAYERS, + || (0usize, vec![0u64; n_genomes]), + |layer| { + let mut counts = vec![0u64; n_genomes]; + for (g, &w) in layer.col_weights().iter().enumerate() { + if g < n_genomes { + counts[g] = w; + } } - } - } + (layer.n(), counts) + }, + |(ka, mut ca), (kb, cb)| { + for (x, y) in ca.iter_mut().zip(cb.iter()) { + *x += y; + } + (ka + kb, ca) + }, + ); Ok((total_kmers, total_counts)) } diff --git a/src/obipipeline/src/throttle.rs b/src/obipipeline/src/throttle.rs index d83b4a67..ed03640c 100644 --- a/src/obipipeline/src/throttle.rs +++ b/src/obipipeline/src/throttle.rs @@ -35,6 +35,19 @@ impl Throttle { } } +impl Throttle { + /// Acquire one slot and return a guard that releases it on drop — + /// for callers with no natural "source iterator" to wrap via + /// [`throttle`] (e.g. bounding concurrency inside a `rayon` parallel + /// iterator's own closure, one call per item, rather than up front in + /// a single source thread). Blocks the calling thread until a slot is + /// free, same as [`throttle`]'s own per-item acquisition. + pub fn acquire_guard(self: &Arc) -> ThrottleGuard { + self.acquire(); + ThrottleGuard(Arc::clone(self)) + } +} + // ── ThrottleGuard ───────────────────────────────────────────────────────────── /// RAII guard: releases one slot in the `Throttle` when dropped.