diff --git a/src/obikindex/src/siblings/alignment.rs b/src/obikindex/src/siblings/alignment.rs index a0409617..93f478d3 100644 --- a/src/obikindex/src/siblings/alignment.rs +++ b/src/obikindex/src/siblings/alignment.rs @@ -78,30 +78,24 @@ impl KmerIndex { // lookups by partition internally, but interleave those sweeps // across layers at the OS level, scattering page-cache access over // every partition at once again and defeating the whole point of - // the grouping. `Vec>` per layer, one entry (column) per - // variable family, appended in layer order for a single - // deterministic column order across the whole index. - let mut partials: Vec>> = Vec::with_capacity(layer_dirs.len()); + // the grouping. Columns appended straight into `sequences` as each + // family comes back from `scan_layer_families` (bounded-batch, not + // whole-layer) — no intermediate per-layer buffer, layer order + // giving a single deterministic column order across the whole index. + let mut sequences: Vec> = vec![Vec::new(); n_genomes]; for layer_dir in &layer_dirs { - let families = scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache)?; - let columns = families - .into_iter() - .filter(|f| f.mask.family_size() >= 2) // monomorphic family — no signal, skip - .map(|f| f.genome_mask.iter().map(|&m| iupac_code(m)).collect()) - .collect(); - partials.push(columns); + scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, |mask, genome_mask| { + if mask.family_size() < 2 { + return; // monomorphic family — no signal, skip + } + for (g, &m) in genome_mask.iter().enumerate() { + sequences[g].push(iupac_code(m)); + } + })?; pb.inc(1); } pb.finish_and_clear(); - let mut sequences: Vec> = vec![Vec::new(); n_genomes]; - for layer_columns in partials { - for column in layer_columns { - for (g, &code) in column.iter().enumerate() { - sequences[g].push(code); - } - } - } Ok(SnpAlignment { sequences }) } } diff --git a/src/obikindex/src/siblings/cardinality.rs b/src/obikindex/src/siblings/cardinality.rs index 652cc09b..5921be7e 100644 --- a/src/obikindex/src/siblings/cardinality.rs +++ b/src/obikindex/src/siblings/cardinality.rs @@ -78,14 +78,12 @@ impl KmerIndex { let pb = progress_bar("cardinality_tally", layer_dirs.len() as u64, "layers"); // One layer at a time — see `snp_pseudo_alignment`'s comment for why // `par_iter()` over layers would defeat `scan_layer_families`'s - // partition-grouped locality. - let mut partials: Vec<[[u64; 5]; 5]> = Vec::with_capacity(layer_dirs.len()); + // partition-grouped locality. Tallied straight into `total` as each + // family comes back — no per-layer buffer. + let mut total = [[0u64; 5]; 5]; for layer_dir in &layer_dirs { - let families = scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache)?; - let mut counts = [[0u64; 5]; 5]; - - for family in &families { - if family.mask.family_size() < 2 { + scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, |mask, genome_mask| { + if mask.family_size() < 2 { // Fully invariant family (never varies anywhere in // the index) — genome-wide background, not // SNP-adjacent signal; would otherwise swamp the @@ -94,10 +92,9 @@ impl KmerIndex { // `+ASC`-corrected alignment/likelihood actually // models. See `base_pair_tally`'s `variable` gate // on its own `same` diagonal for the matching fix. - continue; + return; } - let genome_mask = &family.genome_mask; for i in 0..n_genomes { let card_i = genome_mask[i].count_ones() as usize; for j in (i + 1)..n_genomes { @@ -105,27 +102,17 @@ impl KmerIndex { continue; } let card_j = genome_mask[j].count_ones() as usize; - counts[card_i][card_j] += 1; + total[card_i][card_j] += 1; if card_i != card_j { - counts[card_j][card_i] += 1; + total[card_j][card_i] += 1; } } } - } - - partials.push(counts); + })?; pb.inc(1); } pb.finish_and_clear(); - let mut total = [[0u64; 5]; 5]; - for partial in partials { - for a in 0..5 { - for b in 0..5 { - total[a][b] += partial[a][b]; - } - } - } Ok(CardinalityTally { counts: total }) } } diff --git a/src/obikindex/src/siblings/distance.rs b/src/obikindex/src/siblings/distance.rs index 746bb062..00965acd 100644 --- a/src/obikindex/src/siblings/distance.rs +++ b/src/obikindex/src/siblings/distance.rs @@ -88,12 +88,10 @@ impl KmerIndex { // partition-grouped locality. let mut total = zero(); for layer_dir in &layer_dirs { - let families = scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache)?; let mut acc = zero(); - for family in &families { - let variable = family.mask.family_size() >= 2; - let genome_mask = &family.genome_mask; + scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, |mask, genome_mask| { + let variable = mask.family_size() >= 2; // Per genome: which single form (if exactly one) it // carries — `None` (a `popcount != 1` mask) once a @@ -111,7 +109,7 @@ impl KmerIndex { on_pair(&mut acc, i, j, bi, bj, variable); } } - } + })?; total = combine(total, acc); pb.inc(1); diff --git a/src/obikindex/src/siblings/family_scan.rs b/src/obikindex/src/siblings/family_scan.rs index 4ad7da1f..2d11f18c 100644 --- a/src/obikindex/src/siblings/family_scan.rs +++ b/src/obikindex/src/siblings/family_scan.rs @@ -5,11 +5,22 @@ //! `build_sibling_annex` (see `cache.rs`/`build.rs`'s docs): cross-partition //! lookups are grouped by destination partition and resolved in one //! contiguous sweep per partition, instead of one lookup at a time jumping -//! between partitions in family order. The naive, per-family inline lookup -//! was the shape all four consumers used to have independently; sampling a -//! real ~170 GB index against 137 GB of RAM showed it thrashing the page -//! cache (near-continuous ~2 GB/s pagein), the same failure mode -//! `build_sibling_annex` had already been fixed for. +//! between partitions in family order. +//! +//! Processed in bounded batches of [`FAMILY_BATCH`] families, not the whole +//! layer at once — same rationale and constant as `build_layer_sibling_annex`'s +//! `BATCH_SIZE`. An earlier version of this function materialised every +//! family of the whole layer up front (`Vec`, each with its own +//! heap-allocated `Vec`) before handing anything back to the caller: +//! O(layer's family count × genome count) extra memory, on top of a second +//! full copy for the owned per-family vectors — measured pushing a real run +//! into heavy VM compression/swap even though the per-lookup I/O pattern was +//! already fixed. A batch of a few thousand families needs the same +//! partition-grouping to keep lookups local; it does not need the entire +//! layer resident at once. Callers get each family's `genome_mask` through +//! `on_family`, backed by one reused scratch buffer — no per-family +//! allocation at all, matching the O(genome count) working set the original, +//! pre-locality-fix per-family code had. use std::path::{Path, PathBuf}; use std::sync::atomic::{AtomicU8, Ordering}; @@ -29,14 +40,12 @@ use super::cache::{Mat, PartitionCache}; use super::helpers::{central_base, is_minorant, partition_of}; use super::{olm_to_ok, ANNEX_FILE_NAME, INDEX_SUBDIR}; -/// One layer's minorant family, with its own annex mask and, per genome, -/// the 4-bit base-presence mask resolved across the whole index (bit `b` -/// set iff that genome carries the family member whose own canonical -/// central base is `b`). -pub(super) struct FamilyRow { - pub mask: FamilyMask, - pub genome_mask: Vec, -} +/// Families per batch — bounds `scan_layer_families`'s extra memory to +/// `FAMILY_BATCH * n_genomes` bytes regardless of layer size. Same value and +/// rationale as `build_layer_sibling_annex`'s `BATCH_SIZE`: large enough to +/// amortise the per-batch partition-grouping/sync cost, small enough to keep +/// working memory a rounding error next to the index itself. +const FAMILY_BATCH: usize = 4096; impl KmerIndex { /// Every (partition, layer) directory carrying a sibling annex, checked @@ -66,17 +75,13 @@ impl KmerIndex { } } -/// Every minorant family of one layer, each with its per-genome -/// base-presence mask resolved against the whole (already-open) partition -/// cache. -/// -/// Two passes, not one inline pass: the first is purely local (already-open -/// annex/mphf/matrix, no lookups) and both enumerates the layer's families -/// and collects every cross-partition query they need, grouped by -/// destination partition; the second resolves each partition's whole batch -/// in one contiguous sweep, keeping that partition's mmap'd pages hot for -/// its entire batch instead of faulting them in and out as lookups jump -/// between partitions in family order. +/// Visits every minorant family of one layer, in slot order, batched +/// [`FAMILY_BATCH`] at a time — see the module docs for why. `on_family` is +/// called once per family with its own annex mask and its per-genome +/// base-presence (`genome_mask[g]`: bit `b` set iff genome `g` carries the +/// member whose own canonical central base is `b`), backed by a scratch +/// buffer reused across every call — callers that need to keep data past +/// the call must copy it themselves. pub(super) fn scan_layer_families( layer_dir: &Path, n_parts: usize, @@ -84,7 +89,8 @@ pub(super) fn scan_layer_families( with_counts: bool, k: usize, cache: &PartitionCache, -) -> OKIResult> { + mut on_family: impl FnMut(FamilyMask, &[u8]), +) -> OKIResult<()> { let index_dir = layer_dir.parent().expect("layer_dir has a parent index dir"); let meta = PartitionMeta::load(index_dir).map_err(olm_to_ok)?; let annex = SiblingAnnex::open(&layer_dir.join(ANNEX_FILE_NAME))?; @@ -106,73 +112,77 @@ pub(super) fn scan_layer_families( }; let n_cols = mat.n_cols().min(n_genomes); - // ── Pass 1: local only — enumerate minorant families and collect every - // cross-partition query they need, grouped by destination partition. - struct Family { - slot: usize, - mask: FamilyMask, - } - let mut families: Vec = Vec::new(); - let mut outgoing: Vec> = (0..n_parts).map(|_| Vec::new()).collect(); + let n_slots = annex.len(); + let mut slot = 0usize; + let mut batch_slots: Vec = Vec::with_capacity(FAMILY_BATCH); + let mut scratch = vec![0u8; n_genomes]; // reused across every `on_family` call — no per-family allocation - for slot in 0..annex.len() { - let Some(mask) = annex.get(slot) else { continue }; - let Some(kmer) = slot_kmer[slot] else { continue }; - if !is_minorant(kmer, mask, k) { - continue; // family tallied once, at its minorant - } - let family_idx = families.len(); - families.push(Family { slot, mask }); - - for other in kmer.central_canonical_neighbors() { - if other == kmer { - continue; // local — resolved below straight from `mat`, no lookup + while slot < n_slots { + batch_slots.clear(); + while slot < n_slots && batch_slots.len() < FAMILY_BATCH { + let s = slot; + slot += 1; + let Some(mask) = annex.get(s) else { continue }; + let Some(kmer) = slot_kmer[s] else { continue }; + if !is_minorant(kmer, mask, k) { + continue; // family tallied once, at its minorant } - let base = central_base(other, k); - if !mask.has(base) { - continue; - } - let dest = partition_of(other, n_parts); - outgoing[dest].push((other, family_idx, base)); + batch_slots.push(s); } - } - - let genome_mask: Vec = (0..families.len() * n_genomes).map(|_| AtomicU8::new(0)).collect(); - - // Local presence (the family's own minorant slot) — already open, no - // lookup, so no locality concern either way. - for (family_idx, family) in families.iter().enumerate() { - let kmer = slot_kmer[family.slot].expect("family slot has a kmer"); - let base = central_base(kmer, k); - for g in 0..n_cols { - if mat.carries(g, family.slot) { - genome_mask[family_idx * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed); - } + if batch_slots.is_empty() { + continue; // ran out of minorant slots before filling a batch — loop condition ends it } - } - // ── Pass 2: resolve each destination partition's whole batch in one - // contiguous sweep — same rationale as `build_sibling_annex`'s - // `outgoing` grouping. - outgoing.par_iter().enumerate().filter(|(_, q)| !q.is_empty()).for_each(|(dest, queries)| { - for &(variant, family_idx, base) in queries { - let Some(presence) = cache.find_presence(dest, variant, n_genomes) else { continue }; - for (g, &present) in presence.iter().enumerate() { - if present { + // ── Local only: seed each family's own minorant presence (already + // open, no lookup) and collect this batch's cross-partition + // queries, grouped by destination partition. + let mut outgoing: Vec> = (0..n_parts).map(|_| Vec::new()).collect(); + let genome_mask: Vec = (0..batch_slots.len() * n_genomes).map(|_| AtomicU8::new(0)).collect(); + + for (family_idx, &fslot) in batch_slots.iter().enumerate() { + let kmer = slot_kmer[fslot].expect("batch slot has a kmer"); + let mask = annex.get(fslot).expect("batch slot has a mask"); + let base = central_base(kmer, k); + for g in 0..n_cols { + if mat.carries(g, fslot) { genome_mask[family_idx * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed); } } + for other in kmer.central_canonical_neighbors() { + if other == kmer { + continue; // local — resolved above straight from `mat`, no lookup + } + let b = central_base(other, k); + if !mask.has(b) { + continue; + } + let dest = partition_of(other, n_parts); + outgoing[dest].push((other, family_idx, b)); + } } - }); - Ok(families - .into_iter() - .enumerate() - .map(|(family_idx, family)| FamilyRow { - mask: family.mask, - genome_mask: (0..n_genomes) - .map(|g| genome_mask[family_idx * n_genomes + g].load(Ordering::Relaxed)) - .collect(), - }) - .collect()) + // ── Resolve this batch's queries, one contiguous sweep per + // destination partition — same rationale as `build_sibling_annex`'s + // `outgoing` grouping. + outgoing.par_iter().enumerate().filter(|(_, q)| !q.is_empty()).for_each(|(dest, queries)| { + for &(variant, family_idx, base) in queries { + let Some(presence) = cache.find_presence(dest, variant, n_genomes) else { continue }; + for (g, &present) in presence.iter().enumerate() { + if present { + genome_mask[family_idx * n_genomes + g].fetch_or(1 << base, Ordering::Relaxed); + } + } + } + }); + + for (family_idx, &fslot) in batch_slots.iter().enumerate() { + let mask = annex.get(fslot).expect("batch slot has a mask"); + for (g, dst) in scratch.iter_mut().enumerate() { + *dst = genome_mask[family_idx * n_genomes + g].load(Ordering::Relaxed); + } + on_family(mask, &scratch); + } + } + + Ok(()) } diff --git a/src/obikindex/src/siblings/stats.rs b/src/obikindex/src/siblings/stats.rs index 9d8c7b71..2e6009e0 100644 --- a/src/obikindex/src/siblings/stats.rs +++ b/src/obikindex/src/siblings/stats.rs @@ -65,19 +65,18 @@ impl KmerIndex { ..Default::default() }; for layer_dir in &layer_dirs { - let families = scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache)?; - for family in &families { + scan_layer_families(layer_dir, n_parts, n_genomes, with_counts, k, &cache, |mask, genome_mask| { // "Genome g represents this family" means g carries // *any* of its members, not just the minorant's own — // `genome_mask[g] != 0` is exactly that. - let s = family.mask.siblings() as usize; + let s = mask.siblings() as usize; stats.counts[s] += 1; - for (g, &m) in family.genome_mask.iter().enumerate() { + for (g, &m) in genome_mask.iter().enumerate() { if m != 0 { stats.per_genome[g][s] += 1; } } - } + })?; pb.inc(1); } pb.finish_and_clear();