perf(align): stitch into the transcript in place instead of cloning per node - #264
Open
BenjaminDEMAILLE wants to merge 5 commits into
Open
perf(align): stitch into the transcript in place instead of cloning per node#264BenjaminDEMAILLE wants to merge 5 commits into
BenjaminDEMAILLE wants to merge 5 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>
…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>
…er node `stitchWindowAligns`' include/exclude recursion cloned a `WorkingTranscript` at every node it explored, because the include branch needed its own copy while the exclude branch still needed the original. With a 100k-node budget per window that is a clone, four vector allocations and a copy, per node. Pass the transcript by `&mut` and undo the attempt instead. The mutation surface is small enough to invert exactly: `stitch_align_to_transcript` and the base-case extensions only append to the four vectors and rewrite the *first* and *last* exons plus the scalars, never a middle exon and never a removal or reorder. `WtMark` records the four lengths, copies of the first and last exons, and the scalars; `restore` truncates and writes those back. Vector capacity survives the restore, so sibling branches reuse the same allocations rather than reallocating. An owned clone is now taken only where one is actually needed: when a completed transcript is accepted into the result set, which is bounded by `--alignTranscriptsPerWindowNmax` rather than by the recursion. This supersedes `clone_with_headroom`, which existed to make the per-node clone cheaper and is removed: there is no per-node clone left to soften. Measured, user CPU at 1 thread, `--outSAMtype None`: | dataset | before | after | |---|---|---| | nfcore PE | 22.29s | 21.72s (-2.6%) | | yeast PE | 15.24s | 15.24s (neutral) | The gain tracks stitch depth: nfcore's reads drive a deeper recursion, so more clones disappear. Yeast is unchanged rather than slower, which is the point of checking both. Beyond the timing, this is what a batched or GPU-side stitcher would need: per-branch state is now a mark and an undo rather than a fresh allocation. Output is byte-identical on both datasets: `Aligned.out.sam` and `SJ.out.tab` compare equal against the pre-change binary. 593 tests pass, 0 clippy warnings, fmt clean. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Builds on #257 (contains its commits). Review only the last commit; merge #257 first and this rebases down to it.
What
stitchWindowAligns' include/exclude recursion cloned aWorkingTranscriptat every node, because the include branch needed its own copy while the exclude branch still needed the original. With a 100k-node budget per window that is a clone, four vector allocations and a copy, per node.This passes the transcript by
&mutand undoes the attempt instead.Why the undo is exact
The mutation surface is small enough to invert, and I checked it rather than assumed it.
stitch_align_to_transcriptand the base-case extensions only:Never a middle exon, never a removal, never a reorder. So
WtMark(four lengths + copies of the first and last exons + the scalars) is a complete inverse, andrestoretruncates and writes those back. Capacity survives the restore, so sibling branches reuse allocations instead of reallocating.I got this wrong on the first pass: my initial
WtMarksaved only the last exon, and the base case turned out to mutateexons.first_mut()too. That is why the record carries both.An owned clone is now taken only where one is genuinely needed: when a completed transcript is accepted into the result set, bounded by
--alignTranscriptsPerWindowNmaxrather than by the recursion.Supersedes the headroom commit
clone_with_headroom(the last commit of #257) existed to make the per-node clone cheaper. There is no per-node clone left, so it is removed here. If you would rather not carry that commit only to delete it, I can rebase #257 without it.Measured
User CPU at 1 thread,
--outSAMtype None, 4 interleaved rounds:Ranges do not overlap on nfcore (21.71-21.77 vs 22.26-22.47). The gain tracks stitch depth: nfcore drives a deeper recursion so more clones disappear. Yeast is unchanged rather than slower, which is why both are reported.
Correctness
Aligned.out.samandSJ.out.tabbyte-identical on yeast 50k pairs and the nfcore pair.cargo fmt --checkclean.Why it may be worth more than 2.6%
Per-branch state is now a mark and an undo rather than a fresh allocation, which is the shape a batched or GPU-side stitcher needs. On its own the CPU gain is modest and I am not overselling it; the structural change is the larger part of the value, and it stands or falls on your view of whether that direction is worth having.
🤖 Generated with Claude Code