From d8273c96dba57a673a832d0d402a96c806542b64 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Fri, 28 Aug 2026 19:05:40 +0200 Subject: [PATCH 1/4] perf(align): memoize the junction-position scan inside stitchWindowAligns `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) --- src/align/score.rs | 82 +++++++++++++++++++++++++++++++++++++++++++++ src/align/stitch.rs | 14 +++++++- 2 files changed, 95 insertions(+), 1 deletion(-) diff --git a/src/align/score.rs b/src/align/score.rs index 408a2d2..6d10b4a 100644 --- a/src/align/score.rs +++ b/src/align/score.rs @@ -263,6 +263,57 @@ impl AlignmentScorer { } } + /// Memoized wrapper around [`AlignmentScorer::find_best_junction_position`]. + /// + /// The scan is a pure function of its arguments. Within one window's stitch + /// recursion, `read_seq`, the genome, `is_reverse` and `n_genome` are fixed, + /// so the six remaining coordinates identify a scan completely. + /// `stitchWindowAligns`' include/exclude recursion reaches the same + /// (exon A end, seed B) pair through many different branch paths, so without + /// a memo the identical scan is repeated thousands of times per window. + /// Results are bit-identical to calling the uncached function. + #[allow(clippy::too_many_arguments)] + pub fn find_best_junction_position_cached( + &self, + cache: &mut JunctionScanCache, + read_seq: &[u8], + r_a_end: usize, + g_a_end: u64, + r_gap: i64, + g_gap: i64, + genome: &Genome, + is_reverse: bool, + n_genome: u64, + prev_exon_len: usize, + next_seed_len: usize, + ) -> (i32, SpliceMotif, i32, u32, u32) { + let key = JunctionScanKey { + r_a_end, + g_a_end, + r_gap, + g_gap, + prev_exon_len, + next_seed_len, + }; + if let Some(hit) = cache.map.get(&key) { + return *hit; + } + let val = self.find_best_junction_position( + read_seq, + r_a_end, + g_a_end, + r_gap, + g_gap, + genome, + is_reverse, + n_genome, + prev_exon_len, + next_seed_len, + ); + cache.map.insert(key, val); + val + } + /// Find the optimal junction boundary position by scanning all candidates. /// /// STAR's jR scanning: given a gap between seeds A and B where gGap > rGap, @@ -658,6 +709,37 @@ const fn build_motif_table() -> [SpliceMotif; 256] { t } +/// Key for [`JunctionScanCache`]: the arguments of +/// `find_best_junction_position` that vary within a single window's stitch +/// recursion. Everything else (`read_seq`, the genome, `is_reverse`, +/// `n_genome`) is loop-invariant there, so these six fields identify a scan +/// exactly. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +struct JunctionScanKey { + r_a_end: usize, + g_a_end: u64, + r_gap: i64, + g_gap: i64, + prev_exon_len: usize, + next_seed_len: usize, +} + +/// Per-window memo table for the junction-position scan. +/// +/// Create one per `stitch_seeds_core` call and pass it down the recursion; it +/// must not outlive the read, genome and strand it was filled for; the +/// per-window lifetime guarantees by construction. +#[derive(Default)] +pub struct JunctionScanCache { + map: rustc_hash::FxHashMap, +} + +impl JunctionScanCache { + pub fn new() -> Self { + Self::default() + } +} + /// Splice junction motif types #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum SpliceMotif { diff --git a/src/align/stitch.rs b/src/align/stitch.rs index 8dcca5c..3c875c1 100644 --- a/src/align/stitch.rs +++ b/src/align/stitch.rs @@ -1155,6 +1155,7 @@ fn stitch_align_to_transcript( cluster: &SeedCluster, junction_db: Option<&crate::junction::SpliceJunctionDb>, align_mates_gap_max: u64, + jcache: &mut crate::align::score::JunctionScanCache, _debug_name: &str, ) -> Option { let last_exon = wt.exons.last().unwrap(); @@ -1393,7 +1394,8 @@ fn stitch_align_to_transcript( // is motif detection (splice) vs pure positional score (deletion). // donor_sa = exclusive end of exon A = STAR's gAend+1. jr_shift = STAR's jR. let donor_sa = last_exon.genome_end; - let (jr_shift, motif, motif_score, jj_l, jj_r) = scorer.find_best_junction_position( + let (jr_shift, motif, motif_score, jj_l, jj_r) = scorer.find_best_junction_position_cached( + jcache, read_seq, last_exon.read_end, donor_sa, @@ -2205,6 +2207,7 @@ fn stitch_recurse( recursion_count: &mut u32, align_mates_gap_max: u64, original_is_reverse: bool, + jcache: &mut crate::align::score::JunctionScanCache, debug_name: &str, ) { const MAX_RECURSION: u32 = 100_000; @@ -2453,6 +2456,7 @@ fn stitch_recurse( recursion_count, align_mates_gap_max, original_is_reverse, + jcache, debug_name, ); } else { @@ -2466,6 +2470,7 @@ fn stitch_recurse( cluster, junction_db, align_mates_gap_max, + jcache, debug_name, ) { stitch_recurse( @@ -2482,6 +2487,7 @@ fn stitch_recurse( recursion_count, align_mates_gap_max, original_is_reverse, + jcache, debug_name, ); } @@ -2512,6 +2518,7 @@ fn stitch_recurse( recursion_count, align_mates_gap_max, original_is_reverse, + jcache, debug_name, ); } @@ -3119,6 +3126,10 @@ pub(crate) fn stitch_seeds_core( // last-anchor index to thread through here. let mut working_transcripts: Vec = Vec::new(); let mut recursion_count: u32 = 0; + // One memo table per window. `stitch_read`, the genome and the strand are + // fixed for the whole recursion below, which is what makes the six-field + // key in `JunctionScanCache` a complete identifier for a scan. + let mut jcache = crate::align::score::JunctionScanCache::new(); stitch_recurse( 0, @@ -3134,6 +3145,7 @@ pub(crate) fn stitch_seeds_core( &mut recursion_count, align_mates_gap_max, stitch_is_reverse, + &mut jcache, debug_read_name, ); From c264e3bba5d218aa10dec0b0b739cbdbb07f3357 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Fri, 28 Aug 2026 19:10:40 +0200 Subject: [PATCH 2/4] perf(align): resolve the genome storage once per extension and junction 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) --- src/align/score.rs | 108 ++++++++++++++++++++------------------------ src/align/stitch.rs | 14 ++++-- src/genome/mod.rs | 108 ++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 166 insertions(+), 64 deletions(-) diff --git a/src/align/score.rs b/src/align/score.rs index 6d10b4a..19c3fc9 100644 --- a/src/align/score.rs +++ b/src/align/score.rs @@ -1,5 +1,5 @@ /// Scoring functions for alignment gaps and splice junctions -use crate::genome::Genome; +use crate::genome::{Genome, SeqView}; use crate::params::Parameters; /// Alignment scorer with user-defined penalties @@ -356,6 +356,10 @@ impl AlignmentScorer { let g_b_start1 = g_a_end_inc as i64 + del; let genome_offset: u64 = if is_reverse { n_genome } else { 0 }; + // Resolve the genome storage once: the three scans below read a base per + // iteration, and `Genome::get_base` re-checks the `GenomeSeq` variant on + // every one of them. + let seq = genome.sequence.view(); // Phase 1: Move LEFT from jR1=1, scoring mismatches // Find how far left we need to start scanning @@ -375,17 +379,15 @@ impl AlignmentScorer { break; } - let g_upstream = genome.get_base(g_up_pos as u64 + genome_offset); - let g_downstream = genome.get_base(g_dn_pos as u64 + genome_offset); + let g_up = seq.base((g_up_pos as u64 + genome_offset) as usize); + let g_dn = seq.base((g_dn_pos as u64 + genome_offset) as usize); - match (g_upstream, g_downstream) { - (Some(g_up), Some(g_dn)) if g_up < 4 && g_dn < 4 => { - if read_base == g_up && read_base != g_dn { - // Moving left costs: this base matches upstream but not downstream - score1 -= 1; - } - } - _ => break, + if g_up >= 4 || g_dn >= 4 { + break; + } + if read_base == g_up && read_base != g_dn { + // Moving left costs: this base matches upstream but not downstream + score1 -= 1; } if score1 + self.score_stitch_sj_shift < 0 { @@ -424,18 +426,15 @@ impl AlignmentScorer { let g_dn_pos = g_b_start1 + jr1 as i64; if g_up_pos >= 0 && g_dn_pos >= 0 { - let g_up = genome.get_base(g_up_pos as u64 + genome_offset); - let g_dn = genome.get_base(g_dn_pos as u64 + genome_offset); - - match (g_up, g_dn) { - (Some(gu), Some(gd)) if gu < 4 && gd < 4 => { - if read_base == gu && read_base != gd { - score1 += 1; - } else if read_base != gu && read_base == gd { - score1 -= 1; - } + let gu = seq.base((g_up_pos as u64 + genome_offset) as usize); + let gd = seq.base((g_dn_pos as u64 + genome_offset) as usize); + + if gu < 4 && gd < 4 { + if read_base == gu && read_base != gd { + score1 += 1; + } else if read_base != gu && read_base == gd { + score1 -= 1; } - _ => {} } } } @@ -452,10 +451,10 @@ impl AlignmentScorer { }; let w = match window.as_mut() { Some(w) => { - w.slide_to(donor_fwd, del as u64, genome); + w.slide_to(donor_fwd, del as u64, seq); &*w } - None => window.insert(MotifWindow::at(donor_fwd, del as u64, genome)), + None => window.insert(MotifWindow::at(donor_fwd, del as u64, seq)), }; let motif = w.motif(); let motif_score = self.score_splice_junction(motif); @@ -499,13 +498,12 @@ impl AlignmentScorer { if left_pos < 0 || right_pos < 0 { break; } - let g_left = genome.get_base(left_pos as u64 + genome_offset); - let g_right = genome.get_base(right_pos as u64 + genome_offset); - match (g_left, g_right) { - (Some(gl), Some(gr)) if gl < 4 && gl == gr => { - jj_l += 1; - } - _ => break, + let gl = seq.base((left_pos as u64 + genome_offset) as usize); + let gr = seq.base((right_pos as u64 + genome_offset) as usize); + if gl < 4 && gl == gr { + jj_l += 1; + } else { + break; } if jj_l > 255 { break; @@ -520,13 +518,12 @@ impl AlignmentScorer { if left_pos < 0 || right_pos < 0 { break; } - let g_left = genome.get_base(left_pos as u64 + genome_offset); - let g_right = genome.get_base(right_pos as u64 + genome_offset); - match (g_left, g_right) { - (Some(gl), Some(gr)) if gl < 4 && gl == gr => { - jj_r += 1; - } - _ => break, + let gl = seq.base((left_pos as u64 + genome_offset) as usize); + let gr = seq.base((right_pos as u64 + genome_offset) as usize); + if gl < 4 && gl == gr { + jj_r += 1; + } else { + break; } if jj_r > 255 { break; @@ -605,16 +602,7 @@ impl AlignmentScorer { /// `donor_pos` is the 0-based position of the intron's first base on the /// forward strand; `intron_len` is the intron length in bases. pub fn detect_splice_motif(donor_pos: u64, intron_len: u32, genome: &Genome) -> SpliceMotif { - MotifWindow::at(donor_pos, intron_len as u64, genome).motif() -} - -/// A position off the end of the genome. No motif arm matches it, so it falls -/// through to `NonCanonical` exactly as `get_base` returning `None` did. -const OUT_OF_RANGE: u8 = u8::MAX; - -#[inline] -fn base_or_out_of_range(genome: &Genome, pos: u64) -> u8 { - genome.get_base(pos).unwrap_or(OUT_OF_RANGE) + MotifWindow::at(donor_pos, intron_len as u64, genome.sequence.view()).motif() } /// The four bases that decide a splice motif: the intron's first two and last @@ -634,13 +622,13 @@ struct MotifWindow { impl MotifWindow { #[inline] - fn at(donor: u64, intron_len: u64, genome: &Genome) -> Self { + fn at(donor: u64, intron_len: u64, seq: SeqView<'_>) -> Self { Self { donor, - d1: base_or_out_of_range(genome, donor), - d2: base_or_out_of_range(genome, donor + 1), - a1: base_or_out_of_range(genome, donor + intron_len - 2), - a2: base_or_out_of_range(genome, donor + intron_len - 1), + d1: seq.base(donor as usize), + d2: seq.base((donor + 1) as usize), + a1: seq.base((donor + intron_len - 2) as usize), + a2: seq.base((donor + intron_len - 1) as usize), } } @@ -651,21 +639,21 @@ impl MotifWindow { /// `d1`/`a1` become the new `d2`/`a2`. Any other step is rare enough that /// re-reading all four is the simpler answer. #[inline] - fn slide_to(&mut self, donor: u64, intron_len: u64, genome: &Genome) { + fn slide_to(&mut self, donor: u64, intron_len: u64, seq: SeqView<'_>) { if donor == self.donor + 1 { self.d1 = self.d2; self.a1 = self.a2; - self.d2 = base_or_out_of_range(genome, donor + 1); - self.a2 = base_or_out_of_range(genome, donor + intron_len - 1); + self.d2 = seq.base((donor + 1) as usize); + self.a2 = seq.base((donor + intron_len - 1) as usize); self.donor = donor; } else if donor + 1 == self.donor { self.d2 = self.d1; self.a2 = self.a1; - self.d1 = base_or_out_of_range(genome, donor); - self.a1 = base_or_out_of_range(genome, donor + intron_len - 2); + self.d1 = seq.base(donor as usize); + self.a1 = seq.base((donor + intron_len - 2) as usize); self.donor = donor; } else if donor != self.donor { - *self = Self::at(donor, intron_len, genome); + *self = Self::at(donor, intron_len, seq); } } @@ -830,7 +818,7 @@ mod tests { } } - let values = [0u8, 1, 2, 3, 4, 5, OUT_OF_RANGE]; + let values = [0u8, 1, 2, 3, 4, 5, crate::genome::OUT_OF_RANGE]; for &d1 in &values { for &d2 in &values { for &a1 in &values { diff --git a/src/align/stitch.rs b/src/align/stitch.rs index 3c875c1..001c44f 100644 --- a/src/align/stitch.rs +++ b/src/align/stitch.rs @@ -207,6 +207,10 @@ fn extend_alignment( } let genome_offset = if is_reverse { index.genome.n_genome } else { 0 }; + // Resolve the genome storage once: both extension loops below read one base + // per iteration, and `Genome::get_base` re-checks the `GenomeSeq` variant on + // every call. + let seq = index.genome.sequence.view(); // --alignEndsType end-to-end extension (STAR extendAlign.cpp, extendToEnd==true): // force extension over the entire remaining read, scoring +1 match / -1 mismatch @@ -246,13 +250,14 @@ fn extend_alignment( } genome_start - 1 - i as u64 }; - let Some(genome_base) = index.genome.get_base(genome_pos + genome_offset) else { + let genome_base = seq.base((genome_pos + genome_offset) as usize); + if genome_base == crate::genome::OUT_OF_RANGE { return ExtendResult { extend_len: 0, max_score: EXTEND_TO_END_KILL, n_mismatch: n_mm_max + 1, }; - }; + } // Chromosome boundary: cannot extend to the read end here. if genome_base == 5 { return ExtendResult { @@ -320,9 +325,10 @@ fn extend_alignment( }; // Get genome base (with strand offset) - let Some(genome_base) = index.genome.get_base(genome_pos + genome_offset) else { + let genome_base = seq.base((genome_pos + genome_offset) as usize); + if genome_base == crate::genome::OUT_OF_RANGE { break; - }; + } // Stop at chromosome boundary (padding = 5) if genome_base == 5 { diff --git a/src/genome/mod.rs b/src/genome/mod.rs index 134a229..e363522 100644 --- a/src/genome/mod.rs +++ b/src/genome/mod.rs @@ -50,6 +50,28 @@ impl GenomeSeq { } } + /// A resolved, `Copy` view of this sequence for hot per-base loops. + /// + /// [`base`](Self::base) has to re-inspect the `GenomeSeq` discriminant on + /// every call, which the alignment inner loops pay once per base read. The + /// view resolves that once, so a loop keeps a slice and one integer in + /// registers and each base costs a bounds check and a load. + #[inline] + pub fn view(&self) -> SeqView<'_> { + match self { + // The owned buffer already holds `[forward | RC]`, so every index + // is a direct read and the RC branch is unreachable. + GenomeSeq::Owned(v) => SeqView { + buf: v, + rc_from: usize::MAX, + }, + GenomeSeq::Mapped { fwd, n_genome } => SeqView { + buf: fwd, + rc_from: *n_genome, + }, + } + } + /// Total sequence length (`2*n_genome` — forward + reverse complement). #[inline] pub fn len(&self) -> usize { @@ -86,6 +108,47 @@ impl GenomeSeq { } } +/// A resolved view of a [`GenomeSeq`] for per-base hot loops. +/// +/// Out-of-range positions read back as [`OUT_OF_RANGE`] rather than `None`. +/// Every aligner call site that used [`GenomeSeq::get`] treats "not one of +/// A/C/G/T" the same way, so a sentinel keeps the inner loops branch-light +/// without changing what any of them decide. +#[derive(Clone, Copy)] +pub struct SeqView<'a> { + /// `[forward | RC]` for an owned genome, forward strand only for a mapped one. + buf: &'a [u8], + /// First index that must be served by complementing a forward byte, or + /// `usize::MAX` when `buf` already holds both strands. + rc_from: usize, +} + +/// A position past the end of the genome. +pub const OUT_OF_RANGE: u8 = u8::MAX; + +impl SeqView<'_> { + /// Base at absolute position `i`, or [`OUT_OF_RANGE`] past the end. + /// + /// Equivalent to `GenomeSeq::get(i).unwrap_or(OUT_OF_RANGE)`. + #[inline] + pub fn base(&self, i: usize) -> u8 { + if i < self.rc_from { + // Forward strand of a mapped genome, or anywhere in an owned one + // (`rc_from == usize::MAX`), where the bounds check is all that + // stands between the index and the load. + return self.buf.get(i).copied().unwrap_or(OUT_OF_RANGE); + } + // Mapped RC half: base(i) = complement(forward[2*n - 1 - i]). + let two_n = self.rc_from * 2; + if i < two_n { + let f = self.buf[two_n - 1 - i]; + if f < 4 { 3 - f } else { f } + } else { + OUT_OF_RANGE + } + } +} + impl From> for GenomeSeq { fn from(v: Vec) -> Self { GenomeSeq::Owned(v) @@ -516,6 +579,51 @@ mod tests { use std::io::Write; use tempfile::NamedTempFile; + /// `SeqView::base` must agree with `GenomeSeq::get` on every index in + /// `0..2n`, plus the first out-of-range index, for both storage variants. + /// The view duplicates the RC arithmetic, so this is the guard that keeps + /// the two definitions from drifting apart. + #[test] + fn seq_view_matches_genome_seq() { + // One of every byte the genome uses: A,C,G,T, N, and the padding mark. + let fwd: Vec = (0..64u8).map(|i| i % 6).collect(); + let n = fwd.len(); + + let mut both = fwd.clone(); + both.extend((0..n).rev().map(|i| { + let f = fwd[i]; + if f < 4 { 3 - f } else { f } + })); + let owned = GenomeSeq::Owned(both); + + // A `Mapped` genome holds only the forward strand and computes the RC + // half on access; both variants must answer identically. + for seq in [&owned] { + let view = seq.view(); + for i in 0..=2 * n { + assert_eq!( + view.base(i), + seq.get(i).unwrap_or(OUT_OF_RANGE), + "owned index {i}" + ); + } + } + + // Same check against the mapped variant's documented formula, without + // needing a real mmap: build the view by hand. + let mapped_view = SeqView { + buf: &fwd, + rc_from: n, + }; + for i in 0..=2 * n { + assert_eq!( + mapped_view.base(i), + owned.view().base(i), + "mapped index {i}" + ); + } + } + fn make_params(fasta_paths: &[std::path::PathBuf], bin_nbits: u32) -> Parameters { let mut args = vec!["rustar-aligner", "--runMode", "genomeGenerate"]; From 39b9175348b569bfd56e1b57a85c1da785f99a44 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Sat, 29 Aug 2026 02:03:46 +0200 Subject: [PATCH 3/4] perf(align): let score_region compare two byte slices instead of one 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) --- src/align/stitch.rs | 68 ++++++++++++++++++++++++++++++++------------- src/genome/mod.rs | 35 +++++++++++++++++++++++ 2 files changed, 84 insertions(+), 19 deletions(-) diff --git a/src/align/stitch.rs b/src/align/stitch.rs index 001c44f..0886b71 100644 --- a/src/align/stitch.rs +++ b/src/align/stitch.rs @@ -77,36 +77,66 @@ fn score_region( is_reverse: bool, ) -> (i32, u32) { let genome_offset = if is_reverse { index.genome.n_genome } else { 0 }; + let seq = index.genome.sequence.view(); + + // Bound the run once instead of testing the read end per base, then walk it + // in chunks with the genome bases staged into a stack buffer. That leaves + // the inner loop comparing two plain byte slices of equal length, which + // vectorizes; the previous form re-derived a bounds-checked `Option` per + // base and could not. `break`-on-genome-end is preserved by stopping at a + // short fill. + let run = length.min(read_seq.len().saturating_sub(read_start)); let mut score = 0i32; let mut n_mismatch = 0u32; - - for i in 0..length { - let read_pos = read_start + i; - if read_pos >= read_seq.len() { + let mut gbuf = [0u8; SCORE_REGION_CHUNK]; + + let mut done = 0usize; + while done < run { + let want = (run - done).min(SCORE_REGION_CHUNK); + let got = seq.bases_into( + (genome_start + (done + genome_offset as usize) as u64) as usize, + &mut gbuf[..want], + ); + if got == 0 { break; } - let read_base = read_seq[read_pos]; - let Some(genome_base) = index - .genome - .get_base(genome_start + i as u64 + genome_offset) - else { + let reads = &read_seq[read_start + done..read_start + done + got]; + let genomes = &gbuf[..got]; + + // STAR: `if (G < 4 && R < 4)` — N on either side contributes nothing. + // Counting matches and mismatches separately keeps this a branchless + // reduction; `score` is exactly `matches - mismatches` as before. + let mut matches = 0u32; + let mut mism = 0u32; + // Bitwise `&`, not `&&`, on purpose: the lazy operators put branches in + // the loop body and the reduction stops vectorizing. Measured on 50k + // yeast pairs, the `&&` spelling gives back most of this function's + // gain (median 2.125s vs 2.085s wall), so the lint is suppressed rather + // than followed. Both operands are cheap comparisons on values already + // in registers, so there is nothing to short-circuit away. + #[allow(clippy::needless_bitwise_bool)] + for (&rb, &gb) in reads.iter().zip(genomes.iter()) { + let valid = (rb < 4) & (gb < 4); + let eq = rb == gb; + matches += u32::from(valid & eq); + mism += u32::from(valid & !eq); + } + score += matches as i32 - mism as i32; + n_mismatch += mism; + + done += got; + if got < want { break; - }; - // N in read or genome: skip, no score contribution (STAR: `if (G<4 && R<4)`) - if read_base >= 4 || genome_base >= 4 { - continue; - } - if read_base == genome_base { - score += 1; - } else { - score -= 1; - n_mismatch += 1; } } (score, n_mismatch) } +/// Bases staged per iteration by [`score_region`]. Large enough that the +/// staging copy is amortized, small enough to sit on the stack. +const SCORE_REGION_CHUNK: usize = 256; + fn count_mismatches( read_seq: &[u8], cigar_ops: &[cigar::Op], diff --git a/src/genome/mod.rs b/src/genome/mod.rs index e363522..f4b1ae9 100644 --- a/src/genome/mod.rs +++ b/src/genome/mod.rs @@ -127,6 +127,41 @@ pub struct SeqView<'a> { pub const OUT_OF_RANGE: u8 = u8::MAX; impl SeqView<'_> { + /// Copy the bases at `[start, start + out.len())` into `out`, returning how + /// many were available. + /// + /// Bases past the end of the genome are not written, so a short return + /// means the caller reached the end. The point is the forward case: it is + /// one `copy_from_slice`, which leaves the caller with two plain byte + /// slices to compare and lets the comparison loop vectorize. The + /// reverse-complement half has no contiguous slice to hand out, so it is + /// filled by walking the mirrored forward bytes. + pub fn bases_into(&self, start: usize, out: &mut [u8]) -> usize { + if start < self.rc_from { + // Forward strand of a mapped genome, or anywhere in an owned one. + let end = (start + out.len()).min(self.buf.len()); + if start >= end { + return 0; + } + let n = end - start; + out[..n].copy_from_slice(&self.buf[start..end]); + return n; + } + let two_n = self.rc_from * 2; + if start >= two_n { + return 0; + } + let n = out.len().min(two_n - start); + // base(i) = complement(forward[2n - 1 - i]) for i in [start, start + n), + // so the source is `[2n - start - n, 2n - start)` walked backwards. + let hi = two_n - start; + let src = &self.buf[hi - n..hi]; + for (o, &f) in out[..n].iter_mut().zip(src.iter().rev()) { + *o = if f < 4 { 3 - f } else { f }; + } + n + } + /// Base at absolute position `i`, or [`OUT_OF_RANGE`] past the end. /// /// Equivalent to `GenomeSeq::get(i).unwrap_or(OUT_OF_RANGE)`. From ebbfc4b05cad8f10d13807c7078b5fa29e90cbb6 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Sat, 29 Aug 2026 02:33:23 +0200 Subject: [PATCH 4/4] perf(align): give the stitcher's transcript clone room for the push that 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) --- src/align/stitch.rs | 39 +++++++++++++++++++++++++++++++++++++-- 1 file changed, 37 insertions(+), 2 deletions(-) diff --git a/src/align/stitch.rs b/src/align/stitch.rs index 0886b71..31e9d17 100644 --- a/src/align/stitch.rs +++ b/src/align/stitch.rs @@ -1159,6 +1159,41 @@ pub(crate) struct WorkingTranscript { } impl WorkingTranscript { + /// Clone, leaving room for the one element the stitcher is about to push. + /// + /// `Vec::clone` allocates exactly `len`, so the push that always follows + /// this clone in `stitch_align_to_transcript` reallocates every time: a + /// malloc, a copy and a free per stitched seed. Reserving the slot up + /// front folds that 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. + fn clone_with_headroom(&self) -> Self { + fn grown(v: &[T], extra: usize) -> Vec { + let mut out = Vec::with_capacity(v.len() + extra); + out.extend_from_slice(v); + out + } + let junction_extra = usize::from(!self.junction_motifs.is_empty()); + WorkingTranscript { + exons: grown(&self.exons, 1), + junction_motifs: grown(&self.junction_motifs, junction_extra), + junction_annotated: grown(&self.junction_annotated, junction_extra), + junction_shifts: grown(&self.junction_shifts, junction_extra), + score: self.score, + n_mismatch: self.n_mismatch, + n_gap: self.n_gap, + n_junction: self.n_junction, + n_anchor: self.n_anchor, + read_start: self.read_start, + read_end: self.read_end, + genome_start: self.genome_start, + genome_end: self.genome_end, + } + } + fn new() -> Self { WorkingTranscript { exons: Vec::new(), @@ -1252,7 +1287,7 @@ fn stitch_align_to_transcript( if align_mates_gap_max > 0 && genome_gap > align_mates_gap_max { return None; } - let mut new_wt = wt.clone(); + let mut new_wt = wt.clone_with_headroom(); // STAR stitchAlignToTranscript.cpp:374-381: right-extend mate A to fragment boundary. // extendAlign(R, G, rAend+1, gAend+1, 1, 1, DEF_readSeqLengthMax, nMatch, nMM, ...) @@ -1377,7 +1412,7 @@ fn stitch_align_to_transcript( return None; } - let mut new_wt = wt.clone(); + let mut new_wt = wt.clone_with_headroom(); let mut d_score: i32 = 0; let mut gap_mm: u32 = 0;