From 1c04d1bc42828bbebdfe1617e4faee3dddf7c677 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 22:23:32 +0200 Subject: [PATCH] wip: parallel cascade levels, extracted onto main (no standalone gain) --- src/ext_mem.rs | 129 ++++++++++++++++++++++++++++++++++--------------- 1 file changed, 90 insertions(+), 39 deletions(-) diff --git a/src/ext_mem.rs b/src/ext_mem.rs index 621a3ab..ef89bdd 100644 --- a/src/ext_mem.rs +++ b/src/ext_mem.rs @@ -1074,9 +1074,38 @@ impl<'a> PositionSource<'a> { /// Target subarray size used by [`effective_subproblem_count`] when /// auto-picking `p`. Smaller means more (smaller) subarrays — lower /// per-task phase-1 scratch, at the cost of more phase-3 distribute -/// work (which scales as `O(p² · log(n/p))`, sequentially) and a -/// higher temp-file count. -const PHASE1_TARGET_CHUNK: usize = 65_536; +/// work (which scales as `O(p² · log(n/p))`) and a higher temp-file +/// count. +/// +/// Raised from 65 536 after measuring the trade-off directly. Total work +/// is `n log n` either way, since a smaller `p` moves levels out of +/// phase 4's per-partition cascade and into phase 1's merge sort, but +/// the constants are not equal: phase 3 shrinks quadratically in `p`, +/// and phase 4's cascade does one full pass over its partition per +/// level. Peak RSS is set by phase 4 holding `4 × threads` partitions of +/// `n / p` records at once, so it only starts growing once `p` is small +/// enough for that product to rival the text itself. +/// +/// Measured on chr21 forward ++ revcomp (80 MB), 12 threads: +/// +/// ```text +/// p total peak RSS +/// 48 2.58 s 987 MB +/// 96 2.49 s 538 MB +/// 306 2.85 s 282 MB +/// 612 2.80 s 202 MB <- 131 072 +/// 1224 3.05 s 205 MB <- 65 536 (previous default) +/// ``` +/// +/// 131 072 is the largest step that costs nothing in memory: same peak +/// RSS as before, ~8% less wall time. Going further trades real memory +/// for speed, which is the opposite of what this path is for, so it is +/// left to the caller via `ExtMemOpts::subproblem_count`. +/// +/// At genome scale this changes nothing: `PHASE1_MAX_PARTITIONS` already +/// binds for any `n` above ~1 GB, so GRCh38 still gets `p = 8192`. +const PHASE1_TARGET_CHUNK: usize = 131_072; + /// Hard cap on the number of subarrays. Matches upstream CaPS-SA's /// default of 8192 — phase 3 is now parallelised across rayon /// workers (each subarray distributes independently into per-partition @@ -1807,44 +1836,66 @@ impl CascadeWorkspace { ) }; - let mut new_lens = Vec::with_capacity(run_lens.len().div_ceil(2)); - let mut src_off = 0usize; - let mut dst_off = 0usize; - let mut i = 0; - while i < run_lens.len() { - let l1 = run_lens[i]; - if i + 1 < run_lens.len() { - let l2 = run_lens[i + 1]; - let x_end = src_off + l1; - let xy_end = x_end + l2; - let dst_end = dst_off + l1 + l2; - sample_sort::merge( - text, - lp, - &src_sa[src_off..x_end], - &src_sa[x_end..xy_end], - &src_lcp[src_off..x_end], - &src_lcp[x_end..xy_end], - &mut dst_sa[dst_off..dst_end], - &mut dst_lcp[dst_off..dst_end], - max_ctx, - dispatch, - ); - new_lens.push(l1 + l2); - src_off = xy_end; - dst_off = dst_end; - i += 2; - } else { - // Odd run carries over unchanged. - let end = dst_off + l1; - dst_sa[dst_off..end].copy_from_slice(&src_sa[src_off..src_off + l1]); - dst_lcp[dst_off..end].copy_from_slice(&src_lcp[src_off..src_off + l1]); - new_lens.push(l1); - src_off += l1; - dst_off = end; - i += 1; + // The pairs at one level are independent and write to disjoint + // destination ranges, so the only thing that made this sequential was + // the running `src_off` / `dst_off`. Both are prefix sums, so compute + // them up front and hand each pair its own sub-slices. + // + // This matters because the cascade's last level is a single merge over + // the whole partition. With `p` well above the thread count there is + // enough partition-level parallelism to hide that most of the time, + // but it is what caps phase 4's efficiency: it was running at ~8x on + // 12 threads. + let n_pairs = run_lens.len() / 2; + let mut new_lens: Vec = (0..n_pairs) + .map(|j| run_lens[2 * j] + run_lens[2 * j + 1]) + .collect(); + if run_lens.len() % 2 == 1 { + new_lens.push(run_lens[run_lens.len() - 1]); + } + + // `dst` ranges are exactly `new_lens`; `src` ranges are the pairs. + let mut jobs: Vec<(usize, usize, &mut [I], &mut [I])> = Vec::with_capacity(new_lens.len()); + { + let mut sa_rest: &mut [I] = dst_sa; + let mut lcp_rest: &mut [I] = dst_lcp; + let mut src_off = 0usize; + for (j, &out_len) in new_lens.iter().enumerate() { + let (sa_head, sa_tail) = sa_rest.split_at_mut(out_len); + let (lcp_head, lcp_tail) = lcp_rest.split_at_mut(out_len); + jobs.push((j, src_off, sa_head, lcp_head)); + sa_rest = sa_tail; + lcp_rest = lcp_tail; + src_off += out_len; } } + + jobs.into_par_iter() + .for_each(|(j, src_off, out_sa, out_lcp)| { + if 2 * j + 1 < run_lens.len() { + let l1 = run_lens[2 * j]; + let l2 = run_lens[2 * j + 1]; + let x_end = src_off + l1; + let xy_end = x_end + l2; + sample_sort::merge( + text, + lp, + &src_sa[src_off..x_end], + &src_sa[x_end..xy_end], + &src_lcp[src_off..x_end], + &src_lcp[x_end..xy_end], + out_sa, + out_lcp, + max_ctx, + dispatch, + ); + } else { + // Odd run carries over unchanged. + let l1 = run_lens[2 * j]; + out_sa.copy_from_slice(&src_sa[src_off..src_off + l1]); + out_lcp.copy_from_slice(&src_lcp[src_off..src_off + l1]); + } + }); new_lens } }