Skip to content

perf: merge cascade run pairs in parallel (no measurable effect at 12 cores — needs your machine) #13

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
BenjaminDEMAILLE wants to merge 1 commit into from
Closed

perf: merge cascade run pairs in parallel (no measurable effect at 12 cores — needs your machine) #13

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
129 changes: 90 additions & 39 deletions src/ext_mem.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -1807,44 +1836,66 @@ impl<I: Index> CascadeWorkspace<I> {
)
};

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<usize> = (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
}
}
Expand Down