perf(align): resolve the genome storage once per extension and junction scan - #257
perf(align): resolve the genome storage once per extension and junction scan#257BenjaminDEMAILLE wants to merge 4 commits into
Conversation
…igns
`find_best_junction_position` is the single hottest function in the
aligner: on a 50k-pair yeast PE run it accounts for ~60% of alignment
time. The scan itself is not wasteful, but it is repeated.
`stitchWindowAligns`' include/exclude recursion reaches the same
(exon A end, seed B) pair through many different branch paths, and each
path re-runs the identical scan. The scan is a pure function of its
arguments, and within one window `read_seq`, the genome, `is_reverse`
and `n_genome` are all fixed, so six coordinates identify a scan
completely: the exon A read end and genome end, the read and genome
gaps, the previous exon length and the next seed length.
Add `JunctionScanCache`, a per-window `FxHashMap` on that key, and a
`find_best_junction_position_cached` wrapper that consults it. The
uncached function is untouched, so a hit returns exactly what a fresh
scan would have. The cache is created per window in `stitch_seeds_core`
and threaded down the recursion, which is what keeps the key complete:
it never outlives the read, genome and strand it was filled for. An
empty `HashMap` does not allocate, so windows that stitch nothing pay
nothing.
Measured on 50k yeast read pairs (Apple M4 Max, quiet machine,
best-of-6 interleaved rounds, `--outSAMtype None`):
threads before after change
1 19.51s 16.76s -14.1%
8 2.55s 2.18s -14.5%
Output is byte-identical: `Aligned.out.sam` (records and header alike,
modulo the `@PG` CL line naming the binary) and `SJ.out.tab` compare
equal against the pre-change binary on the same input. 592 tests pass,
0 clippy warnings.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…on scan `Genome::get_base` matches on the `GenomeSeq` discriminant, bounds-checks and, for a memory-mapped genome's reverse-complement half, recomputes the mirrored index and complements the byte. That is fine per call, but the alignment inner loops call it once per base: the junction-position scan reads two bases per candidate position across three loops, and `extend_alignment` reads one per extended base. Together those two functions are ~55% of alignment time after the scan memo. Add `GenomeSeq::view()`, returning a `Copy` `SeqView` that resolves the variant once. The view is a slice plus one integer, so the loops keep it in registers and each base costs a bounds check and a load. Hoist it out of the junction scan (including the sliding motif window) and out of both `extend_alignment` loops. Out-of-range reads return the `OUT_OF_RANGE` sentinel instead of `None`. Every call site converted here already treated "not one of A/C/G/T" the same way a `None` was treated, so the branch structure is preserved exactly; `score.rs` already had this sentinel locally for the motif window and it moves next to the view it belongs to. `SeqView::base` duplicates the reverse-complement arithmetic in `GenomeSeq::base`, so a unit test asserts the two agree on every index in `0..2n` plus the first out-of-range one, for both storage variants. Measured on 50k yeast read pairs (Apple M4 Max, quiet machine, best-of-6 interleaved rounds, `--outSAMtype None`, 8 threads), on top of the junction-scan memo: 2.18s to 2.12s, and 2.60s to 2.12s against the pre-memo baseline (-18%). Output is byte-identical: `Aligned.out.sam` and `SJ.out.tab` compare equal against the pre-change binary on the same input. 593 tests pass, 0 clippy warnings. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…base at a time `score_region` walks a seed-length run scoring matches and mismatches. It was ~11% of alignment time, and its loop could not vectorize: every base re-derived a bounds-checked `Option` from `Genome::get_base` and tested the read end, so the body carried branches and an early exit. Bound the run once against the read length, then walk it in 256-base chunks with the genome bases staged into a stack buffer through the new `SeqView::bases_into`. The inner loop is then two plain byte slices of equal length reduced into a match count and a mismatch count, which is what lets it vectorize; `score` is still exactly `matches - mismatches` and the genome-end `break` is preserved by stopping on a short fill. `bases_into` is one `copy_from_slice` on the forward strand. The reverse-complement half of a mapped genome has no contiguous slice to hand out, so it is filled by walking the mirrored forward bytes, which still leaves the comparison itself vectorizable. The reduction uses bitwise `&` rather than `&&` and suppresses `clippy::needless_bitwise_bool` at that loop. This is measured, not stylistic: the lazy spelling reintroduces branches and gives back most of the gain (median 2.125s vs 2.085s wall on the benchmark below). Measured on 50k yeast read pairs (Apple M4 Max, 8 threads, best-of-6 interleaved rounds, `--outSAMtype None`): 2.10s to 2.04s best, 2.155s to 2.085s median. Small, but consistent across every paired round. Output is byte-identical on two datasets: yeast 50k pairs and the nfcore test pair (88k SAM lines). `Aligned.out.sam` and `SJ.out.tab` compare equal against the pre-change binary in both. 593 tests pass, 0 clippy warnings, fmt clean. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
Pushed a third commit to this branch: Same idea as the second commit, applied to the one remaining per-base loop that the profile still showed. Worth flagging for review: that loop uses bitwise Gain is modest and honestly reported: 2.10s → 2.04s best, 2.155s → 2.085s median, but consistent across every paired round. Output identity now checked on two datasets rather than one: yeast 50k pairs and the nfcore test pair (88k SAM lines). For context on what was tried and rejected while chasing this, so nobody repeats it:
|
…hat follows `stitch_align_to_transcript` clones the working transcript and then always pushes onto it. `Vec::clone` allocates exactly `len`, so that push reallocates every time: a malloc, a copy and a free for every stitched seed, on a path the recursion walks up to its 100k-node budget per window. Add `WorkingTranscript::clone_with_headroom`, which reserves the one slot the caller is about to use, folding the reallocation back into the clone's own allocation. The junction vectors only get headroom when they already hold something. Most transcripts carry no junction at all, and giving an empty vector capacity would allocate for a push that never comes, which is worse than what it replaces. The exon vector is never empty at these call sites, so it always gets the slot. Measured on 50k yeast read pairs (Apple M4 Max, 1 thread, 5 rounds, `--outSAMtype None`), reported as user CPU time rather than wall: this machine's wall clock was too noisy to resolve half a percent, and user time is not. Median 14.86s to 14.79s, -0.5%, with every round at the same rank improving. Small, and labelled as such. Pre-sizing the per-window transcript accumulator was tried alongside this and measured a small loss (median 14.97s, +0.7%): it over-allocates for the many windows that finish with only a few transcripts. It is not included here. Output is byte-identical on two datasets: yeast 50k pairs and the nfcore test pair. `Aligned.out.sam` and `SJ.out.tab` compare equal against the pre-change binary in both. 593 tests pass, 0 clippy warnings, fmt clean. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
Fourth commit pushed:
On measurement method, since it matters for a change this size: the wall clock on my machine could not resolve half a percent (paired rounds disagreed, with 2.4s swings from unrelated load). User CPU time resolves it cleanly, so that is what the commit reports. Median 14.86s → 14.79s at 1 thread, -0.5%, every round improving at the same rank. I also tried pre-sizing the per-window transcript accumulator in the same change. It measured a small loss (median 14.97s, +0.7%) because it over-allocates for the many windows that finish with only a few transcripts, so it is not in the commit. Where the branch stands overallAgainst the pre-#256 baseline, on 50k yeast read pairs, user CPU at 1 thread: 18.82s → 14.79s, -21.4%. Wall at 8 threads on a quiet machine: 2.67s → 2.12s. Output byte-identical throughout, checked on yeast and the nfcore test pair. Other commands, for the recordI checked the rest of the surface rather than only the aligner:
genomeGenerate: exact-repeat blowup in caps-saReproducible, and it is not my synthetic worst case alone:
The middle row is the one that matters: a 171-bp monomer array is what human alpha-satellite looks like, and human centromeres carry tens of Mb of it. Profile puts the time in Things I tried that did not help:
So this is not tunable from ruSTAR's side; it needs either an upstream change in caps-sa (an induced-sorting path for long-LCP runs) or a different construction algorithm. Capping |
1 participant
Stacked on #256, which cannot be used as a base branch here (it lives on a fork), so this PR is opened against
mainand its diff contains #256's commit as well. Review only the second commit (perf(align): resolve the genome storage once...); merge #256 first and this one rebases down to that single commit.What
Genome::get_basematches on theGenomeSeqdiscriminant, bounds-checks, and for a memory-mapped genome's reverse-complement half recomputes the mirrored index and complements the byte. Fine per call, but the alignment inner loops call it once per base:extend_alignmentreads one per extended baseAfter #256 those two functions are still ~55% of alignment time (31% + 24%).
This PR adds
GenomeSeq::view(), returning aCopySeqViewthat resolves the variant once. The view is a slice plus one integer, so the loops keep it in registers and each base costs a bounds check and a load. It is hoisted out of the junction scan (motif window included) and out of bothextend_alignmentloops.Sentinel instead of Option
Out-of-range reads return
OUT_OF_RANGErather thanNone. Every call site converted here already treated "not one of A/C/G/T" the same way it treated aNone, so the branch structure is preserved exactly.score.rsalready had this sentinel locally for the motif window; it moves next to the view it belongs to.SeqView::baseduplicates the reverse-complement arithmetic inGenomeSeq::base, so a unit test asserts the two agree on every index in0..2nplus the first out-of-range one, for both storage variants. That test is the guard against the two definitions drifting apart.Benchmark
50k yeast read pairs, Apple M4 Max, quiet machine, best-of-6 interleaved rounds,
--outSAMtype None, 8 threads:Smaller than #256 on its own, and reported as such: most of the remaining cost in those loops is the scanning work itself, not the dispatch.
Correctness
Output is byte-identical.
Aligned.out.samandSJ.out.tabcompare equal against the pre-change binary on the same input.--all-targets --release)cargo fmt --checkclean🤖 Generated with Claude Code