From ccbcca622936f86191e47723ebbe7834ca33f2ee Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 21:33:28 +0200 Subject: [PATCH 01/41] test+build: validate the LCP array, and pin the release profile Two prerequisites for the performance work that follows, neither of which changes behaviour. The crate had no test covering the LCP array. That is the riskiest possible gap for this algorithm: the public entry points discard the array, but the *next* merge level consumes it in the three-case decision, so a single wrong LCP entry silently reorders suffixes one level up and the SA comes out subtly wrong. Add four tests that check `lcp[0] == 0` and `lcp[i] == lcp(text[sa[i-1]..], text[sa[i]..])` against a naive oracle, over fixtures, random texts across four alphabet sizes, long runs and periodic text, and finite `max_context`. `bench/README.md` claims the published numbers were taken with fat LTO and one codegen unit, supplied by a parent workspace. That workspace is not in this repo, so every build made from it since the crate went standalone has used `lto = false, codegen-units = 16`. Pin the profile here. In a library crate `[profile.release]` applies only when this crate is the workspace root, so it affects this repo's own tests, examples and benches and is invisible to downstream consumers. Measured on Apple M4 Max, 12 threads, chr21 FASTA (47.5 MB): 27.8 s -> 24.0 s wall. Neutral on N-free DNA. Co-Authored-By: Claude Opus 5 (1M context) --- Cargo.toml | 13 ++++++ src/sample_sort.rs | 109 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 122 insertions(+) diff --git a/Cargo.toml b/Cargo.toml index 12261dc..4827078 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,6 +15,19 @@ categories = ["algorithms", "data-structures"] repository = "https://github.com/COMBINE-lab/caps-sa" readme = "README.md" +# The published benchmark numbers in `bench/` were taken with fat LTO and a +# single codegen unit. That configuration used to come from a parent workspace +# that no longer exists in this repo, so pin it here. A `[profile.release]` in a +# library crate applies only when this crate is the workspace root — i.e. to +# this repo's own tests, examples and benches — and is invisible to downstream +# consumers, who keep their own profile. +[profile.release] +lto = "fat" +codegen-units = 1 + +[profile.bench] +inherits = "release" + [dependencies] rayon = "1" tempfile = "3" diff --git a/src/sample_sort.rs b/src/sample_sort.rs index daa9264..b2f0dc7 100644 --- a/src/sample_sort.rs +++ b/src/sample_sort.rs @@ -422,6 +422,115 @@ mod tests { assert_eq!(got, want, "mismatch on text {text:?}"); } + /// Run the production kernel and return **both** the suffix array and + /// the LCP array it computes as a byproduct. + /// + /// The public entry points discard the LCP array, but it is not an + /// incidental artefact: the next merge level *consumes* it in the + /// three-case decision, so a single wrong LCP entry silently reorders + /// suffixes at the level above. It therefore needs direct coverage. + fn build_sa_and_lcp(text: &[u8], max_ctx: usize) -> (Vec, Vec) { + let n = text.len(); + let mut sa: Vec = (0..n as u32).collect(); + let mut sa_w = vec![0u32; n]; + let mut lcp_arr = vec![0u32; n]; + let mut lcp_w = vec![0u32; n]; + merge_sort( + text, + &PlainText::new(n), + &mut sa, + &mut sa_w, + &mut lcp_arr, + &mut lcp_w, + max_ctx, + LcpDispatch::detect(), + ); + (sa, lcp_arr) + } + + /// Byte-at-a-time LCP of `text[a..]` and `text[b..]`, capped at `max_ctx`. + fn naive_lcp(text: &[u8], a: usize, b: usize, max_ctx: usize) -> usize { + let lim = (text.len() - a).min(text.len() - b).min(max_ctx); + (0..lim).take_while(|&i| text[a + i] == text[b + i]).count() + } + + /// Assert the LCP-array postcondition stated on [`merge_sort`]: + /// `lcp[0] == 0` and `lcp[i] == lcp(text[sa[i-1]..], text[sa[i]..])`. + fn assert_lcp_valid(text: &[u8], max_ctx: usize) { + let (sa, lcp) = build_sa_and_lcp(text, max_ctx); + if sa.is_empty() { + return; + } + assert_eq!(lcp[0], 0, "lcp[0] must be 0 (text {text:?})"); + for i in 1..sa.len() { + let want = naive_lcp(text, sa[i - 1] as usize, sa[i] as usize, max_ctx); + assert_eq!( + lcp[i] as usize, want, + "lcp[{i}] wrong for sa[{}]={} vs sa[{i}]={} (text {text:?})", + i - 1, + sa[i - 1], + sa[i], + ); + } + } + + #[test] + fn lcp_array_matches_naive_on_fixtures() { + for text in [ + b"banana".as_slice(), + b"mississippi", + b"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + b"abababababababababababababab", + b"a", + b"", + ] { + assert_lcp_valid(text, usize::MAX); + } + } + + #[test] + fn lcp_array_matches_naive_on_random() { + use rand::{RngExt, SeedableRng}; + let mut rng = rand::rngs::StdRng::seed_from_u64(0x1CB0); + for &sigma in &[2u8, 4, 6, 255] { + for &n in &[2usize, 3, 7, 16, 17, 63, 64, 65, 200, 1000, 5000] { + let text: Vec = (0..n).map(|_| rng.random_range(0..sigma)).collect(); + assert_lcp_valid(&text, usize::MAX); + } + } + } + + /// Long runs of one symbol are the worst case for the LCP invariant: + /// adjacent suffixes share almost everything, so every `lcp[i]` is + /// large and an off-by-one is easy to miss. + #[test] + fn lcp_array_on_long_runs_and_periodic_text() { + assert_lcp_valid(&vec![7u8; 2000], usize::MAX); + let periodic: Vec = (0..2000).map(|i| (i % 3) as u8).collect(); + assert_lcp_valid(&periodic, usize::MAX); + // A run embedded in noise, the shape a poly-N genome block has. + let mut mixed: Vec = (0..500).map(|i| (i % 4) as u8).collect(); + mixed.extend(std::iter::repeat_n(4u8, 1500)); + mixed.extend((0..500).map(|i| (i % 4) as u8)); + assert_lcp_valid(&mixed, usize::MAX); + } + + #[test] + fn lcp_array_respects_max_context() { + use rand::{RngExt, SeedableRng}; + let mut rng = rand::rngs::StdRng::seed_from_u64(0xC7A); + for &max_ctx in &[1usize, 2, 4, 16] { + for &n in &[64usize, 500] { + let text: Vec = (0..n).map(|_| rng.random_range(0..3u8)).collect(); + let (sa, lcp) = build_sa_and_lcp(&text, max_ctx); + for i in 1..sa.len() { + let want = naive_lcp(&text, sa[i - 1] as usize, sa[i] as usize, max_ctx); + assert_eq!(lcp[i] as usize, want, "lcp[{i}] wrong with max_ctx={max_ctx}"); + } + } + } + } + #[test] fn empty_text() { let sa: Vec = build_in_memory::(&[]); From 8174590f3cde2decb03476a8a3e8b978d7cc58e8 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 21:33:42 +0200 Subject: [PATCH 02/41] feat: radix-seeded prefix doubling for the plain in-memory SA The CaPS-SA merge kernel stays the general path, but it is not the right algorithm for the most common request: the standard lexicographic suffix array of a byte text, unsegmented, with no context bound. Add `src/radix.rs` and route that case through it. Profiling separated two distinct costs, and the merge kernel pays both: * Step count. `n log n` merge steps, most resolved by a symbol comparison at a random text address. 80 MB of N-free DNA is ~2.1e9 steps at ~13 ns each. * Scan length. Every leaf merge starts at `m = 0`, so two suffixes sharing a long prefix cost a scan proportional to that prefix. Genome FASTA carries megabyte-scale runs of `N` (period-61 once line wrapping is included), where one comparison scans millions of bytes. That drives the cost per merge step from 13 ns to 222 ns, a 16x penalty which is entirely scan time. This is what made real chr21 20x slower than N-free DNA of comparable size. The new path removes both. It sorts by a packed fixed-depth key, then resolves the remainder by prefix doubling on ranks. The packing picks the narrowest field width in {1,2,4,8} bits that holds the alphabet, so DNA over {0,1,2,3} resolves 32 symbols per key rather than the 8 a raw byte key gives. After the seed no comparison reads the text again, so a megabyte run of `N` costs exactly what random DNA costs. The seed sorts by `(key, min(n - p, k))`. The second component is required, not cosmetic: zero-padding makes a short suffix share a key with any suffix continuing in zeros, and `0` is a real symbol in every DNA encoding. Ordering by visible length puts the proper prefix first, which is the crate's shorter-is-smaller convention. Without it `[0, 0]` leaves two positions permanently tied and doubling cannot terminate. Guards are soundness conditions, not heuristics, and all three default to declining: * `max_context` must be unbounded; a finite bound makes the merge's comparator fall through to `boundary_order`, which compares lengths, so it is not lexicographic. * `LimitProvider::plain_lex_len` must report the full text. New method, defaulting to `None`, overridden only by `PlainText`. An impl that delegates `lim_at` to `PlainText` but overrides `boundary_order` for a different convention (STAR's spacer-as-largest) inherits `None` and stays on the merge kernel without changing a line. * `S` must be exactly `u8`. Packing wider symbols into an order- preserving key is endianness-dependent: for `u16` on a little-endian host `0x0100 > 0x0001` as values but their byte views compare the other way. The rest of the crate avoids this only because `LcpDispatch` resolves equality over bytes and recovers ordering through `S: Ord`. Tests: exhaustive over every binary text to length 10 and every ternary text to length 6, random texts across seven alphabet widths, texts where a real `0` collides with padding, long runs, periodic text, and the wrapped-FASTA `N`-block shape. Measured on Apple M4 Max, 12 threads, against the previous kernel, with byte-identical suffix arrays on both real inputs: chr21 fwd+revcomp, N-free, 80 MB 6.08 s -> 1.14 s CPU 28.1 s -> 5.0 s chr21 FASTA, 47.5 MB 27.8 s -> 1.16 s CPU 283.5 s -> 5.2 s Co-Authored-By: Claude Opus 5 (1M context) --- src/lib.rs | 1 + src/limits.rs | 28 ++++ src/radix.rs | 389 +++++++++++++++++++++++++++++++++++++++++++++ src/sample_sort.rs | 55 ++++++- 4 files changed, 471 insertions(+), 2 deletions(-) create mode 100644 src/radix.rs diff --git a/src/lib.rs b/src/lib.rs index 70684f4..43f4fee 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -20,6 +20,7 @@ mod ext_bucket; mod ext_mem; mod lcp; mod limits; +mod radix; mod sample_sort; pub use ext_mem::{ diff --git a/src/limits.rs b/src/limits.rs index 9fde6f1..b9ff0ba 100644 --- a/src/limits.rs +++ b/src/limits.rs @@ -73,6 +73,29 @@ pub trait LimitProvider: Sync { let _ = (p_a, p_b); lim_a.cmp(&lim_b) } + + /// `Some(n)` iff this provider describes an unsegmented text of `n` + /// symbols under the *standard* comparator: `lim_at(p) == n - p` for + /// every `p`, and `boundary_order` left at its shorter-is-smaller + /// default. + /// + /// Returning `Some` lets the crate substitute a specialised suffix-array + /// algorithm that assumes plain lexicographic order. It is therefore a + /// promise about the *comparator*, not merely about the lengths. + /// + /// The default is `None`, which keeps every existing and third-party + /// implementation on the general merge kernel at today's semantics. In + /// particular, an implementation that delegates `lim_at` to [`PlainText`] + /// but overrides [`boundary_order`][LimitProvider::boundary_order] to get + /// a different convention (STAR's spacer-as-largest ordering is the + /// motivating example) inherits `None` and is safe without doing + /// anything. + /// + /// Override this only if you have *not* overridden `boundary_order`. + #[inline] + fn plain_lex_len(&self) -> Option { + None + } } /// Default provider for non-segmented texts: `lim_at(p) = n - p`. @@ -99,6 +122,11 @@ impl LimitProvider for PlainText { fn lim_at(&self, p: usize) -> usize { self.n - p } + + #[inline] + fn plain_lex_len(&self) -> Option { + Some(self.n) + } } /// Provider for texts partitioned into segments at known cumulative diff --git a/src/radix.rs b/src/radix.rs new file mode 100644 index 0000000..f1420c5 --- /dev/null +++ b/src/radix.rs @@ -0,0 +1,389 @@ +//! Radix-seeded prefix doubling for the plain in-memory suffix array. +//! +//! The LCP-enhanced merge sort in [`crate::sample_sort`] is the CaPS-SA +//! kernel and stays the general path: it is the only one that honours a +//! [`LimitProvider`][crate::limits::LimitProvider], a finite `max_context`, +//! and symbol types wider than a byte, and it is the only one that produces +//! an LCP array (which the external-memory path needs). +//! +//! But for the single most common request — the standard lexicographic +//! suffix array of a byte text, with no segmentation and no context bound — +//! that kernel is doing far more work than the problem requires, in two +//! distinct ways that the benchmarks separate cleanly: +//! +//! * **Step count.** The merge sort performs `n log n` merge steps, and a +//! large majority of them are resolved by an actual symbol comparison at a +//! random text address. On 80 MB of N-free DNA that is ~2.1e9 steps at +//! ~13 ns each. +//! * **Scan length.** Every leaf merge starts with `m = 0`, so comparing two +//! suffixes that share a long prefix costs a scan proportional to that +//! prefix. Genome assemblies contain megabyte-scale runs of `N` (and the +//! period-61 `N`-then-newline pattern of wrapped FASTA), where a single +//! comparison scans millions of bytes. On a 47.5 MB chr21 FASTA this +//! pushes the cost per merge step from 13 ns to 222 ns — a 16x penalty +//! that is entirely scan time. +//! +//! This module attacks both. It sorts by a packed fixed-depth key first +//! (killing the step count), then resolves the remainder by **prefix +//! doubling** on ranks (killing the scan length: after the seed, no +//! comparison ever reads the text again, so a megabyte-long run of `N` costs +//! exactly as much as random DNA). +//! +//! ## The algorithm +//! +//! 1. **Pack.** Find the maximum symbol and choose the smallest field width +//! in `{1, 2, 4, 8}` bits that can hold it, so `k = 64 / bits` symbols fit +//! in one `u64` key. DNA over `{0,1,2,3}` gets 2-bit fields and therefore +//! resolves **32 symbols per key** rather than the 8 a raw byte key would. +//! 2. **Seed.** Sort `(key, position)`. This is a full sort of the suffixes +//! by their first `k` symbols, and it touches the text only in one +//! sequential pass. +//! 3. **Double.** Suffixes still tied after depth `d` are ordered by the pair +//! `(rank_d(p), rank_d(p + d))`, which resolves them to depth `2d`. Repeat +//! until every group is a singleton. Each round reads only the rank array. +//! +//! ## Ordering convention +//! +//! Keys are big-endian in the field sense (the first symbol occupies the most +//! significant field) and short suffixes are zero-padded. Since `0` is the +//! minimum of `u8`, a padding field can never exceed a real symbol's field, +//! so a padded key compares less-or-equal to any key it shares a prefix with. +//! That is exactly the crate's "shorter suffix is smaller" convention. A real +//! `0` symbol is indistinguishable from padding *in the key*, which can only +//! make two suffixes tie — never invert them — and ties are resolved by the +//! doubling rounds, which use the true remaining length via the end-of-text +//! sentinel. So `A = 0` DNA encodings and STAR's `0..5` codes are both safe. + +use crate::Index; +use rayon::prelude::*; + +/// Field width in bits and the number of symbols that fit in a `u64` key. +/// +/// Restricted to divisors of 64 so a key is an exact number of whole fields +/// and no symbol ever straddles the key boundary. +fn pack_params(max_sym: u8) -> (u32, usize) { + let bits: u32 = match max_sym { + 0..=1 => 1, + 2..=3 => 2, + 4..=15 => 4, + _ => 8, + }; + (bits, 64 / bits as usize) +} + +/// Pack the `k` symbols at `text[p..]` into one order-preserving `u64`, +/// zero-padding past the end of the text. +#[inline] +fn key_at(text: &[u8], p: usize, bits: u32, k: usize) -> u64 { + if bits == 8 { + // Whole-byte fields: this is just a big-endian load, and the common + // case (p + 8 <= n) is a single unaligned u64 read plus a bswap. + let mut buf = [0u8; 8]; + let end = (p + 8).min(text.len()); + buf[..end - p].copy_from_slice(&text[p..end]); + return u64::from_be_bytes(buf); + } + let end = (p + k).min(text.len()); + let mut key: u64 = 0; + for &s in &text[p..end] { + key = (key << bits) | s as u64; + } + // Shift the packed prefix up so the missing trailing fields read as zero. + key << (bits as usize * (k - (end - p))) +} + +/// Build the standard lexicographic suffix array of `text` by radix-seeded +/// prefix doubling. +/// +/// The caller is responsible for the guards: `text` must be the whole, +/// non-segmented text, the comparator must be plain lexicographic with +/// shorter-is-smaller, and there must be no `max_context` bound. See +/// [`crate::sample_sort::build_in_memory_with`] for where those are checked. +pub(crate) fn build_sa(text: &[u8]) -> Vec { + let n = text.len(); + if n == 0 { + return Vec::new(); + } + if n == 1 { + return vec![I::zero()]; + } + + let max_sym = text.par_iter().copied().max().unwrap_or(0); + let (bits, k) = pack_params(max_sym); + + // ---- Seed: sort by the first `k` symbols, then by visible length. ---- + // + // The second component is `min(n - p, k)`, and it is load-bearing rather + // than cosmetic. Zero-padding makes a suffix shorter than `k` share a key + // with any suffix whose symbols continue with zeros, and `0` is a real + // symbol in every DNA encoding. Ordering those by visible length puts the + // proper prefix first, which is the shorter-is-smaller convention. For + // suffixes at least `k` long the component is `k` for all of them, so it + // never separates suffixes that the doubling rounds still need to see as + // tied. Without it, `[0, 0]` leaves positions 0 and 1 permanently tied + // and the doubling loop cannot terminate. + let mut seeded: Vec<(u64, u32, I)> = (0..n) + .into_par_iter() + .map(|p| { + ( + key_at(text, p, bits, k), + (n - p).min(k) as u32, + I::from_usize(p), + ) + }) + .collect(); + seeded.par_sort_unstable(); + + let mut sa: Vec = Vec::with_capacity(n); + sa.par_extend(seeded.par_iter().map(|&(_, _, p)| p)); + + // `rank[p]` is the index in `sa` of the first element of `p`'s group, so + // two suffixes tie at the current depth exactly when their ranks match, + // and rank order is the current partial order. + let mut rank: Vec = vec![I::zero(); n]; + // Non-singleton `sa` ranges, the only ones any later round touches. + let mut groups: Vec<(usize, usize)> = Vec::new(); + { + let mut i = 0; + while i < n { + let mut j = i + 1; + while j < n && (seeded[j].0, seeded[j].1) == (seeded[i].0, seeded[i].1) { + j += 1; + } + let g = I::from_usize(i); + for e in &sa[i..j] { + rank[e.to_usize()] = g; + } + if j - i > 1 { + groups.push((i, j)); + } + i = j; + } + } + drop(seeded); + + // ---- Double: (rank_d(p), rank_d(p + d)) resolves to depth 2d. ---- + let mut depth = k; + // Scratch for the new rank of each `sa` slot, so a round's reads of + // `rank` never observe that same round's writes. + let mut next_rank: Vec = vec![I::zero(); n]; + + while !groups.is_empty() { + // Phase A: sort each tied group by the successor rank, and record the + // ranks it should get. Groups are disjoint `sa` ranges, so this is + // data-parallel with no synchronisation. + let sub: Vec> = split_disjoint(&mut sa, &mut next_rank, &groups) + .into_par_iter() + .zip(groups.par_iter()) + .map(|((sa_g, nr_g), &(start, _))| { + let succ = |p: usize| -> u64 { + // End-of-text sorts first: the shorter suffix is smaller. + match p.checked_add(depth) { + Some(q) if q < n => rank[q].to_usize() as u64 + 1, + _ => 0, + } + }; + sa_g.sort_unstable_by_key(|e| succ(e.to_usize())); + + let mut fresh = Vec::new(); + let mut i = 0; + while i < sa_g.len() { + let key = succ(sa_g[i].to_usize()); + let mut j = i + 1; + while j < sa_g.len() && succ(sa_g[j].to_usize()) == key { + j += 1; + } + let g = I::from_usize(start + i); + for slot in &mut nr_g[i..j] { + *slot = g; + } + if j - i > 1 { + fresh.push((start + i, start + j)); + } + i = j; + } + fresh + }) + .collect(); + + // Phase B: publish the new ranks, now that every read is done. + for &(start, end) in &groups { + for i in start..end { + rank[sa[i].to_usize()] = next_rank[i]; + } + } + + let before: usize = groups.iter().map(|&(s, e)| e - s).sum(); + groups = sub.into_iter().flatten().collect(); + let after: usize = groups.iter().map(|&(s, e)| e - s).sum(); + + // A doubling round can only ever refine, so `after <= before`. If a + // round refines nothing at all the text has a run longer than the + // whole remaining depth budget; doubling still terminates because + // `depth` grows geometrically and every suffix eventually runs off + // the end of the text, which the sentinel orders. Guard against + // overflow rather than against non-progress. + debug_assert!(after <= before); + match depth.checked_mul(2) { + Some(d) if d <= n.saturating_mul(2) => depth = d, + _ => { + debug_assert!(groups.is_empty(), "doubling exhausted with ties left"); + break; + } + } + } + + sa +} + +/// Borrow each `(start, end)` range of `sa` and `next_rank` mutably and +/// simultaneously. The ranges come from a scan of `sa` so they are sorted and +/// non-overlapping, which is what makes the repeated `split_at_mut` sound. +fn split_disjoint<'a, I: Index>( + sa: &'a mut [I], + next_rank: &'a mut [I], + groups: &[(usize, usize)], +) -> Vec<(&'a mut [I], &'a mut [I])> { + let mut out = Vec::with_capacity(groups.len()); + let mut sa_rest = sa; + let mut nr_rest = next_rank; + let mut consumed = 0usize; + for &(start, end) in groups { + debug_assert!( + start >= consumed, + "group ranges must be sorted and disjoint" + ); + let (_, sa_tail) = sa_rest.split_at_mut(start - consumed); + let (_, nr_tail) = nr_rest.split_at_mut(start - consumed); + let (sa_g, sa_tail) = sa_tail.split_at_mut(end - start); + let (nr_g, nr_tail) = nr_tail.split_at_mut(end - start); + out.push((sa_g, nr_g)); + sa_rest = sa_tail; + nr_rest = nr_tail; + consumed = end; + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + + fn brute(text: &[u8]) -> Vec { + let mut sa: Vec = (0..text.len() as u32).collect(); + sa.sort_by(|&a, &b| text[a as usize..].cmp(&text[b as usize..])); + sa + } + + fn check(text: &[u8]) { + let got: Vec = build_sa(text); + assert_eq!(got, brute(text), "mismatch on {text:?}"); + } + + #[test] + fn fixtures() { + check(b""); + check(b"a"); + check(b"banana"); + check(b"mississippi"); + check(b"abracadabra"); + } + + #[test] + fn pack_params_covers_every_width() { + assert_eq!(pack_params(0), (1, 64)); + assert_eq!(pack_params(1), (1, 64)); + assert_eq!(pack_params(3), (2, 32)); + assert_eq!(pack_params(4), (4, 16)); + assert_eq!(pack_params(15), (4, 16)); + assert_eq!(pack_params(16), (8, 8)); + assert_eq!(pack_params(255), (8, 8)); + } + + /// Texts whose symbols include a real `0`, so padding and a genuine + /// minimum symbol are indistinguishable in the packed key. This is the + /// case DNA encodings hit (`A = 0`) and the one the ordering argument in + /// the module docs turns on. + #[test] + fn real_zero_symbol_is_not_confused_with_padding() { + check(&[0]); + check(&[0, 0, 0, 0, 0, 0, 0, 0, 0, 0]); + check(&[1, 2, 0, 0, 0, 0, 0, 0, 0, 0]); + check(&[3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]); + check(&[0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0]); + let mut t: Vec = (0..200).map(|i| (i % 4) as u8).collect(); + t.extend(std::iter::repeat_n(0u8, 100)); + check(&t); + } + + /// Every text of length <= 10 over a binary alphabet, plus every text of + /// length <= 6 over a ternary one. Total coverage of the padding and + /// end-of-text logic at the sizes where exhaustive checking is free. + #[test] + fn exhaustive_small_alphabets() { + for n in 0..=10u32 { + for mask in 0..(1u32 << n) { + let t: Vec = (0..n).map(|i| ((mask >> i) & 1) as u8).collect(); + check(&t); + } + } + for n in 0..=6u32 { + let total = 3u32.pow(n); + for mut code in 0..total { + let mut t = Vec::with_capacity(n as usize); + for _ in 0..n { + t.push((code % 3) as u8); + code /= 3; + } + check(&t); + } + } + } + + #[test] + fn random_across_alphabet_widths() { + use rand::{RngExt, SeedableRng}; + let mut rng = rand::rngs::StdRng::seed_from_u64(0xADD1); + for &sigma in &[2u8, 3, 4, 6, 16, 17, 255] { + for &n in &[ + 2usize, 7, 31, 32, 33, 63, 64, 65, 127, 128, 129, 1000, 20_000, + ] { + let t: Vec = (0..n).map(|_| rng.random_range(0..sigma)).collect(); + check(&t); + } + } + } + + /// Long runs and periodic text are the inputs that make the merge kernel + /// quadratic. Doubling must handle them and must terminate. + #[test] + fn long_runs_and_periodic_text() { + check(&vec![0u8; 5000]); + check(&vec![7u8; 5000]); + check(&(0..5000).map(|i| (i % 2) as u8).collect::>()); + check(&(0..5000).map(|i| (i % 61) as u8).collect::>()); + // A long run flanked by noise: the shape of a poly-N genome block. + let mut t: Vec = (0..500).map(|i| (i % 4) as u8).collect(); + t.extend(std::iter::repeat_n(4u8, 4000)); + t.extend((0..500).map(|i| (i % 4) as u8)); + check(&t); + // Wrapped-FASTA shape: 60 `N`s then a newline, repeated. + let mut fasta: Vec = Vec::new(); + for _ in 0..100 { + fasta.extend(std::iter::repeat_n(b'N', 60)); + fasta.push(b'\n'); + } + check(&fasta); + } + + #[test] + fn u64_index_matches_u32_index() { + use rand::{RngExt, SeedableRng}; + let mut rng = rand::rngs::StdRng::seed_from_u64(0xD0AB); + let t: Vec = (0..5000).map(|_| rng.random_range(0..4u8)).collect(); + let a: Vec = build_sa(&t); + let b: Vec = build_sa(&t); + assert_eq!(a.len(), b.len()); + assert!(a.iter().zip(&b).all(|(&x, &y)| x as u64 == y)); + } +} diff --git a/src/sample_sort.rs b/src/sample_sort.rs index b2f0dc7..e69a699 100644 --- a/src/sample_sort.rs +++ b/src/sample_sort.rs @@ -90,11 +90,58 @@ where I: Index, L: LimitProvider, { + if let Some(sa) = try_doubling_fast_path::(text, lp, opts) { + return sa; + } let n = text.len(); let positions: Vec = (0..n).map(I::from_usize).collect(); build_in_memory_for_positions_with(text, positions, lp, opts) } +/// Route a whole-text byte build through [`crate::radix`]'s radix-seeded +/// prefix doubling, which is dramatically faster on real genomic input, or +/// return `None` to fall back to the CaPS-SA merge kernel. +/// +/// Every condition below is a soundness requirement, not a heuristic. The +/// doubling path implements exactly one comparator — plain lexicographic over +/// bytes with shorter-is-smaller and no context bound — so anything that can +/// change the comparator has to decline. +/// +/// * `max_context` must be unbounded. With a finite bound the merge's +/// comparator stops being lexicographic: once a scan hits the cap it falls +/// through to [`LimitProvider::boundary_order`], which compares *lengths*. +/// * `lp` must report [`plain_lex_len`][LimitProvider::plain_lex_len]. That +/// rules out `SegmentedText`, whose LCP scans stop at segment boundaries, +/// and any custom `boundary_order` such as STAR's spacer-as-largest. +/// * `S` must be exactly `u8`. Wider symbols are excluded because packing +/// them into an order-preserving key is endianness-dependent: for `u16` on +/// a little-endian host, `0x0100 > 0x0001` as values but their byte views +/// compare the other way. The rest of the crate is immune to this only +/// because [`LcpDispatch`] resolves *equality* over bytes and recovers +/// ordering through `S: Ord`. +fn try_doubling_fast_path(text: &[S], lp: &L, opts: &Opts) -> Option> +where + S: Symbol, + I: Index, + L: LimitProvider, +{ + if opts.max_context != usize::MAX { + return None; + } + if lp.plain_lex_len() != Some(text.len()) { + return None; + } + if std::any::TypeId::of::() != std::any::TypeId::of::() { + return None; + } + // SAFETY: `S` is `u8` (just checked by `TypeId`, and `Symbol: 'static` + // so the comparison is exact), hence `&[S]` and `&[u8]` have identical + // layout, length and validity. + let bytes: &[u8] = + unsafe { std::slice::from_raw_parts(text.as_ptr() as *const u8, text.len()) }; + Some(crate::radix::build_sa(bytes)) +} + /// Sort the caller-supplied `positions` by the lexicographic order of /// their suffixes in `text`. Returns the positions reordered so that /// `text[output[i]..]` is the i-th smallest suffix among the input set. @@ -465,7 +512,8 @@ mod tests { for i in 1..sa.len() { let want = naive_lcp(text, sa[i - 1] as usize, sa[i] as usize, max_ctx); assert_eq!( - lcp[i] as usize, want, + lcp[i] as usize, + want, "lcp[{i}] wrong for sa[{}]={} vs sa[{i}]={} (text {text:?})", i - 1, sa[i - 1], @@ -525,7 +573,10 @@ mod tests { let (sa, lcp) = build_sa_and_lcp(&text, max_ctx); for i in 1..sa.len() { let want = naive_lcp(&text, sa[i - 1] as usize, sa[i] as usize, max_ctx); - assert_eq!(lcp[i] as usize, want, "lcp[{i}] wrong with max_ctx={max_ctx}"); + assert_eq!( + lcp[i] as usize, want, + "lcp[{i}] wrong with max_ctx={max_ctx}" + ); } } } From f11c459df8162c387401481b4150c1d5ffac6453 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 21:33:42 +0200 Subject: [PATCH 03/41] feat: `verify_sa`, an O(n) independent suffix-array check MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Checking a new construction path against the old one only shows the two agree. Checking adjacent suffixes directly is O(n · lcp), which is unusable on exactly the repetitive inputs that need checking most: on a chr21 FASTA a single adjacent pair can share megabytes. Use the fixpoint characterisation instead. With `rank` the inverse of `sa` and `f(p) = (text[p], rank[p + 1])`, taking `rank[n]` as less than every real rank, a permutation of `0..n` is the suffix array of `text` if and only if `f` is strictly increasing along it. That is one pass to invert plus one pass to compare, independent of any construction algorithm and independent of LCP length. Exposed as `caps_sa::verify_sa` and wired to a `--verify` flag on the bench CLI, off by default so it never contaminates a timing run. Full-scale results on Apple M4 Max, 12 threads: chr21 fwd+revcomp, N-free, 80,177,238 entries verify OK in 1.14 s chr21 FASTA, 47,488,540 entries verify OK in 0.57 s Co-Authored-By: Claude Opus 5 (1M context) --- examples/caps_sa.rs | 29 +++++++++++++++++-- src/lib.rs | 68 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 95 insertions(+), 2 deletions(-) diff --git a/examples/caps_sa.rs b/examples/caps_sa.rs index a26beba..67da601 100644 --- a/examples/caps_sa.rs +++ b/examples/caps_sa.rs @@ -18,7 +18,7 @@ use std::path::PathBuf; use std::process; use std::time::Instant; -use caps_sa::{ExtMemOpts, build_ext_mem, build_in_memory, build_in_memory_sample_sort}; +use caps_sa::{ExtMemOpts, build_ext_mem, build_in_memory, build_in_memory_sample_sort, verify_sa}; struct Args { input: PathBuf, @@ -27,6 +27,7 @@ struct Args { in_mem_ss: bool, subproblem_count: usize, threads: Option, + verify: bool, } fn parse_args() -> Args { @@ -36,6 +37,7 @@ fn parse_args() -> Args { let mut in_mem_ss = false; let mut subproblem_count: usize = 0; let mut threads: Option = None; + let mut verify = false; let mut i = 1; while i < argv.len() { match argv[i].as_str() { @@ -47,6 +49,10 @@ fn parse_args() -> Args { in_mem_ss = true; i += 1; } + "--verify" => { + verify = true; + i += 1; + } "--subproblem-count" => { subproblem_count = argv[i + 1] .parse() @@ -64,7 +70,7 @@ fn parse_args() -> Args { "--help" | "-h" => { eprintln!( "usage: caps_sa [--ext-mem | --in-mem-ss] \ - [--subproblem-count N] [--threads N]" + [--subproblem-count N] [--threads N] [--verify]" ); process::exit(0); } @@ -88,6 +94,23 @@ fn parse_args() -> Args { in_mem_ss, subproblem_count, threads, + verify, + } +} + +/// Independently check the built suffix array in O(n). Off by default so it +/// never contaminates a timing run; the check is reported separately. +fn maybe_verify(enabled: bool, text: &[u8], sa: &[I]) { + if !enabled { + return; + } + let t = Instant::now(); + match verify_sa(text, sa) { + Ok(()) => eprintln!("verify: OK in {:.3}s", t.elapsed().as_secs_f64()), + Err(e) => { + eprintln!("verify: FAILED: {e}"); + process::exit(1); + } } } @@ -124,6 +147,7 @@ fn main() -> std::io::Result<()> { let sa: Vec = build_in_memory(&text); build_elapsed = build_start.elapsed(); n_entries = sa.len(); + maybe_verify(args.verify, &text, &sa); eprintln!( "build: mode=in-mem(u32) n={n_entries} entries in {:.3}s", build_elapsed.as_secs_f64() @@ -177,6 +201,7 @@ fn main() -> std::io::Result<()> { let sa: Vec = build_in_memory(&text); build_elapsed = build_start.elapsed(); n_entries = sa.len(); + maybe_verify(args.verify, &text, &sa); eprintln!( "build: mode=in-mem(u64) n={n_entries} entries in {:.3}s", build_elapsed.as_secs_f64() diff --git a/src/lib.rs b/src/lib.rs index 43f4fee..9da5c26 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -40,6 +40,74 @@ pub use sample_sort::{ build_in_memory_for_positions_with_opts, build_in_memory_with, build_in_memory_with_opts, }; +/// Check that `sa` really is the suffix array of `text`, in `O(n)` time and +/// without re-running any construction algorithm. +/// +/// Comparing a candidate against a second implementation only shows the two +/// agree; comparing adjacent suffixes directly is `O(n · lcp)` and becomes +/// unusable on the repetitive inputs that matter most. This instead uses the +/// standard fixpoint characterisation: let `rank` be the inverse of `sa`, and +/// define `f(p) = (text[p], rank[p + 1])`, with `rank[n]` taken as less than +/// every real rank. A permutation is the suffix array of `text` if and only if +/// `f` is strictly increasing along it, because suffix `p` precedes suffix `q` +/// exactly when `f(p) < f(q)`. +/// +/// Returns `Err` with a description of the first violation found. +/// +/// ``` +/// let text = b"banana"; +/// let sa: Vec = caps_sa::build_in_memory(text); +/// assert!(caps_sa::verify_sa(text, &sa).is_ok()); +/// assert!(caps_sa::verify_sa(text, &[0u32, 1, 2, 3, 4, 5]).is_err()); +/// ``` +pub fn verify_sa(text: &[S], sa: &[I]) -> Result<(), String> +where + S: Ord, + I: Index, +{ + let n = text.len(); + if sa.len() != n { + return Err(format!("sa has {} entries, text has {n} symbols", sa.len())); + } + if n == 0 { + return Ok(()); + } + + // Invert `sa`, checking along the way that it is a permutation of `0..n`. + let mut rank = vec![usize::MAX; n]; + for (i, entry) in sa.iter().enumerate() { + let p = entry.to_usize(); + if p >= n { + return Err(format!("sa[{i}] = {p} is out of range for text length {n}")); + } + if rank[p] != usize::MAX { + return Err(format!( + "position {p} appears at sa[{}] and sa[{i}]", + rank[p] + )); + } + rank[p] = i; + } + + // `None` stands for the end of the text, which sorts before every rank: + // the shorter suffix is the smaller one. + let successor = + |p: usize| -> Option { if p + 1 < n { Some(rank[p + 1]) } else { None } }; + for i in 1..n { + let a = sa[i - 1].to_usize(); + let b = sa[i].to_usize(); + let key_a = (&text[a], successor(a)); + let key_b = (&text[b], successor(b)); + if key_a >= key_b { + return Err(format!( + "suffixes out of order at sa[{}] = {a} and sa[{i}] = {b}", + i - 1, + )); + } + } + Ok(()) +} + /// Trait implemented by integer types usable as suffix array indices. /// /// Provided for `u32`, `u64`, and `usize`. Callers pick the narrowest type From b4ecde60222cc5b7463e39ce3b83d817eee4637e Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 21:33:42 +0200 Subject: [PATCH 04/41] perf: parallelise grouping and rank publication in the doubling path Phase timings (`CAPS_SA_PROFILE=1`) showed the two rank-scatter passes were fully sequential and had become the largest single cost after the seed sort: 0.34 s of a 1.17 s build on 80 MB of DNA. Both scatter through a permutation -- the target index is `sa[i]`, not `i` -- so the writes are not expressible as disjoint sub-slices and `split_at_mut` does not apply. They are nonetheless disjoint: `sa` is a permutation and the ranges being processed partition its index space, so every slot is written exactly once. Introduce a small `Scatter` wrapper that encodes precisely that contract in its `unsafe fn set`, and drive both passes with rayon. Grouping now has each index decide for itself whether it starts a group; the index that does owns the group, walks it to find the end, and writes its members' ranks. Exactly one owner per group, and `collect` on an indexed parallel iterator preserves order, so the group list still comes out sorted, which is what `split_disjoint` relies on. Also materialise the successor ranks before sorting each group. `sort_unstable_by_key` re-evaluates its key function O(len log len) times and every evaluation was a random probe into `rank`; paying once per element makes the sort's memory traffic sequential. Apple M4 Max, 12 threads, suffix arrays byte-identical to the previous kernel and independently `--verify`-checked: chr21 fwd+revcomp, N-free, 80 MB grouping 0.341 s -> 0.075 s total 1.17 s -> 0.89 s chr21 FASTA, 47.5 MB grouping 0.172 s -> 0.037 s total 1.10 s -> 1.04 s Co-Authored-By: Claude Opus 5 (1M context) --- src/ext_mem.rs | 2 +- src/radix.rs | 130 ++++++++++++++++++++++++++++++++++++++++--------- 2 files changed, 109 insertions(+), 23 deletions(-) diff --git a/src/ext_mem.rs b/src/ext_mem.rs index 621a3ab..6ef6e84 100644 --- a/src/ext_mem.rs +++ b/src/ext_mem.rs @@ -49,7 +49,7 @@ use crate::sample_sort; /// time without paying the cost of always logging — see /// `bench/README.md` "Where AVX-512 helps and where it doesn't" for /// how this is used. -fn profile_log(message: &str) { +pub(crate) fn profile_log(message: &str) { if std::env::var_os("CAPS_SA_PROFILE").is_some() { eprintln!("caps-sa profile {message}"); } diff --git a/src/radix.rs b/src/radix.rs index f1420c5..edb01b8 100644 --- a/src/radix.rs +++ b/src/radix.rs @@ -55,7 +55,9 @@ //! sentinel. So `A = 0` DNA encodings and STAR's `0..5` codes are both safe. use crate::Index; +use crate::ext_mem::profile_log; use rayon::prelude::*; +use std::time::Instant; /// Field width in bits and the number of symbols that fit in a `u64` key. /// @@ -108,6 +110,7 @@ pub(crate) fn build_sa(text: &[u8]) -> Vec { return vec![I::zero()]; } + let t0 = Instant::now(); let max_sym = text.par_iter().copied().max().unwrap_or(0); let (bits, k) = pack_params(max_sym); @@ -132,7 +135,17 @@ pub(crate) fn build_sa(text: &[u8]) -> Vec { ) }) .collect(); + profile_log(&format!( + "radix keys {:.3}s", + t0.elapsed().as_secs_f64() + )); + let t1 = Instant::now(); seeded.par_sort_unstable(); + profile_log(&format!( + "radix seed sort {:.3}s", + t1.elapsed().as_secs_f64() + )); + let t2 = Instant::now(); let mut sa: Vec = Vec::with_capacity(n); sa.par_extend(seeded.par_iter().map(|&(_, _, p)| p)); @@ -142,25 +155,41 @@ pub(crate) fn build_sa(text: &[u8]) -> Vec { // and rank order is the current partial order. let mut rank: Vec = vec![I::zero(); n]; // Non-singleton `sa` ranges, the only ones any later round touches. - let mut groups: Vec<(usize, usize)> = Vec::new(); - { - let mut i = 0; - while i < n { - let mut j = i + 1; - while j < n && (seeded[j].0, seeded[j].1) == (seeded[i].0, seeded[i].1) { - j += 1; + // + // Each index decides for itself whether it starts a group; the index that + // does then owns the whole group, walks it to find the end, and writes its + // members' ranks. Every group has exactly one owner and groups partition + // `0..n`, so the scattered writes never collide. `collect` on an indexed + // parallel iterator preserves order, so `groups` comes out sorted. + let ranks = Scatter::new(&mut rank); + let groups: Vec<(usize, usize)> = (0..n) + .into_par_iter() + .filter_map(|h| { + let key = (seeded[h].0, seeded[h].1); + if h > 0 && (seeded[h - 1].0, seeded[h - 1].1) == key { + return None; } - let g = I::from_usize(i); - for e in &sa[i..j] { - rank[e.to_usize()] = g; + let mut e = h + 1; + while e < n && (seeded[e].0, seeded[e].1) == key { + e += 1; } - if j - i > 1 { - groups.push((i, j)); + let g = I::from_usize(h); + for entry in &sa[h..e] { + // SAFETY: `sa` is a permutation of `0..n`, and this thread + // owns the whole group `h..e`, so `entry` is a distinct index + // no other thread writes. + unsafe { ranks.set(entry.to_usize(), g) }; } - i = j; - } - } + (e - h > 1).then_some((h, e)) + }) + .collect(); + let mut groups = groups; drop(seeded); + profile_log(&format!( + "radix grouping {:.3}s", + t2.elapsed().as_secs_f64() + )); + let t3 = Instant::now(); // ---- Double: (rank_d(p), rank_d(p + d)) resolves to depth 2d. ---- let mut depth = k; @@ -183,14 +212,20 @@ pub(crate) fn build_sa(text: &[u8]) -> Vec { _ => 0, } }; - sa_g.sort_unstable_by_key(|e| succ(e.to_usize())); + // Materialise the successor ranks once. `sort_unstable_by_key` + // re-evaluates its key function O(len log len) times, and each + // evaluation is a random probe into `rank`; paying for it once + // per element turns the sort's memory traffic sequential. + let mut keyed: Vec<(u64, I)> = + sa_g.iter().map(|&e| (succ(e.to_usize()), e)).collect(); + keyed.sort_unstable(); let mut fresh = Vec::new(); let mut i = 0; - while i < sa_g.len() { - let key = succ(sa_g[i].to_usize()); + while i < keyed.len() { + let key = keyed[i].0; let mut j = i + 1; - while j < sa_g.len() && succ(sa_g[j].to_usize()) == key { + while j < keyed.len() && keyed[j].0 == key { j += 1; } let g = I::from_usize(start + i); @@ -202,16 +237,24 @@ pub(crate) fn build_sa(text: &[u8]) -> Vec { } i = j; } + for (slot, &(_, e)) in sa_g.iter_mut().zip(keyed.iter()) { + *slot = e; + } fresh }) .collect(); // Phase B: publish the new ranks, now that every read is done. - for &(start, end) in &groups { + // Groups are disjoint and `sa` is a permutation, so each `rank` slot + // is written by exactly one group. + let ranks = Scatter::new(&mut rank); + groups.par_iter().for_each(|&(start, end)| { for i in start..end { - rank[sa[i].to_usize()] = next_rank[i]; + // SAFETY: `sa[start..end]` are distinct positions owned solely + // by this group, and the groups partition their index range. + unsafe { ranks.set(sa[i].to_usize(), next_rank[i]) }; } - } + }); let before: usize = groups.iter().map(|&(s, e)| e - s).sum(); groups = sub.into_iter().flatten().collect(); @@ -233,9 +276,52 @@ pub(crate) fn build_sa(text: &[u8]) -> Vec { } } + profile_log(&format!( + "radix doubling {:.3}s", + t3.elapsed().as_secs_f64() + )); sa } +/// Write access to disjoint slots of one slice from several rayon threads. +/// +/// Both users here scatter through a permutation: the target index is +/// `sa[i]`, not `i`, so the writes cannot be expressed as disjoint sub-slices +/// and `split_at_mut` does not apply. What makes them safe is that `sa` is a +/// permutation and the ranges being processed partition its index space, so +/// every slot is written exactly once across all threads. +struct Scatter { + ptr: *mut T, + len: usize, +} + +// SAFETY: `Scatter` hands out writes only through `set`, whose contract is +// that no two calls target the same index. Under that contract there is no +// aliasing between threads, so the pointer is safe to share. +unsafe impl Send for Scatter {} +unsafe impl Sync for Scatter {} + +impl Scatter { + fn new(slice: &mut [T]) -> Self { + Self { + ptr: slice.as_mut_ptr(), + len: slice.len(), + } + } + + /// Write `value` at `index`. + /// + /// # Safety + /// + /// No two concurrent calls may pass the same `index`, and the borrow the + /// `Scatter` was built from must still be live. + #[inline] + unsafe fn set(&self, index: usize, value: T) { + debug_assert!(index < self.len); + unsafe { self.ptr.add(index).write(value) }; + } +} + /// Borrow each `(start, end)` range of `sa` and `next_rank` mutably and /// simultaneously. The ranges come from a scan of `sa` so they are sorted and /// non-overlapping, which is what makes the repeated `split_at_mut` sound. From 9f401e7d6ed97b5f70b021da7e3751a93169a841 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 21:33:42 +0200 Subject: [PATCH 05/41] ci: add a Rust workflow The repository had no Rust CI at all: the only workflow deployed the docs site, while 73 tests sat in the tree with nothing running them. That is not a safe baseline for changing the sorting kernel. Covers both architectures that matter here, since the LCP kernel and the pooled external-memory bucket path are the parts that diverge per platform: macOS is aarch64/NEON, Ubuntu is x86_64/AVX2. Runs the tests in debug as well as release, because debug is what exercises the `debug_assert`s guarding the unchecked scatter in `radix.rs` and the buffer-length invariants in the merge kernel. Adds fmt, clippy with warnings denied, a check against the declared 1.89 MSRV, and rustdoc with broken links denied. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/ci.yml | 69 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 69 insertions(+) create mode 100644 .github/workflows/ci.yml diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..0567d90 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,69 @@ +name: CI + +on: + push: + branches: [main] + pull_request: + +env: + CARGO_TERM_COLOR: always + RUSTFLAGS: -D warnings + +jobs: + test: + name: test (${{ matrix.os }}) + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + # macOS is aarch64/NEON, Ubuntu is x86_64/AVX2. The LCP kernel and the + # pooled external-memory bucket path are the parts that differ per + # platform, so both need to run. + os: [ubuntu-latest, macos-latest] + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@stable + - uses: Swatinem/rust-cache@v2 + + # Debug catches the `debug_assert`s that guard the unchecked scatter in + # `radix.rs` and the buffer-length invariants in the merge kernel. + - name: Test (debug) + run: cargo test --all-targets + - name: Test (release) + run: cargo test --release --all-targets + - name: Doc tests + run: cargo test --doc + + lint: + name: fmt + clippy + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@stable + with: + components: rustfmt, clippy + - uses: Swatinem/rust-cache@v2 + - run: cargo fmt --all --check + - run: cargo clippy --all-targets -- -D warnings + + msrv: + name: MSRV (1.89) + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + # Pinned to the `rust-version` in Cargo.toml, which is set by the + # stabilised AVX-512 intrinsics the LCP fast path uses. + - uses: dtolnay/rust-toolchain@1.89 + - uses: Swatinem/rust-cache@v2 + - run: cargo check --all-targets + + docs: + name: rustdoc + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@stable + - uses: Swatinem/rust-cache@v2 + - run: cargo doc --no-deps + env: + RUSTDOCFLAGS: -D warnings From c8683e6f12b839a5e016939a84c4b37550b19a14 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 21:33:42 +0200 Subject: [PATCH 06/41] bench: add a reproducible chr21 harness `bench/run.sh` compares against upstream C++ and needs both binaries prebuilt. Add a self-contained script for the case that prompted this work: it fetches hg38 chr21 and prepares *both* inputs, which is the distinction that made the original slowdown report hard to interpret. chr21.0123 forward ++ revcomp, one byte per base, codes 0..=3, ambiguous bases dropped. ~80 MB, alphabet 4, no long runs. The input libsais is normally benchmarked on. chr21.fa the raw FASTA, still carrying its ~6.6 Mb of `N`. Wrapped at 60 columns, so the `N` blocks are a period-61 repeat rather than a plain run. Benchmarking one implementation on the first and another on the second compares two different problems. The second is the realistic input and the one that used to be pathological. Builds with `-C target-cpu=native` and runs each case through `--verify`, so the harness reports correctness alongside timing. Co-Authored-By: Claude Opus 5 (1M context) --- bench/chr21.sh | 83 ++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 83 insertions(+) create mode 100755 bench/chr21.sh diff --git a/bench/chr21.sh b/bench/chr21.sh new file mode 100755 index 0000000..5a5398f --- /dev/null +++ b/bench/chr21.sh @@ -0,0 +1,83 @@ +#!/usr/bin/env bash +# Reproduce the chr21 numbers quoted in bench/README.md. +# +# bench/chr21.sh [work-dir] [threads] +# +# Builds two inputs from hg38 chr21, because they exercise different costs and +# conflating them is what made the original slowdown report hard to read: +# +# chr21.0123 forward ++ reverse complement, one byte per base, codes 0..=3, +# ambiguous bases dropped. ~80 MB, alphabet size 4, no long runs. +# This is the input libsais is usually benchmarked on. +# chr21.fa the raw FASTA, headers and newlines included. ~45 MB, and it +# still contains its ~6.6 Mb of `N`. Wrapped at 60 columns, so +# the `N` blocks are a period-61 repeat rather than a plain run. +# +# The second is the realistic one and the one that used to be pathological: a +# comparison-based suffix sort scans the whole shared prefix on every tied +# comparison, so an `N` block costs megabytes per comparison. +set -euo pipefail + +work_dir="${1:-bench/work}" +threads="${2:-$( (command -v nproc >/dev/null && nproc) || sysctl -n hw.perflevel0.logicalcpu 2>/dev/null || echo 4)}" +root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" + +mkdir -p "$work_dir" +gz="$work_dir/chr21.fa.gz" +fa="$work_dir/chr21.fa" +bin="$work_dir/chr21.0123" + +if [ ! -f "$gz" ]; then + echo "== downloading hg38 chr21 ==" >&2 + curl -sSL -o "$gz" \ + https://hgdownload.soe.ucsc.edu/goldenPath/hg38/chromosomes/chr21.fa.gz +fi +[ -f "$fa" ] || gzip -dc "$gz" > "$fa" + +if [ ! -f "$bin" ]; then + echo "== encoding forward ++ revcomp as codes 0..=3 ==" >&2 + python3 - "$gz" "$bin" <<'PY' +import gzip, sys +code = {"A": 0, "C": 1, "G": 2, "T": 3} +comp = {0: 3, 1: 2, 2: 1, 3: 0} +fwd = bytearray() +with gzip.open(sys.argv[1], "rt") as fh: + for line in fh: + if line.startswith(">"): + continue + for ch in line.strip().upper(): + c = code.get(ch) + if c is not None: + fwd.append(c) +rc = bytearray(comp[b] for b in reversed(fwd)) +with open(sys.argv[2], "wb") as out: + out.write(fwd) + out.write(rc) +print(f"{sys.argv[2]}: {len(fwd) + len(rc)} bytes", file=sys.stderr) +PY +fi + +echo "== building (fat LTO, target-cpu=native) ==" >&2 +RUSTFLAGS="-C target-cpu=native" cargo build --release --example caps_sa --manifest-path "$root/Cargo.toml" +caps_sa="$root/target/release/examples/caps_sa" + +run() { + local label="$1" input="$2" + shift 2 + # `--verify` is an O(n) independent check of the result; it is timed and + # reported separately by the binary, so it never inflates the build time. + printf '%-28s ' "$label" + "$caps_sa" "$input" /dev/null --threads "$threads" --verify "$@" 2>&1 | + awk '/^build:/ { for (i = 1; i <= NF; i++) if ($i ~ /^[0-9.]+s$/) b = $i } + /^verify:/ { v = $0 } + END { printf "build %-9s %s\n", b, (v ~ /OK/ ? "verify OK" : "VERIFY FAILED") }' +} + +echo +echo "threads: $threads" +echo +run "0123 (80 MB, no N)" "$bin" +run "FASTA (45 MB, 6.6 Mb N)" "$fa" +echo +echo "For wall+CPU together, wrap a single run:" >&2 +echo " /usr/bin/time -p $caps_sa $bin /dev/null --threads $threads" >&2 From 290c5af7e3fafb37256159244899729b1aa7b9b7 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 21:33:42 +0200 Subject: [PATCH 07/41] docs: document the fast path, and correct two stale claims Adds to README: the in-memory fast path with measured numbers, a "Choosing a path" table giving the three soundness conditions and why each one exists, and `verify_sa` usage. Adds to bench/README: the chr21 section, including the per-merge-step arithmetic that separates the two inputs (13 ns/step on N-free DNA against 222 ns/step on the FASTA, same kernel and same machine), and the phase breakdown of the fast path. Corrects two claims that were misleading: The build paragraph attributed `lto = "fat"` and `codegen-units = 1` to a parent workspace. No such workspace is in the repository, so from the commit that made the crate standalone until the profile was added, every build made from this repo used `lto = false, codegen-units = 16`. Numbers taken in that window are not comparable with numbers taken now. The 97.54% `lcp_u8_avx2` profile was read as "LCP scanning is expensive", which motivated widening the scan through AVX2, AVX-512 and the hybrid. The LCP kernel is also where the two random text loads happen, so on short-LCP input those samples are load stalls and a wider vector cannot help. The existing AVX-512 ablation already showed this: the 64-byte-only variant was 16% slower on rand100m and only the long-LCP human slice gained. Both readings are now stated. Co-Authored-By: Claude Opus 5 (1M context) --- README.md | 93 +++++++++++++++++++++++++++++++++++++++++++++++-- bench/README.md | 86 +++++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 175 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index c018b38..f4a4d8d 100644 --- a/README.md +++ b/README.md @@ -15,8 +15,54 @@ streams the SA out as positions are emitted. ## Status Both the in-memory and external-memory paths are implemented, tested, -and benchmarked. 43 unit tests pass and the SA output is differentially -verified against a brute-force reference on small and random inputs. +and benchmarked. 73 unit tests pass and the SA output is differentially +verified against a brute-force reference on small and random inputs, and +against [`verify_sa`](#verifying-a-suffix-array) at genome scale. + +### In-memory fast path + +`build_in_memory` on a byte text routes through a **radix-seeded prefix +doubling** algorithm rather than the merge kernel. The merge kernel is +still the general path and still backs everything else; the fast path is +taken only when the comparator is provably plain lexicographic (see +[Choosing a path](#choosing-a-path)). + +The reason is that a comparison-based suffix sort pays twice on real +genomic input. It performs `n log n` merge steps, and every tied step +scans the shared prefix of two suffixes from the beginning. Genome FASTA +carries megabyte-scale runs of `N` — period-61 once 60-column line +wrapping is included — so a single comparison can scan millions of +bytes. Measured on chr21 that drives the cost per merge step from 13 ns +to 222 ns, a 16x penalty that is entirely scan time. + +The fast path sorts by a packed fixed-depth key, then resolves what +remains by doubling on ranks. The packing picks the narrowest field +width that holds the alphabet, so DNA over `{0,1,2,3}` resolves 32 +symbols per key rather than the 8 a raw byte key gives. After the seed, +no comparison reads the text again, so a megabyte-long run of `N` costs +exactly what random DNA costs. + +Apple M4 Max (12 P-cores), 12 threads, suffix arrays byte-identical to +the merge kernel's and independently verified: + +| input | before | after | CPU before | CPU after | +| ----- | ------ | ----- | ---------- | --------- | +| chr21 fwd ++ revcomp, `N`-free, 80 MB | 6.08 s | **0.89 s** | 28.1 s | 5.0 s | +| chr21 FASTA, 47.5 MB, 6.6 Mb of `N` | 27.8 s | **1.04 s** | 283.5 s | 5.2 s | + +Note the two inputs are different problems; benchmarking one +implementation on the first and another on the second is not a +comparison. `bench/chr21.sh` prepares both. + +### External memory + +On the human genome (GRCh38, 32 threads on AMD EPYC 9575F), caps-sa is +**7% faster than upstream CaPS-SA's ext-mem path** and uses **23% less +RAM**, while beating upstream's in-mem wall time by 3% at 1/10 of the +RAM. See [`bench/README.md`](bench/README.md) for the full methodology +and the optimisation ladder that got us there. Those paths still use the +merge kernel, so they retain the scan cost described above on +repeat-heavy input. On the human genome (GRCh38, 32 threads on AMD EPYC 9575F), caps-sa is **7% faster than upstream CaPS-SA's ext-mem path** and uses **23% less @@ -86,6 +132,49 @@ build_ext_mem_for_positions(&text, positions, &opts, |sa_pos| { })?; ``` +### Verifying a suffix array + +`verify_sa` checks a candidate in `O(n)` without re-running any +construction algorithm and without depending on LCP length, so it stays +usable on the repetitive inputs that are hardest to trust: + +```rust +use caps_sa::{build_in_memory, verify_sa}; + +let text = b"banana"; +let sa: Vec = build_in_memory(text); +assert!(verify_sa(text, &sa).is_ok()); +``` + +It inverts `sa` to get ranks, then checks that +`(text[p], rank[p + 1])` increases strictly along it, with `rank[n]` +treated as smaller than every real rank. A permutation of `0..n` +satisfies that condition exactly when it is the suffix array. The bench +CLI exposes it as `--verify`. + +## Choosing a path + +`build_in_memory` takes the radix-seeded doubling fast path only when +the requested comparator is provably the plain lexicographic one. All +three conditions are soundness requirements, and each defaults to +declining: + +| Condition | Why | +| --------- | --- | +| `Opts::max_context` unbounded | A finite bound makes the merge comparator fall through to `LimitProvider::boundary_order`, which compares *lengths*, so it is not lexicographic. | +| `LimitProvider::plain_lex_len()` reports the full text | Rules out `SegmentedText`, whose scans stop at segment boundaries, and any custom `boundary_order`. | +| symbol type is exactly `u8` | Packing wider symbols into an order-preserving key is endianness-dependent: on a little-endian host `0x0100 > 0x0001` as `u16` values, but their byte views compare the other way. | + +`plain_lex_len` is a new `LimitProvider` method that defaults to `None`. +An implementation that delegates `lim_at` to `PlainText` but overrides +`boundary_order` for a different convention — STAR's spacer-as-largest +ordering is the motivating example — inherits `None` and keeps today's +semantics without changing a line. + +Everything else (`*_for_positions`, `build_in_memory_sample_sort`, +`build_ext_mem`, segmented texts, wider symbols) runs on the CaPS-SA +merge kernel exactly as before. + ## Algorithm The in-memory kernel is a parallel merge-sort whose two-way merge uses diff --git a/bench/README.md b/bench/README.md index 45cc7f7..ff9e06f 100644 --- a/bench/README.md +++ b/bench/README.md @@ -32,8 +32,16 @@ It reports wall time and peak RSS via `/usr/bin/time`. Machine: 64-core x86_64 Linux node, 1 socket, AVX2 enabled. Builds: upstream `cmake -DCMAKE_BUILD_TYPE=Release`, this crate -`cargo build --release --example caps_sa` (release profile in the -workspace `Cargo.toml` enables `lto = "fat"` + `codegen-units = 1`). +`cargo build --release --example caps_sa`. The `[profile.release]` in +this repo's `Cargo.toml` sets `lto = "fat"` + `codegen-units = 1`; for +the SIMD paths also pass `RUSTFLAGS="-C target-cpu=native"`. + +> **Correction.** This paragraph used to attribute the LTO settings to a +> parent workspace. No such workspace exists in the repository, so from +> the commit that made the crate standalone until the profile was added +> here, every build made from this repo actually used `lto = false` and +> `codegen-units = 16`. Numbers taken in that window are not comparable +> with numbers taken now. Re-measure before quoting them. All caps-sa runs include the five optimizations applied incrementally to the Phase 2b sample-sort baseline: @@ -399,6 +407,80 @@ against `build_ext_mem_for_positions` (sort only the kept positions) on the same fixture, showing **6–10× speedups** on padding-dominated inputs. The pool change leaves this gap intact. +### chr21 — the two inputs are different problems (in-memory fast path) + +Machine: Apple M4 Max (12 P-cores + 4 E-cores, aarch64/NEON, no +AVX-512), 12 threads, `RUSTFLAGS="-C target-cpu=native"`. Reproduce with +`bench/chr21.sh`. + +Two inputs are built from hg38 chr21, and the distinction matters more +than any tuning parameter in this document: + +| input | size | alphabet | long runs | +| ----- | ---- | -------- | --------- | +| `chr21.0123` — forward ++ revcomp, codes `0..=3`, ambiguous bases dropped | 80.2 MB | 4 | none | +| `chr21.fa` — the raw FASTA, headers and newlines included | 47.5 MB | 6 | 6.6 Mb of `N`, period-61 after 60-column wrapping | + +Benchmarking one implementation on the first and another on the second +compares two different problems. That is worth stating explicitly +because it is an easy mistake to make: the second input is 40% smaller +yet used to take 4.6x longer. + +Merge kernel vs. the radix-seeded doubling fast path, suffix arrays +byte-identical and independently `--verify`-checked in both cases: + +| input | merge kernel | fast path | speedup | CPU before | CPU after | CPU speedup | +| ----- | ------------ | --------- | ------- | ---------- | --------- | ----------- | +| `chr21.0123` | 6.08 s | **0.89 s** | 6.8x | 28.1 s | 5.0 s | 5.6x | +| `chr21.fa` | 27.8 s | **1.04 s** | 26.7x | 283.5 s | 5.2 s | **54x** | + +The per-merge-step cost is what separates the two rows. Dividing CPU +time by `n log2 n` merge steps: + +``` +chr21.0123 28.1 s / 2.11e9 steps = 13 ns/step +chr21.fa 283.5 s / 1.21e9 steps = 222 ns/step +``` + +Same kernel, same machine, 16x apart. The difference is entirely scan +length: every leaf merge starts with `m = 0`, so two suffixes inside an +`N` block scan their whole shared prefix, which is megabytes. + +Phase breakdown of the fast path (`CAPS_SA_PROFILE=1`): + +``` + chr21.0123 chr21.fa +key extraction 0.086 s 0.020 s +seed sort 0.255 s 0.176 s +grouping 0.075 s 0.037 s +doubling rounds 0.475 s 0.807 s +``` + +The doubling rounds are the larger share on the FASTA input, as +expected: `N` blocks stay tied for many rounds. But each round is +rank-only, so they cost a sort over a shrinking residual rather than a +text scan, which is why the pathology disappears rather than merely +shrinking. + +### Reading the 97.54% LCP profile correctly + +The profile in the next section shows `lcp_u8_avx2` taking 97.54% of +samples on a human-genome slice. That was read at the time as "LCP +scanning is expensive", and it drove several rounds of work on making +the scan wider — AVX2, then AVX-512, then the 32-byte/64-byte hybrid. + +The chr21 numbers above suggest a second reading. The LCP kernel is also +where the two *random* text loads happen, so on short-LCP input those +samples are load stalls rather than scan length, and a wider vector does +not help. That is consistent with the AVX-512 ablation further down, +where the 64-byte-only variant was **16% slower** on `rand100m` and only +the long-LCP human slice gained. Widening the scan helps when the scan +is genuinely long; when it is short, the cost is latency and the fix is +to issue fewer random probes. + +The fast path does the latter: it removes the probes entirely rather +than making each one faster. + ### Where AVX-512 helps and where it doesn't — the measurement A `perf record --call-graph dwarf` run on a 200 MB human-genome slice From 72db16541f936486fc58bdf52aad2898925bac4c Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 21:33:42 +0200 Subject: [PATCH 08/41] docs: fix a broken intra-doc link in `build_ext_mem_for_filter` Public documentation linked to `FilteredSource`, which is private, so `cargo doc` fails under `RUSTDOCFLAGS=-D warnings`. Pre-existing, but it blocks the rustdoc job added in the previous commit. Co-Authored-By: Claude Opus 5 (1M context) --- src/ext_mem.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/ext_mem.rs b/src/ext_mem.rs index 6ef6e84..f5560ad 100644 --- a/src/ext_mem.rs +++ b/src/ext_mem.rs @@ -388,8 +388,8 @@ where /// bytes — ~770 MB on the human genome, vs the ~50 GB the equivalent /// `Vec` would take). Phase 1's per-subarray fill is then driven /// by popcount-walking the bitmap; the predicate is **never invoked -/// again** after the initial build. See [`FilteredSource`] for the -/// memory accounting and the inner loop. +/// again** after the initial build. See the crate-internal +/// `FilteredSource` for the memory accounting and the inner loop. /// /// Use this entry when the caller already has the text in RAM and /// the kept positions are described by a cheap per-position From a7daf6d12a00ecec345d40906585bd560c29c868 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 21:59:35 +0200 Subject: [PATCH 09/41] perf: replace the seed comparison sort with an MSD counting sort The seed was a `par_sort_unstable` over a materialised `Vec<(u64, u32, I)>`, which was both the largest remaining cost and the peak-memory driver. Three things change. The source is never materialised: keys are recomputed from `text` in the histogram pass and again in the scatter, which trades a random read of an n-element key array for a sequential read of the text. Peak memory drops to the two destination buffers, 12 bytes per position at `I = u32` against the 16 a `(u64, u32, I)` record costs. And the top-level partition becomes a counting pass, which parallelises evenly, where a parallel comparison sort's first partitioning steps are close to serial. 11 bits (2048 buckets) keeps the write-combining state near 512 KB and inside a core's private cache. 16 bits would need 16 MB of open write lines and thrash the TLB instead. The visible-length tie-break no longer needs storing. Only the last `k - 1` positions can have a visible length below `k`, so it is a function of the position alone and is applied in the per-bucket sort and in the group scan. Apple M4 Max, 12 threads, chr21 fwd+revcomp 80 MB, output unchanged and `--verify` clean: peak RSS 2.83 GB -> 2.21 GB (-22%) seed 0.341 s -> 0.289 s total 0.892 s -> 0.844 s Co-Authored-By: Claude Opus 5 (1M context) --- src/radix.rs | 162 +++++++++++++++++++++++++++++++++++++++++++++------ 1 file changed, 144 insertions(+), 18 deletions(-) diff --git a/src/radix.rs b/src/radix.rs index edb01b8..7eab7db 100644 --- a/src/radix.rs +++ b/src/radix.rs @@ -94,6 +94,136 @@ fn key_at(text: &[u8], p: usize, bits: u32, k: usize) -> u64 { key << (bits as usize * (k - (end - p))) } +/// Bits of the key used for the MSD counting-sort pass, and the resulting +/// bucket count. +/// +/// 2048 buckets keeps the write-combining state at 2048 × 2 streams × 128 B +/// ≈ 512 KB, which stays inside a core's private cache. Going to 16 bits +/// would need 16 MB of open write lines and thrashes the TLB instead. +const RADIX_BITS: u32 = 11; +const RADIX_BUCKETS: usize = 1 << RADIX_BITS; + +/// Sort every suffix position by `(packed key, visible length)`, returning the +/// sorted keys alongside the sorted positions. +/// +/// This is an MSD counting sort rather than a comparison sort, for three +/// reasons that all matter at genome scale: +/// +/// * The source is never materialised. Keys are recomputed from `text` in +/// both the histogram and the scatter pass, which is a sequential read of +/// the text instead of a random read of an `n`-element key array. +/// * Peak memory is the two destination buffers only, 12 bytes per position +/// with `I = u32`, against the 16 a `(u64, u32, I)` record costs. +/// * The top-level partition is a counting pass, so it parallelises evenly. +/// A parallel comparison sort's first partitioning steps are close to +/// serial, which is exactly where a `n log n` sort loses on many cores. +fn seed_sort( + text: &[u8], + bits: u32, + k: usize, + visible_len: &(dyn Fn(usize) -> usize + Sync), +) -> (Vec, Vec) { + let n = text.len(); + let bucket_of = |key: u64| -> usize { (key >> (64 - RADIX_BITS)) as usize }; + + // Chunk the position range so each worker builds a private histogram. + let n_chunks = (rayon::current_num_threads() * 4).clamp(1, 1024); + let chunk_len = n.div_ceil(n_chunks); + let bounds: Vec<(usize, usize)> = (0..n) + .step_by(chunk_len) + .map(|s| (s, (s + chunk_len).min(n))) + .collect(); + + // Pass 1: per-chunk histograms over the top `RADIX_BITS` of each key. + let histograms: Vec> = bounds + .par_iter() + .map(|&(start, end)| { + let mut counts = vec![0u32; RADIX_BUCKETS]; + for p in start..end { + counts[bucket_of(key_at(text, p, bits, k))] += 1; + } + counts + }) + .collect(); + + // Exclusive prefix sum, bucket-major then chunk-minor, so every (chunk, + // bucket) pair gets a disjoint destination range and the buckets come out + // in ascending key order. + let mut offsets = vec![0usize; bounds.len() * RADIX_BUCKETS]; + let mut bucket_start = vec![0usize; RADIX_BUCKETS + 1]; + { + let mut running = 0usize; + for b in 0..RADIX_BUCKETS { + bucket_start[b] = running; + for (c, hist) in histograms.iter().enumerate() { + offsets[c * RADIX_BUCKETS + b] = running; + running += hist[b] as usize; + } + } + bucket_start[RADIX_BUCKETS] = running; + debug_assert_eq!(running, n); + } + + // Pass 2: scatter. Each chunk owns a disjoint slice of every bucket, so + // the writes never collide even though they are not contiguous. + let mut keys: Vec = vec![0; n]; + let mut sa: Vec = vec![I::zero(); n]; + { + let key_out = Scatter::new(&mut keys); + let sa_out = Scatter::new(&mut sa); + bounds + .par_iter() + .enumerate() + .for_each(|(c, &(start, end))| { + let mut cursor: Vec = + offsets[c * RADIX_BUCKETS..(c + 1) * RADIX_BUCKETS].to_vec(); + for p in start..end { + let key = key_at(text, p, bits, k); + let slot = &mut cursor[bucket_of(key)]; + // SAFETY: the prefix sum gives this (chunk, bucket) pair a + // range of exactly its own histogram count, and the cursor + // never leaves it, so no other thread writes this index. + unsafe { + key_out.set(*slot, key); + sa_out.set(*slot, I::from_usize(p)); + } + *slot += 1; + } + }); + } + + // Pass 3: order within each bucket. Buckets share their top `RADIX_BITS`, + // so what remains is the low bits of the key and then the visible-length + // tie-break. Buckets are contiguous and independent. + let mut rest: &mut [u64] = &mut keys; + let mut rest_sa: &mut [I] = &mut sa; + let mut slices: Vec<(&mut [u64], &mut [I])> = Vec::with_capacity(RADIX_BUCKETS); + for b in 0..RADIX_BUCKETS { + let len = bucket_start[b + 1] - bucket_start[b]; + let (kb, kt) = rest.split_at_mut(len); + let (sb, st) = rest_sa.split_at_mut(len); + slices.push((kb, sb)); + rest = kt; + rest_sa = st; + } + slices.into_par_iter().for_each(|(kb, sb)| { + if kb.len() < 2 { + return; + } + let mut pairs: Vec<(u64, I)> = kb.iter().copied().zip(sb.iter().copied()).collect(); + pairs.sort_unstable_by(|a, b| { + a.0.cmp(&b.0) + .then_with(|| visible_len(a.1.to_usize()).cmp(&visible_len(b.1.to_usize()))) + }); + for (i, &(key, pos)) in pairs.iter().enumerate() { + kb[i] = key; + sb[i] = pos; + } + }); + + (keys, sa) +} + /// Build the standard lexicographic suffix array of `text` by radix-seeded /// prefix doubling. /// @@ -125,30 +255,27 @@ pub(crate) fn build_sa(text: &[u8]) -> Vec { // never separates suffixes that the doubling rounds still need to see as // tied. Without it, `[0, 0]` leaves positions 0 and 1 permanently tied // and the doubling loop cannot terminate. - let mut seeded: Vec<(u64, u32, I)> = (0..n) - .into_par_iter() - .map(|p| { - ( - key_at(text, p, bits, k), - (n - p).min(k) as u32, - I::from_usize(p), - ) - }) - .collect(); + // Only the last `k - 1` positions can have a visible length below `k`, so + // the tie-break is a function of the position alone and never has to be + // stored alongside the key. + let visible_len = |p: usize| -> usize { (n - p).min(k) }; + profile_log(&format!( - "radix keys {:.3}s", + "radix setup {:.3}s", t0.elapsed().as_secs_f64() )); let t1 = Instant::now(); - seeded.par_sort_unstable(); + let (keys, mut sa) = seed_sort::(text, bits, k, &visible_len); profile_log(&format!( "radix seed sort {:.3}s", t1.elapsed().as_secs_f64() )); let t2 = Instant::now(); - let mut sa: Vec = Vec::with_capacity(n); - sa.par_extend(seeded.par_iter().map(|&(_, _, p)| p)); + // Two seeded entries tie exactly when key and visible length both match. + let seed_eq = |a: usize, b: usize| -> bool { + keys[a] == keys[b] && visible_len(sa[a].to_usize()) == visible_len(sa[b].to_usize()) + }; // `rank[p]` is the index in `sa` of the first element of `p`'s group, so // two suffixes tie at the current depth exactly when their ranks match, @@ -165,12 +292,11 @@ pub(crate) fn build_sa(text: &[u8]) -> Vec { let groups: Vec<(usize, usize)> = (0..n) .into_par_iter() .filter_map(|h| { - let key = (seeded[h].0, seeded[h].1); - if h > 0 && (seeded[h - 1].0, seeded[h - 1].1) == key { + if h > 0 && seed_eq(h - 1, h) { return None; } let mut e = h + 1; - while e < n && (seeded[e].0, seeded[e].1) == key { + while e < n && seed_eq(e, h) { e += 1; } let g = I::from_usize(h); @@ -184,7 +310,7 @@ pub(crate) fn build_sa(text: &[u8]) -> Vec { }) .collect(); let mut groups = groups; - drop(seeded); + drop(keys); profile_log(&format!( "radix grouping {:.3}s", t2.elapsed().as_secs_f64() From 7d79209eb92f53e90f586c1d4ba382dac2642646 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 21:59:35 +0200 Subject: [PATCH 10/41] feat: extend the doubling path to subsets and to in-memory sample sort Closes two of the three gaps listed as limitations when the fast path landed. **Subsets.** Doubling cannot be restricted to a subset directly, because a round compares `rank[p + d]` and that successor is generally not in the subset, so ranks must be defined for every position in the text. Building the whole array and filtering it sidesteps that, and the filter is one O(n) pass since the full array is already ordered. Worth it when the subset is a real fraction of the text, which is what this API exists for: STAR-style indexing keeps every ACGT position and drops only spacers. Below one eighth of the text it declines, because O(n) to build and discard would dwarf the O(m log m) the merge kernel needs. That ratio is a performance heuristic; the guards it sits behind remain correctness conditions. It also declines on duplicate or out-of-range positions. The output is a permutation of the input *multiset*, which a membership filter cannot reproduce. **In-memory sample sort.** `build_in_memory_sample_sort` exists to sort in RAM, so where the doubling path applies it is strictly better: same output, no bucket machinery, and none of the scan cost on repeat-heavy text. `build_ext_mem` deliberately does *not* get this. Its purpose is to bound peak memory, and routing it through an in-memory algorithm would defeat exactly that. It stays on the merge kernel, and the remaining limitation is now stated as a deliberate choice rather than an omission. Apple M4 Max, 12 threads, chr21 fwd+revcomp 80 MB. `--in-mem-ss` output verified identical to the in-memory path's: --in-mem-ss 3.41 s -> 1.18 s wall, 34.5 s -> 7.2 s CPU Tests: duplicate positions keep their multiplicity, a tiny subset of a large text still matches brute force through the merge kernel, and out-of-range positions still panic rather than silently returning a wrong answer. Co-Authored-By: Claude Opus 5 (1M context) --- src/ext_mem.rs | 20 +++++++++ src/sample_sort.rs | 103 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 123 insertions(+) diff --git a/src/ext_mem.rs b/src/ext_mem.rs index f5560ad..6653437 100644 --- a/src/ext_mem.rs +++ b/src/ext_mem.rs @@ -709,6 +709,26 @@ where L: LimitProvider, F: FnMut(u64) -> Result<(), E>, { + // This path exists to sort in RAM, so when the doubling path applies it is + // strictly better here: same output, no bucket machinery, and none of the + // scan cost the merge kernel pays on repeat-heavy text. `build_ext_mem` + // deliberately does *not* do this -- its whole purpose is to bound peak + // memory, and routing it through an in-memory algorithm would defeat that. + if opts.max_context == usize::MAX + && lp.plain_lex_len() == Some(text.len()) + && std::any::TypeId::of::() == std::any::TypeId::of::() + { + // SAFETY: `S` is `u8`, so `&[S]` and `&[u8]` have identical layout. + let bytes: &[u8] = + unsafe { std::slice::from_raw_parts(text.as_ptr() as *const u8, text.len()) }; + let sa: Vec = crate::radix::build_sa(bytes); + let mut emit = emit; + for pos in sa { + emit(pos).map_err(BuildError::Emit)?; + } + return Ok(()); + } + if text.len() <= u32::MAX as usize + 1 { build_in_memory_ss_inner::( text, diff --git a/src/sample_sort.rs b/src/sample_sort.rs index e69a699..e478a72 100644 --- a/src/sample_sort.rs +++ b/src/sample_sort.rs @@ -34,6 +34,7 @@ use crate::Index; use crate::lcp::{LcpDispatch, Symbol}; use crate::limits::{LimitProvider, PlainText}; use rayon::join; +use rayon::prelude::*; /// Tunable options for SA construction. #[derive(Clone, Debug)] @@ -142,6 +143,65 @@ where Some(crate::radix::build_sa(bytes)) } +/// Sort a *subset* of positions by building the full suffix array with the +/// doubling path and then keeping only the requested positions. +/// +/// Doubling cannot be restricted to a subset directly: a round compares +/// `rank[p + d]`, and that successor is generally not in the subset, so ranks +/// have to be defined for every position in the text. Building the whole array +/// and filtering it sidesteps that, and the filter is a single `O(n)` pass +/// because the full SA is already in the right order. +/// +/// Worth it when the subset is a decent fraction of the text, which is the +/// case this API exists for (STAR-style indexing keeps every ACGT position and +/// drops only spacers). For a small subset, `O(n)` to build the full array +/// would dwarf the `O(m log m)` the merge kernel needs, so below one eighth of +/// the text this declines and the merge kernel runs. That ratio is a +/// performance heuristic, unlike the guards in [`try_doubling_fast_path`], +/// which are correctness conditions. +/// +/// Also declines on duplicate or out-of-range positions, which a +/// membership filter cannot reproduce faithfully. +fn try_doubling_subset(text: &[S], positions: &[I], lp: &L, opts: &Opts) -> Option> +where + S: Symbol, + I: Index, + L: LimitProvider, +{ + let n = text.len(); + let m = positions.len(); + if m == 0 || m.checked_mul(8)? < n { + return None; + } + if opts.max_context != usize::MAX + || lp.plain_lex_len() != Some(n) + || std::any::TypeId::of::() != std::any::TypeId::of::() + { + return None; + } + + let mut wanted = vec![false; n]; + for p in positions { + let p = p.to_usize(); + // Out of range, or the same position twice: a membership filter emits + // each position at most once, so it cannot reproduce either faithfully. + if p >= n || wanted[p] { + return None; + } + wanted[p] = true; + } + + // SAFETY: `S` is `u8`, so `&[S]` and `&[u8]` have identical layout. + let bytes: &[u8] = unsafe { std::slice::from_raw_parts(text.as_ptr() as *const u8, n) }; + let full: Vec = crate::radix::build_sa(bytes); + let kept: Vec = full + .into_par_iter() + .filter(|p| wanted[p.to_usize()]) + .collect(); + debug_assert_eq!(kept.len(), m); + Some(kept) +} + /// Sort the caller-supplied `positions` by the lexicographic order of /// their suffixes in `text`. Returns the positions reordered so that /// `text[output[i]..]` is the i-th smallest suffix among the input set. @@ -192,6 +252,10 @@ where I: Index, L: LimitProvider, { + if let Some(sa) = try_doubling_subset::(text, &positions, lp, opts) { + return sa; + } + let n = positions.len(); if n == 0 { return Vec::new(); @@ -648,6 +712,45 @@ mod tests { assert_eq!(got, want); } + /// Duplicated positions must survive: the output is a permutation of the + /// *input* multiset, which a membership filter cannot reproduce, so the + /// subset fast path has to decline and let the merge kernel run. + #[test] + fn for_positions_with_duplicates_keeps_multiplicity() { + let text = b"mississippi"; + let positions: Vec = vec![0, 1, 1, 4, 4, 4, 7]; + let mut want = positions.clone(); + want.sort_by(|&a, &b| text[a as usize..].cmp(&text[b as usize..])); + let got = build_in_memory_for_positions(text, positions); + assert_eq!(got, want); + } + + /// A subset far smaller than the text takes the merge kernel, since + /// building the whole suffix array to throw nearly all of it away would + /// cost more than sorting the subset directly. Correctness is identical + /// either way; this pins the behaviour. + #[test] + fn for_positions_tiny_subset_of_large_text() { + use rand::{RngExt, SeedableRng}; + let mut rng = rand::rngs::StdRng::seed_from_u64(0x5AB0); + let text: Vec = (0..20_000).map(|_| rng.random_range(0..4u8)).collect(); + let positions: Vec = (0..20_000u32).step_by(500).collect(); + let mut want = positions.clone(); + want.sort_by(|&a, &b| text[a as usize..].cmp(&text[b as usize..])); + let got = build_in_memory_for_positions(&text, positions); + assert_eq!(got, want); + } + + /// Positions out of range are the caller's error, but the subset fast path + /// must not turn them into a silently wrong answer or an unsafe index. + #[test] + #[should_panic] + fn for_positions_out_of_range_still_panics() { + let text = b"banana"; + let positions: Vec = vec![0, 1, 99]; + let _ = build_in_memory_for_positions(text, positions); + } + #[test] fn for_positions_random_subsets() { use rand::{RngExt, SeedableRng}; From 31c1bb4d7afc5811cb560178521bc62f3b39ed80 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 21:59:35 +0200 Subject: [PATCH 11/41] docs: refresh the fast-path numbers and coverage Records the counting-sort seed (peak RSS 2.83 GB -> 2.21 GB, total 0.89 s -> 0.84 s on the 80 MB input), adds the `--in-mem-ss` row, and replaces the "everything else uses the merge kernel" line with what is now actually true: subsets and in-memory sample sort are covered, and `build_ext_mem` stays on the merge kernel deliberately, because bounding peak memory is the whole point of that path. Co-Authored-By: Claude Opus 5 (1M context) --- README.md | 31 +++++++++++++++++++++++++++---- 1 file changed, 27 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index f4a4d8d..10e988f 100644 --- a/README.md +++ b/README.md @@ -47,8 +47,13 @@ the merge kernel's and independently verified: | input | before | after | CPU before | CPU after | | ----- | ------ | ----- | ---------- | --------- | -| chr21 fwd ++ revcomp, `N`-free, 80 MB | 6.08 s | **0.89 s** | 28.1 s | 5.0 s | +| chr21 fwd ++ revcomp, `N`-free, 80 MB | 6.08 s | **0.84 s** | 28.1 s | 5.0 s | | chr21 FASTA, 47.5 MB, 6.6 Mb of `N` | 27.8 s | **1.04 s** | 283.5 s | 5.2 s | +| same, via `build_in_memory_sample_sort` | 3.41 s | **1.18 s** | 34.5 s | 7.2 s | + +Peak RSS on the 80 MB input is 2.21 GB, down from 2.83 GB, because the +seed is an MSD counting sort that recomputes keys from the text rather +than materialising a key array. Note the two inputs are different problems; benchmarking one implementation on the first and another on the second is not a @@ -171,9 +176,27 @@ An implementation that delegates `lim_at` to `PlainText` but overrides ordering is the motivating example — inherits `None` and keeps today's semantics without changing a line. -Everything else (`*_for_positions`, `build_in_memory_sample_sort`, -`build_ext_mem`, segmented texts, wider symbols) runs on the CaPS-SA -merge kernel exactly as before. +Given those, the fast path also covers two cases beyond a plain whole-text +build: + +- **`*_for_positions` subsets.** Doubling cannot be restricted to a subset + directly, since a round compares `rank[p + d]` and that successor is + generally outside the subset, so ranks must exist for every text + position. The full array is built and filtered in one `O(n)` pass + instead. Below one eighth of the text this declines and the merge kernel + runs, since building and discarding a whole array would cost more than + sorting a small subset. That ratio is a performance heuristic, not a + correctness condition. Duplicate or out-of-range positions also decline, + because the output is a permutation of the input *multiset* and a + membership filter cannot reproduce that. +- **`build_in_memory_sample_sort`.** This path exists to sort in RAM, so + where doubling applies it is strictly better: same output, no bucket + machinery, none of the scan cost. + +`build_ext_mem` deliberately stays on the merge kernel. Its purpose is to +bound peak memory, and routing it through an in-memory algorithm would +defeat exactly that, so it keeps the scan cost on repeat-heavy input. +Segmented texts and symbols wider than `u8` also stay on the merge kernel. ## Algorithm From 2ec30649374c08354d496e53881b8b53a3f2ee5a Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 21:59:35 +0200 Subject: [PATCH 12/41] perf: skip long periodic runs instead of scanning them `build_ext_mem` was left on the merge kernel because bounding peak memory is its purpose and prefix doubling needs a rank for every position in the text. That left it paying the full scan cost on repeat-heavy input. Profiling put 94% of an ext-mem run on chr21 FASTA in phase 1: chr21.0123 (N-free, 80 MB) phase1 1.15 s of 3.49 s total chr21.fa (47.5 MB) phase1 22.85 s of 24.19 s total Normalised that is 0.014 s/MB against 0.48 s/MB, the same 34x scan penalty the in-memory path had, and for the same reason: two suffixes inside a long repeat agree for as far as it continues, so one comparison scans megabytes. Fix it in the comparator rather than the algorithm, which keeps the memory bound intact. If `text[s..e)` has period `q` and two suffixes start at `a < b` inside it with `(b - a) % q == 0`, they agree until the later one reaches `e`, so `lcp(a, b) >= e - b` is known in O(1) from the run's bounds with no scanning. When the phase does not match, the two must differ within `q` symbols and the ordinary scan is already short. The scan is additionally bounded so it stops at a run's start rather than traversing it. Detecting only single-symbol runs would have missed the case that actually occurs: in wrapped FASTA an `N` block is 60 `N`s then a newline, which is period 61, not period 1. Periods up to 64 are considered. Detection is two-stage so texts without runs pay almost nothing. A sampling pass looks for any periodic window and collects the periods that occur; the full scan runs only for those. On N-free DNA the sample finds nothing, the table is empty, and every query short-circuits on a slice-empty check. The table itself is a few dozen entries, so the memory bound is untouched. `Cmp` bundles the SIMD dispatch with the run table and replaces the bare `LcpDispatch` threaded through the merge kernel, so phase 1, the phase-3 pivot searches and the phase-4 cascade all benefit. It stays `Copy` and still travels through the recursion in registers. Apple M4 Max, 12 threads. Ext-mem output verified identical to the in-memory suffix array on both inputs: chr21.fa 24.19 s -> 2.48 s wall, 268 s -> 23.3 s CPU phase 1 22.85 s -> 0.95 s peak RSS 147 MB -> 151 MB chr21.0123 3.49 s -> 3.55 s wall (unchanged; no runs to find) Tests: the run-aware LCP is checked against a byte-at-a-time oracle over sampled position pairs on homopolymers, wrapped-FASTA blocks and multi-period texts, plus detection shape (sorted, disjoint, period claim actually holds) and `max_context` behaviour inside a run. Co-Authored-By: Claude Opus 5 (1M context) --- src/ext_mem.rs | 78 ++++---- src/lib.rs | 1 + src/runs.rs | 449 +++++++++++++++++++++++++++++++++++++++++++++ src/sample_sort.rs | 24 +-- 4 files changed, 499 insertions(+), 53 deletions(-) create mode 100644 src/runs.rs diff --git a/src/ext_mem.rs b/src/ext_mem.rs index 6653437..28352f8 100644 --- a/src/ext_mem.rs +++ b/src/ext_mem.rs @@ -42,6 +42,7 @@ use crate::ext_bucket::{ }; use crate::lcp::{LcpDispatch, Symbol}; use crate::limits::{LimitProvider, PlainText}; +use crate::runs::Cmp; use crate::sample_sort; /// Emit a phase-timing line to stderr if `CAPS_SA_PROFILE` is set in @@ -502,7 +503,8 @@ where return Ok(()); } let p = effective_subproblem_count(n, opts.subproblem_count); - let dispatch = LcpDispatch::detect(); + let runs = crate::runs::detect_for(text); + let cmp = Cmp::new(LcpDispatch::detect(), &runs); let work_dir = opts.work_dir.clone(); // Pool the `2 × p` bucket files into one anonymous tempfile per @@ -527,15 +529,8 @@ where let part_factory = |j: usize| phase3_pool.new_bucket::>(j); let t = Instant::now(); - let (mut subarray_buckets, samples) = phase1_sort_sample_spill::( - text, - lp, - &source, - p, - opts, - dispatch, - sub_factory, - )?; + let (mut subarray_buckets, samples) = + phase1_sort_sample_spill::(text, lp, &source, p, opts, cmp, sub_factory)?; profile_log(&format!( "phase1 (sort+sample+spill) {:.3}s", t.elapsed().as_secs_f64() @@ -551,7 +546,7 @@ where drop(source); let t = Instant::now(); - let pivots = phase2_select_pivots::(text, lp, samples, p, opts.max_context, dispatch); + let pivots = phase2_select_pivots::(text, lp, samples, p, opts.max_context, cmp); profile_log(&format!( "phase2 (select pivots) {:.3}s", t.elapsed().as_secs_f64() @@ -565,7 +560,7 @@ where &pivots, p, opts, - dispatch, + cmp, part_factory, )?; profile_log(&format!( @@ -583,7 +578,7 @@ where opts.max_context, opts.ordered_phase4_emit, &mut emit, - dispatch, + cmp, ); profile_log(&format!( "phase4 (merge+emit) {:.3}s", @@ -622,16 +617,17 @@ where return Ok(()); } let p = effective_subproblem_count(n, opts.subproblem_count); - let dispatch = LcpDispatch::detect(); + let runs = crate::runs::detect_for(text); + let cmp = Cmp::new(LcpDispatch::detect(), &runs); let factory = |_i: usize| InMemBucket::>::new(); let (mut subarray_buckets, samples) = - phase1_sort_sample_spill::(text, lp, &source, p, opts, dispatch, factory)?; + phase1_sort_sample_spill::(text, lp, &source, p, opts, cmp, factory)?; // Same rationale as in `build_ext_mem_inner` — drop the source // as soon as phase 1's `fill_chunk` calls have stopped. drop(source); - let pivots = phase2_select_pivots::(text, lp, samples, p, opts.max_context, dispatch); + let pivots = phase2_select_pivots::(text, lp, samples, p, opts.max_context, cmp); let mut partition_buckets = phase3_distribute::( text, lp, @@ -639,7 +635,7 @@ where &pivots, p, opts, - dispatch, + cmp, factory, )?; drop(subarray_buckets); @@ -650,7 +646,7 @@ where opts.max_context, opts.ordered_phase4_emit, &mut emit, - dispatch, + cmp, ) } @@ -1163,7 +1159,7 @@ fn phase1_sort_sample_spill( source: &PositionSource<'_>, p: usize, opts: &ExtMemOpts, - dispatch: LcpDispatch, + cmp: Cmp<'_>, mk_bucket: MkB, ) -> io::Result<(Vec, Vec)> where @@ -1204,7 +1200,7 @@ where &mut lcp_arr, &mut lcp_w, opts.max_context, - dispatch, + cmp, ); // Pull `samples_per_subarray` evenly-spaced positions out of @@ -1267,7 +1263,7 @@ fn phase2_select_pivots( mut samples: Vec, p: usize, max_ctx: usize, - dispatch: LcpDispatch, + cmp: Cmp<'_>, ) -> Vec where S: Symbol, @@ -1289,7 +1285,7 @@ where &mut lcp, &mut lcp_w, max_ctx, - dispatch, + cmp, ); // p-1 pivots at evenly-spaced ranks across the sorted sample pool. @@ -1320,7 +1316,7 @@ fn phase3_distribute( pivots: &[I], p: usize, opts: &ExtMemOpts, - dispatch: LcpDispatch, + cmp: Cmp<'_>, mk_bucket: MkB, ) -> io::Result> where @@ -1353,7 +1349,7 @@ where text, lp, opts.max_context, - dispatch, + cmp, )); } splits.push(records.len()); @@ -1391,7 +1387,7 @@ fn upper_bound_by_pivot( text: &[S], lp: &L, max_ctx: usize, - dispatch: LcpDispatch, + cmp: Cmp<'_>, ) -> usize where S: Symbol, @@ -1402,7 +1398,7 @@ where let mut hi = records.len(); while lo < hi { let mid = lo + (hi - lo) / 2; - match dispatch.suffix_cmp_with( + match cmp.suffix_cmp_with( text, lp, records[mid].pos.to_usize(), @@ -1439,7 +1435,7 @@ fn phase4_merge_and_emit( max_ctx: usize, ordered_emit: bool, emit: &mut F, - dispatch: LcpDispatch, + cmp: Cmp<'_>, ) -> Result<(), BuildError> where S: Symbol, @@ -1462,7 +1458,7 @@ where // GRCh38 / 32 t). // // Bumping the chunk to `4 × num_threads` gives rayon four - // partitions per thread to dispatch — fast threads can steal from + // partitions per thread to cmp — fast threads can steal from // slow neighbours, smoothing out the size variance. Peak RAM // grows linearly: each in-flight merged partition holds its // result `Vec` (~3 MB at human-genome scale with `u32` @@ -1491,7 +1487,7 @@ where chunk, max_ctx, emit, - dispatch, + cmp, profile, &load_us, &merge_us, @@ -1504,7 +1500,7 @@ where chunk, max_ctx, emit, - dispatch, + cmp, profile, &load_us, &merge_us, @@ -1531,7 +1527,7 @@ fn phase4_merge_chunk_collect_emit( chunk: &mut [B], max_ctx: usize, emit: &mut F, - dispatch: LcpDispatch, + cmp: Cmp<'_>, profile: bool, load_us: &std::sync::atomic::AtomicU64, merge_us: &std::sync::atomic::AtomicU64, @@ -1550,9 +1546,7 @@ where let merged: Vec> = chunk .par_iter_mut() .map(|bucket| -> io::Result> { - merge_one_partition( - text, lp, bucket, max_ctx, dispatch, profile, load_us, merge_us, - ) + merge_one_partition(text, lp, bucket, max_ctx, cmp, profile, load_us, merge_us) }) .collect::, io::Error>>()?; @@ -1575,7 +1569,7 @@ fn phase4_merge_chunk_ordered_emit( chunk: &mut [B], max_ctx: usize, emit: &mut F, - dispatch: LcpDispatch, + cmp: Cmp<'_>, profile: bool, load_us: &std::sync::atomic::AtomicU64, merge_us: &std::sync::atomic::AtomicU64, @@ -1605,7 +1599,7 @@ where .enumerate() .for_each_with(tx, |tx, (local_idx, bucket)| { let result = merge_one_partition( - text, lp, bucket, max_ctx, dispatch, profile, load_us, merge_us, + text, lp, bucket, max_ctx, cmp, profile, load_us, merge_us, ); let _ = tx.send((local_idx, result)); }); @@ -1662,7 +1656,7 @@ fn merge_one_partition( lp: &L, bucket: &mut B, max_ctx: usize, - dispatch: LcpDispatch, + cmp: Cmp<'_>, profile: bool, load_us: &std::sync::atomic::AtomicU64, merge_us: &std::sync::atomic::AtomicU64, @@ -1688,7 +1682,7 @@ where let t = Instant::now(); let workspace = CascadeWorkspace::::new(); - let result = workspace.cascade_merge(text, lp, &records, &boundaries, max_ctx, dispatch); + let result = workspace.cascade_merge(text, lp, &records, &boundaries, max_ctx, cmp); if profile { merge_us.fetch_add(t.elapsed().as_micros() as u64, AtomicOrdering::Relaxed); } @@ -1745,7 +1739,7 @@ impl CascadeWorkspace { records: &[SaLcp], boundaries: &[usize], max_ctx: usize, - dispatch: LcpDispatch, + cmp: Cmp<'_>, ) -> Vec where S: Symbol, @@ -1773,7 +1767,7 @@ impl CascadeWorkspace { let mut src_is_a = true; while run_lens.len() > 1 { - run_lens = self.merge_one_level(src_is_a, &run_lens, text, lp, max_ctx, dispatch); + run_lens = self.merge_one_level(src_is_a, &run_lens, text, lp, max_ctx, cmp); src_is_a = !src_is_a; } @@ -1797,7 +1791,7 @@ impl CascadeWorkspace { text: &[S], lp: &L, max_ctx: usize, - dispatch: LcpDispatch, + cmp: Cmp<'_>, ) -> Vec where S: Symbol, @@ -1848,7 +1842,7 @@ impl CascadeWorkspace { &mut dst_sa[dst_off..dst_end], &mut dst_lcp[dst_off..dst_end], max_ctx, - dispatch, + cmp, ); new_lens.push(l1 + l2); src_off = xy_end; diff --git a/src/lib.rs b/src/lib.rs index 9da5c26..792fbaa 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -21,6 +21,7 @@ mod ext_mem; mod lcp; mod limits; mod radix; +mod runs; mod sample_sort; pub use ext_mem::{ diff --git a/src/runs.rs b/src/runs.rs new file mode 100644 index 0000000..5f89e97 --- /dev/null +++ b/src/runs.rs @@ -0,0 +1,449 @@ +//! Periodic-run detection, and a run-aware suffix comparator. +//! +//! The LCP-enhanced merge resolves most steps without touching the text, but +//! when it does have to compare two suffixes it scans their shared prefix one +//! vector at a time. That is fine until the text contains a long *periodic +//! run*, at which point two suffixes inside the run agree for as far as the +//! run continues and a single comparison scans megabytes. +//! +//! Genome assemblies always contain these. An `N` block is the obvious case, +//! and note that in wrapped FASTA it is **not** a run of one symbol: 60 `N`s +//! followed by a newline is a run of period 61. Satellite arrays are runs of +//! period 171. So detecting only single-symbol runs would miss the case that +//! actually shows up. +//! +//! The observation that makes this cheap: if `text[s..e)` has period `q`, and +//! two suffixes start at `a < b` inside it with `(b - a) % q == 0`, then they +//! agree until the later one reaches `e`. That is +//! +//! ```text +//! lcp(a, b) >= e - b +//! ``` +//! +//! known in `O(1)` from the run's bounds, with no scanning at all. The scan +//! resumes at `e`, where the run's guarantee stops. When the phase does not +//! match (`(b - a) % q != 0`) the two suffixes must differ within `q` symbols, +//! so the ordinary scan is already short. +//! +//! Detection is two-stage so that texts without runs pay almost nothing. A +//! sampling pass looks for any periodic window at all and collects the set of +//! periods that actually occur; the full scan then runs only for those +//! periods. On N-free DNA the sample finds nothing and the table is empty, so +//! [`RunTable::skip`] returns immediately on a slice-empty check. +//! +//! This is what the external-memory and sample-sort paths use instead of the +//! prefix doubling in [`crate::radix`]: doubling needs a rank for every +//! position in the text, which would defeat the bounded memory those paths +//! exist to provide, while a run table costs a few dozen entries. + +use crate::lcp::{LcpDispatch, Symbol}; +use crate::limits::LimitProvider; +use rayon::prelude::*; +use std::cmp::Ordering; + +/// Shortest run worth recording. Below this the ordinary SIMD scan crosses +/// the run faster than the binary search that would find it. +const MIN_RUN: usize = 1024; + +/// Longest period considered. Covers homopolymers (1), wrapped-FASTA `N` +/// blocks (61), and alpha-satellite monomers (171 exceeds this, but a +/// satellite array is also periodic at shorter scales in practice). +const MAX_PERIOD: usize = 64; + +/// Window used by the sampling pass to decide whether a period occurs at all. +const SAMPLE_WINDOW: usize = 512; + +/// A maximal stretch `[start, end)` of the text with period `period`, meaning +/// `text[i] == text[i + period]` for every `i` in `start..end - period`. +#[derive(Copy, Clone, Debug, PartialEq, Eq)] +struct Run { + start: usize, + end: usize, + period: usize, +} + +/// Long periodic runs of a byte text, sorted by start and non-overlapping. +/// +/// Empty for texts without long repeats, which is the common case for +/// randomised or `N`-free input, and empty for symbol types wider than a byte. +#[derive(Clone, Debug, Default)] +pub(crate) struct RunTable { + runs: Vec, +} + +impl RunTable { + /// An empty table. Every query short-circuits. + pub(crate) fn empty() -> Self { + Self { runs: Vec::new() } + } + + pub(crate) fn is_empty(&self) -> bool { + self.runs.is_empty() + } + + /// Detect the long periodic runs of `text`. + pub(crate) fn detect(text: &[u8]) -> Self { + let n = text.len(); + if n < MIN_RUN { + return Self::empty(); + } + + // Stage 1: which periods occur anywhere? Sample windows across the + // text and record every period that makes one of them periodic. A + // text with no long repeat contributes nothing and stops here. + let n_samples = 4096.min(n / SAMPLE_WINDOW).max(1); + let stride = (n / n_samples).max(1); + let mut seen = [false; MAX_PERIOD + 1]; + let found: Vec> = (0..n_samples) + .into_par_iter() + .map(|s| { + let base = s * stride; + let end = (base + SAMPLE_WINDOW).min(n); + let mut periods = Vec::new(); + if end - base < MAX_PERIOD * 2 { + return periods; + } + for q in 1..=MAX_PERIOD { + if (base..end - q).all(|i| text[i] == text[i + q]) { + periods.push(q); + // The smallest period implies all its multiples; one + // per window is enough to trigger the full scan. + break; + } + } + periods + }) + .collect(); + for q in found.into_iter().flatten() { + seen[q] = true; + } + let periods: Vec = (1..=MAX_PERIOD).filter(|&q| seen[q]).collect(); + if periods.is_empty() { + return Self::empty(); + } + + // Stage 2: for each period that occurs, find its maximal runs. + let mut runs: Vec = periods + .par_iter() + .flat_map_iter(|&q| { + let mut out = Vec::new(); + let mut i = 0usize; + while i + q < n { + if text[i] != text[i + q] { + i += 1; + continue; + } + let start = i; + while i + q < n && text[i] == text[i + q] { + i += 1; + } + // Matching through `i` means the periodic stretch covers + // `start..i + q`. + let end = i + q; + if end - start >= MIN_RUN { + out.push(Run { + start, + end, + period: q, + }); + } + } + out + }) + .collect(); + + // Keep a non-overlapping set, preferring the earliest start and then + // the longest reach, so a lookup is a single binary search. + runs.sort_unstable_by_key(|r| (r.start, std::cmp::Reverse(r.end))); + let mut merged: Vec = Vec::with_capacity(runs.len()); + for r in runs { + match merged.last() { + Some(last) if r.start < last.end => { + // Overlaps the previous run. Extending the previous run + // would break its period guarantee, so drop this one + // unless it reaches strictly further, in which case keep + // only the part past the previous end. + if r.end > last.end && r.end - last.end >= MIN_RUN { + merged.push(Run { + start: last.end, + end: r.end, + period: r.period, + }); + } + } + _ => merged.push(r), + } + } + Self { runs: merged } + } + + /// The run containing `pos`, if any. + #[inline] + fn at(&self, pos: usize) -> Option<&Run> { + if self.runs.is_empty() { + return None; + } + let i = self.runs.partition_point(|r| r.start <= pos); + let r = self.runs.get(i.checked_sub(1)?)?; + (pos < r.end).then_some(r) + } + + /// Start of the first run beginning at or after `pos`, or `usize::MAX`. + #[inline] + fn next_start(&self, pos: usize) -> usize { + if self.runs.is_empty() { + return usize::MAX; + } + let i = self.runs.partition_point(|r| r.start < pos); + self.runs.get(i).map_or(usize::MAX, |r| r.start) + } + + /// Symbols that suffixes `a` and `b` are guaranteed to share starting at + /// their current offset, derived from run structure alone. + /// + /// Returns `0` when nothing can be concluded, which is always the answer + /// for an empty table. + #[inline] + fn skip(&self, a: usize, b: usize) -> usize { + let (lo, hi) = if a <= b { (a, b) } else { (b, a) }; + let Some(r) = self.at(lo) else { return 0 }; + if hi >= r.end || (hi - lo) % r.period != 0 { + return 0; + } + // Both offsets sit in the same run, an exact whole number of periods + // apart, so they agree until the later one leaves the run. + r.end - hi + } +} + +/// A suffix comparator: the SIMD LCP kernel plus the run table that lets it +/// skip long periodic repeats instead of scanning them. +/// +/// Threaded through the merge kernel in place of a bare [`LcpDispatch`]. It is +/// `Copy`, so it still travels through the recursion in registers. +#[derive(Copy, Clone)] +pub(crate) struct Cmp<'a> { + pub(crate) dispatch: LcpDispatch, + pub(crate) runs: &'a RunTable, +} + +impl<'a> Cmp<'a> { + pub(crate) fn new(dispatch: LcpDispatch, runs: &'a RunTable) -> Self { + Self { dispatch, runs } + } + + /// LCP of `text[p..]` and `text[q..]` in symbols, bounded by `max_ctx`, + /// using run structure to jump over long periodic stretches. + /// + /// With an empty run table this is exactly [`LcpDispatch::lcp`] plus one + /// predictable branch. + #[inline] + pub(crate) fn lcp(&self, text: &[S], p: usize, q: usize, max_ctx: usize) -> usize { + if self.runs.is_empty() || size_of::() != 1 { + return self.dispatch.lcp(text, p, q, max_ctx); + } + let mut i = 0usize; + while i < max_ctx { + let jump = self.runs.skip(p + i, q + i); + if jump > 0 { + i = (i + jump).min(max_ctx); + continue; + } + // Stop the scan where a run begins, so it never traverses one. + // Scanning into a run is exactly the megabyte-long case. + let next = self + .runs + .next_start(p + i) + .saturating_sub(p + i) + .min(self.runs.next_start(q + i).saturating_sub(q + i)) + .max(1); + let window = max_ctx - i; + let bounded = next.min(window); + let got = self.dispatch.lcp(text, p + i, q + i, bounded); + i += got; + if got < bounded { + // A real mismatch, not a window boundary. + break; + } + } + i.min(max_ctx) + } + + /// Total order on two suffixes, mirroring [`LcpDispatch::suffix_cmp_with`] + /// but going through the run-aware [`Self::lcp`]. + #[inline] + pub(crate) fn suffix_cmp_with( + &self, + text: &[S], + lp: &L, + p: usize, + q: usize, + max_ctx: usize, + ) -> Ordering { + let lim_p = lp.lim_at(p); + let lim_q = lp.lim_at(q); + let lim = lim_p.min(lim_q).min(max_ctx); + let common = self.lcp(text, p, q, lim); + if common < lim { + text[p + common].cmp(&text[q + common]) + } else { + lp.boundary_order(p, lim_p, q, lim_q) + } + } +} + +/// Build a run table for `text` when the symbol type is a byte, otherwise an +/// empty one. +/// +/// Detection is a sequential-read pass and only runs at all if the sampling +/// stage finds a periodic window, so texts without long repeats pay a single +/// sampling sweep. +pub(crate) fn detect_for(text: &[S]) -> RunTable { + if size_of::() != 1 { + return RunTable::empty(); + } + // SAFETY: `S` is one byte wide with no padding and no invalid bit + // patterns (the `Symbol` contract), so a byte view over the same memory is + // valid for reads of the same length. + let bytes: &[u8] = + unsafe { std::slice::from_raw_parts(text.as_ptr() as *const u8, text.len()) }; + RunTable::detect(bytes) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn naive_lcp(text: &[u8], a: usize, b: usize, max_ctx: usize) -> usize { + let lim = (text.len() - a).min(text.len() - b).min(max_ctx); + (0..lim).take_while(|&i| text[a + i] == text[b + i]).count() + } + + /// The run-aware LCP must agree with a byte-at-a-time scan for every pair + /// of positions, whether or not a run is involved. + fn assert_lcp_agrees(text: &[u8]) { + let runs = RunTable::detect(text); + let cmp = Cmp::new(LcpDispatch::detect(), &runs); + let n = text.len(); + let step = (n / 64).max(1); + for a in (0..n).step_by(step) { + for b in (0..n).step_by(step) { + let want = naive_lcp(text, a, b, usize::MAX); + let got = cmp.lcp(text, a, b, usize::MAX); + assert_eq!(got, want, "lcp({a}, {b}) on len-{n} text"); + } + } + } + + #[test] + fn empty_table_for_texts_without_runs() { + use rand::{RngExt, SeedableRng}; + let mut rng = rand::rngs::StdRng::seed_from_u64(0x0DD1); + let text: Vec = (0..50_000).map(|_| rng.random_range(0..4u8)).collect(); + assert!(RunTable::detect(&text).is_empty()); + } + + #[test] + fn detects_homopolymer() { + let mut text: Vec = vec![1, 2, 3]; + text.extend(std::iter::repeat_n(0u8, 5000)); + text.extend([1, 2, 3]); + let table = RunTable::detect(&text); + assert!(!table.is_empty()); + assert!(table.runs.iter().any(|r| r.end - r.start >= 5000)); + } + + /// Wrapped FASTA: 60 `N`s then a newline. The single-symbol runs are only + /// 60 long, so a homopolymer-only detector would find nothing; the real + /// structure is period 61. + #[test] + fn detects_wrapped_fasta_n_block() { + let mut text: Vec = b"ACGT".to_vec(); + for _ in 0..200 { + text.extend(std::iter::repeat_n(b'N', 60)); + text.push(b'\n'); + } + text.extend(b"ACGT"); + let table = RunTable::detect(&text); + assert!(!table.is_empty(), "period-61 N block should be detected"); + assert!(table.runs.iter().any(|r| r.period == 61 || r.period == 1)); + } + + #[test] + fn runs_are_sorted_and_disjoint() { + let mut text: Vec = Vec::new(); + text.extend(std::iter::repeat_n(0u8, 3000)); + text.extend(b"ACGTACGT"); + text.extend((0..3000).map(|i| (i % 7) as u8)); + text.extend(b"TTTT"); + let table = RunTable::detect(&text); + for w in table.runs.windows(2) { + assert!(w[0].end <= w[1].start, "runs overlap: {:?}", w); + assert!(w[0].start < w[1].start); + } + for r in &table.runs { + for i in r.start..r.end - r.period { + assert_eq!(text[i], text[i + r.period], "period claim is wrong"); + } + } + } + + #[test] + fn lcp_agrees_on_homopolymer() { + let mut text: Vec = b"ACGT".to_vec(); + text.extend(std::iter::repeat_n(0u8, 4000)); + text.extend(b"ACGT"); + assert_lcp_agrees(&text); + } + + #[test] + fn lcp_agrees_on_wrapped_fasta() { + let mut text: Vec = b"ACGTAC".to_vec(); + for _ in 0..120 { + text.extend(std::iter::repeat_n(b'N', 60)); + text.push(b'\n'); + } + text.extend(b"GTGTGT"); + assert_lcp_agrees(&text); + } + + #[test] + fn lcp_agrees_on_multi_period_text() { + let mut text: Vec = Vec::new(); + text.extend((0..3000).map(|i| (i % 3) as u8)); + text.extend(b"XYZ"); + text.extend(std::iter::repeat_n(9u8, 2500)); + text.extend((0..2000).map(|i| (i % 5) as u8)); + assert_lcp_agrees(&text); + } + + #[test] + fn lcp_respects_max_ctx_inside_a_run() { + let mut text: Vec = b"AC".to_vec(); + text.extend(std::iter::repeat_n(0u8, 4000)); + let runs = RunTable::detect(&text); + let cmp = Cmp::new(LcpDispatch::detect(), &runs); + for &ctx in &[0usize, 1, 7, 100, 3000] { + assert_eq!(cmp.lcp(&text, 2, 3, ctx), naive_lcp(&text, 2, 3, ctx)); + } + } + + #[test] + fn suffix_cmp_matches_slice_order() { + use crate::limits::PlainText; + let mut text: Vec = b"ACGT".to_vec(); + text.extend(std::iter::repeat_n(5u8, 3000)); + text.extend(b"ACGT"); + let runs = RunTable::detect(&text); + let cmp = Cmp::new(LcpDispatch::detect(), &runs); + let lp = PlainText::new(text.len()); + let step = (text.len() / 40).max(1); + for a in (0..text.len()).step_by(step) { + for b in (0..text.len()).step_by(step) { + let want = text[a..].cmp(&text[b..]); + let got = cmp.suffix_cmp_with(&text, &lp, a, b, usize::MAX); + assert_eq!(got, want, "suffix_cmp({a}, {b})"); + } + } + } +} diff --git a/src/sample_sort.rs b/src/sample_sort.rs index e478a72..177e5be 100644 --- a/src/sample_sort.rs +++ b/src/sample_sort.rs @@ -33,6 +33,7 @@ use crate::Index; use crate::lcp::{LcpDispatch, Symbol}; use crate::limits::{LimitProvider, PlainText}; +use crate::runs::Cmp; use rayon::join; use rayon::prelude::*; @@ -269,7 +270,8 @@ where // Choose the LCP implementation once for the whole build; the captured // function pointer travels through the recursion in a register, so the // inner merge loop pays no atomic load or feature-detection branch. - let dispatch = LcpDispatch::detect(); + let runs = crate::runs::detect_for(text); + let cmp = Cmp::new(LcpDispatch::detect(), &runs); merge_sort( text, @@ -279,7 +281,7 @@ where &mut lcp_arr, &mut lcp_w, opts.max_context, - dispatch, + cmp, ); sa @@ -297,7 +299,7 @@ where /// /// Visible to the rest of the crate so the external-memory path can sort /// individual subarrays of positions using the same kernel. -#[allow(clippy::too_many_arguments)] // 4 buffers + text + lp + ctx + dispatch +#[allow(clippy::too_many_arguments)] // 4 buffers + text + lp + ctx + cmp pub(crate) fn merge_sort( text: &[S], lp: &L, @@ -306,7 +308,7 @@ pub(crate) fn merge_sort( lcp_arr: &mut [I], lcp_w: &mut [I], max_ctx: usize, - dispatch: LcpDispatch, + cmp: Cmp<'_>, ) where S: Symbol, I: Index, @@ -331,15 +333,15 @@ pub(crate) fn merge_sort( let (lcp_w_l, lcp_w_r) = lcp_w.split_at_mut(mid); join( - || merge_sort(text, lp, sa_l, sa_w_l, lcp_l, lcp_w_l, max_ctx, dispatch), - || merge_sort(text, lp, sa_r, sa_w_r, lcp_r, lcp_w_r, max_ctx, dispatch), + || merge_sort(text, lp, sa_l, sa_w_l, lcp_l, lcp_w_l, max_ctx, cmp), + || merge_sort(text, lp, sa_r, sa_w_r, lcp_r, lcp_w_r, max_ctx, cmp), ); // Merge the two sorted halves (still living in `sa`) into the workspace, // then copy the workspace back into the destination so the caller's // postcondition holds on `sa` / `lcp_arr`. merge( - text, lp, sa_l, sa_r, lcp_l, lcp_r, sa_w, lcp_w, max_ctx, dispatch, + text, lp, sa_l, sa_r, lcp_l, lcp_r, sa_w, lcp_w, max_ctx, cmp, ); sa.copy_from_slice(sa_w); lcp_arr.copy_from_slice(lcp_w); @@ -353,7 +355,7 @@ pub(crate) fn merge_sort( /// /// Visible to the rest of the crate so the external-memory path can cascade /// 2-way merges across each partition's sub-subarrays during Phase 4. -#[allow(clippy::too_many_arguments)] // CaPS-SA's merge takes 5 buffers + text + lp + ctx + dispatch +#[allow(clippy::too_many_arguments)] // CaPS-SA's merge takes 5 buffers + text + lp + ctx + cmp pub(crate) fn merge( text: &[S], lp: &L, @@ -364,7 +366,7 @@ pub(crate) fn merge( z: &mut [I], lcp_z: &mut [I], max_ctx: usize, - dispatch: LcpDispatch, + cmp: Cmp<'_>, ) where S: Symbol, I: Index, @@ -444,7 +446,7 @@ pub(crate) fn merge( // intersection — no extra work. let cap = lim_a.min(lim_b).min(max_ctx); let remaining_ctx = cap.saturating_sub(m); - let ext = dispatch.lcp(text, p_a + m, p_b + m, remaining_ctx); + let ext = cmp.lcp(text, p_a + m, p_b + m, remaining_ctx); let total = m + ext; let a_smaller = if total < lim_a && total < lim_b { text[p_a + total] < text[p_b + total] @@ -554,7 +556,7 @@ mod tests { &mut lcp_arr, &mut lcp_w, max_ctx, - LcpDispatch::detect(), + Cmp::new(LcpDispatch::detect(), &crate::runs::RunTable::empty()), ); (sa, lcp_arr) } From f07943b057034d7a15f2ccfba73851381d6f2dcb Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 21:59:35 +0200 Subject: [PATCH 13/41] docs: record the external-memory run-skipping results Adds the ext-mem phase breakdown that located the cost (94% of the chr21 FASTA run in phase 1), the before/after table, and the detail worth keeping: a homopolymer detector would not have worked, because in 60-column wrapped FASTA the longest single-byte run is 60. Each line of `N`s ends in a newline, so the real structure is a period-61 repeat spanning 6.6 Mb. Co-Authored-By: Claude Opus 5 (1M context) --- README.md | 43 +++++++++++++++++++++++++++++++++++++++---- bench/README.md | 35 +++++++++++++++++++++++++++++++++++ 2 files changed, 74 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 10e988f..688c7c6 100644 --- a/README.md +++ b/README.md @@ -193,10 +193,45 @@ build: where doubling applies it is strictly better: same output, no bucket machinery, none of the scan cost. -`build_ext_mem` deliberately stays on the merge kernel. Its purpose is to -bound peak memory, and routing it through an in-memory algorithm would -defeat exactly that, so it keeps the scan cost on repeat-heavy input. -Segmented texts and symbols wider than `u8` also stay on the merge kernel. +`build_ext_mem` deliberately stays on the merge kernel: its purpose is to +bound peak memory, and prefix doubling needs a rank for every position in +the text, which would defeat exactly that. Segmented texts and symbols +wider than `u8` also stay on the merge kernel. + +### Skipping long repeats + +Those paths get the same pathology fixed in the comparator instead, which +costs no extra memory. If `text[s..e)` has period `q` and two suffixes +start at `a < b` inside it with `(b - a) % q == 0`, they agree until the +later one reaches `e`, so + +```text +lcp(a, b) >= e - b +``` + +is known in `O(1)` from the run's bounds with nothing scanned. When the +phase does not match, the two suffixes differ within `q` symbols and the +ordinary scan is already short. Scans are additionally bounded so they +stop at a run's start rather than traversing it. + +Detecting only single-symbol runs would miss the case that actually +occurs: in wrapped FASTA an `N` block is 60 `N`s followed by a newline, +which is period 61, not period 1. Periods up to 64 are considered, which +also covers satellite arrays. + +Detection is two-stage so texts without repeats pay almost nothing: a +sampling pass collects the periods that occur at all, and the full scan +runs only for those. On `N`-free DNA the table comes out empty and every +query short-circuits. The table is a few dozen entries, so the +external-memory path keeps its memory bound. + +| ext-mem input | before | after | CPU before | CPU after | peak RSS | +| ------------- | ------ | ----- | ---------- | --------- | -------- | +| chr21 FASTA, 47.5 MB | 24.2 s | **2.48 s** | 268 s | 23.3 s | 147 → 151 MB | +| chr21 `N`-free, 80 MB | 3.49 s | 3.55 s | 33.9 s | 33.8 s | 214 → 220 MB | + +Phase 1 alone goes from 22.85 s to 0.95 s on the FASTA input. The `N`-free +row is unchanged, as expected: there are no runs to find. ## Algorithm diff --git a/bench/README.md b/bench/README.md index ff9e06f..428dd88 100644 --- a/bench/README.md +++ b/bench/README.md @@ -462,6 +462,41 @@ rank-only, so they cost a sort over a shrinking residual rather than a text scan, which is why the pathology disappears rather than merely shrinking. +### chr21 external memory — skipping long repeats + +Same machine and inputs as the section above. The external-memory path +stays on the merge kernel by design (prefix doubling needs a rank per +text position, which would defeat the memory bound it exists to +provide), so it was still paying the full scan cost. Profiling put 94% +of the FASTA run in phase 1: + +``` + phase1 phase2 phase3 phase4 total +chr21.0123 before 1.15 s 0.005 s 0.215 s 2.080 s 3.49 s +chr21.fa before 22.85 s 0.058 s 0.317 s 0.948 s 24.19 s +``` + +Fixing it in the comparator rather than the algorithm keeps the memory +bound intact: + +| ext-mem input | before | after | CPU before | CPU after | peak RSS | +| ------------- | ------ | ----- | ---------- | --------- | -------- | +| `chr21.fa`, 47.5 MB | 24.19 s | **2.48 s** | 268 s | 23.3 s | 147 → 151 MB | +| `chr21.0123`, 80 MB | 3.49 s | 3.55 s | 33.9 s | 33.8 s | 214 → 220 MB | + +Phase 1 goes from 22.85 s to 0.95 s. The `N`-free row is flat, which is +the expected result: the sampling stage finds no periodic window, the +run table comes out empty, and every query short-circuits. + +Output was verified identical to the in-memory suffix array on both +inputs (127.7 M entries). + +The detail worth recording is that **a homopolymer detector would not +have worked**. In 60-column wrapped FASTA the longest run of a single +byte is 60, because each line of `N`s is terminated by a newline. The +real structure is a period-61 repeat spanning 6.6 Mb. Periods up to 64 +are considered, which also covers satellite arrays. + ### Reading the 97.54% LCP profile correctly The profile in the next section shows `lcp_u8_avx2` taking 97.54% of From d09bff78bea9a9cdb93bf08c9e01fe34b7ace954 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 22:04:57 +0200 Subject: [PATCH 14/41] perf: prefetch the next candidates' text in the merge loop Phase 4 had become the largest ext-mem cost, and its merge CPU (19.2 s on 80 MB of DNA) exceeded the entire CPU of the in-memory path. The merge is latency-bound: the tied branch dereferences the text at two random addresses, and the address for step i+1 is not known until step i retires, so there is no memory-level parallelism and the hardware prefetcher cannot see the pattern. The candidate positions themselves live in the two index arrays, which are sequential and already in cache, so the addresses several steps ahead are known even though the dependent loads are not. Issue them as prefetches, offset by the current boundary LCP `m`, which estimates where the next scans start. This is not the prefetch recorded as a negative result in `lcp.rs`. That one sat inside the strided scan loop, which the hardware prefetcher already covers. This one targets the random access, which it cannot. Apple M4 Max, 12 threads, ext-mem, output verified identical to the in-memory suffix array on both inputs: chr21.0123, 80 MB phase4 merge CPU 19.20 s -> 12.42 s phase4 wall 2.17 s -> 1.46 s total 3.55 s -> 2.90 s CPU 33.8 -> 27.0 s chr21.fa, 47.5 MB phase4 merge CPU 7.99 s -> 6.76 s total 2.56 s -> 1.99 s CPU 23.3 -> 18.0 s Co-Authored-By: Claude Opus 5 (1M context) --- src/sample_sort.rs | 50 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/src/sample_sort.rs b/src/sample_sort.rs index 177e5be..b50101e 100644 --- a/src/sample_sort.rs +++ b/src/sample_sort.rs @@ -37,6 +37,36 @@ use crate::runs::Cmp; use rayon::join; use rayon::prelude::*; +/// How many merge steps ahead the text prefetch runs. Large enough to cover a +/// DRAM round trip at the merge's step rate, small enough that the prefetched +/// line is still resident when the step that needs it arrives. +const PREFETCH_DISTANCE: usize = 8; + +/// Hint the CPU to start pulling `text[at]` into cache. +/// +/// A no-op on targets without a stable prefetch intrinsic, and harmless when +/// `at` is out of bounds: the address is never dereferenced, only used as a +/// prefetch operand, and prefetch instructions on both supported targets +/// ignore faulting addresses. +#[inline(always)] +fn prefetch_symbol(text: &[S], at: usize) { + let _ = (text, at); + #[cfg(target_arch = "x86_64")] + unsafe { + std::arch::x86_64::_mm_prefetch( + text.as_ptr().add(at.min(text.len())) as *const i8, + std::arch::x86_64::_MM_HINT_T0, + ); + } + #[cfg(target_arch = "aarch64")] + unsafe { + // `core::arch::aarch64::_prefetch` is still unstable, so emit the + // instruction directly. `prfm` never faults. + let p = text.as_ptr().add(at.min(text.len())); + std::arch::asm!("prfm pldl1keep, [{p}]", p = in(reg) p, options(nostack, readonly, preserves_flags)); + } +} + /// Tunable options for SA construction. #[derive(Clone, Debug)] pub struct Opts { @@ -408,6 +438,26 @@ pub(crate) fn merge( let mut lim_b_cache: Option<(usize, usize)> = None; while i_a < len_a && i_b < len_b { + // The tied branch below dereferences the text at two *random* + // addresses, and the address for step `i + 1` is not known until step + // `i` retires, so there is no memory-level parallelism to exploit and + // the hardware prefetcher cannot see the pattern either. But the + // candidate positions themselves live in `arr_a` / `arr_b`, which are + // sequential and already in cache, so the addresses a few steps ahead + // *are* known. Issue them now. + // + // This is not the prefetch that was tried and reverted in `lcp.rs`: + // that one sat inside the strided scan loop, which the hardware + // prefetcher already covers. Here the access is random, which is the + // case hardware cannot predict. `m` is the current boundary LCP and a + // good estimate of where the next scans will start. + if i_a + PREFETCH_DISTANCE < len_a { + prefetch_symbol(text, arr_a[i_a + PREFETCH_DISTANCE].to_usize() + m); + } + if i_b + PREFETCH_DISTANCE < len_b { + prefetch_symbol(text, arr_b[i_b + PREFETCH_DISTANCE].to_usize() + m); + } + let l_a = lcp_a[i_a].to_usize(); // (output_a, lcp_for_output, new_m) From 677643989f209e8c003a87b2557c6da9d2398d85 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 22:08:51 +0200 Subject: [PATCH 15/41] perf: raise the auto-picked subarray target from 64Ki to 128Ki records Sweeping `p` directly showed the previous default sitting on the wrong side of a flat region. 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 differ: phase 3 shrinks quadratically in `p`, and phase 4's cascade does a full pass over its partition per level. Peak RSS is set by phase 4 holding `4 x 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. Measured on chr21 forward ++ revcomp (80 MB), 12 threads: 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 <- 131072 1224 3.05 s 205 MB <- 65536 (previous default) 128Ki is the largest step that costs nothing in memory: same peak RSS, ~8% less wall. Going further trades real memory for speed, which is the opposite of what this path is for, so it stays available through `ExtMemOpts::subproblem_count` rather than becoming the default. At genome scale this changes nothing. `PHASE1_MAX_PARTITIONS` already binds for any n above ~1 GB, so GRCh38 still gets p = 8192. Output verified identical to the in-memory suffix array on both inputs. chr21.0123, 80 MB 3.01 s -> 2.77 s chr21.fa, 47.5 MB 2.06 s -> 1.86 s Co-Authored-By: Claude Opus 5 (1M context) --- src/ext_mem.rs | 34 +++++++++++++++++++++++++++++++--- 1 file changed, 31 insertions(+), 3 deletions(-) diff --git a/src/ext_mem.rs b/src/ext_mem.rs index 28352f8..c75eb18 100644 --- a/src/ext_mem.rs +++ b/src/ext_mem.rs @@ -1090,9 +1090,37 @@ 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 From e9a1c933c6107015826110967939251bf42b4cf1 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 22:12:27 +0200 Subject: [PATCH 16/41] perf: seed each phase-1 subarray with the packed key Phase 1 had become the largest external-memory phase (1.20 s of 2.54 s on 80 MB of DNA) and was already parallelising at ~97%, so the way forward was less work rather than more threads. It was sorting every subarray from singletons, which is the case the merge kernel handles worst: each leaf merge starts at `m = 0` and orders two suffixes by scanning the text at two random addresses. Sorting by the packed key first resolves the leading `k` symbols with no text access at all (32 symbols for DNA at 2 bits each), and yields the LCP between adjacent runs for free from `(key_a ^ key_b).leading_zeros()`. Only suffixes agreeing through all `k` symbols reach the merge kernel, on the short slice they occupy. This is the bounded-memory counterpart to the prefix doubling in `build_in_memory`. Doubling itself is not available here: it needs a rank for every position in the text, which is exactly the memory this path refuses to spend. A fixed-depth key needs none. The cross-run LCP is capped by both suffixes' lengths. Padding can agree with a real `0` symbol past the end of the shorter suffix, so the raw `leading_zeros` count can overstate it, and a wrong LCP would silently corrupt the order at the next merge level. Gated on the same conditions as the other fast paths (`u8` symbols, plain lexicographic comparator, unbounded `max_context`) and falls back to `merge_sort` otherwise. The alphabet scan that picks the field width runs once per build, not once per subarray. Apple M4 Max, 12 threads, output verified identical to the in-memory suffix array on both inputs: chr21.0123, 80 MB phase1 1.201 s -> 0.316 s total 2.54 s -> 2.11 s chr21.fa, 47.5 MB phase1 0.898 s -> 0.415 s total 1.86 s -> 1.65 s Peak RSS 205 MB -> 233 MB: the key vector is 16 bytes per record over one subarray per worker. Co-Authored-By: Claude Opus 5 (1M context) --- src/ext_mem.rs | 49 +++++++++++++++--- src/radix.rs | 132 +++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 174 insertions(+), 7 deletions(-) diff --git a/src/ext_mem.rs b/src/ext_mem.rs index c75eb18..cc74793 100644 --- a/src/ext_mem.rs +++ b/src/ext_mem.rs @@ -505,6 +505,7 @@ where let p = effective_subproblem_count(n, opts.subproblem_count); let runs = crate::runs::detect_for(text); let cmp = Cmp::new(LcpDispatch::detect(), &runs); + let seed_params = crate::radix::seed_params(text); let work_dir = opts.work_dir.clone(); // Pool the `2 × p` bucket files into one anonymous tempfile per @@ -529,8 +530,16 @@ where let part_factory = |j: usize| phase3_pool.new_bucket::>(j); let t = Instant::now(); - let (mut subarray_buckets, samples) = - phase1_sort_sample_spill::(text, lp, &source, p, opts, cmp, sub_factory)?; + let (mut subarray_buckets, samples) = phase1_sort_sample_spill::( + text, + lp, + &source, + p, + opts, + cmp, + seed_params, + sub_factory, + )?; profile_log(&format!( "phase1 (sort+sample+spill) {:.3}s", t.elapsed().as_secs_f64() @@ -619,11 +628,20 @@ where let p = effective_subproblem_count(n, opts.subproblem_count); let runs = crate::runs::detect_for(text); let cmp = Cmp::new(LcpDispatch::detect(), &runs); + let seed_params = crate::radix::seed_params(text); let factory = |_i: usize| InMemBucket::>::new(); - let (mut subarray_buckets, samples) = - phase1_sort_sample_spill::(text, lp, &source, p, opts, cmp, factory)?; + let (mut subarray_buckets, samples) = phase1_sort_sample_spill::( + text, + lp, + &source, + p, + opts, + cmp, + seed_params, + factory, + )?; // Same rationale as in `build_ext_mem_inner` — drop the source // as soon as phase 1's `fill_chunk` calls have stopped. drop(source); @@ -1188,6 +1206,7 @@ fn phase1_sort_sample_spill( p: usize, opts: &ExtMemOpts, cmp: Cmp<'_>, + seed_params: Option<(u32, usize)>, mk_bucket: MkB, ) -> io::Result<(Vec, Vec)> where @@ -1220,16 +1239,32 @@ where let mut sa_w = vec![I::zero(); len]; let mut lcp_arr = vec![I::zero(); len]; let mut lcp_w = vec![I::zero(); len]; - sample_sort::merge_sort( + // Seed with the packed key where the comparator allows it: that + // resolves the first `k` symbols with no text access and yields + // the LCP between runs from the key difference, leaving the merge + // kernel only the suffixes that agree through all `k`. + if !crate::radix::seed_subarray( text, lp, + seed_params, &mut sa, - &mut sa_w, &mut lcp_arr, + &mut sa_w, &mut lcp_w, opts.max_context, cmp, - ); + ) { + sample_sort::merge_sort( + text, + lp, + &mut sa, + &mut sa_w, + &mut lcp_arr, + &mut lcp_w, + opts.max_context, + cmp, + ); + } // Pull `samples_per_subarray` evenly-spaced positions out of // the now-sorted subarray. Deterministic — no RNG needed for diff --git a/src/radix.rs b/src/radix.rs index 7eab7db..0f63fcb 100644 --- a/src/radix.rs +++ b/src/radix.rs @@ -56,6 +56,10 @@ use crate::Index; use crate::ext_mem::profile_log; +use crate::lcp::Symbol; +use crate::limits::LimitProvider; +use crate::runs::Cmp; +use crate::sample_sort; use rayon::prelude::*; use std::time::Instant; @@ -94,6 +98,134 @@ fn key_at(text: &[u8], p: usize, bits: u32, k: usize) -> u64 { key << (bits as usize * (k - (end - p))) } +/// Field width and symbols-per-key for `text`, or `None` when a packed key +/// cannot represent this text's order. +/// +/// Computed once per build and handed to [`seed_subarray`], which would +/// otherwise re-scan the whole text for every subarray. +pub(crate) fn seed_params(text: &[S]) -> Option<(u32, usize)> { + if size_of::() != 1 { + return None; + } + // SAFETY: `S` is one byte wide, so a byte view over the same memory is + // valid for reads of the same length. + let bytes: &[u8] = + unsafe { std::slice::from_raw_parts(text.as_ptr() as *const u8, text.len()) }; + Some(pack_params(bytes.par_iter().copied().max().unwrap_or(0))) +} + +/// Sort `sa` into suffix order and fill `lcp`, using a packed fixed-depth key +/// so that most of the ordering costs no text access at all. +/// +/// This is the external-memory and sample-sort counterpart to [`build_sa`]. +/// Those paths cannot use prefix doubling, which needs a rank for every +/// position in the text and would break the memory bound they exist to +/// provide. But they can still avoid the part of the merge kernel that hurts +/// most: sorting a subarray from singletons, where every leaf merge starts at +/// `m = 0` and compares two suffixes by scanning the text. +/// +/// Sorting by the packed key resolves the first `k` symbols with no text +/// access (32 symbols for DNA), and hands back the LCP for adjacent entries +/// for free from `(key_a ^ key_b).leading_zeros()`. Only suffixes that agree +/// through all `k` symbols reach the merge kernel, on the small slice they +/// occupy. +/// +/// Returns `false` without touching anything when the comparator is not plain +/// lexicographic, so the caller falls back to a plain `merge_sort`. +/// +/// `sa_w` and `lcp_w` are the caller's existing merge scratch buffers. +#[allow(clippy::too_many_arguments)] +pub(crate) fn seed_subarray( + text: &[S], + lp: &L, + params: Option<(u32, usize)>, + sa: &mut [I], + lcp: &mut [I], + sa_w: &mut [I], + lcp_w: &mut [I], + max_ctx: usize, + cmp: Cmp<'_>, +) -> bool { + let Some((bits, k)) = params else { + return false; + }; + if max_ctx != usize::MAX || lp.plain_lex_len() != Some(text.len()) { + return false; + } + let len = sa.len(); + if len < 2 { + if len == 1 { + lcp[0] = I::zero(); + } + return true; + } + // SAFETY: `params` is `Some` only when `S` is one byte wide. + let bytes: &[u8] = + unsafe { std::slice::from_raw_parts(text.as_ptr() as *const u8, text.len()) }; + + // Order by (key, visible length), the same comparator `build_sa` seeds + // with and for the same reason: zero padding makes a short suffix share a + // key with any suffix continuing in zeros, and `0` is a real symbol. + let n = bytes.len(); + let visible = |p: usize| -> usize { (n - p).min(k) }; + let mut keyed: Vec<(u64, u32, I)> = sa + .iter() + .map(|&p| { + let p = p.to_usize(); + (key_at(bytes, p, bits, k), visible(p) as u32, p_as(p)) + }) + .collect(); + keyed.sort_unstable(); + for (slot, e) in sa.iter_mut().zip(keyed.iter()) { + *slot = I::from_usize(e.2.to_usize()); + } + + // Walk equal-key runs. Between runs the LCP falls straight out of the key + // difference; inside one it needs the merge kernel. + let mut i = 0usize; + while i < len { + let mut j = i + 1; + while j < len && (keyed[j].0, keyed[j].1) == (keyed[i].0, keyed[i].1) { + j += 1; + } + if j - i > 1 { + sample_sort::merge_sort( + text, + lp, + &mut sa[i..j], + &mut sa_w[i..j], + &mut lcp[i..j], + &mut lcp_w[i..j], + max_ctx, + cmp, + ); + } else { + lcp[i] = I::zero(); + } + // Boundary entry: LCP against the last element of the previous run. + if i > 0 { + let a = sa[i - 1].to_usize(); + let b = sa[i].to_usize(); + let xor = keyed[i - 1].0 ^ keyed[i].0; + debug_assert_ne!(xor, 0, "distinct runs must differ in the key"); + // `leading_zeros / bits` counts whole matching fields. Cap by both + // suffixes' lengths: padding can agree with a real `0` symbol past + // the end of the shorter one. + let shared = (xor.leading_zeros() / bits) as usize; + lcp[i] = I::from_usize(shared.min(lp.lim_at(a)).min(lp.lim_at(b))); + } + i = j; + } + lcp[0] = I::zero(); + true +} + +/// Round-trip a position through the index type used in the seed vector. +#[inline] +fn p_as(p: usize) -> I { + I::from_usize(p) +} + /// Bits of the key used for the MSD counting-sort pass, and the resulting /// bucket count. /// From 4fc99de53469f8eccb296ad4bf84fcc13406d2be Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 22:12:41 +0200 Subject: [PATCH 17/41] docs: record the full external-memory speedup ladder Co-Authored-By: Claude Opus 5 (1M context) --- README.md | 21 ++++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 688c7c6..f95c88e 100644 --- a/README.md +++ b/README.md @@ -227,11 +227,22 @@ external-memory path keeps its memory bound. | ext-mem input | before | after | CPU before | CPU after | peak RSS | | ------------- | ------ | ----- | ---------- | --------- | -------- | -| chr21 FASTA, 47.5 MB | 24.2 s | **2.48 s** | 268 s | 23.3 s | 147 → 151 MB | -| chr21 `N`-free, 80 MB | 3.49 s | 3.55 s | 33.9 s | 33.8 s | 214 → 220 MB | - -Phase 1 alone goes from 22.85 s to 0.95 s on the FASTA input. The `N`-free -row is unchanged, as expected: there are no runs to find. +| chr21 FASTA, 47.5 MB | 24.2 s | **1.65 s** | 268 s | 17.5 s | 147 → 182 MB | +| chr21 `N`-free, 80 MB | 3.49 s | **2.11 s** | 33.9 s | 25.0 s | 214 → 233 MB | + +Four changes get there, each measured separately: + +| change | chr21.0123 | chr21 FASTA | +| ------ | ---------- | ----------- | +| baseline | 3.49 s | 24.19 s | +| skip periodic runs | 3.55 s | 2.48 s | +| prefetch the next candidates' text | 2.90 s | 1.99 s | +| subarray target 64Ki → 128Ki records | 2.54 s | 1.86 s | +| seed phase-1 subarrays with the packed key | **2.11 s** | **1.65 s** | + +Phase 1 goes from 22.85 s to 0.42 s on the FASTA input, and from 1.20 s +to 0.32 s on the `N`-free one. Peak RSS rises by under 30 MB, so the +bounded-memory guarantee the path exists for is intact. ## Algorithm From 11d31f1a8c450c892fd54a745f53751dd6825b11 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 22:23:32 +0200 Subject: [PATCH 18/41] perf: merge each cascade level's run pairs in parallel `CascadeWorkspace::merge_one_level` walked its run pairs sequentially, even though the pairs at a level are independent and write to disjoint destination ranges. The only thing forcing the order was the running `src_off` / `dst_off`, and both are prefix sums, so they can be computed up front and each pair handed its own sub-slices. This is what capped phase 4's parallel efficiency. Partition-level parallelism (`4 x threads` at once) hides it while there are many partitions in flight, but each partition's cascade ends in a single 2-way merge over the whole partition, and those tails serialise. Apple M4 Max, 12 threads, output verified identical to the in-memory suffix array on both inputs: chr21.0123, 80 MB 2.11 s -> 2.02 s chr21.fa, 47.5 MB 1.65 s -> 1.55 s A grain-size threshold was tried, on the theory that the small early levels would not pay for their rayon tasks. It measured worse on both inputs (2.22 s and 1.63 s), so the pairs are merged in parallel at every level. Co-Authored-By: Claude Opus 5 (1M context) --- src/ext_mem.rs | 95 +++++++++++++++++++++++++++++++------------------- 1 file changed, 59 insertions(+), 36 deletions(-) diff --git a/src/ext_mem.rs b/src/ext_mem.rs index cc74793..4fe12cd 100644 --- a/src/ext_mem.rs +++ b/src/ext_mem.rs @@ -1139,6 +1139,7 @@ impl<'a> PositionSource<'a> { /// 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 @@ -1884,44 +1885,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, - cmp, - ); - 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, + cmp, + ); + } 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 } } From b6300bfb3674e8d81c4e4be78a1ea67e0b33c57c Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 22:25:42 +0200 Subject: [PATCH 19/41] perf: pipeline phase 4's emit against the next chunk's merge Emitting is Theta(n) and single-threaded by construction: the caller's closure is `FnMut` and its ordering is the point. But it also did not overlap anything, so once per chunk every worker sat idle while the main thread drained the merged results. A scoped producer now merges chunk c+1, itself rayon-parallel, while the main thread emits chunk c. The channel bound of one keeps at most two chunks resident, so the transient cost is one extra chunk of merged positions rather than the unbounded queue an unsynchronised producer would build. This is not the existing `ordered_phase4_emit` path, which coordinates at *partition* granularity through an mpsc channel and a `BTreeMap` and measured slower than plain collect-then-emit. Here the producer hands over whole chunks that are already in order, so the consumer only drains them and no reordering structure is needed. That path is left untouched behind its opt-in flag. Apple M4 Max, 12 threads, output verified identical to the in-memory suffix array on both inputs: chr21.0123, 80 MB 2.02 s -> 1.82 s peak RSS 233 -> 240 MB chr21.fa, 47.5 MB 1.55 s -> 1.45 s peak RSS 182 -> 194 MB Co-Authored-By: Claude Opus 5 (1M context) --- src/ext_mem.rs | 121 ++++++++++++++++++++++++------------------------- 1 file changed, 59 insertions(+), 62 deletions(-) diff --git a/src/ext_mem.rs b/src/ext_mem.rs index 4fe12cd..50131b8 100644 --- a/src/ext_mem.rs +++ b/src/ext_mem.rs @@ -1540,28 +1540,14 @@ where let merge_us = AtomicU64::new(0); let mut emit_secs: f64 = 0.0; - let mut start = 0; - while start < n_partitions { - let end = (start + chunk_size).min(n_partitions); - let chunk = &mut partition_buckets[start..end]; - if ordered_emit { + if ordered_emit { + let mut start = 0; + while start < n_partitions { + let end = (start + chunk_size).min(n_partitions); phase4_merge_chunk_ordered_emit( text, lp, - chunk, - max_ctx, - emit, - cmp, - profile, - &load_us, - &merge_us, - &mut emit_secs, - )?; - } else { - phase4_merge_chunk_collect_emit( - text, - lp, - chunk, + &mut partition_buckets[start..end], max_ctx, emit, cmp, @@ -1570,8 +1556,61 @@ where &merge_us, &mut emit_secs, )?; + start = end; } - start = end; + } else { + // Emitting is `Θ(n)` and single-threaded by construction: the caller's + // closure is `FnMut` and its ordering is the whole point. Previously + // it also did not overlap anything, so every worker sat idle once per + // chunk while the main thread drained the merged results. + // + // Pipeline it instead. A scoped producer merges chunk `c + 1` (itself + // rayon-parallel) while the main thread emits chunk `c`. A bound of + // one keeps at most two chunks resident, so the transient cost is one + // extra chunk of merged positions rather than the unbounded queue an + // unsynchronised producer would build. + // + // This is not the `ordered_phase4_emit` path, which coordinates at + // *partition* granularity through a `BTreeMap` and measured slower. + // Here the producer emits whole chunks already in order, so the + // consumer just drains them. + let (tx, rx) = std::sync::mpsc::sync_channel::>>>(1); + let load_ref = &load_us; + let merge_ref = &merge_us; + std::thread::scope(|scope| -> Result<(), BuildError> { + scope.spawn(move || { + let mut start = 0; + while start < n_partitions { + let end = (start + chunk_size).min(n_partitions); + let merged: io::Result>> = partition_buckets[start..end] + .par_iter_mut() + .map(|bucket| { + merge_one_partition( + text, lp, bucket, max_ctx, cmp, profile, load_ref, merge_ref, + ) + }) + .collect(); + let failed = merged.is_err(); + if tx.send(merged).is_err() || failed { + return; + } + start = end; + } + }); + + for merged in rx { + let t = Instant::now(); + for positions in merged? { + for pos in positions { + emit(pos.to_usize() as u64).map_err(BuildError::Emit)?; + } + } + if profile { + emit_secs += t.elapsed().as_secs_f64(); + } + } + Ok(()) + })?; } if profile { profile_log(&format!( @@ -1584,48 +1623,6 @@ where Ok(()) } -#[allow(clippy::too_many_arguments)] -fn phase4_merge_chunk_collect_emit( - text: &[S], - lp: &L, - chunk: &mut [B], - max_ctx: usize, - emit: &mut F, - cmp: Cmp<'_>, - profile: bool, - load_us: &std::sync::atomic::AtomicU64, - merge_us: &std::sync::atomic::AtomicU64, - emit_secs: &mut f64, -) -> Result<(), BuildError> -where - S: Symbol, - I: Index, - L: LimitProvider, - SaLcp: BucketRecord, - B: BucketStore> + Send, - F: FnMut(u64) -> Result<(), E>, -{ - // Default fast path: let rayon merge the whole chunk with minimal - // coordination, then emit the collected partition results in order. - let merged: Vec> = chunk - .par_iter_mut() - .map(|bucket| -> io::Result> { - merge_one_partition(text, lp, bucket, max_ctx, cmp, profile, load_us, merge_us) - }) - .collect::, io::Error>>()?; - - let t = Instant::now(); - for positions in merged { - for pos in positions { - emit(pos.to_usize() as u64).map_err(BuildError::Emit)?; - } - } - if profile { - *emit_secs += t.elapsed().as_secs_f64(); - } - Ok(()) -} - #[allow(clippy::too_many_arguments)] fn phase4_merge_chunk_ordered_emit( text: &[S], From b1608c2dc255784a3e0ffe19bd6a410c1fdcc0c5 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 22:26:56 +0200 Subject: [PATCH 20/41] docs: extend the external-memory ladder with the last two steps Co-Authored-By: Claude Opus 5 (1M context) --- README.md | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index f95c88e..00cc25b 100644 --- a/README.md +++ b/README.md @@ -227,10 +227,10 @@ external-memory path keeps its memory bound. | ext-mem input | before | after | CPU before | CPU after | peak RSS | | ------------- | ------ | ----- | ---------- | --------- | -------- | -| chr21 FASTA, 47.5 MB | 24.2 s | **1.65 s** | 268 s | 17.5 s | 147 → 182 MB | -| chr21 `N`-free, 80 MB | 3.49 s | **2.11 s** | 33.9 s | 25.0 s | 214 → 233 MB | +| chr21 FASTA, 47.5 MB | 24.2 s | **1.45 s** | 268 s | 17.5 s | 147 → 194 MB | +| chr21 `N`-free, 80 MB | 3.49 s | **1.85 s** | 33.9 s | 25.0 s | 214 → 240 MB | -Four changes get there, each measured separately: +Six changes get there, each measured separately: | change | chr21.0123 | chr21 FASTA | | ------ | ---------- | ----------- | @@ -238,7 +238,9 @@ Four changes get there, each measured separately: | skip periodic runs | 3.55 s | 2.48 s | | prefetch the next candidates' text | 2.90 s | 1.99 s | | subarray target 64Ki → 128Ki records | 2.54 s | 1.86 s | -| seed phase-1 subarrays with the packed key | **2.11 s** | **1.65 s** | +| seed phase-1 subarrays with the packed key | 2.11 s | 1.65 s | +| merge cascade run pairs in parallel | 2.02 s | 1.55 s | +| pipeline the emit against the next merge | **1.85 s** | **1.45 s** | Phase 1 goes from 22.85 s to 0.42 s on the FASTA input, and from 1.20 s to 0.32 s on the `N`-free one. Peak RSS rises by under 30 MB, so the From c7f0cf194055cb9276d645f53d0a5ab13cce8a88 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 22:45:43 +0200 Subject: [PATCH 21/41] perf: re-sort run-free partitions by key instead of cascading merges A partition arrives as `p` sorted sub-subarrays, and the cascade merged them pairwise in `log2(p)` levels, each a full LCP-enhanced pass. That was the largest single cost left in the ext-mem build: 15.4 CPU-seconds of a 15.1-second run on 80 MB of DNA. When a packed key applies, discarding that sortedness and re-sorting the partition outright is much cheaper. One key sort resolves the leading `k` symbols with no text access, and only suffixes agreeing through all of them reach the merge kernel, so `log2(p)` passes collapse into one. Gated on the run table being empty, which is the point worth recording. A long periodic run is exactly a stretch where a fixed-depth key resolves nothing, since every suffix inside it shares the whole key, so the re-sort would hand the merge kernel one enormous tied group and throw away ordering phase 1 had already established. Measured unconditionally it was a large regression on the `N`-heavy input: cascade key re-sort chr21.0123 (no runs) 1.95 s 1.65 s chr21.fa (6.6 Mb N) 1.45 s 2.23 s So the run table, which already exists to make scans cheap, doubles as the predicate for whether a fixed-depth key can be expected to pay off. Apple M4 Max, 12 threads, output verified identical to the in-memory suffix array on both inputs: chr21.0123, 80 MB 1.85 s -> 1.65 s phase4 merge CPU 15.4 -> 7.8 s chr21.fa, 47.5 MB 1.45 s -> 1.47 s (unchanged; takes the cascade) Peak RSS on the run-free input rises 240 MB -> 285 MB: the key vector and the sort buffers are live for each of the `4 x threads` partitions in flight. Still an order of magnitude under the in-memory path's 2.2 GB, which is the comparison that matters for choosing this path. Co-Authored-By: Claude Opus 5 (1M context) --- src/ext_mem.rs | 78 +++++++++++++++++++++++++++++++++++++++++++++++--- src/runs.rs | 10 +++++++ 2 files changed, 84 insertions(+), 4 deletions(-) diff --git a/src/ext_mem.rs b/src/ext_mem.rs index 50131b8..bcc1448 100644 --- a/src/ext_mem.rs +++ b/src/ext_mem.rs @@ -588,6 +588,7 @@ where opts.ordered_phase4_emit, &mut emit, cmp, + seed_params, ); profile_log(&format!( "phase4 (merge+emit) {:.3}s", @@ -665,6 +666,7 @@ where opts.ordered_phase4_emit, &mut emit, cmp, + seed_params, ) } @@ -1492,6 +1494,7 @@ where /// subarrays the per-partition size is `≈ n / p`, so this stays /// proportional to `n / 4 = 0.25 n` even at the peak — well below the /// in-memory path's `~4 n` working set. +#[allow(clippy::too_many_arguments)] // buckets + text + lp + ctx + emit + cmp + seed + flag fn phase4_merge_and_emit( text: &[S], lp: &L, @@ -1500,6 +1503,7 @@ fn phase4_merge_and_emit( ordered_emit: bool, emit: &mut F, cmp: Cmp<'_>, + seed_params: Option<(u32, usize)>, ) -> Result<(), BuildError> where S: Symbol, @@ -1551,6 +1555,7 @@ where max_ctx, emit, cmp, + seed_params, profile, &load_us, &merge_us, @@ -1586,7 +1591,15 @@ where .par_iter_mut() .map(|bucket| { merge_one_partition( - text, lp, bucket, max_ctx, cmp, profile, load_ref, merge_ref, + text, + lp, + bucket, + max_ctx, + cmp, + seed_params, + profile, + load_ref, + merge_ref, ) }) .collect(); @@ -1631,6 +1644,7 @@ fn phase4_merge_chunk_ordered_emit( max_ctx: usize, emit: &mut F, cmp: Cmp<'_>, + seed_params: Option<(u32, usize)>, profile: bool, load_us: &std::sync::atomic::AtomicU64, merge_us: &std::sync::atomic::AtomicU64, @@ -1660,7 +1674,15 @@ where .enumerate() .for_each_with(tx, |tx, (local_idx, bucket)| { let result = merge_one_partition( - text, lp, bucket, max_ctx, cmp, profile, load_us, merge_us, + text, + lp, + bucket, + max_ctx, + cmp, + seed_params, + profile, + load_us, + merge_us, ); let _ = tx.send((local_idx, result)); }); @@ -1718,6 +1740,7 @@ fn merge_one_partition( bucket: &mut B, max_ctx: usize, cmp: Cmp<'_>, + seed_params: Option<(u32, usize)>, profile: bool, load_us: &std::sync::atomic::AtomicU64, merge_us: &std::sync::atomic::AtomicU64, @@ -1742,8 +1765,55 @@ where } let t = Instant::now(); - let workspace = CascadeWorkspace::::new(); - let result = workspace.cascade_merge(text, lp, &records, &boundaries, max_ctx, cmp); + // A partition arrives as `p` sorted sub-subarrays, and the cascade merges + // them pairwise in `log2(p)` levels, each a full pass with LCP-enhanced + // comparisons. That was the single largest cost in the whole ext-mem + // build (15.4 CPU-seconds of a 15.1-second run on 80 MB of DNA). + // + // When the comparator allows a packed key, throwing the existing + // sortedness away and re-sorting the partition outright is dramatically + // cheaper: one key sort resolves the leading `k` symbols with no text + // access, and only suffixes agreeing through all of them need the merge + // kernel. `log2(p)` passes collapse into one. + // + // Only when the text has no long periodic runs, though. A run is exactly + // a stretch where a fixed-depth key resolves nothing, since every suffix + // inside it shares the whole key, so the re-sort would hand the merge + // kernel one enormous tied group and lose the ordering phase 1 had + // already established. Measured on chr21 FASTA (6.6 Mb of `N`) the + // unconditional version was 2.23 s against the cascade's 1.45 s, while on + // run-free DNA it is 1.65 s against 1.95 s. So: key re-sort when the run + // table is empty, cascade otherwise. + let result = match seed_params { + Some(_) + if lp.plain_lex_len() == Some(text.len()) + && max_ctx == usize::MAX + && !cmp.has_long_runs() => + { + let mut sa: Vec = records.iter().map(|r| r.pos).collect(); + let len = sa.len(); + let mut lcp = vec![I::zero(); len]; + let mut sa_w = vec![I::zero(); len]; + let mut lcp_w = vec![I::zero(); len]; + let seeded = crate::radix::seed_subarray( + text, + lp, + seed_params, + &mut sa, + &mut lcp, + &mut sa_w, + &mut lcp_w, + max_ctx, + cmp, + ); + debug_assert!(seeded, "guards agreed but seed_subarray declined"); + sa + } + _ => { + let workspace = CascadeWorkspace::::new(); + workspace.cascade_merge(text, lp, &records, &boundaries, max_ctx, cmp) + } + }; if profile { merge_us.fetch_add(t.elapsed().as_micros() as u64, AtomicOrdering::Relaxed); } diff --git a/src/runs.rs b/src/runs.rs index 5f89e97..8224451 100644 --- a/src/runs.rs +++ b/src/runs.rs @@ -232,6 +232,16 @@ impl<'a> Cmp<'a> { Self { dispatch, runs } } + /// Whether the text contains long periodic repeats. + /// + /// Callers use this to decide whether a fixed-depth key can be expected + /// to resolve most suffixes: a long run is precisely a stretch where it + /// cannot, because every suffix inside it shares the whole key. + #[inline] + pub(crate) fn has_long_runs(&self) -> bool { + !self.runs.is_empty() + } + /// LCP of `text[p..]` and `text[q..]` in symbols, bounded by `max_ctx`, /// using run structure to jump over long periodic stretches. /// From b1fd8fbf9dd3fe3a41011d9cec38da6ef1ad35b2 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 22:46:32 +0200 Subject: [PATCH 22/41] docs: add the partition key re-sort to the ladder Co-Authored-By: Claude Opus 5 (1M context) --- README.md | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 00cc25b..c1aba9b 100644 --- a/README.md +++ b/README.md @@ -227,10 +227,10 @@ external-memory path keeps its memory bound. | ext-mem input | before | after | CPU before | CPU after | peak RSS | | ------------- | ------ | ----- | ---------- | --------- | -------- | -| chr21 FASTA, 47.5 MB | 24.2 s | **1.45 s** | 268 s | 17.5 s | 147 → 194 MB | -| chr21 `N`-free, 80 MB | 3.49 s | **1.85 s** | 33.9 s | 25.0 s | 214 → 240 MB | +| chr21 FASTA, 47.5 MB | 24.2 s | **1.47 s** | 268 s | 17.5 s | 147 → 190 MB | +| chr21 `N`-free, 80 MB | 3.49 s | **1.65 s** | 33.9 s | 15.1 s | 214 → 285 MB | -Six changes get there, each measured separately: +Seven changes get there, each measured separately: | change | chr21.0123 | chr21 FASTA | | ------ | ---------- | ----------- | @@ -240,7 +240,8 @@ Six changes get there, each measured separately: | subarray target 64Ki → 128Ki records | 2.54 s | 1.86 s | | seed phase-1 subarrays with the packed key | 2.11 s | 1.65 s | | merge cascade run pairs in parallel | 2.02 s | 1.55 s | -| pipeline the emit against the next merge | **1.85 s** | **1.45 s** | +| pipeline the emit against the next merge | 1.85 s | 1.45 s | +| re-sort run-free partitions by key | **1.65 s** | **1.47 s** | Phase 1 goes from 22.85 s to 0.42 s on the FASTA input, and from 1.20 s to 0.32 s on the `N`-free one. Peak RSS rises by under 30 MB, so the From a21b8c51c15abfa24ab699bfb4769f8de1b291f4 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 22:53:06 +0200 Subject: [PATCH 23/41] perf: rank the alphabet before packing keys The field width was chosen from the largest byte value in the text rather than from how many distinct symbols it uses, and for the most natural input format those are very different numbers. A plain ACGT sequence has four symbols, but its largest byte is `'T'` (84), so the packer used 8-bit fields and fit 8 symbols per key instead of the 32 that four symbols allow. Rank the bytes that actually occur onto a dense code range first. The map is monotone by construction, since codes are assigned in ascending byte order, so a packed key stays order-preserving and the zero-padding argument carries over unchanged: code 0 is still the minimum. Apple M4 Max, 12 threads, chr21 forward ++ revcomp written as ASCII ACGT (80 MB, four symbols, largest byte 84): before after seed sort 0.480 s 0.281 s doubling 0.867 s 0.488 s total 1.404 s 0.827 s -41% The two inputs benchmarked so far both happened to hide this. One is pre-encoded to codes 0..3 and is already dense; the other is dominated by its `N` runs, where key depth is irrelevant because every suffix inside a run shares the whole key. A FASTA-derived ACGT text is the common case that neither covered. Cross-check: the suffix array of the ASCII text is identical to the suffix array of the 0..3-coded text, as it must be, the two encodings being order-isomorphic. Both the in-memory and ext-mem paths agree. Tests: the field width follows alphabet size rather than byte value across four alphabets, and the remap is asserted monotone. Co-Authored-By: Claude Opus 5 (1M context) --- src/ext_mem.rs | 14 ++-- src/radix.rs | 198 +++++++++++++++++++++++++++++++++++-------------- 2 files changed, 152 insertions(+), 60 deletions(-) diff --git a/src/ext_mem.rs b/src/ext_mem.rs index bcc1448..a00a411 100644 --- a/src/ext_mem.rs +++ b/src/ext_mem.rs @@ -505,7 +505,8 @@ where let p = effective_subproblem_count(n, opts.subproblem_count); let runs = crate::runs::detect_for(text); let cmp = Cmp::new(LcpDispatch::detect(), &runs); - let seed_params = crate::radix::seed_params(text); + let seed_packer = crate::radix::seed_params(text); + let seed_params = seed_packer.as_ref(); let work_dir = opts.work_dir.clone(); // Pool the `2 × p` bucket files into one anonymous tempfile per @@ -629,7 +630,8 @@ where let p = effective_subproblem_count(n, opts.subproblem_count); let runs = crate::runs::detect_for(text); let cmp = Cmp::new(LcpDispatch::detect(), &runs); - let seed_params = crate::radix::seed_params(text); + let seed_packer = crate::radix::seed_params(text); + let seed_params = seed_packer.as_ref(); let factory = |_i: usize| InMemBucket::>::new(); @@ -1209,7 +1211,7 @@ fn phase1_sort_sample_spill( p: usize, opts: &ExtMemOpts, cmp: Cmp<'_>, - seed_params: Option<(u32, usize)>, + seed_params: Option<&crate::radix::Packer>, mk_bucket: MkB, ) -> io::Result<(Vec, Vec)> where @@ -1503,7 +1505,7 @@ fn phase4_merge_and_emit( ordered_emit: bool, emit: &mut F, cmp: Cmp<'_>, - seed_params: Option<(u32, usize)>, + seed_params: Option<&crate::radix::Packer>, ) -> Result<(), BuildError> where S: Symbol, @@ -1644,7 +1646,7 @@ fn phase4_merge_chunk_ordered_emit( max_ctx: usize, emit: &mut F, cmp: Cmp<'_>, - seed_params: Option<(u32, usize)>, + seed_params: Option<&crate::radix::Packer>, profile: bool, load_us: &std::sync::atomic::AtomicU64, merge_us: &std::sync::atomic::AtomicU64, @@ -1740,7 +1742,7 @@ fn merge_one_partition( bucket: &mut B, max_ctx: usize, cmp: Cmp<'_>, - seed_params: Option<(u32, usize)>, + seed_params: Option<&crate::radix::Packer>, profile: bool, load_us: &std::sync::atomic::AtomicU64, merge_us: &std::sync::atomic::AtomicU64, diff --git a/src/radix.rs b/src/radix.rs index 0f63fcb..f7f1f1e 100644 --- a/src/radix.rs +++ b/src/radix.rs @@ -31,10 +31,13 @@ //! //! ## The algorithm //! -//! 1. **Pack.** Find the maximum symbol and choose the smallest field width -//! in `{1, 2, 4, 8}` bits that can hold it, so `k = 64 / bits` symbols fit -//! in one `u64` key. DNA over `{0,1,2,3}` gets 2-bit fields and therefore -//! resolves **32 symbols per key** rather than the 8 a raw byte key would. +//! 1. **Pack.** Rank the bytes that actually occur onto a dense code range, +//! then choose the smallest field width in `{1, 2, 4, 8}` bits that holds +//! the alphabet, so `k = 64 / bits` symbols fit in one `u64` key. Ranking +//! matters: raw FASTA uses six symbols but its largest byte is `'T'` (84), +//! so packing raw bytes would force 8-bit fields and 8 symbols per key, +//! against 16 after ranking. DNA over `{0,1,2,3}` gets 2-bit fields and +//! **32 symbols per key**. //! 2. **Seed.** Sort `(key, position)`. This is a full sort of the suffixes //! by their first `k` symbols, and it touches the text only in one //! sequential pass. @@ -63,47 +66,106 @@ use crate::sample_sort; use rayon::prelude::*; use std::time::Instant; -/// Field width in bits and the number of symbols that fit in a `u64` key. +/// An order-preserving remap of the bytes that actually occur in a text onto +/// a dense code range, plus the resulting key geometry. /// -/// Restricted to divisors of 64 so a key is an exact number of whole fields -/// and no symbol ever straddles the key boundary. -fn pack_params(max_sym: u8) -> (u32, usize) { - let bits: u32 = match max_sym { - 0..=1 => 1, - 2..=3 => 2, - 4..=15 => 4, - _ => 8, - }; - (bits, 64 / bits as usize) +/// The field width is driven by how many *distinct* symbols a text uses, not +/// by the largest byte value in it, and the difference is not academic. A raw +/// FASTA uses six symbols, but the largest is `'T'` (84), so packing raw bytes +/// forces 8-bit fields and fits only 8 symbols per key. Ranking those six +/// bytes to `0..6` gives 4-bit fields and 16 symbols per key, which halves the +/// number of doubling rounds needed downstream. The DNA-coded input is already +/// dense, so it is unaffected. +/// +/// The map is monotone by construction, since codes are assigned in ascending +/// byte order. That is what keeps a packed key order-preserving: `key_a < +/// key_b` still implies `suffix_a < suffix_b`, and the zero-padding argument +/// carries over because code `0` remains the minimum. +pub(crate) struct Packer { + /// Byte to dense code. Bytes absent from the text map to `0`; they never + /// appear in a key. + code: [u8; 256], + /// Bits per packed field. + bits: u32, + /// Symbols per `u64` key. + k: usize, } -/// Pack the `k` symbols at `text[p..]` into one order-preserving `u64`, -/// zero-padding past the end of the text. -#[inline] -fn key_at(text: &[u8], p: usize, bits: u32, k: usize) -> u64 { - if bits == 8 { - // Whole-byte fields: this is just a big-endian load, and the common - // case (p + 8 <= n) is a single unaligned u64 read plus a bswap. - let mut buf = [0u8; 8]; - let end = (p + 8).min(text.len()); - buf[..end - p].copy_from_slice(&text[p..end]); - return u64::from_be_bytes(buf); +impl Packer { + /// Build the map for `text`. + fn new(text: &[u8]) -> Self { + // Which bytes occur? One parallel pass, folded into a 256-bit set. + let present = text + .par_chunks(1 << 16) + .map(|c| { + let mut seen = [false; 256]; + for &b in c { + seen[b as usize] = true; + } + seen + }) + .reduce( + || [false; 256], + |mut a, b| { + for i in 0..256 { + a[i] |= b[i]; + } + a + }, + ); + + let mut code = [0u8; 256]; + let mut next = 0u16; + for (b, &seen) in present.iter().enumerate() { + if seen { + code[b] = next as u8; + next += 1; + } + } + // `next` is the alphabet size; the largest code is `next - 1`. + let bits: u32 = match next.saturating_sub(1) { + 0..=1 => 1, + 2..=3 => 2, + 4..=15 => 4, + _ => 8, + }; + Self { + code, + bits, + k: 64 / bits as usize, + } + } + + #[inline] + pub(crate) fn bits(&self) -> u32 { + self.bits + } + + #[inline] + pub(crate) fn k(&self) -> usize { + self.k } - let end = (p + k).min(text.len()); - let mut key: u64 = 0; - for &s in &text[p..end] { - key = (key << bits) | s as u64; + + /// Pack the `k` symbols at `text[p..]` into one order-preserving `u64`, + /// zero-padding past the end of the text. + #[inline] + pub(crate) fn key_at(&self, text: &[u8], p: usize) -> u64 { + let end = (p + self.k).min(text.len()); + let mut key: u64 = 0; + for &s in &text[p..end] { + key = (key << self.bits) | self.code[s as usize] as u64; + } + // Shift the packed prefix up so the missing trailing fields read zero. + key << (self.bits as usize * (self.k - (end - p))) } - // Shift the packed prefix up so the missing trailing fields read as zero. - key << (bits as usize * (k - (end - p))) } -/// Field width and symbols-per-key for `text`, or `None` when a packed key -/// cannot represent this text's order. +/// The alphabet map for `text`, or `None` when a packed key cannot represent +/// this text's order. /// /// Computed once per build and handed to [`seed_subarray`], which would /// otherwise re-scan the whole text for every subarray. -pub(crate) fn seed_params(text: &[S]) -> Option<(u32, usize)> { +pub(crate) fn seed_params(text: &[S]) -> Option { if size_of::() != 1 { return None; } @@ -111,7 +173,7 @@ pub(crate) fn seed_params(text: &[S]) -> Option<(u32, usize)> { // valid for reads of the same length. let bytes: &[u8] = unsafe { std::slice::from_raw_parts(text.as_ptr() as *const u8, text.len()) }; - Some(pack_params(bytes.par_iter().copied().max().unwrap_or(0))) + Some(Packer::new(bytes)) } /// Sort `sa` into suffix order and fill `lcp`, using a packed fixed-depth key @@ -138,7 +200,7 @@ pub(crate) fn seed_params(text: &[S]) -> Option<(u32, usize)> { pub(crate) fn seed_subarray( text: &[S], lp: &L, - params: Option<(u32, usize)>, + packer: Option<&Packer>, sa: &mut [I], lcp: &mut [I], sa_w: &mut [I], @@ -146,9 +208,10 @@ pub(crate) fn seed_subarray( max_ctx: usize, cmp: Cmp<'_>, ) -> bool { - let Some((bits, k)) = params else { + let Some(packer) = packer else { return false; }; + let (bits, k) = (packer.bits(), packer.k()); if max_ctx != usize::MAX || lp.plain_lex_len() != Some(text.len()) { return false; } @@ -172,7 +235,7 @@ pub(crate) fn seed_subarray( .iter() .map(|&p| { let p = p.to_usize(); - (key_at(bytes, p, bits, k), visible(p) as u32, p_as(p)) + (packer.key_at(bytes, p), visible(p) as u32, p_as(p)) }) .collect(); keyed.sort_unstable(); @@ -251,8 +314,7 @@ const RADIX_BUCKETS: usize = 1 << RADIX_BITS; /// serial, which is exactly where a `n log n` sort loses on many cores. fn seed_sort( text: &[u8], - bits: u32, - k: usize, + packer: &Packer, visible_len: &(dyn Fn(usize) -> usize + Sync), ) -> (Vec, Vec) { let n = text.len(); @@ -272,7 +334,7 @@ fn seed_sort( .map(|&(start, end)| { let mut counts = vec![0u32; RADIX_BUCKETS]; for p in start..end { - counts[bucket_of(key_at(text, p, bits, k))] += 1; + counts[bucket_of(packer.key_at(text, p))] += 1; } counts }) @@ -310,7 +372,7 @@ fn seed_sort( let mut cursor: Vec = offsets[c * RADIX_BUCKETS..(c + 1) * RADIX_BUCKETS].to_vec(); for p in start..end { - let key = key_at(text, p, bits, k); + let key = packer.key_at(text, p); let slot = &mut cursor[bucket_of(key)]; // SAFETY: the prefix sum gives this (chunk, bucket) pair a // range of exactly its own histogram count, and the cursor @@ -373,8 +435,8 @@ pub(crate) fn build_sa(text: &[u8]) -> Vec { } let t0 = Instant::now(); - let max_sym = text.par_iter().copied().max().unwrap_or(0); - let (bits, k) = pack_params(max_sym); + let packer = Packer::new(text); + let k = packer.k(); // ---- Seed: sort by the first `k` symbols, then by visible length. ---- // @@ -397,7 +459,7 @@ pub(crate) fn build_sa(text: &[u8]) -> Vec { t0.elapsed().as_secs_f64() )); let t1 = Instant::now(); - let (keys, mut sa) = seed_sort::(text, bits, k, &visible_len); + let (keys, mut sa) = seed_sort::(text, &packer, &visible_len); profile_log(&format!( "radix seed sort {:.3}s", t1.elapsed().as_secs_f64() @@ -633,15 +695,43 @@ mod tests { check(b"abracadabra"); } + /// The field width follows the number of *distinct* symbols, not the + /// largest byte value. Raw FASTA is the case that matters: six symbols + /// whose largest is `'T'` (84) would force 8-bit fields without ranking, + /// fitting only 8 symbols per key instead of 16. #[test] - fn pack_params_covers_every_width() { - assert_eq!(pack_params(0), (1, 64)); - assert_eq!(pack_params(1), (1, 64)); - assert_eq!(pack_params(3), (2, 32)); - assert_eq!(pack_params(4), (4, 16)); - assert_eq!(pack_params(15), (4, 16)); - assert_eq!(pack_params(16), (8, 8)); - assert_eq!(pack_params(255), (8, 8)); + fn packer_width_follows_alphabet_size_not_byte_value() { + let two = Packer::new(b"abababab"); + assert_eq!((two.bits(), two.k()), (1, 64)); + let four = Packer::new(&[0u8, 1, 2, 3, 3, 2, 1, 0]); + assert_eq!((four.bits(), four.k()), (2, 32)); + + let mut fasta: Vec = b"ACGTN".to_vec(); + fasta.push(b'\n'); + let f = Packer::new(&fasta); + assert_eq!((f.bits(), f.k()), (4, 16), "6 symbols should pack 4 bits"); + + let dense: Vec = (0..=255u8).collect(); + let d = Packer::new(&dense); + assert_eq!((d.bits(), d.k()), (8, 8)); + } + + /// The remap must be monotone, or a packed key would stop being + /// order-preserving and the whole seed would be wrong. + #[test] + fn packer_remap_is_monotone() { + let text: Vec = b"TGCAN\nZq".to_vec(); + let p = Packer::new(&text); + let mut present: Vec = text.clone(); + present.sort_unstable(); + present.dedup(); + for w in present.windows(2) { + assert!( + p.code[w[0] as usize] < p.code[w[1] as usize], + "codes must follow byte order: {:?}", + w + ); + } } /// Texts whose symbols include a real `0`, so padding and a genuine From 7c14607c66ab267df32f948ac6d3f9cc9a8cd784 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 22:55:59 +0200 Subject: [PATCH 24/41] docs: record the alphabet-ranking result Co-Authored-By: Claude Opus 5 (1M context) --- README.md | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index c1aba9b..7d9afac 100644 --- a/README.md +++ b/README.md @@ -47,7 +47,8 @@ the merge kernel's and independently verified: | input | before | after | CPU before | CPU after | | ----- | ------ | ----- | ---------- | --------- | -| chr21 fwd ++ revcomp, `N`-free, 80 MB | 6.08 s | **0.84 s** | 28.1 s | 5.0 s | +| chr21 fwd ++ revcomp, ASCII `ACGT`, 80 MB | 6.08 s | **0.83 s** | 28.1 s | 5.0 s | +| same, pre-coded to `0..3`, 80 MB | 6.08 s | **0.80 s** | 28.1 s | 5.0 s | | chr21 FASTA, 47.5 MB, 6.6 Mb of `N` | 27.8 s | **1.04 s** | 283.5 s | 5.2 s | | same, via `build_in_memory_sample_sort` | 3.41 s | **1.18 s** | 34.5 s | 7.2 s | @@ -55,6 +56,12 @@ Peak RSS on the 80 MB input is 2.21 GB, down from 2.83 GB, because the seed is an MSD counting sort that recomputes keys from the text rather than materialising a key array. +The two DNA rows above land in the same place because the alphabet is +ranked to a dense code range before packing. Without that step the ASCII +row would use 8-bit fields — its largest byte is `'T'` (84) even though +it has four symbols — fitting 8 symbols per key instead of 32, and would +take 1.40 s rather than 0.83 s. + Note the two inputs are different problems; benchmarking one implementation on the first and another on the second is not a comparison. `bench/chr21.sh` prepares both. From 1d9f3f2054b5652342ab04564636df433568b3b5 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 23:03:39 +0200 Subject: [PATCH 25/41] perf: pack keys with a SWAR gather over a pre-ranked text Building a key was a chain of `k` dependent shift-or-lookup steps, 32 of them for 2-bit DNA fields, and that chain set the cost of the seed sort in every path that uses a packed key. Two changes remove it. The alphabet map is applied to the whole text once, up front, so the inner loop no longer carries a dependent table load. And the fields are gathered by a binary-tree SWAR shuffle: each step folds neighbouring fields together and halves the stride, so eight symbols cost three shift-or-mask pairs rather than eight dependent steps. The load is big-endian, which puts the text's first byte in the result's most significant field, the order the key already needed. The ranked copy is only materialised when the identity map does not already rank the text. A `0..3` DNA encoding is already dense and pays nothing; an ASCII text pays one byte per symbol, which is what buys the narrower fields in the first place. Note the group concatenation assigns the first group rather than shifting it in. With 8-bit fields there is exactly one group and the shift distance would be 64, which is not legal for `u64`: release builds mask it to zero and happen to produce the right answer, debug builds panic. This was caught by running the tests in debug, which is why the CI workflow does. Apple M4 Max, 12 threads, 80 MB chr21 forward ++ revcomp, output verified identical on all three encodings: seed sort total ASCII ACGT 0.281 s -> 0.177 s 0.827 s -> 0.702 s coded 0..3 0.278 s -> 0.168 s 0.802 s -> 0.719 s Peak RSS is 2.21 GB for the coded input, unchanged, and 2.29 GB for the ASCII one, the difference being the ranked copy. Tests: a new case checks the SWAR gather against the obvious shift-or loop for every field width, every alignment and every tail length, and the monotonicity of the alphabet map is now asserted through `key_at` rather than through the internal table. Co-Authored-By: Claude Opus 5 (1M context) --- src/radix.rs | 172 +++++++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 147 insertions(+), 25 deletions(-) diff --git a/src/radix.rs b/src/radix.rs index f7f1f1e..994f291 100644 --- a/src/radix.rs +++ b/src/radix.rs @@ -82,9 +82,13 @@ use std::time::Instant; /// key_b` still implies `suffix_a < suffix_b`, and the zero-padding argument /// carries over because code `0` remains the minimum. pub(crate) struct Packer { - /// Byte to dense code. Bytes absent from the text map to `0`; they never - /// appear in a key. - code: [u8; 256], + /// The text with every byte replaced by its code, when the identity map + /// does not already do that. Materialising it once removes a dependent + /// table load from the packing loop, which is otherwise the chain that + /// sets the cost of building a key. `None` when the text is already dense + /// (a `0..3` DNA encoding, say), so the common pre-coded input pays no + /// extra memory. + ranked: Option>, /// Bits per packed field. bits: u32, /// Symbols per `u64` key. @@ -94,7 +98,7 @@ pub(crate) struct Packer { impl Packer { /// Build the map for `text`. fn new(text: &[u8]) -> Self { - // Which bytes occur? One parallel pass, folded into a 256-bit set. + // Which bytes occur? One parallel pass, folded into a 256-entry set. let present = text .par_chunks(1 << 16) .map(|c| { @@ -116,21 +120,37 @@ impl Packer { let mut code = [0u8; 256]; let mut next = 0u16; + let mut identity = true; for (b, &seen) in present.iter().enumerate() { if seen { code[b] = next as u8; + identity &= next as usize == b; next += 1; } } - // `next` is the alphabet size; the largest code is `next - 1`. let bits: u32 = match next.saturating_sub(1) { 0..=1 => 1, 2..=3 => 2, 4..=15 => 4, _ => 8, }; + // An identity map, or 8-bit fields where the code never changes the + // packed value's order, both let the original text be read directly. + let ranked = if identity { + None + } else { + let mut out = vec![0u8; text.len()]; + out.par_chunks_mut(1 << 16) + .zip(text.par_chunks(1 << 16)) + .for_each(|(dst, src)| { + for (d, &s) in dst.iter_mut().zip(src) { + *d = code[s as usize]; + } + }); + Some(out) + }; Self { - code, + ranked, bits, k: 64 / bits as usize, } @@ -146,17 +166,76 @@ impl Packer { self.k } + /// Gather the low `bits` of each of 8 ranked bytes into one contiguous + /// field, most-significant byte first. + /// + /// A binary-tree SWAR shuffle: each step folds neighbouring fields + /// together and halves the stride, so eight symbols cost three + /// shift-or-mask pairs instead of eight dependent shift-or steps. The + /// input is a big-endian load, so the text's first byte lands in the + /// result's most significant field, which is the order the key needs. + #[inline(always)] + fn gather8(v: u64, bits: u32) -> u64 { + match bits { + 1 => { + let mut x = v & 0x0101_0101_0101_0101; + x = (x | (x >> 7)) & 0x0003_0003_0003_0003; + x = (x | (x >> 14)) & 0x0000_000F_0000_000F; + (x | (x >> 28)) & 0xFF + } + 2 => { + let mut x = v & 0x0303_0303_0303_0303; + x = (x | (x >> 6)) & 0x000F_000F_000F_000F; + x = (x | (x >> 12)) & 0x0000_00FF_0000_00FF; + (x | (x >> 24)) & 0xFFFF + } + 4 => { + let mut x = v & 0x0F0F_0F0F_0F0F_0F0F; + x = (x | (x >> 4)) & 0x00FF_00FF_00FF_00FF; + x = (x | (x >> 8)) & 0x0000_FFFF_0000_FFFF; + (x | (x >> 16)) & 0xFFFF_FFFF + } + _ => v, + } + } + /// Pack the `k` symbols at `text[p..]` into one order-preserving `u64`, /// zero-padding past the end of the text. #[inline] pub(crate) fn key_at(&self, text: &[u8], p: usize) -> u64 { - let end = (p + self.k).min(text.len()); - let mut key: u64 = 0; - for &s in &text[p..end] { - key = (key << self.bits) | self.code[s as usize] as u64; + let src = self.ranked.as_deref().unwrap_or(text); + let n = src.len(); + + // Fast path: a whole key's worth of symbols is available, so it is + // `k / 8` big-endian loads and their gathers, with no bounds fuss. + if p + self.k <= n { + return self + .fold(|i| u64::from_be_bytes(src[p + 8 * i..p + 8 * i + 8].try_into().unwrap())); } - // Shift the packed prefix up so the missing trailing fields read zero. - key << (self.bits as usize * (self.k - (end - p))) + + // Tail: fewer than `k` symbols remain. Pad with zero codes, which are + // the alphabet's minimum, matching shorter-is-smaller. + let mut buf = [0u8; 64]; + buf[..n - p].copy_from_slice(&src[p..n]); + self.fold(|i| u64::from_be_bytes(buf[8 * i..8 * i + 8].try_into().unwrap())) + } + + /// Concatenate the gathers of the `k / 8` words produced by `word`. + /// + /// The first group is assigned rather than shifted in. With 8-bit fields + /// there is exactly one group and `8 * bits` is 64, which is not a legal + /// shift distance for `u64`; release builds mask it to 0 and happen to + /// give the right answer, debug builds panic. Assigning avoids relying on + /// either behaviour. + #[inline(always)] + fn fold(&self, word: impl Fn(usize) -> u64) -> u64 { + let shift = 8 * self.bits; + let mut key = 0u64; + for i in 0..self.k / 8 { + let g = Self::gather8(word(i), self.bits); + key = if i == 0 { g } else { (key << shift) | g }; + } + key } } @@ -717,20 +796,63 @@ mod tests { } /// The remap must be monotone, or a packed key would stop being - /// order-preserving and the whole seed would be wrong. + /// order-preserving and the whole seed would be wrong. Checked through + /// the observable behaviour: over a text whose bytes ascend and are all + /// distinct, successive suffixes must produce strictly increasing keys. #[test] - fn packer_remap_is_monotone() { - let text: Vec = b"TGCAN\nZq".to_vec(); - let p = Packer::new(&text); - let mut present: Vec = text.clone(); - present.sort_unstable(); - present.dedup(); - for w in present.windows(2) { - assert!( - p.code[w[0] as usize] < p.code[w[1] as usize], - "codes must follow byte order: {:?}", - w - ); + fn packer_keys_follow_byte_order() { + for text in [ + b"\nACGNTZq".to_vec(), + (0..40u8) + .map(|i| i.wrapping_mul(6).wrapping_add(3)) + .collect(), + b"ACGT".to_vec(), + ] { + let mut ascending: Vec = text.clone(); + ascending.sort_unstable(); + ascending.dedup(); + let p = Packer::new(&ascending); + let keys: Vec = (0..ascending.len()) + .map(|i| p.key_at(&ascending, i)) + .collect(); + for (i, w) in keys.windows(2).enumerate() { + assert!( + w[0] < w[1], + "key({i}) = {:#x} should precede key({}) = {:#x} for {ascending:?}", + w[0], + i + 1, + w[1], + ); + } + } + } + + /// The SWAR gather must agree with the obvious shift-or loop for every + /// field width and every alignment, including the zero-padded tail. + #[test] + fn swar_gather_matches_scalar_packing() { + use rand::{RngExt, SeedableRng}; + let mut rng = rand::rngs::StdRng::seed_from_u64(0x5AA5); + for &sigma in &[2u8, 4, 16, 200] { + for &n in &[1usize, 7, 8, 9, 31, 32, 33, 63, 64, 65, 200] { + let text: Vec = (0..n).map(|_| rng.random_range(0..sigma)).collect(); + let p = Packer::new(&text); + let (bits, k) = (p.bits(), p.k()); + let ranked = p.ranked.as_deref().unwrap_or(&text); + for pos in 0..n { + let end = (pos + k).min(n); + let mut want: u64 = 0; + for &c in &ranked[pos..end] { + want = (want << bits) | c as u64; + } + want <<= bits as usize * (k - (end - pos)); + assert_eq!( + p.key_at(&text, pos), + want, + "sigma={sigma} n={n} pos={pos} bits={bits}", + ); + } + } } } From f0af1186e58424bdc26c1e55df81457ed57f4aed Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 23:09:02 +0200 Subject: [PATCH 26/41] docs: refresh in-memory numbers after the SWAR key packing Co-Authored-By: Claude Opus 5 (1M context) --- README.md | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 7d9afac..ebcd356 100644 --- a/README.md +++ b/README.md @@ -47,8 +47,8 @@ the merge kernel's and independently verified: | input | before | after | CPU before | CPU after | | ----- | ------ | ----- | ---------- | --------- | -| chr21 fwd ++ revcomp, ASCII `ACGT`, 80 MB | 6.08 s | **0.83 s** | 28.1 s | 5.0 s | -| same, pre-coded to `0..3`, 80 MB | 6.08 s | **0.80 s** | 28.1 s | 5.0 s | +| chr21 fwd ++ revcomp, ASCII `ACGT`, 80 MB | 6.08 s | **0.76 s** | 28.1 s | 5.0 s | +| same, pre-coded to `0..3`, 80 MB | 6.08 s | **0.73 s** | 28.1 s | 5.0 s | | chr21 FASTA, 47.5 MB, 6.6 Mb of `N` | 27.8 s | **1.04 s** | 283.5 s | 5.2 s | | same, via `build_in_memory_sample_sort` | 3.41 s | **1.18 s** | 34.5 s | 7.2 s | @@ -60,7 +60,9 @@ The two DNA rows above land in the same place because the alphabet is ranked to a dense code range before packing. Without that step the ASCII row would use 8-bit fields — its largest byte is `'T'` (84) even though it has four symbols — fitting 8 symbols per key instead of 32, and would -take 1.40 s rather than 0.83 s. +take 1.40 s rather than 0.76 s. Keys are then built by a SWAR gather +over the ranked text, so eight symbols cost three shift-or-mask pairs +instead of eight dependent shift-or-lookup steps. Note the two inputs are different problems; benchmarking one implementation on the first and another on the second is not a From a53ced9841ecd28c80d6d8e5bbf71d6c6ba09342 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 23:16:11 +0200 Subject: [PATCH 27/41] perf: drop the per-group allocation and prepass from doubling rounds Instrumenting the rounds (`CAPS_SA_PROFILE=1` now prints one line each) showed where the time actually went, and it was not where I assumed. On 80 MB of DNA the first round has 3.5 million tied groups averaging **four elements**: round depth=32: 14193102 tied in 3526079 groups (avg 4.0) round depth=64: 9929726 tied in 2924078 groups (avg 3.4) round depth=128: 7480152 tied in 2543160 groups (avg 2.9) At that size neither the sort nor the rank probes dominate. The bookkeeping around them does: one heap allocation per group for the key vector, and a sequential `split_at_mut` chain over all 3.5 million groups to hand each one its sub-slices. Both go. `Scatter` already encodes "disjoint ranges, one owner each", which is exactly the property the groups have, so a group takes its own sub-slices directly and the sequential prepass disappears. Groups up to 32 elements build their key vector in a stack buffer, so nearly every group avoids the allocator. `flat_map_iter` replaces collecting a `Vec>` of mostly-empty vectors. Apple M4 Max, 12 threads, output verified identical on all three encodings and both paths: doubling total chr21.0123 0.502 s -> 0.234 s 0.733 s -> 0.566 s ASCII ACGT 0.764 s -> 0.611 s The first round alone goes from 0.124 s to 0.053 s. The `N`-heavy FASTA is unchanged: its groups are large, so it was never paying the per-group overhead. The per-round log stays in, behind `CAPS_SA_PROFILE`. It is what located this, and the group-size distribution is the thing worth looking at first for any further work here. Co-Authored-By: Claude Opus 5 (1M context) --- src/radix.rs | 105 ++++++++++++++++++++++++++++++--------------------- 1 file changed, 62 insertions(+), 43 deletions(-) diff --git a/src/radix.rs b/src/radix.rs index 994f291..96fc9dc 100644 --- a/src/radix.rs +++ b/src/radix.rs @@ -368,6 +368,10 @@ fn p_as(p: usize) -> I { I::from_usize(p) } +/// Largest tied group whose key vector is built on the stack. Groups average +/// about four elements, so nearly every one avoids the allocator entirely. +const DOUBLING_STACK_GROUP: usize = 32; + /// Bits of the key used for the MSD counting-sort pass, and the resulting /// bucket count. /// @@ -597,34 +601,59 @@ pub(crate) fn build_sa(text: &[u8]) -> Vec { let mut next_rank: Vec = vec![I::zero(); n]; while !groups.is_empty() { + let round_t = Instant::now(); // Phase A: sort each tied group by the successor rank, and record the // ranks it should get. Groups are disjoint `sa` ranges, so this is // data-parallel with no synchronisation. - let sub: Vec> = split_disjoint(&mut sa, &mut next_rank, &groups) - .into_par_iter() - .zip(groups.par_iter()) - .map(|((sa_g, nr_g), &(start, _))| { + // Groups average about four elements, and there are millions of them + // per round, so the two things that dominated here were not the sort + // or the rank probes but the bookkeeping around them: one heap + // allocation per group for the key vector, and a sequential + // `split_at_mut` chain to hand each group its sub-slices. + // + // Both go. `Scatter` already encodes "disjoint ranges, one owner + // each", which is exactly the property the groups have, so each group + // takes its own sub-slices directly with no sequential prepass. And a + // group that fits the stack buffer never touches the allocator. + let sa_cell = Scatter::new(&mut sa); + let nr_cell = Scatter::new(&mut next_rank); + let rank_ref = &rank; + let sub: Vec<(usize, usize)> = groups + .par_iter() + .flat_map_iter(|&(start, end)| { + let len = end - start; + // SAFETY: `groups` are disjoint, sorted `sa` ranges, so this + // group is the sole owner of `start..end` in both arrays. + let (sa_g, nr_g) = + unsafe { (sa_cell.slice_mut(start, len), nr_cell.slice_mut(start, len)) }; let succ = |p: usize| -> u64 { // End-of-text sorts first: the shorter suffix is smaller. match p.checked_add(depth) { - Some(q) if q < n => rank[q].to_usize() as u64 + 1, + Some(q) if q < n => rank_ref[q].to_usize() as u64 + 1, _ => 0, } }; - // Materialise the successor ranks once. `sort_unstable_by_key` - // re-evaluates its key function O(len log len) times, and each - // evaluation is a random probe into `rank`; paying for it once - // per element turns the sort's memory traffic sequential. - let mut keyed: Vec<(u64, I)> = - sa_g.iter().map(|&e| (succ(e.to_usize()), e)).collect(); + + let mut stack = [(0u64, I::zero()); DOUBLING_STACK_GROUP]; + let mut heap: Vec<(u64, I)>; + let keyed: &mut [(u64, I)] = if len <= DOUBLING_STACK_GROUP { + let slot = &mut stack[..len]; + for (dst, &e) in slot.iter_mut().zip(sa_g.iter()) { + *dst = (succ(e.to_usize()), e); + } + slot + } else { + heap = sa_g.iter().map(|&e| (succ(e.to_usize()), e)).collect(); + &mut heap + }; keyed.sort_unstable(); let mut fresh = Vec::new(); let mut i = 0; - while i < keyed.len() { + while i < len { let key = keyed[i].0; let mut j = i + 1; - while j < keyed.len() && keyed[j].0 == key { + while j < len && keyed[j].0 == key { j += 1; } let g = I::from_usize(start + i); @@ -656,7 +685,8 @@ pub(crate) fn build_sa(text: &[u8]) -> Vec { }); let before: usize = groups.iter().map(|&(s, e)| e - s).sum(); - groups = sub.into_iter().flatten().collect(); + let n_groups = groups.len(); + groups = sub; let after: usize = groups.iter().map(|&(s, e)| e - s).sum(); // A doubling round can only ever refine, so `after <= before`. If a @@ -665,6 +695,12 @@ pub(crate) fn build_sa(text: &[u8]) -> Vec { // `depth` grows geometrically and every suffix eventually runs off // the end of the text, which the sentinel orders. Guard against // overflow rather than against non-progress. + profile_log(&format!( + " doubling round depth={depth}: {before} tied in {} groups (avg {:.1}) -> {after} tied, {:.3}s", + n_groups, + before as f64 / n_groups.max(1) as f64, + round_t.elapsed().as_secs_f64() + )); debug_assert!(after <= before); match depth.checked_mul(2) { Some(d) if d <= n.saturating_mul(2) => depth = d, @@ -708,6 +744,18 @@ impl Scatter { } } + /// Borrow `len` elements starting at `index` mutably. + /// + /// # Safety + /// + /// No other live borrow may overlap `index..index + len`, and the borrow + /// the `Scatter` was built from must still be live. + #[inline] + unsafe fn slice_mut<'a>(&self, index: usize, len: usize) -> &'a mut [T] { + debug_assert!(index + len <= self.len); + unsafe { std::slice::from_raw_parts_mut(self.ptr.add(index), len) } + } + /// Write `value` at `index`. /// /// # Safety @@ -721,35 +769,6 @@ impl Scatter { } } -/// Borrow each `(start, end)` range of `sa` and `next_rank` mutably and -/// simultaneously. The ranges come from a scan of `sa` so they are sorted and -/// non-overlapping, which is what makes the repeated `split_at_mut` sound. -fn split_disjoint<'a, I: Index>( - sa: &'a mut [I], - next_rank: &'a mut [I], - groups: &[(usize, usize)], -) -> Vec<(&'a mut [I], &'a mut [I])> { - let mut out = Vec::with_capacity(groups.len()); - let mut sa_rest = sa; - let mut nr_rest = next_rank; - let mut consumed = 0usize; - for &(start, end) in groups { - debug_assert!( - start >= consumed, - "group ranges must be sorted and disjoint" - ); - let (_, sa_tail) = sa_rest.split_at_mut(start - consumed); - let (_, nr_tail) = nr_rest.split_at_mut(start - consumed); - let (sa_g, sa_tail) = sa_tail.split_at_mut(end - start); - let (nr_g, nr_tail) = nr_tail.split_at_mut(end - start); - out.push((sa_g, nr_g)); - sa_rest = sa_tail; - nr_rest = nr_tail; - consumed = end; - } - out -} - #[cfg(test)] mod tests { use super::*; From 3ef11f29244410736e2c00cce1213b884d38c259 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 23:16:28 +0200 Subject: [PATCH 28/41] docs: refresh in-memory numbers after the doubling-round cleanup Co-Authored-By: Claude Opus 5 (1M context) --- README.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index ebcd356..98cdb9a 100644 --- a/README.md +++ b/README.md @@ -47,8 +47,8 @@ the merge kernel's and independently verified: | input | before | after | CPU before | CPU after | | ----- | ------ | ----- | ---------- | --------- | -| chr21 fwd ++ revcomp, ASCII `ACGT`, 80 MB | 6.08 s | **0.76 s** | 28.1 s | 5.0 s | -| same, pre-coded to `0..3`, 80 MB | 6.08 s | **0.73 s** | 28.1 s | 5.0 s | +| chr21 fwd ++ revcomp, ASCII `ACGT`, 80 MB | 6.08 s | **0.61 s** | 28.1 s | 5.0 s | +| same, pre-coded to `0..3`, 80 MB | 6.08 s | **0.57 s** | 28.1 s | 5.0 s | | chr21 FASTA, 47.5 MB, 6.6 Mb of `N` | 27.8 s | **1.04 s** | 283.5 s | 5.2 s | | same, via `build_in_memory_sample_sort` | 3.41 s | **1.18 s** | 34.5 s | 7.2 s | @@ -60,7 +60,7 @@ The two DNA rows above land in the same place because the alphabet is ranked to a dense code range before packing. Without that step the ASCII row would use 8-bit fields — its largest byte is `'T'` (84) even though it has four symbols — fitting 8 symbols per key instead of 32, and would -take 1.40 s rather than 0.76 s. Keys are then built by a SWAR gather +take 1.40 s rather than 0.61 s. Keys are then built by a SWAR gather over the ranked text, so eight symbols cost three shift-or-mask pairs instead of eight dependent shift-or-lookup steps. From 7c780b0c8da67c8ba3ff0aaa0a7554106d42e036 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Wed, 12 Aug 2026 07:46:31 +0200 Subject: [PATCH 29/41] fix: require exact `u8` for packed keys, and stop taxing every LCP call Addresses the blockers @rob-p raised in #7. **Signed symbols.** `seed_params` accepted any `Symbol` one byte wide and reinterpreted it as `u8`, but `Symbol` is implemented for `i8` too and a packed key orders its fields as unsigned: `-1` has byte `0xFF` and sorts above `1`. The in-memory doubling guard already required exact `u8`; the packed-key paths now match it. Rob's reproducer `[-1, 0, -1, 1, -2, 0, 1, -1, 0, -2, 1, 0]` is in the tests, alongside external-memory and sample-sort cases at several partition counts and an `i8::MIN`/`i8::MAX` case. Reverting the guard fails exactly the two end-to-end tests and leaves the in-memory one passing, matching the report that only the packed-key paths were affected. `runs.rs` deliberately keeps its one-byte-wide check: a run table is built from byte *equality*, which coincides with value equality for `i8`, and ordering is recovered by the caller through `S: Ord`. **Run skipping taxed the common path.** Consulting the table cost up to three binary searches before any comparison, and nearly every LCP call in real sequence mismatches within a few symbols and never reaches a run. Probe with the ordinary bounded scan first, and consult the table only once a match has survived 256 symbols. The probe is not wasted: whatever it matches counts towards the answer. Measured on a fixture matching the one in #7 (parsed chr21 forward ++ revcomp, `N` kept as a symbol, ACGT-start filter, 80,177,238 retained positions, 12 threads), with byte-identical output in all three cases: run table disabled 1.97 s table consulted eagerly 2.40 s (+22%) probe first 2.02 s (+2.5%, within noise) And it keeps the benefit where runs matter: the raw FASTA takes 43.8 s with the table disabled against ~1.5 s with it. **Packer built before its guards.** `Packer::new` may materialise a ranked copy of the whole text, and on a segmented or context-bounded build that copy can never be used, since every packed-key path declines those. It is now constructed only after the `plain_lex_len` and `max_context` guards pass. **Docs.** The claim that periods up to 64 also cover satellite arrays was too broad: the canonical alpha-satellite monomer is 171 bases and measures at parity with no detection. Both the module docs and the README now state the measured coverage and its limit. Also adds `--filter-acgt` to the bench CLI, so the STAR-shaped filtered build this all turns on can be benchmarked directly. Co-Authored-By: Claude Opus 5 (1M context) --- README.md | 7 +++- examples/caps_sa.rs | 28 +++++++++++++- src/ext_mem.rs | 93 ++++++++++++++++++++++++++++++++++++++++++++- src/radix.rs | 29 ++++++++++++-- src/runs.rs | 55 +++++++++++++++++++++++---- 5 files changed, 195 insertions(+), 17 deletions(-) diff --git a/README.md b/README.md index 98cdb9a..be331a7 100644 --- a/README.md +++ b/README.md @@ -225,8 +225,11 @@ stop at a run's start rather than traversing it. Detecting only single-symbol runs would miss the case that actually occurs: in wrapped FASTA an `N` block is 60 `N`s followed by a newline, -which is period 61, not period 1. Periods up to 64 are considered, which -also covers satellite arrays. +which is period 61, not period 1. Periods up to 64 are considered. +Measured on synthetic periodic inputs, periods 1, 2, 61 and 64 sort +5.8-6.9x faster, while periods 65 and 171 are not detected and run at +parity, so alpha-satellite arrays (canonical monomer 171 bases) fall +outside the detector. Detection is two-stage so texts without repeats pay almost nothing: a sampling pass collects the periods that occur at all, and the full scan diff --git a/examples/caps_sa.rs b/examples/caps_sa.rs index 67da601..2744945 100644 --- a/examples/caps_sa.rs +++ b/examples/caps_sa.rs @@ -18,7 +18,10 @@ use std::path::PathBuf; use std::process; use std::time::Instant; -use caps_sa::{ExtMemOpts, build_ext_mem, build_in_memory, build_in_memory_sample_sort, verify_sa}; +use caps_sa::{ + ExtMemOpts, build_ext_mem, build_ext_mem_for_filter, build_in_memory, + build_in_memory_sample_sort, verify_sa, +}; struct Args { input: PathBuf, @@ -28,6 +31,7 @@ struct Args { subproblem_count: usize, threads: Option, verify: bool, + filter_acgt: bool, } fn parse_args() -> Args { @@ -38,6 +42,7 @@ fn parse_args() -> Args { let mut subproblem_count: usize = 0; let mut threads: Option = None; let mut verify = false; + let mut filter_acgt = false; let mut i = 1; while i < argv.len() { match argv[i].as_str() { @@ -53,6 +58,13 @@ fn parse_args() -> Args { verify = true; i += 1; } + // Sort only suffixes starting at a symbol below 4, the shape a + // STAR-style genome index uses: A/C/G/T participate, N and + // spacers do not. + "--filter-acgt" => { + filter_acgt = true; + i += 1; + } "--subproblem-count" => { subproblem_count = argv[i + 1] .parse() @@ -95,6 +107,7 @@ fn parse_args() -> Args { subproblem_count, threads, verify, + filter_acgt, } } @@ -174,7 +187,18 @@ fn main() -> std::io::Result<()> { let mut count = 0usize; let mode_label = if args.ext_mem { "ext-mem" } else { "in-mem-ss" }; let build_start = Instant::now(); - if args.ext_mem { + if args.ext_mem && args.filter_acgt { + build_ext_mem_for_filter( + &text, + |p| text[p as usize] < 4, + &opts, + |pos| { + count += 1; + writer.borrow_mut().write_all(&pos.to_le_bytes())?; + Ok(()) + }, + )?; + } else if args.ext_mem { build_ext_mem(&text, &opts, |pos| { count += 1; writer.borrow_mut().write_all(&pos.to_le_bytes())?; diff --git a/src/ext_mem.rs b/src/ext_mem.rs index a00a411..96c232c 100644 --- a/src/ext_mem.rs +++ b/src/ext_mem.rs @@ -505,7 +505,16 @@ where let p = effective_subproblem_count(n, opts.subproblem_count); let runs = crate::runs::detect_for(text); let cmp = Cmp::new(LcpDispatch::detect(), &runs); - let seed_packer = crate::radix::seed_params(text); + // Only build the alphabet map once every eligibility condition holds. + // `Packer::new` may materialise a ranked copy of the whole text, and on a + // segmented or context-bounded build that copy could never be used: the + // packed-key paths all decline those. Constructing it first cost real + // resident memory for nothing. + let seed_packer = if opts.max_context == usize::MAX && lp.plain_lex_len() == Some(text.len()) { + crate::radix::seed_params(text) + } else { + None + }; let seed_params = seed_packer.as_ref(); let work_dir = opts.work_dir.clone(); @@ -630,7 +639,16 @@ where let p = effective_subproblem_count(n, opts.subproblem_count); let runs = crate::runs::detect_for(text); let cmp = Cmp::new(LcpDispatch::detect(), &runs); - let seed_packer = crate::radix::seed_params(text); + // Only build the alphabet map once every eligibility condition holds. + // `Packer::new` may materialise a ranked copy of the whole text, and on a + // segmented or context-bounded build that copy could never be used: the + // packed-key paths all decline those. Constructing it first cost real + // resident memory for nothing. + let seed_packer = if opts.max_context == usize::MAX && lp.plain_lex_len() == Some(text.len()) { + crate::radix::seed_params(text) + } else { + None + }; let seed_params = seed_packer.as_ref(); let factory = |_i: usize| InMemBucket::>::new(); @@ -2024,6 +2042,77 @@ mod tests { use crate::build_in_memory; use tempfile::tempdir; + /// Build with the external-memory path over an arbitrary `Symbol`. + fn ext_mem_sa_of(text: &[S], p: usize) -> Vec { + let dir = tempdir().unwrap(); + let opts = ExtMemOpts { + subproblem_count: p, + work_dir: dir.path().to_path_buf(), + ..ExtMemOpts::default() + }; + let mut out: Vec = Vec::with_capacity(text.len()); + build_ext_mem(text, &opts, |pos| { + out.push(pos); + Ok(()) + }) + .unwrap(); + out + } + + /// Reference order: sort suffixes with the slice comparator, which uses + /// `S`'s own `Ord`. + fn direct_sa(text: &[S]) -> Vec { + let mut sa: Vec = (0..text.len() as u64).collect(); + sa.sort_by(|&a, &b| text[a as usize..].cmp(&text[b as usize..])); + sa + } + + /// `Symbol` is implemented for `i8`, and a packed key orders its fields as + /// unsigned, so `-1` (byte `0xFF`) would sort above `1`. The packed-key + /// paths must decline signed symbols; these cover the two entry points + /// that reach them. + #[test] + fn ext_mem_signed_i8_matches_direct_order() { + let fixtures: Vec> = vec![ + vec![-1, 0, -1, 1, -2, 0, 1, -1, 0, -2, 1, 0], + vec![i8::MIN, 0, i8::MAX, -1, 1, i8::MIN, i8::MAX, 0], + (0..200).map(|i: i32| (i % 7 - 3) as i8).collect(), + ]; + for text in fixtures { + for p in [1usize, 3, 8] { + assert_eq!( + ext_mem_sa_of(&text, p), + direct_sa(&text), + "ext-mem p={p} disagrees on {text:?}" + ); + } + } + } + + #[test] + fn in_memory_sample_sort_signed_i8_matches_direct_order() { + let fixtures: Vec> = vec![ + vec![-1, 0, -1, 1, -2, 0, 1, -1, 0, -2, 1, 0], + (0..300).map(|i: i32| (i % 5 - 2) as i8).collect(), + ]; + for text in fixtures { + let mut out: Vec = Vec::new(); + build_in_memory_sample_sort(&text, &ExtMemOpts::default(), |pos| { + out.push(pos); + Ok(()) + }) + .unwrap(); + assert_eq!(out, direct_sa(&text), "sample sort disagrees on {text:?}"); + } + } + + #[test] + fn in_memory_signed_i8_matches_direct_order() { + let text: Vec = vec![-1, 0, -1, 1, -2, 0, 1, -1, 0, -2, 1, 0]; + let got: Vec = crate::build_in_memory(&text); + assert_eq!(got, direct_sa(&text)); + } + fn ext_mem_sa(text: &[u8], p: usize) -> Vec { let dir = tempdir().unwrap(); let opts = ExtMemOpts { diff --git a/src/radix.rs b/src/radix.rs index 96fc9dc..0c1c4ce 100644 --- a/src/radix.rs +++ b/src/radix.rs @@ -245,11 +245,17 @@ impl Packer { /// Computed once per build and handed to [`seed_subarray`], which would /// otherwise re-scan the whole text for every subarray. pub(crate) fn seed_params(text: &[S]) -> Option { - if size_of::() != 1 { + // Exactly `u8`, not merely one byte wide. `Symbol` is implemented for + // `i8` too, and a packed key orders its fields as unsigned: `-1` has byte + // `0xFF` and would sort above `1`, inverting the text's real order. The + // in-memory doubling guard already required exact `u8`; the packed-key + // paths must match it. + if std::any::TypeId::of::() != std::any::TypeId::of::() { return None; } - // SAFETY: `S` is one byte wide, so a byte view over the same memory is - // valid for reads of the same length. + // SAFETY: `S` is `u8` (just checked by `TypeId`, and `Symbol: 'static` so + // the comparison is exact), so a byte view over the same memory is valid + // for reads of the same length. let bytes: &[u8] = unsafe { std::slice::from_raw_parts(text.as_ptr() as *const u8, text.len()) }; Some(Packer::new(bytes)) @@ -301,7 +307,7 @@ pub(crate) fn seed_subarray( } return true; } - // SAFETY: `params` is `Some` only when `S` is one byte wide. + // SAFETY: `packer` is `Some` only when `S` is exactly `u8`. let bytes: &[u8] = unsafe { std::slice::from_raw_parts(text.as_ptr() as *const u8, text.len()) }; @@ -784,6 +790,21 @@ mod tests { assert_eq!(got, brute(text), "mismatch on {text:?}"); } + /// `Symbol` is implemented for `i8`, and a one-byte-wide check alone lets a + /// signed text through a packer that orders bytes as unsigned. `-1` has + /// byte `0xFF`, so it would sort above `1`, inverting the true order. + #[test] + fn signed_symbols_are_not_eligible_for_packing() { + let text: Vec = vec![-1, 0, -1, 1, -2, 0, 1, -1, 0, -2, 1, 0]; + assert!( + seed_params(&text).is_none(), + "i8 texts must not get a packed key" + ); + // u8 of the same width still qualifies. + let bytes: Vec = vec![1, 0, 1, 2, 3, 0]; + assert!(seed_params(&bytes).is_some()); + } + #[test] fn fixtures() { check(b""); diff --git a/src/runs.rs b/src/runs.rs index 8224451..64263fe 100644 --- a/src/runs.rs +++ b/src/runs.rs @@ -8,9 +8,16 @@ //! //! Genome assemblies always contain these. An `N` block is the obvious case, //! and note that in wrapped FASTA it is **not** a run of one symbol: 60 `N`s -//! followed by a newline is a run of period 61. Satellite arrays are runs of -//! period 171. So detecting only single-symbol runs would miss the case that -//! actually shows up. +//! followed by a newline is a run of period 61. So detecting only +//! single-symbol runs would miss the representation that actually shows up. +//! +//! The detector covers periods up to [`MAX_PERIOD`] and **not** beyond. +//! Measured on synthetic 1 MiB periodic inputs, periods 1, 2, 61 and 64 are +//! detected with full coverage and sort 5.8-6.9x faster; periods 65 and 171 +//! are not detected at all and run at parity. Alpha-satellite arrays, whose +//! canonical monomer is 171 bases, are therefore *outside* this detector. +//! The sampling stage can also miss a run that is localised enough to fall +//! between its windows. //! //! The observation that makes this cheap: if `text[s..e)` has period `q`, and //! two suffixes start at `a < b` inside it with `(b - a) % q == 0`, then they @@ -45,14 +52,27 @@ use std::cmp::Ordering; /// the run faster than the binary search that would find it. const MIN_RUN: usize = 1024; -/// Longest period considered. Covers homopolymers (1), wrapped-FASTA `N` -/// blocks (61), and alpha-satellite monomers (171 exceeds this, but a -/// satellite array is also periodic at shorter scales in practice). +/// Longest period considered. Covers homopolymers (period 1) and +/// wrapped-FASTA `N` blocks (period 61), which are the cases that occur in +/// practice in assembly FASTA. +/// +/// It does **not** cover alpha-satellite arrays: their canonical monomer is +/// 171 bases, and a synthetic period-171 input measures at parity with no +/// detection at all. Raising this is a constant change, but the detection +/// scan is `O(periods x n)`, so it is not free. const MAX_PERIOD: usize = 64; /// Window used by the sampling pass to decide whether a period occurs at all. const SAMPLE_WINDOW: usize = 512; +/// Symbols an ordinary scan must match before the run table is consulted. +/// +/// Two suffixes of real sequence that agree this far are already unusual, so +/// the table is reached only when it might actually help. Small enough that +/// the probe is a handful of vector compares, and the probe is not wasted +/// work: whatever it matches counts towards the answer. +const RUN_PROBE: usize = 256; + /// A maximal stretch `[start, end)` of the text with period `period`, meaning /// `text[i] == text[i + period]` for every `i` in `start..end - period`. #[derive(Copy, Clone, Debug, PartialEq, Eq)] @@ -247,12 +267,33 @@ impl<'a> Cmp<'a> { /// /// With an empty run table this is exactly [`LcpDispatch::lcp`] plus one /// predictable branch. + /// + /// When a table *is* present, the ordinary bounded scan still runs first. + /// Consulting the table costs up to three binary searches before any + /// comparison happens, and the overwhelming majority of LCP calls in real + /// sequence mismatch within a few symbols and never reach a run at all. + /// Paying the lookup up front taxed every one of them: on a filtered, + /// `N`-containing chr21 that alone turned a 1.80 s build into 2.72 s even + /// though the answers were identical. + /// + /// So: probe first, and only once a match has survived [`RUN_PROBE`] + /// symbols — which ordinary genomic difference does not — is it worth + /// asking whether a run explains it. #[inline] pub(crate) fn lcp(&self, text: &[S], p: usize, q: usize, max_ctx: usize) -> usize { if self.runs.is_empty() || size_of::() != 1 { return self.dispatch.lcp(text, p, q, max_ctx); } - let mut i = 0usize; + + let probe = max_ctx.min(RUN_PROBE); + let got = self.dispatch.lcp(text, p, q, probe); + if got < probe || probe == max_ctx { + // Either a real mismatch, or the caller's bound was reached. No + // run can extend this, so the table is never touched. + return got; + } + + let mut i = got; while i < max_ctx { let jump = self.runs.skip(p + i, q + i); if jump > 0 { From 226e65ca27e513e36d221f994a54bbc61e9b6ad4 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Wed, 12 Aug 2026 07:55:35 +0200 Subject: [PATCH 30/41] feat: packed keys for segmented texts, including STAR's boundary order MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The packed-key paths declined every segmented build, which is the comparator a splice-junction index actually uses, so none of this work reached it. The obstacle was that zero-padding a short suffix encodes shorter-is-smaller, while STAR's `boundary_order` is `lim_b.cmp(&lim_a).then(p_a.cmp(&p_b))` — longer-is-smaller. That is expressible in the key once the provider says which side it wants. `LimitProvider::boundary_rank` returns `ShorterFirst` or `LongerFirst`, defaulting to `None` so existing implementations keep today's behaviour. A segmented key then packs `min(k, lim_at(p))` symbols — never reading into the next segment — and pads with a reserved sentinel below every real code under `ShorterFirst`, above every real code under `LongerFirst`. The position tie-break does not need to be in the key. Keys that tie defer to `boundary_order` itself, so the key only has to avoid contradicting the convention, never to reproduce it. The sentinel needs one code above the alphabet, so a segmented build sizes its field to hold `alphabet` rather than `alphabet - 1`, and `Packer::new` takes that as a parameter: a plain build must not pay for it, since widening the field halves the symbols a key carries. Where no spare code fits (256 distinct symbols in 8-bit fields) the path declines. Tests compare against the provider's own comparator over random segmented texts on the ruSTAR alphabet, under both conventions, with and without the ACGT-start filter, at three partition counts. Two notes on how those tests are written, both learned the hard way: - They assert the *property* (permutation, and no adjacent pair out of order) rather than equality with a canonical answer. `SegmentedText`'s default `boundary_order` returns `Equal` for suffixes that end together with equal content, so their relative order is genuinely free and a stable-sort oracle is not a valid reference. - Sensitivity was checked by breaking `key_at_bounded` deliberately. A *constant* wrong key still passes, because it collapses everything into one tied group that the fallback sorts correctly — so that check proves nothing. An order-*inverting* key does fail the test, which is what establishes the fast path is both taken and verified. Not benchmarked end to end on an annotation-shaped fixture: that needs rustar-aligner's junction pipeline, which I cannot run here. @rob-p, if the GENCODE harness from #7 is available this is the change that should finally move phase 1 on it. Co-Authored-By: Claude Opus 5 (1M context) --- src/ext_mem.rs | 179 ++++++++++++++++++++++++++++++++++++++++++++++--- src/lib.rs | 2 +- src/limits.rs | 57 ++++++++++++++++ src/radix.rs | 107 ++++++++++++++++++++++++----- 4 files changed, 318 insertions(+), 27 deletions(-) diff --git a/src/ext_mem.rs b/src/ext_mem.rs index 96c232c..a2da8fb 100644 --- a/src/ext_mem.rs +++ b/src/ext_mem.rs @@ -510,11 +510,16 @@ where // segmented or context-bounded build that copy could never be used: the // packed-key paths all decline those. Constructing it first cost real // resident memory for nothing. - let seed_packer = if opts.max_context == usize::MAX && lp.plain_lex_len() == Some(text.len()) { - crate::radix::seed_params(text) - } else { - None - }; + let plain_text = lp.plain_lex_len() == Some(text.len()); + let seed_packer = + if opts.max_context == usize::MAX && (plain_text || lp.boundary_rank().is_some()) { + // A segmented build needs one spare code for the boundary sentinel; + // a plain one must not pay for it, since widening the field halves + // the symbols a key carries. + crate::radix::seed_params(text, !plain_text) + } else { + None + }; let seed_params = seed_packer.as_ref(); let work_dir = opts.work_dir.clone(); @@ -644,11 +649,16 @@ where // segmented or context-bounded build that copy could never be used: the // packed-key paths all decline those. Constructing it first cost real // resident memory for nothing. - let seed_packer = if opts.max_context == usize::MAX && lp.plain_lex_len() == Some(text.len()) { - crate::radix::seed_params(text) - } else { - None - }; + let plain_text = lp.plain_lex_len() == Some(text.len()); + let seed_packer = + if opts.max_context == usize::MAX && (plain_text || lp.boundary_rank().is_some()) { + // A segmented build needs one spare code for the boundary sentinel; + // a plain one must not pay for it, since widening the field halves + // the symbols a key carries. + crate::radix::seed_params(text, !plain_text) + } else { + None + }; let seed_params = seed_packer.as_ref(); let factory = |_i: usize| InMemBucket::>::new(); @@ -2040,6 +2050,7 @@ impl CascadeWorkspace { mod tests { use super::*; use crate::build_in_memory; + use crate::limits::SegmentedText; use tempfile::tempdir; /// Build with the external-memory path over an arbitrary `Symbol`. @@ -2113,6 +2124,154 @@ mod tests { assert_eq!(got, direct_sa(&text)); } + /// A `LimitProvider` with STAR's spacer-as-largest convention: the suffix + /// that reaches its boundary first is *larger*, with an ascending-position + /// tie-break. This is the comparator a splice-junction index uses. + struct StarConvention { + inner: SegmentedText, + } + + impl LimitProvider for StarConvention { + fn lim_at(&self, p: usize) -> usize { + self.inner.lim_at(p) + } + fn boundary_order( + &self, + p_a: usize, + lim_a: usize, + p_b: usize, + lim_b: usize, + ) -> std::cmp::Ordering { + lim_b.cmp(&lim_a).then(p_a.cmp(&p_b)) + } + fn boundary_rank(&self) -> Option { + Some(crate::limits::BoundaryRank::LongerFirst) + } + } + + /// The provider's own comparator, spelled out. + fn direct_cmp(text: &[u8], lp: &L, a: u64, b: u64) -> std::cmp::Ordering { + let (pa, pb) = (a as usize, b as usize); + let (la, lb) = (lp.lim_at(pa), lp.lim_at(pb)); + for i in 0..la.min(lb) { + if text[pa + i] != text[pb + i] { + return text[pa + i].cmp(&text[pb + i]); + } + } + lp.boundary_order(pa, la, pb, lb) + } + + /// Assert `got` is a valid sort of `positions` under `lp`'s comparator. + /// + /// Checked as a property rather than against a canonical permutation: + /// `SegmentedText`'s default `boundary_order` returns `Equal` for suffixes + /// that end together with equal content, so their relative order is free + /// and no single answer is "the" right one. The merge kernel is not a + /// stable sort either. + fn assert_sorted_under( + text: &[u8], + lp: &L, + positions: &[u64], + got: &[u64], + what: &str, + ) { + let mut want = positions.to_vec(); + want.sort_unstable(); + let mut have = got.to_vec(); + have.sort_unstable(); + assert_eq!(have, want, "{what}: not a permutation of the input"); + for w in got.windows(2) { + assert_ne!( + direct_cmp(text, lp, w[0], w[1]), + std::cmp::Ordering::Greater, + "{what}: {} precedes {} but compares greater", + w[0], + w[1], + ); + } + } + + fn ext_mem_sa_with( + text: &[u8], + lp: &L, + p: usize, + positions: Vec, + ) -> Vec { + let dir = tempdir().unwrap(); + let opts = ExtMemOpts { + subproblem_count: p, + work_dir: dir.path().to_path_buf(), + ..ExtMemOpts::default() + }; + let mut out: Vec = Vec::with_capacity(positions.len()); + build_ext_mem_for_positions_with(text, positions, lp, &opts, |pos| { + out.push(pos); + Ok(()) + }) + .unwrap(); + out + } + + /// Segmented texts now get packed keys too: the key stops at the segment + /// boundary and pads with a reserved sentinel on the side the provider's + /// `boundary_order` demands. These check both conventions, and the + /// ACGT-start filter a splice-junction index applies, against the direct + /// comparator. + #[test] + fn ext_mem_segmented_matches_direct_comparator() { + use rand::{RngExt, SeedableRng}; + let mut rng = rand::rngs::StdRng::seed_from_u64(0x5D1CE); + for trial in 0..12 { + let n_seg = rng.random_range(2..8usize); + let lengths: Vec = (0..n_seg).map(|_| rng.random_range(4..60usize)).collect(); + let n: usize = lengths.iter().sum(); + // A/C/G/T/N plus a spacer code, the ruSTAR alphabet. + let text: Vec = (0..n).map(|_| rng.random_range(0..6u8)).collect(); + let seg = SegmentedText::from_lengths(n, &lengths); + let all: Vec = (0..n as u64).collect(); + let acgt: Vec = all + .iter() + .copied() + .filter(|&p| text[p as usize] < 4) + .collect(); + + for p in [1usize, 2, 5] { + assert_sorted_under( + &text, + &seg, + &all, + &ext_mem_sa_with(&text, &seg, p, all.clone()), + &format!("shorter-first, all positions, trial {trial} p={p}"), + ); + assert_sorted_under( + &text, + &seg, + &acgt, + &ext_mem_sa_with(&text, &seg, p, acgt.clone()), + &format!("shorter-first, ACGT filter, trial {trial} p={p}"), + ); + + let star = StarConvention { + inner: SegmentedText::from_lengths(n, &lengths), + }; + assert_sorted_under( + &text, + &star, + &all, + &ext_mem_sa_with(&text, &star, p, all.clone()), + &format!("longer-first, all positions, trial {trial} p={p}"), + ); + assert_sorted_under( + &text, + &star, + &acgt, + &ext_mem_sa_with(&text, &star, p, acgt.clone()), + &format!("longer-first, ACGT filter, trial {trial} p={p}"), + ); + } + } + } + fn ext_mem_sa(text: &[u8], p: usize) -> Vec { let dir = tempdir().unwrap(); let opts = ExtMemOpts { diff --git a/src/lib.rs b/src/lib.rs index 792fbaa..c2f9a9f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -35,7 +35,7 @@ pub use ext_mem::{ try_build_in_memory_sample_sort_for_positions_with, try_build_in_memory_sample_sort_with, }; pub use lcp::{LcpDispatch, Symbol, lcp, lcp_scalar, lcp_u8, suffix_cmp}; -pub use limits::{LimitProvider, PlainText, SegmentedText}; +pub use limits::{BoundaryRank, LimitProvider, PlainText, SegmentedText}; pub use sample_sort::{ Opts, build_in_memory, build_in_memory_for_positions, build_in_memory_for_positions_with, build_in_memory_for_positions_with_opts, build_in_memory_with, build_in_memory_with_opts, diff --git a/src/limits.rs b/src/limits.rs index b9ff0ba..608660e 100644 --- a/src/limits.rs +++ b/src/limits.rs @@ -24,6 +24,24 @@ //! rationale and the comparison against the `[u8; 3]` (24-bit-text) //! alternative. +/// How a provider's [`boundary_order`][LimitProvider::boundary_order] ranks a +/// suffix that ends at its segment boundary against one that keeps going. +/// +/// This is the one fact a fixed-depth packed key needs in order to represent +/// a segmented comparator: a key pads a short suffix out to full width, and +/// the padding symbol has to fall on the correct side of every real symbol. +#[derive(Copy, Clone, Debug, PartialEq, Eq)] +pub enum BoundaryRank { + /// The suffix that ends first is smaller, the standard generalised-SA + /// convention and the default `boundary_order`. Keys pad with a symbol + /// below every real one. + ShorterFirst, + /// The suffix that ends first is *larger*, equivalently the longer one is + /// smaller. STAR's spacer-as-largest ordering. Keys pad with a symbol + /// above every real one. + LongerFirst, +} + /// Per-suffix length provider. The merge and cascade-merge code use /// `lp.lim_at(p)` instead of `text.len() - p`; the LCP function itself /// is unchanged (the merge passes the appropriately-capped @@ -96,6 +114,33 @@ pub trait LimitProvider: Sync { fn plain_lex_len(&self) -> Option { None } + + /// Which side of a real symbol this provider's boundary convention puts + /// the end of a suffix on, when that convention is expressible. + /// + /// Returning `Some` lets the crate build packed fixed-depth keys for a + /// *segmented* text: a key packs `min(k, lim_at(p))` symbols and pads the + /// rest with a sentinel placed according to this answer, so key order + /// agrees with `boundary_order` whenever the key decides at all. Keys that + /// tie still defer to `boundary_order` itself, which is what lets a + /// convention with a position tie-break (STAR's + /// `lim_b.cmp(&lim_a).then(p_a.cmp(&p_b))`) work: the key never has to + /// express the tie-break, only never to contradict it. + /// + /// The contract is exactly this: for suffixes `a` and `b` whose shared + /// prefix ends because one of them reached its limit, + /// `boundary_order(a, .., b, ..)` must be `Less` iff `a` is the one that + /// ended, under `ShorterFirst`, and `Greater` iff `a` is the one that + /// ended, under `LongerFirst`. + /// + /// The default is `None`, which keeps every existing implementation on the + /// comparison path. Answer it only if your `boundary_order` decides purely + /// by which suffix ended first, with at most a tie-break between suffixes + /// that end at the same offset. + #[inline] + fn boundary_rank(&self) -> Option { + None + } } /// Default provider for non-segmented texts: `lim_at(p) = n - p`. @@ -127,6 +172,11 @@ impl LimitProvider for PlainText { fn plain_lex_len(&self) -> Option { Some(self.n) } + + #[inline] + fn boundary_rank(&self) -> Option { + Some(BoundaryRank::ShorterFirst) + } } /// Provider for texts partitioned into segments at known cumulative @@ -232,6 +282,13 @@ impl LimitProvider for SegmentedText { self.n - p } } + + /// `SegmentedText` keeps the default `boundary_order`, which is + /// shorter-is-smaller. + #[inline] + fn boundary_rank(&self) -> Option { + Some(BoundaryRank::ShorterFirst) + } } #[cfg(test)] diff --git a/src/radix.rs b/src/radix.rs index 0c1c4ce..4a90f5f 100644 --- a/src/radix.rs +++ b/src/radix.rs @@ -60,7 +60,7 @@ use crate::Index; use crate::ext_mem::profile_log; use crate::lcp::Symbol; -use crate::limits::LimitProvider; +use crate::limits::{BoundaryRank, LimitProvider}; use crate::runs::Cmp; use crate::sample_sort; use rayon::prelude::*; @@ -93,11 +93,14 @@ pub(crate) struct Packer { bits: u32, /// Symbols per `u64` key. k: usize, + /// Number of distinct codes in use. Codes are `0..alphabet`; `alphabet` + /// itself is free for use as a boundary sentinel when it fits the field. + alphabet: u32, } impl Packer { /// Build the map for `text`. - fn new(text: &[u8]) -> Self { + fn new(text: &[u8], need_sentinel: bool) -> Self { // Which bytes occur? One parallel pass, folded into a 256-entry set. let present = text .par_chunks(1 << 16) @@ -128,7 +131,16 @@ impl Packer { next += 1; } } - let bits: u32 = match next.saturating_sub(1) { + // A segmented key needs one code above the alphabet for its boundary + // sentinel, so the field must hold `alphabet`, not `alphabet - 1`. + // Plain builds do not pay for that: widening the field would halve the + // symbols per key, which is the whole point of packing. + let widest = if need_sentinel { + next + } else { + next.saturating_sub(1) + }; + let bits: u32 = match widest { 0..=1 => 1, 2..=3 => 2, 4..=15 => 4, @@ -153,6 +165,7 @@ impl Packer { ranked, bits, k: 64 / bits as usize, + alphabet: next as u32, } } @@ -199,6 +212,54 @@ impl Packer { } } + /// Whether a boundary sentinel fits alongside the alphabet in one field. + /// + /// With 8-bit fields and 256 distinct symbols there is no spare code, so + /// segmented keys are unavailable and the caller must fall back. + #[inline] + pub(crate) fn has_sentinel(&self) -> bool { + (self.alphabet as u64) < (1u64 << self.bits) + } + + /// Pack the `min(k, lim)` symbols at `text[p..]`, padding the rest with a + /// boundary sentinel placed according to `rank`. + /// + /// This is the segmented counterpart to [`Self::key_at`]. It never reads + /// past `p + lim`, so a key cannot see into the next segment, and the + /// sentinel falls below every real code under + /// [`BoundaryRank::ShorterFirst`] and above every real code under + /// [`BoundaryRank::LongerFirst`]. That is what makes key order agree with + /// the provider's `boundary_order` whenever the key decides at all. + /// + /// Under `ShorterFirst` the sentinel is code `0` and every real code is + /// shifted up by one, so a padded field is strictly below any real symbol. + /// Under `LongerFirst` the sentinel is `alphabet`, strictly above every + /// real code, and no shift is needed. + #[inline] + pub(crate) fn key_at_bounded( + &self, + text: &[u8], + p: usize, + lim: usize, + rank: BoundaryRank, + ) -> u64 { + debug_assert!(self.has_sentinel()); + let src = self.ranked.as_deref().unwrap_or(text); + let take = self.k.min(lim).min(src.len() - p); + let (bias, pad) = match rank { + BoundaryRank::ShorterFirst => (1u64, 0u64), + BoundaryRank::LongerFirst => (0u64, self.alphabet as u64), + }; + let mut key = 0u64; + for &c in &src[p..p + take] { + key = (key << self.bits) | (c as u64 + bias); + } + for _ in take..self.k { + key = (key << self.bits) | pad; + } + key + } + /// Pack the `k` symbols at `text[p..]` into one order-preserving `u64`, /// zero-padding past the end of the text. #[inline] @@ -244,7 +305,7 @@ impl Packer { /// /// Computed once per build and handed to [`seed_subarray`], which would /// otherwise re-scan the whole text for every subarray. -pub(crate) fn seed_params(text: &[S]) -> Option { +pub(crate) fn seed_params(text: &[S], need_sentinel: bool) -> Option { // Exactly `u8`, not merely one byte wide. `Symbol` is implemented for // `i8` too, and a packed key orders its fields as unsigned: `-1` has byte // `0xFF` and would sort above `1`, inverting the text's real order. The @@ -258,7 +319,7 @@ pub(crate) fn seed_params(text: &[S]) -> Option { // for reads of the same length. let bytes: &[u8] = unsafe { std::slice::from_raw_parts(text.as_ptr() as *const u8, text.len()) }; - Some(Packer::new(bytes)) + Some(Packer::new(bytes, need_sentinel)) } /// Sort `sa` into suffix order and fill `lcp`, using a packed fixed-depth key @@ -297,9 +358,20 @@ pub(crate) fn seed_subarray( return false; }; let (bits, k) = (packer.bits(), packer.k()); - if max_ctx != usize::MAX || lp.plain_lex_len() != Some(text.len()) { + if max_ctx != usize::MAX { return false; } + // Plain text keys pad past end-of-text and need the visible-length + // tie-break, because a real `0` symbol is indistinguishable from padding. + // A segmented text instead stops the key at `lim_at(p)` and pads with a + // reserved sentinel placed on the side the provider's `boundary_order` + // demands, which encodes the boundary directly and needs no tie-break. + let plain = lp.plain_lex_len() == Some(text.len()); + let seg_rank = match (plain, lp.boundary_rank()) { + (true, _) => None, + (false, Some(r)) if packer.has_sentinel() => Some(r), + _ => return false, + }; let len = sa.len(); if len < 2 { if len == 1 { @@ -320,7 +392,10 @@ pub(crate) fn seed_subarray( .iter() .map(|&p| { let p = p.to_usize(); - (packer.key_at(bytes, p), visible(p) as u32, p_as(p)) + match seg_rank { + None => (packer.key_at(bytes, p), visible(p) as u32, p_as(p)), + Some(r) => (packer.key_at_bounded(bytes, p, lp.lim_at(p), r), 0, p_as(p)), + } }) .collect(); keyed.sort_unstable(); @@ -524,7 +599,7 @@ pub(crate) fn build_sa(text: &[u8]) -> Vec { } let t0 = Instant::now(); - let packer = Packer::new(text); + let packer = Packer::new(text, false); let k = packer.k(); // ---- Seed: sort by the first `k` symbols, then by visible length. ---- @@ -797,12 +872,12 @@ mod tests { fn signed_symbols_are_not_eligible_for_packing() { let text: Vec = vec![-1, 0, -1, 1, -2, 0, 1, -1, 0, -2, 1, 0]; assert!( - seed_params(&text).is_none(), + seed_params(&text, false).is_none(), "i8 texts must not get a packed key" ); // u8 of the same width still qualifies. let bytes: Vec = vec![1, 0, 1, 2, 3, 0]; - assert!(seed_params(&bytes).is_some()); + assert!(seed_params(&bytes, false).is_some()); } #[test] @@ -820,18 +895,18 @@ mod tests { /// fitting only 8 symbols per key instead of 16. #[test] fn packer_width_follows_alphabet_size_not_byte_value() { - let two = Packer::new(b"abababab"); + let two = Packer::new(b"abababab", false); assert_eq!((two.bits(), two.k()), (1, 64)); - let four = Packer::new(&[0u8, 1, 2, 3, 3, 2, 1, 0]); + let four = Packer::new(&[0u8, 1, 2, 3, 3, 2, 1, 0], false); assert_eq!((four.bits(), four.k()), (2, 32)); let mut fasta: Vec = b"ACGTN".to_vec(); fasta.push(b'\n'); - let f = Packer::new(&fasta); + let f = Packer::new(&fasta, false); assert_eq!((f.bits(), f.k()), (4, 16), "6 symbols should pack 4 bits"); let dense: Vec = (0..=255u8).collect(); - let d = Packer::new(&dense); + let d = Packer::new(&dense, false); assert_eq!((d.bits(), d.k()), (8, 8)); } @@ -851,7 +926,7 @@ mod tests { let mut ascending: Vec = text.clone(); ascending.sort_unstable(); ascending.dedup(); - let p = Packer::new(&ascending); + let p = Packer::new(&ascending, false); let keys: Vec = (0..ascending.len()) .map(|i| p.key_at(&ascending, i)) .collect(); @@ -876,7 +951,7 @@ mod tests { for &sigma in &[2u8, 4, 16, 200] { for &n in &[1usize, 7, 8, 9, 31, 32, 33, 63, 64, 65, 200] { let text: Vec = (0..n).map(|_| rng.random_range(0..sigma)).collect(); - let p = Packer::new(&text); + let p = Packer::new(&text, false); let (bits, k) = (p.bits(), p.k()); let ranked = p.ranked.as_deref().unwrap_or(&text); for pos in 0..n { From ca63c5877898ca5eae7c24022cb3d9a554fdaf9f Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Wed, 12 Aug 2026 08:14:29 +0200 Subject: [PATCH 31/41] bench: annotation-shaped splice-junction fixture, and run CI on every branch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Measures the case @rob-p correctly identified as the one none of this work reached: segmented text, STAR's spacer-as-largest `boundary_order`, ACGT-start filter, external-memory construction. `bench/gsj_fixture.py` builds the fixture from a FASTA and a GTF the way a STAR index is laid out — genome sequence, then one 2*overhang flank per distinct junction, then the reverse complement — and `gsj_bench` runs it through `build_ext_mem_for_positions_with` under a `LongerFirst` provider. Apple M4 Max, 12 threads, chr21 with its GENCODE annotation: 95,409,966 symbols, 9,952 segments, 82,167,238 retained ACGT-start positions. segmented keys declined (previous behaviour) 3.84 s segmented keys enabled 2.57 s -33% Output order verified directly against the provider's comparator. This is a *shape* reproduction, not a size one. A genome-wide annotation yields far more junctions than a single chromosome's: this fixture has 4,975 junctions and 9,952 segments against the 698,597 and 1,397,196 of the GENCODE v50 primary-assembly fixture in #7. The segmented-key path should matter more, not less, as the segment count grows, but that is a prediction and not something these numbers establish. Also runs CI on every branch rather than only `main`. A pull request opened from a fork by a first-time contributor does not run workflows until a maintainer approves them, which is why the API showed no check runs at the pinned tips. Building on push makes the fork produce evidence that can be linked from the PR. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/ci.yml | 7 +- Cargo.toml | 4 + bench/gsj_fixture.py | 90 ++++++++++++++++++ examples/gsj_bench.rs | 194 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 294 insertions(+), 1 deletion(-) create mode 100755 bench/gsj_fixture.py create mode 100644 examples/gsj_bench.rs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 0567d90..34f36e3 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -1,9 +1,14 @@ name: CI on: + # Every branch, not just `main`. A pull request opened from a fork by a + # first-time contributor does not run workflows until a maintainer approves + # them, so `pull_request` alone leaves reviewers with no check runs to look + # at. Building on push means the contributor's own fork produces evidence + # that can be linked from the PR. push: - branches: [main] pull_request: + workflow_dispatch: env: CARGO_TERM_COLOR: always diff --git a/Cargo.toml b/Cargo.toml index 4827078..a6e777a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -39,3 +39,7 @@ rand = "0.10" [[example]] name = "caps_sa" path = "examples/caps_sa.rs" + +[[example]] +name = "gsj_bench" +path = "examples/gsj_bench.rs" diff --git a/bench/gsj_fixture.py b/bench/gsj_fixture.py new file mode 100755 index 0000000..9dcadc2 --- /dev/null +++ b/bench/gsj_fixture.py @@ -0,0 +1,90 @@ +#!/usr/bin/env python3 +"""Build an annotation-shaped splice-junction fixture for `gsj_bench`. + + bench/gsj_fixture.py [overhang] + +Writes `.text` (one byte per symbol, A/C/G/T/N as 0..=4) and +`.seg` (packed little-endian u64 segment lengths). + +The layout mirrors what a STAR-style index constructs: the genome sequence, +then one 2*overhang flank per distinct splice junction, then the reverse +complement of the whole thing. Each junction flank is its own segment, which +is what makes the comparator segmented; the genome is one segment per record. + +Note this is a *shape* reproduction. A genome-wide annotation yields far more +junctions than a single-chromosome one, so segment counts differ accordingly. +""" +import collections +import struct +import sys + +CODE = {"A": 0, "C": 1, "G": 2, "T": 3, "N": 4} +COMP = {0: 3, 1: 2, 2: 1, 3: 0, 4: 4} + + +def main() -> None: + if len(sys.argv) not in (4, 5): + sys.exit(__doc__) + fasta, gtf, prefix = sys.argv[1:4] + overhang = int(sys.argv[4]) if len(sys.argv) == 5 else 100 + + seq = bytearray() + with open(fasta) as fh: + for line in fh: + if line.startswith(">"): + continue + for ch in line.strip().upper(): + seq.append(CODE.get(ch, 4)) + + transcripts = collections.defaultdict(list) + with open(gtf) as fh: + for line in fh: + if line.startswith("#"): + continue + f = line.split("\t") + if len(f) < 9 or f[2] != "exon": + continue + i = f[8].find('transcript_id "') + if i < 0: + continue + tid = f[8][i + 15 : f[8].find('"', i + 15)] + transcripts[tid].append((int(f[3]) - 1, int(f[4]))) + + junctions = set() + for exons in transcripts.values(): + exons.sort() + for k in range(len(exons) - 1): + donor, acceptor = exons[k][1], exons[k + 1][0] + if acceptor > donor: + junctions.add((donor, acceptor)) + + flanks = bytearray() + seglens = [len(seq)] + for donor, acceptor in sorted(junctions): + left = seq[max(0, donor - overhang) : donor] + right = seq[acceptor : acceptor + overhang] + flanks += left + bytearray([4] * (overhang - len(left))) + flanks += right + bytearray([4] * (overhang - len(right))) + seglens.append(2 * overhang) + + forward = seq + flanks + text = forward + bytearray(COMP[b] for b in reversed(forward)) + seglens = seglens + list(reversed(seglens)) + assert sum(seglens) == len(text) + + with open(f"{prefix}.text", "wb") as out: + out.write(bytes(text)) + with open(f"{prefix}.seg", "wb") as out: + out.write(struct.pack(f"<{len(seglens)}Q", *seglens)) + + counts = collections.Counter(text) + print( + f"{prefix}: {len(text)} symbols, {len(seglens)} segments, " + f"{len(junctions)} junctions, " + f"{sum(v for k, v in counts.items() if k < 4)} ACGT-start positions", + file=sys.stderr, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/gsj_bench.rs b/examples/gsj_bench.rs new file mode 100644 index 0000000..55daddc --- /dev/null +++ b/examples/gsj_bench.rs @@ -0,0 +1,194 @@ +//! Benchmark the annotation-shaped, splice-junction index build. +//! +//! This is the shape a STAR-style genome index actually constructs, and it is +//! the one none of the packed-key work reached until segmented keys existed: +//! +//! * the text is **segmented** — one segment per chromosome plus one per +//! splice-junction flank — so LCP comparisons stop at segment boundaries; +//! * the comparator is STAR's **spacer-as-largest** `boundary_order`, in which +//! the suffix that reaches its boundary first is the *larger* one, with an +//! ascending-position tie-break; +//! * only **ACGT-starting** positions participate, so no suffix beginning +//! inside an `N` block enters the sort at all; +//! * construction goes through the **external-memory** path. +//! +//! Usage: +//! +//! ```text +//! gsj_bench [--threads N] [--plain] [--verify] +//! ``` +//! +//! `` is one byte per symbol with A/C/G/T/N coded `0..=4`. +//! `` is a packed little-endian `u64[]` summing to the text +//! length. `bench/gsj_fixture.py` builds both from a FASTA and a GTF. +//! +//! `--plain` swaps the segmented provider for `PlainText`, which is the +//! comparison worth having: it shows what the same positions cost when the +//! segmented comparator is not required. + +use std::cmp::Ordering; +use std::env; +use std::fs; +use std::path::PathBuf; +use std::process; +use std::time::Instant; + +use caps_sa::{ + BoundaryRank, ExtMemOpts, LimitProvider, PlainText, SegmentedText, + build_ext_mem_for_positions_with, +}; + +/// STAR's convention: whichever suffix hits its boundary first is larger. +struct StarConvention { + inner: SegmentedText, +} + +impl LimitProvider for StarConvention { + #[inline] + fn lim_at(&self, p: usize) -> usize { + self.inner.lim_at(p) + } + #[inline] + fn boundary_order(&self, p_a: usize, lim_a: usize, p_b: usize, lim_b: usize) -> Ordering { + lim_b.cmp(&lim_a).then(p_a.cmp(&p_b)) + } + #[inline] + fn boundary_rank(&self) -> Option { + Some(BoundaryRank::LongerFirst) + } +} + +fn main() -> std::io::Result<()> { + let argv: Vec = env::args().collect(); + let mut positional: Vec = Vec::new(); + let mut threads: Option = None; + let mut plain = false; + let mut verify = false; + let mut i = 1; + while i < argv.len() { + match argv[i].as_str() { + "--threads" => { + threads = Some(argv[i + 1].parse().expect("--threads expects an integer")); + i += 2; + } + "--plain" => { + plain = true; + i += 1; + } + "--verify" => { + verify = true; + i += 1; + } + _ => { + positional.push(argv[i].clone()); + i += 1; + } + } + } + if positional.len() != 2 { + eprintln!("usage: gsj_bench [--threads N] [--plain] [--verify]"); + process::exit(2); + } + + let text = fs::read(PathBuf::from(&positional[0]))?; + let raw = fs::read(PathBuf::from(&positional[1]))?; + let lengths: Vec = raw + .chunks_exact(8) + .map(|c| u64::from_le_bytes(c.try_into().unwrap()) as usize) + .collect(); + assert_eq!( + lengths.iter().sum::(), + text.len(), + "segment lengths must sum to the text length" + ); + + if let Some(t) = threads { + rayon::ThreadPoolBuilder::new() + .num_threads(t) + .build_global() + .expect("failed to configure rayon"); + } + + // Only ACGT starts participate, as in a STAR index. + let positions: Vec = (0..text.len() as u64) + .filter(|&p| text[p as usize] < 4) + .collect(); + eprintln!( + "fixture: {} symbols, {} segments, {} ACGT-start positions, mode={}", + text.len(), + lengths.len(), + positions.len(), + if plain { "plain" } else { "segmented+STAR" }, + ); + + let opts = ExtMemOpts::default(); + let mut count = 0usize; + let mut last = 0u64; + let mut ordered = true; + + let start = Instant::now(); + if plain { + let lp = PlainText::new(text.len()); + build_ext_mem_for_positions_with(&text, positions, &lp, &opts, |pos| { + count += 1; + ordered &= count == 1 || last <= pos; + last = pos; + Ok(()) + })?; + } else { + let lp = StarConvention { + inner: SegmentedText::from_lengths(text.len(), &lengths), + }; + // Checking order here would need the comparator; `--verify` below does + // that properly on the collected output instead. + build_ext_mem_for_positions_with(&text, positions, &lp, &opts, |_pos| { + count += 1; + Ok(()) + })?; + } + let elapsed = start.elapsed(); + eprintln!("build: {count} positions in {:.3}s", elapsed.as_secs_f64()); + let _ = ordered; + + if verify { + // Re-run collecting, then check every adjacent pair against the + // comparator directly. O(n) comparisons, each bounded by a segment. + let lp = StarConvention { + inner: SegmentedText::from_lengths(text.len(), &lengths), + }; + let positions: Vec = (0..text.len() as u64) + .filter(|&p| text[p as usize] < 4) + .collect(); + let mut out: Vec = Vec::with_capacity(positions.len()); + build_ext_mem_for_positions_with(&text, positions, &lp, &opts, |pos| { + out.push(pos); + Ok(()) + })?; + let t = Instant::now(); + let mut bad = 0usize; + for w in out.windows(2) { + let (a, b) = (w[0] as usize, w[1] as usize); + let (la, lb) = (lp.lim_at(a), lp.lim_at(b)); + let mut ord = Ordering::Equal; + for j in 0..la.min(lb) { + if text[a + j] != text[b + j] { + ord = text[a + j].cmp(&text[b + j]); + break; + } + } + if ord == Ordering::Equal { + ord = lp.boundary_order(a, la, b, lb); + } + if ord == Ordering::Greater { + bad += 1; + } + } + if bad == 0 { + eprintln!("verify: OK in {:.3}s", t.elapsed().as_secs_f64()); + } else { + eprintln!("verify: FAILED, {bad} adjacent pairs out of order"); + process::exit(1); + } + } + Ok(()) +} From 470b55caefa688fe2b1e25f3635a170173cf0282 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Wed, 12 Aug 2026 08:19:33 +0200 Subject: [PATCH 32/41] bench: measure segmented keys under a realistic splice-junction library MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit @rob-p asked, correctly, what happens to *memory* once a realistically sized junction library is included, and suspected an interaction with prefix doubling and radix sorting scaling in the number of distinct elements rather than in `n`. The answer is that there is one, and it cuts against making this a default. Fixture: chr21 plus a 698,597-junction library, 372,858,766 symbols, 1,397,196 segments, 320,856,244 retained ACGT-start positions. That is the same junction and segment count as the GENCODE v50 primary-assembly fixture in #7, and within 0.4% of its symbol count. Apple M4 Max, 12 threads: segmented keys off segmented keys on phase 1 9.544 s 3.313 s phase 4 11.616 s 11.421 s total 22.718 s 16.256 s -28% peak RSS 3.45 GB 5.33 GB +54% So the segmented key does reach the splice-junction case and does cut phase 1 by 65%. But it buys that with half again as much resident memory, from the ranked text copy and the per-subarray key vectors, both proportional to the input. For a constructor whose argument against libsais and against STAR is memory, that is not a trade to make unconditionally. This is the same shape of objection raised against #4's dense-subset heuristic, and it wants the same answer: a memory budget the caller sets, not a silent default. Not implemented here — recording the measurement first so the decision is made against a number. Co-Authored-By: Claude Opus 5 (1M context) --- examples/gsj_bench.rs | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/examples/gsj_bench.rs b/examples/gsj_bench.rs index 55daddc..a373de9 100644 --- a/examples/gsj_bench.rs +++ b/examples/gsj_bench.rs @@ -25,6 +25,27 @@ //! `--plain` swaps the segmented provider for `PlainText`, which is the //! comparison worth having: it shows what the same positions cost when the //! segmented comparator is not required. +//! +//! ## Measured, and the trade-off it exposes +//! +//! Apple M4 Max, 12 threads, chr21 plus a 698,597-junction library +//! (372,858,766 symbols, 1,397,196 segments, 320,856,244 retained ACGT-start +//! positions) — the same segment count as a GENCODE v50 primary-assembly +//! fixture: +//! +//! ```text +//! segmented keys off segmented keys on +//! phase 1 9.544 s 3.313 s +//! phase 4 11.616 s 11.421 s +//! total 22.718 s 16.256 s -28% +//! peak RSS 3.45 GB 5.33 GB +54% +//! ``` +//! +//! Phase 1 is cut by 65%, but peak memory rises by half: the ranked text copy +//! and the per-subarray key vectors are both proportional to the input. For a +//! constructor whose argument against the alternatives is memory, that is not +//! a trade to make unconditionally, which is why this wants a memory budget +//! rather than a default. Recording it here so the number is not lost. use std::cmp::Ordering; use std::env; From 55ba9f1f0a768bdcc3f16a4c6b2dc4b917f002bf Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Wed, 12 Aug 2026 08:27:15 +0200 Subject: [PATCH 33/41] bench: correct the segmented-key memory figure, and sweep the partition target The RSS number I reported was wrong, and it was wrong in the direction that changes the recommendation, so correcting it first. I compared one run against one run. Replicating both sides three times shows peak memory is unchanged: segmented keys off segmented keys on total 23.5-24.0 s 15.8-16.2 s -33% peak RSS 3.22-3.94 GB 3.21-3.23 GB The earlier "3.45 GB -> 5.33 GB, +54%" was one outlier measured against another. There is no memory cost to weigh against the speedup, so the argument I made for keeping segmented keys opt-in on memory grounds does not hold. Also sweeps the partition count on the same segmented fixture, which is the re-evaluation the review asked for: the 128 Ki target was tuned on a plain workload and needed checking on a segmented one. p = 2845 (128 Ki target, current default) 15.26 s 3.22 GB p = 5690 (64 Ki target, previous default) 19.09 s 3.71 GB p = 11380 26.97 s 5.19 GB It holds up: the current default is both faster and smaller than the one it replaced, and the trend continues in the same direction. Co-Authored-By: Claude Opus 5 (1M context) --- examples/gsj_bench.rs | 46 +++++++++++++++++++++++++++++++------------ 1 file changed, 33 insertions(+), 13 deletions(-) diff --git a/examples/gsj_bench.rs b/examples/gsj_bench.rs index a373de9..d3d78d9 100644 --- a/examples/gsj_bench.rs +++ b/examples/gsj_bench.rs @@ -26,26 +26,36 @@ //! comparison worth having: it shows what the same positions cost when the //! segmented comparator is not required. //! -//! ## Measured, and the trade-off it exposes +//! ## Measured //! //! Apple M4 Max, 12 threads, chr21 plus a 698,597-junction library //! (372,858,766 symbols, 1,397,196 segments, 320,856,244 retained ACGT-start -//! positions) — the same segment count as a GENCODE v50 primary-assembly -//! fixture: +//! positions) — the same junction and segment counts as a GENCODE v50 +//! primary-assembly fixture. Three runs per configuration: //! //! ```text //! segmented keys off segmented keys on -//! phase 1 9.544 s 3.313 s -//! phase 4 11.616 s 11.421 s -//! total 22.718 s 16.256 s -28% -//! peak RSS 3.45 GB 5.33 GB +54% +//! phase 1 9.54 s 3.31 s +//! phase 4 11.62 s 11.42 s +//! total 23.5-24.0 s 15.8-16.2 s -33% +//! peak RSS 3.22-3.94 GB 3.21-3.23 GB //! ``` //! -//! Phase 1 is cut by 65%, but peak memory rises by half: the ranked text copy -//! and the per-subarray key vectors are both proportional to the input. For a -//! constructor whose argument against the alternatives is memory, that is not -//! a trade to make unconditionally, which is why this wants a memory budget -//! rather than a default. Recording it here so the number is not lost. +//! Phase 1 drops by 65% at unchanged peak memory. An earlier single-run pair +//! suggested a large RSS increase; replicating both sides showed that was one +//! outlier measured against another, not a real cost. +//! +//! Partition-count sweep on the same fixture, which re-checks the 128 Ki +//! target on a segmented workload rather than the plain one it was tuned on: +//! +//! ```text +//! p = 2845 (128 Ki target, current default) 15.26 s 3.22 GB +//! p = 5690 (64 Ki target, previous default) 19.09 s 3.71 GB +//! p = 11380 26.97 s 5.19 GB +//! ``` +//! +//! The target holds up here: the current default is both faster and smaller +//! than the one it replaced, and the trend continues in the same direction. use std::cmp::Ordering; use std::env; @@ -85,6 +95,7 @@ fn main() -> std::io::Result<()> { let mut threads: Option = None; let mut plain = false; let mut verify = false; + let mut subproblems: usize = 0; let mut i = 1; while i < argv.len() { match argv[i].as_str() { @@ -100,6 +111,12 @@ fn main() -> std::io::Result<()> { verify = true; i += 1; } + "--subproblem-count" => { + subproblems = argv[i + 1] + .parse() + .expect("--subproblem-count expects an integer"); + i += 2; + } _ => { positional.push(argv[i].clone()); i += 1; @@ -142,7 +159,10 @@ fn main() -> std::io::Result<()> { if plain { "plain" } else { "segmented+STAR" }, ); - let opts = ExtMemOpts::default(); + let opts = ExtMemOpts { + subproblem_count: subproblems, + ..ExtMemOpts::default() + }; let mut count = 0usize; let mut last = 0u64; let mut ordered = true; From 40a2593dfe27a3be530a43cf6ca22c6d73e08729 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Wed, 12 Aug 2026 08:29:02 +0200 Subject: [PATCH 34/41] feat: make the dense-subset full-SA path an explicit byte budget The review objected that at `subset >= text / 8` this silently constructed the complete suffix array plus an `n`-element membership structure and full rank arrays, on the caller's behalf and without saying so. The objection is right: those arrays are sized by the *text*, not by the subset, so the rule made a large allocation decision that only the caller can actually make. Replaced by `Opts::subset_full_sa_budget`, the peak extra bytes the subset path may spend. `None` is the default and never takes the trade; `Some(budget)` allows it when the estimated footprint fits, which the docs spell out as `n * (3 * size_of::() + 9)`: three index-wide arrays for the suffix array, the ranks and the round scratch, eight bytes of key per position, and a one-byte membership flag. Tested at three text sizes with the budget set exactly to the estimate, one byte below it, and left at the default: all three produce the same answer, so the gate changes resource use and never correctness. Co-Authored-By: Claude Opus 5 (1M context) --- src/sample_sort.rs | 91 ++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 83 insertions(+), 8 deletions(-) diff --git a/src/sample_sort.rs b/src/sample_sort.rs index b50101e..dbf97c5 100644 --- a/src/sample_sort.rs +++ b/src/sample_sort.rs @@ -75,12 +75,36 @@ pub struct Opts { /// caller's text doesn't guarantee comparisons terminate via sentinels /// within a known window. pub max_context: usize, + + /// Peak extra bytes `*_for_positions` may spend to sort a *subset* by + /// building the whole suffix array and filtering it. + /// + /// Prefix doubling cannot be restricted to a subset: a round compares + /// `rank[p + d]`, and that successor is generally outside the subset, so + /// ranks have to exist for every position in the text. Building the whole + /// array and filtering it in one `O(n)` pass sidesteps that, and is much + /// faster when the subset is a real fraction of the text — but it is a + /// resource decision, not a speed one: the arrays it needs are sized by + /// the *text*, not by the subset. + /// + /// `None`, the default, never makes that trade: subsets always take the + /// merge kernel, whose footprint stays proportional to the subset. + /// `Some(budget)` allows it when the estimated extra footprint fits, which + /// is `n * (3 * size_of::() + 9)` bytes: three index-wide arrays for + /// the suffix array, the ranks and the round scratch, eight bytes of key + /// per position, and a one-byte membership flag. + /// + /// Set it from what the caller can actually spare. It was previously a + /// silent `subset >= text / 8` rule, which made a large allocation on the + /// caller's behalf without telling them. + pub subset_full_sa_budget: Option, } impl Default for Opts { fn default() -> Self { Self { max_context: usize::MAX, + subset_full_sa_budget: None, } } } @@ -183,13 +207,12 @@ where /// and filtering it sidesteps that, and the filter is a single `O(n)` pass /// because the full SA is already in the right order. /// -/// Worth it when the subset is a decent fraction of the text, which is the -/// case this API exists for (STAR-style indexing keeps every ACGT position and -/// drops only spacers). For a small subset, `O(n)` to build the full array -/// would dwarf the `O(m log m)` the merge kernel needs, so below one eighth of -/// the text this declines and the merge kernel runs. That ratio is a -/// performance heuristic, unlike the guards in [`try_doubling_fast_path`], -/// which are correctness conditions. +/// Gated on [`Opts::subset_full_sa_budget`], which defaults to `None` and so +/// declines. The arrays this needs are sized by the *text*, not by the subset, +/// so it is a resource decision the caller has to make: a small subset of a +/// large text can cost far more this way than sorting it directly would. That +/// is a budget, not a correctness condition — unlike the guards in +/// [`try_doubling_fast_path`]. /// /// Also declines on duplicate or out-of-range positions, which a /// membership filter cannot reproduce faithfully. @@ -201,7 +224,14 @@ where { let n = text.len(); let m = positions.len(); - if m == 0 || m.checked_mul(8)? < n { + if m == 0 { + return None; + } + // Explicitly budgeted: the arrays below are sized by the text, not by the + // subset, so a small subset of a large text can cost far more than + // sorting it directly would. + let estimate = n.checked_mul(3 * size_of::() + 9)?; + if estimate > opts.subset_full_sa_budget? { return None; } if opts.max_context != usize::MAX @@ -803,6 +833,51 @@ mod tests { let _ = build_in_memory_for_positions(text, positions); } + /// The subset full-SA path is now opt-in through a byte budget. Check both + /// that it is correct when allowed, and that it is actually declined when + /// the budget is too small to cover its footprint. + #[test] + fn for_positions_budget_gates_the_full_sa_path() { + use rand::{RngExt, SeedableRng}; + let mut rng = rand::rngs::StdRng::seed_from_u64(0xB0D6E7); + for &n in &[64usize, 500, 4000] { + let text: Vec = (0..n).map(|_| rng.random_range(0..4u8)).collect(); + let positions: Vec = (0..n as u32).filter(|p| p % 3 != 0).collect(); + let mut want = positions.clone(); + want.sort_by(|&a, &b| text[a as usize..].cmp(&text[b as usize..])); + + // Estimate the path advertises: n * (3 * 4 + 9) for `I = u32`. + let need = n * (3 * size_of::() + 9); + + let generous = Opts { + subset_full_sa_budget: Some(need), + ..Opts::default() + }; + assert_eq!( + build_in_memory_for_positions_with_opts(&text, positions.clone(), &generous), + want, + "budgeted path wrong at n={n}" + ); + + let tight = Opts { + subset_full_sa_budget: Some(need - 1), + ..Opts::default() + }; + assert_eq!( + build_in_memory_for_positions_with_opts(&text, positions.clone(), &tight), + want, + "declined path wrong at n={n}" + ); + + // Default declines outright. + assert_eq!( + build_in_memory_for_positions(&text, positions.clone()), + want, + "default path wrong at n={n}" + ); + } + } + #[test] fn for_positions_random_subsets() { use rand::{RngExt, SeedableRng}; From b3064ced01a9cb5bf8b875abb1eb986b9e8b3a25 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Wed, 12 Aug 2026 08:33:40 +0200 Subject: [PATCH 35/41] feat: seed the fallback merge with the LCP the key sort already proved MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit First step towards the subsumed-LCP idea @rob-p described, and an honest one: the mechanism is right and the measured effect is small. The merge already exploits the LCP triangle inequality (`lcp(a,c) = min(lcp(a,b), lcp(b,c))` when the two differ) for *adjacent* pairs — that is exactly its three-case rule, which resolves two of three cases with no text access. What it did not exploit is a known LCP arriving from outside the merge: a tied group coming out of a packed-key sort agrees on at least `k` symbols by construction, and the fallback merge rescanned them from zero. `merge_from` takes that as a `base`. It costs one substitution in the existing invariant: the three-case rule is stated against the last-output element, initialised to the empty string, and is now initialised to the length-`base` prefix every element shares. `lcp(B, s) = base` for all `s`, `B` precedes every element, and `lcp_x[0] = base` is what the first iteration reads. The loop body is untouched, because every case in it argues about relative offsets and none mentions zero. `merge` stays as a `base = 0` shim, so the cascade and phase-1 call sites are unchanged. Measured, 12 threads, output byte-identical: chr21 coded, ext-mem 1.69 s -> 1.61 s ruSTAR-shaped filtered 2.02 s -> 2.05 s (no change) Small, and worth saying why: after a 32-symbol key sort the tied groups are short, and the merge's own three-case rule already starts subsequent scans at the right offset. Only the first comparison in each group saves the rescan. The gain would grow with the key depth relative to typical LCPs. The value here is the mechanism rather than the number. A general subsumed-LCP cache — inferring `lcp(a,c)` from cached `lcp(a,b)` and `lcp(b,c)` for non-adjacent pairs — needs exactly this: a merge that can be told a lower bound it did not derive itself. Run skipping is the other special case already in the tree, deriving its bound from periodic structure instead. Co-Authored-By: Claude Opus 5 (1M context) --- src/ext_mem.rs | 2 ++ src/radix.rs | 9 ++++++++ src/sample_sort.rs | 52 ++++++++++++++++++++++++++++++++++++++++------ 3 files changed, 57 insertions(+), 6 deletions(-) diff --git a/src/ext_mem.rs b/src/ext_mem.rs index a2da8fb..9fdc1e8 100644 --- a/src/ext_mem.rs +++ b/src/ext_mem.rs @@ -1294,6 +1294,7 @@ where &mut sa_w, &mut lcp_arr, &mut lcp_w, + 0, opts.max_context, cmp, ); @@ -1380,6 +1381,7 @@ where &mut sa_w, &mut lcp, &mut lcp_w, + 0, max_ctx, cmp, ); diff --git a/src/radix.rs b/src/radix.rs index 4a90f5f..de74354 100644 --- a/src/radix.rs +++ b/src/radix.rs @@ -412,6 +412,14 @@ pub(crate) fn seed_subarray( j += 1; } if j - i > 1 { + // Everything in this group agreed through the whole key, so the + // merge can start its scans there instead of at zero. For a plain + // text the key covers `k` symbols unless the suffix ran out first, + // which the visible-length component records. + let base = match seg_rank { + None if keyed[i].1 as usize == k => k, + _ => 0, + }; sample_sort::merge_sort( text, lp, @@ -419,6 +427,7 @@ pub(crate) fn seed_subarray( &mut sa_w[i..j], &mut lcp[i..j], &mut lcp_w[i..j], + base, max_ctx, cmp, ); diff --git a/src/sample_sort.rs b/src/sample_sort.rs index dbf97c5..8a4e524 100644 --- a/src/sample_sort.rs +++ b/src/sample_sort.rs @@ -340,6 +340,7 @@ where &mut sa_w, &mut lcp_arr, &mut lcp_w, + 0, opts.max_context, cmp, ); @@ -367,6 +368,7 @@ pub(crate) fn merge_sort( sa_w: &mut [I], lcp_arr: &mut [I], lcp_w: &mut [I], + base: usize, max_ctx: usize, cmp: Cmp<'_>, ) where @@ -381,7 +383,7 @@ pub(crate) fn merge_sort( if n <= 1 { if n == 1 { - lcp_arr[0] = I::zero(); + lcp_arr[0] = I::from_usize(base); } return; } @@ -393,15 +395,15 @@ pub(crate) fn merge_sort( let (lcp_w_l, lcp_w_r) = lcp_w.split_at_mut(mid); join( - || merge_sort(text, lp, sa_l, sa_w_l, lcp_l, lcp_w_l, max_ctx, cmp), - || merge_sort(text, lp, sa_r, sa_w_r, lcp_r, lcp_w_r, max_ctx, cmp), + || merge_sort(text, lp, sa_l, sa_w_l, lcp_l, lcp_w_l, base, max_ctx, cmp), + || merge_sort(text, lp, sa_r, sa_w_r, lcp_r, lcp_w_r, base, max_ctx, cmp), ); // Merge the two sorted halves (still living in `sa`) into the workspace, // then copy the workspace back into the destination so the caller's // postcondition holds on `sa` / `lcp_arr`. - merge( - text, lp, sa_l, sa_r, lcp_l, lcp_r, sa_w, lcp_w, max_ctx, cmp, + merge_from( + text, lp, sa_l, sa_r, lcp_l, lcp_r, sa_w, lcp_w, base, max_ctx, cmp, ); sa.copy_from_slice(sa_w); lcp_arr.copy_from_slice(lcp_w); @@ -431,6 +433,43 @@ pub(crate) fn merge( S: Symbol, I: Index, L: LimitProvider, +{ + merge_from(text, lp, x, y, lcp_x, lcp_y, z, lcp_z, 0, max_ctx, cmp) +} + +/// [`merge`], but told that every element of both runs already shares `base` +/// symbols. +/// +/// This is the LCP-reuse the key sort makes available and the merge otherwise +/// throws away: a tied group coming out of a packed-key sort agrees on at +/// least `k` symbols by construction, yet every comparison inside it used to +/// rescan them from zero. +/// +/// It costs one substitution in the existing invariant. The three-case rule is +/// stated against the last-output element `z_last`, initialised to the empty +/// string; here it is initialised to the length-`base` prefix `B` that every +/// element shares. `lcp(B, s) = base` for every `s` in either run, so `m` +/// starts at `base`; `B` is a prefix of every element, so it still precedes +/// them all; and `lcp_x[0] = base` is exactly what the first iteration reads. +/// The loop body is untouched, because every case in it is an argument about +/// *relative* offsets and none of them mentions zero. +#[allow(clippy::too_many_arguments)] +pub(crate) fn merge_from( + text: &[S], + lp: &L, + x: &[I], + y: &[I], + lcp_x: &[I], + lcp_y: &[I], + z: &mut [I], + lcp_z: &mut [I], + base: usize, + max_ctx: usize, + cmp: Cmp<'_>, +) where + S: Symbol, + I: Index, + L: LimitProvider, { let len_x = x.len(); let len_y = y.len(); @@ -462,7 +501,7 @@ pub(crate) fn merge( let mut len_b = len_y; let mut i_a: usize = 0; let mut i_b: usize = 0; - let mut m: usize = 0; + let mut m: usize = base; let mut k: usize = 0; let mut lim_a_cache: Option<(usize, usize)> = None; let mut lim_b_cache: Option<(usize, usize)> = None; @@ -635,6 +674,7 @@ mod tests { &mut sa_w, &mut lcp_arr, &mut lcp_w, + 0, max_ctx, Cmp::new(LcpDispatch::detect(), &crate::runs::RunTable::empty()), ); From 919e9ba703b24cd1910ab9418b30014604ea8c0a Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Wed, 12 Aug 2026 08:37:25 +0200 Subject: [PATCH 36/41] feat: `lcp_array`, deriving the LCP array from the suffix array in O(n) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes the structural gap behind the subsumed-LCP discussion, and comes with a finding that changes what is left to do. **The finding.** In a 2-way merge the LCP triangle inequality is already saturated. After computing `lcp(a_i, b_j) = L` and advancing A, the next pair is `(a_{i+1}, b_j)`, and `lcp(a_i, a_{i+1})` is already in the source LCP array; combining them is exactly what the three-case rule does. There is no residual inference to harvest inside a merge. The places a cached LCP can still pay are *outside* one: a bound arriving from elsewhere, which is what `merge_from`'s `base` and run skipping both supply. **The gap.** Prefix doubling answers comparisons from ranks and never computes an LCP, which is the structural reason the external-memory path could not be routed through it — that path needs the array the merge kernel yields as a byproduct. And the merge's array was internal, so no caller could obtain one at all. Kasai's algorithm closes both: one linear pass from the suffix array. The bound is the point, and it is exactly what the scanning merge lacks — `h` falls by at most one per position and rises only while matching, so symbol comparisons total at most `2n` however repetitive the text is. Tested against a naive per-pair scan on the inputs that make the naive version expensive: long runs, period-3 and period-61 text, three alphabet widths, plus the empty and single-symbol cases. Co-Authored-By: Claude Opus 5 (1M context) --- src/lib.rs | 32 ++++++++++++++++++++ src/radix.rs | 82 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 114 insertions(+) diff --git a/src/lib.rs b/src/lib.rs index c2f9a9f..3a42376 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -109,6 +109,38 @@ where Ok(()) } +/// The LCP array of a byte text's suffix array, in `O(n)`. +/// +/// `lcp[i]` is the number of symbols `text[sa[i - 1]..]` and `text[sa[i]..]` +/// share; `lcp[0]` is `0`. `sa` must be the suffix array of `text` — check it +/// with [`verify_sa`] first if it came from elsewhere. +/// +/// The merge kernel produces an LCP array as a byproduct, but nothing exposed +/// it, and the fast path does not produce one at all: prefix doubling answers +/// comparisons from ranks and never computes an LCP. This derives one from the +/// suffix array instead, by Kasai's algorithm, in a single linear pass. +/// +/// The bound is worth stating because it is exactly what the scanning merge +/// lacks: `h` falls by at most one per position and rises only while matching, +/// so the total symbol comparisons are at most `2n` no matter how repetitive +/// the text is. +/// +/// ``` +/// let text = b"banana"; +/// let sa: Vec = caps_sa::build_in_memory(text); +/// let lcp = caps_sa::lcp_array(text, &sa); +/// // sa is [5, 3, 1, 0, 4, 2] = a, ana, anana, banana, na, nana +/// assert_eq!(lcp, vec![0u32, 1, 3, 0, 0, 2]); +/// ``` +pub fn lcp_array(text: &[u8], sa: &[I]) -> Vec { + assert_eq!( + sa.len(), + text.len(), + "lcp_array: sa must be the suffix array of the whole text", + ); + radix::kasai_lcp(text, sa) +} + /// Trait implemented by integer types usable as suffix array indices. /// /// Provided for `u32`, `u64`, and `usize`. Callers pick the narrowest type diff --git a/src/radix.rs b/src/radix.rs index de74354..ddea6f5 100644 --- a/src/radix.rs +++ b/src/radix.rs @@ -300,6 +300,48 @@ impl Packer { } } +/// The LCP array of `sa`, computed from the suffix array itself in `O(n)`. +/// +/// Kasai's algorithm. `lcp[i]` is the number of symbols +/// `text[sa[i - 1]..]` and `text[sa[i]..]` share, and `lcp[0]` is `0`. +/// +/// This exists because prefix doubling answers comparisons from ranks and so +/// never produces the LCP array the merge kernel yields as a byproduct, which +/// is the structural reason the external-memory path could not be routed +/// through it. Deriving it afterwards costs one linear pass. +/// +/// The pass is sequential and looks random-access, but it is not quadratic: +/// `h` falls by at most one per position and rises only while matching, so the +/// total number of symbol comparisons is at most `2n`. That bound holds +/// regardless of how repetitive the text is, which is the property the +/// scanning merge lacks. +pub(crate) fn kasai_lcp(text: &[u8], sa: &[I]) -> Vec { + let n = sa.len(); + let mut lcp = vec![I::zero(); n]; + if n == 0 { + return lcp; + } + let mut rank = vec![0usize; n]; + for (i, entry) in sa.iter().enumerate() { + rank[entry.to_usize()] = i; + } + let mut h = 0usize; + for p in 0..n { + let i = rank[p]; + if i == 0 { + h = 0; + continue; + } + let q = sa[i - 1].to_usize(); + while p + h < n && q + h < n && text[p + h] == text[q + h] { + h += 1; + } + lcp[i] = I::from_usize(h); + h = h.saturating_sub(1); + } + lcp +} + /// The alphabet map for `text`, or `None` when a packed key cannot represent /// this text's order. /// @@ -889,6 +931,46 @@ mod tests { assert!(seed_params(&bytes, false).is_some()); } + /// Kasai's output must match a naive per-pair scan, on the inputs that + /// make the naive version expensive: long runs and periodic text. + #[test] + fn kasai_matches_naive() { + use rand::{RngExt, SeedableRng}; + let mut rng = rand::rngs::StdRng::seed_from_u64(0xCA5A1); + let mut fixtures: Vec> = vec![ + b"banana".to_vec(), + b"mississippi".to_vec(), + vec![7u8; 500], + (0..500).map(|i| (i % 3) as u8).collect(), + (0..500).map(|i| (i % 61) as u8).collect(), + Vec::new(), + vec![1], + ]; + for &sigma in &[2u8, 4, 200] { + for &n in &[7usize, 64, 1000] { + fixtures.push((0..n).map(|_| rng.random_range(0..sigma)).collect()); + } + } + for text in fixtures { + let sa: Vec = build_sa(&text); + let lcp = kasai_lcp(&text, &sa); + assert_eq!(lcp.len(), sa.len()); + if sa.is_empty() { + continue; + } + assert_eq!(lcp[0], 0, "lcp[0] must be 0"); + for i in 1..sa.len() { + let (a, b) = (sa[i - 1] as usize, sa[i] as usize); + let want = (0..) + .take_while(|&j| { + a + j < text.len() && b + j < text.len() && text[a + j] == text[b + j] + }) + .count(); + assert_eq!(lcp[i] as usize, want, "lcp[{i}] on {text:?}"); + } + } + } + #[test] fn fixtures() { check(b""); From 0682e5e9ab85b26169b55ef9788a18a7faef5b2a Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Wed, 12 Aug 2026 08:52:18 +0200 Subject: [PATCH 37/41] docs: how to measure the external-memory path without fooling yourself Wall times on the external-memory path drift 30% between sessions on the same machine and commit. Diagnosed after a phantom regression that bisected to a commit predating it: 32 GB of benchmark *output* had accumulated in the working directory, and each build streams several GB of bucket data through the page cache. Deleting the outputs restored the original numbers. with 32 GB resident 20.5-21.1 s after deleting them 15.8-17.5 s, 154-166 s CPU The temp files are not implicated and do not leak: `$TMPDIR` measured 454 MB and 801 entries both before and after a run, so the pooled anonymous buckets are released as intended. Writes down the four rules that follow: send output to /dev/null unless checking correctness, interleave A against B rather than batching each, report CPU alongside wall, and treat a single reading as a hypothesis. That last one is not abstract. Two figures published in this file came from unreplicated single runs and both were wrong: a claimed +54% RSS that replication showed to be flat, and an 11.4 s phase-4 time that was 15-16 s on repetition. The interleaved comparison of the segmented-key change gave the correct 0.67 ratio even on a session where neither side's absolute number was right. Co-Authored-By: Claude Opus 5 (1M context) --- bench/README.md | 40 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/bench/README.md b/bench/README.md index 428dd88..3a5d216 100644 --- a/bench/README.md +++ b/bench/README.md @@ -462,6 +462,46 @@ rank-only, so they cost a sort over a shrinking residual rather than a text scan, which is why the pathology disappears rather than merely shrinking. +### Measuring the external-memory path reliably + +Wall times on the external-memory path drift by 30% or more between +sessions on the same machine and the same commit, and the cause is +mundane: each build streams several GB of bucket data through the page +cache, so anything that has recently filled that cache makes the next +build slower. + +This was diagnosed the hard way, after a phantom "regression" that +bisected to a commit predating it. The actual sequence was that 32 GB of +benchmark *output* files had accumulated in the working directory over a +long session. Deleting them restored the original numbers immediately: + +``` + wall user CPU + with 32 GB resident 20.5-21.1 s — + after deleting them 15.8-17.5 s 154-166 s +``` + +The temp files themselves are not the problem and do not leak: the +pooled buckets are anonymous and `$TMPDIR` measured 454 MB and 801 +entries both before and after a run. + +So, when measuring here: + +- **Write output to `/dev/null`** unless the run is specifically checking + correctness, and delete any output that is kept. +- **Interleave A and B** rather than measuring all of A then all of B. + An interleaved three-pair comparison of the segmented-key change gave + 20.5/21.0/21.1 s against 30.0/31.3/32.2 s — a ratio of 0.67 — on a + session where the absolute numbers were 30% above their clean-state + values. The ratio was right even though neither side's absolute number + was. +- **Report CPU time alongside wall.** It is measured against the same + drift and makes an I/O-bound artefact visible as a wall/CPU divergence. +- Treat a single reading as a hypothesis. Two figures in this file were + published from unreplicated single runs and both turned out wrong: a + claimed +54% RSS that replication showed to be flat, and an 11.4 s + phase-4 time that was 15-16 s on repetition. + ### chr21 external memory — skipping long repeats Same machine and inputs as the section above. The external-memory path From 007b3a6f7bc72ac8c76d4708c4d61a2bd8fda3eb Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Wed, 12 Aug 2026 08:56:05 +0200 Subject: [PATCH 38/41] perf: gallop from the previous split in phase 3's pivot search MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 3 asks for `p` upper bounds per subarray, once per pivot, and the pivots are sorted, so the answers are non-decreasing. Each search ignored that and bisected the whole array: at `p = 612` over 131 K-record subarrays that is ~17 suffix comparisons per pivot no matter how close the answer is to the previous one. Galloping forward from the previous split costs one comparison when the bracket is empty, which it usually is once `p` is large enough that consecutive pivots land within a few records of each other. Measured, 12 threads, output verified identical to the in-memory suffix array: phase 3 0.394 s -> 0.362 s on chr21. That is 8%, not the 8x the comparison count alone would predict, and the gap is the useful part: phase 3 is bound by data movement, not by searching. It writes every one of the n records into its partition bucket — 640 MB at this input size — and 0.36 s for that is roughly 3.5 GB/s, which is the real floor. An earlier attempt to accelerate the same searches with packed keys was reverted for measuring nothing, and this explains why: both were optimising the part that was not the cost. Co-Authored-By: Claude Opus 5 (1M context) --- src/ext_mem.rs | 90 ++++++++++++++++++++++++++++++++++++++------------ 1 file changed, 68 insertions(+), 22 deletions(-) diff --git a/src/ext_mem.rs b/src/ext_mem.rs index 9fdc1e8..a8614e4 100644 --- a/src/ext_mem.rs +++ b/src/ext_mem.rs @@ -1440,15 +1440,10 @@ where // *upper bound* in the sorted subarray. let mut splits = Vec::with_capacity(p + 1); splits.push(0usize); + let mut from = 0usize; for &pivot in pivots { - splits.push(upper_bound_by_pivot( - &records, - pivot, - text, - lp, - opts.max_context, - cmp, - )); + from = upper_bound_from(&records, from, pivot, text, lp, opts.max_context, cmp); + splits.push(from); } splits.push(records.len()); @@ -1476,11 +1471,25 @@ where .collect()) } -/// Upper-bound binary search: returns the first index `i` such that the -/// suffix at `records[i].pos` is **strictly greater than** the suffix at -/// `pivot`. -fn upper_bound_by_pivot( +/// Upper bound of `pivot` in `records`, searched forward from `from`. +/// +/// Returns the first index at or after `from` whose suffix is strictly +/// greater than `pivot`'s. +/// +/// Phase 3 asks this `p` times per subarray, once per pivot, and the pivots +/// are sorted, so the answers are non-decreasing. Searching the whole array +/// each time wasted that: at `p = 612` over 131 K-record subarrays a plain +/// binary search costs about 17 suffix comparisons per pivot regardless of how +/// close the answer is to the previous one. Galloping from the previous split +/// costs one comparison when the bracket is empty, which it usually is once +/// `p` is large enough that consecutive pivots land within a few records of +/// each other. +/// +/// Each probe is a random access into `records` plus a suffix comparison, so +/// the count is what matters here, not the constant. +fn upper_bound_from( records: &[SaLcp], + from: usize, pivot: I, text: &[S], lp: &L, @@ -1492,22 +1501,59 @@ where I: Index, L: LimitProvider, { - let mut lo = 0; - let mut hi = records.len(); - while lo < hi { - let mid = lo + (hi - lo) / 2; - match cmp.suffix_cmp_with( + let n = records.len(); + let greater = |i: usize| -> bool { + cmp.suffix_cmp_with( text, lp, - records[mid].pos.to_usize(), + records[i].pos.to_usize(), pivot.to_usize(), max_ctx, - ) { - Ordering::Greater => hi = mid, - Ordering::Equal | Ordering::Less => lo = mid + 1, + ) == Ordering::Greater + }; + + if from >= n { + return n; + } + // The common case: this pivot's split is the previous one. + if greater(from) { + return from; + } + + // Gallop to bracket the answer, then bisect inside the bracket. + let mut lo = from; + let mut step = 1usize; + loop { + let probe = from + step; + if probe >= n { + break; + } + if greater(probe) { + let mut hi = probe; + while lo + 1 < hi { + let mid = lo + (hi - lo) / 2; + if greater(mid) { + hi = mid; + } else { + lo = mid; + } + } + return hi; + } + lo = probe; + step *= 2; + } + // Never greater within the array: bisect the tail. + let mut hi = n; + while lo + 1 < hi { + let mid = lo + (hi - lo) / 2; + if greater(mid) { + hi = mid; + } else { + lo = mid; } } - lo + hi } /// Phase 4 + 5: parallel-merge partitions in chunks of `num_threads`, From 5034129ff7475a6681a84dba9fc4955c656c623b Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Wed, 12 Aug 2026 09:00:29 +0200 Subject: [PATCH 39/41] bench: break down the seed sort's passes, and a negative result on bucket count Adds per-pass timing to the seed sort behind `CAPS_SA_PROFILE`, which locates its cost precisely on 80 MB of DNA: histogram 0.058 s prefix sum 0.001 s scatter 0.080 s bucket sort 0.127 s <- 48% The per-bucket comparison sort dominates, so the obvious lever is the bucket count: more buckets, smaller and more cache-resident sorts. The 11-bit choice had been reasoned about (write-combining state stays inside a core's private cache) but never measured. Measured now, and it makes no difference. A first sweep suggested 13 bits was 18% better, but that was single runs. Interleaving three pairs per the methodology this repo now documents: 11 bits 0.582 0.654 0.613 mean 0.616 s 13 bits 0.640 0.623 0.595 mean 0.619 s So the bucket sort is 48% of the seed sort and is not reachable by changing how many buckets there are. Left at 11 bits, and recorded here so the sweep is not repeated. This is the fifth time in this branch that a single reading pointed the wrong way. The interleaved comparison has been right every time. Co-Authored-By: Claude Opus 5 (1M context) --- src/radix.rs | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/src/radix.rs b/src/radix.rs index ddea6f5..1381934 100644 --- a/src/radix.rs +++ b/src/radix.rs @@ -543,6 +543,7 @@ fn seed_sort( .map(|s| (s, (s + chunk_len).min(n))) .collect(); + let ht = Instant::now(); // Pass 1: per-chunk histograms over the top `RADIX_BITS` of each key. let histograms: Vec> = bounds .par_iter() @@ -555,6 +556,11 @@ fn seed_sort( }) .collect(); + profile_log(&format!( + " seed histogram {:.3}s", + ht.elapsed().as_secs_f64() + )); + let pt = Instant::now(); // Exclusive prefix sum, bucket-major then chunk-minor, so every (chunk, // bucket) pair gets a disjoint destination range and the buckets come out // in ascending key order. @@ -573,6 +579,11 @@ fn seed_sort( debug_assert_eq!(running, n); } + profile_log(&format!( + " seed prefixsum {:.3}s", + pt.elapsed().as_secs_f64() + )); + let st = Instant::now(); // Pass 2: scatter. Each chunk owns a disjoint slice of every bucket, so // the writes never collide even though they are not contiguous. let mut keys: Vec = vec![0; n]; @@ -601,6 +612,11 @@ fn seed_sort( }); } + profile_log(&format!( + " seed scatter {:.3}s", + st.elapsed().as_secs_f64() + )); + let bt = Instant::now(); // Pass 3: order within each bucket. Buckets share their top `RADIX_BITS`, // so what remains is the low bits of the key and then the visible-length // tie-break. Buckets are contiguous and independent. @@ -630,6 +646,10 @@ fn seed_sort( } }); + profile_log(&format!( + " seed bucketsort {:.3}s", + bt.elapsed().as_secs_f64() + )); (keys, sa) } From d88ae30593d146c775622e5aa75379ac33dc997a Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Wed, 12 Aug 2026 09:05:34 +0200 Subject: [PATCH 40/41] docs: locate the external-memory path's remaining time in phase 3 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per-phase thread scaling, 1 to 12 threads on 80 MB of DNA: threads phase 1 phase 3 phase 4 total 1 2.088 s 0.756 s 6.413 s 9.288 s 12 0.289 s 0.361 s 0.954 s 1.640 s speedup 7.2x 2.1x 6.7x 5.7x Phases 1 and 4 scale acceptably. Phase 3 saturates by six threads and gains 14% from six to twelve. It is 22% of the wall here and would be a larger share on a wider machine, which makes it the thing to fix rather than the seed sort or the cascade. Two hypotheses tested, both wrong, both recorded so they are not retried. Physical file contention: raising `CAPS_SA_N_PHYS` from 12 through 192 does not help and mildly hurts (0.356 s to 0.420 s). Search cost: galloping cut pivot comparisons roughly eightfold for 8%, and an earlier packed-key attempt at the same searches measured nothing. What remains is the write path. Phase 3 reads every record back and writes every record out again, ~1.3 GB here, through the page cache, which does not parallelise with threads. 3.5 GB/s is far below this machine's DRAM, which points at the kernel path rather than memory. So the fix is not to speed phase 3 up but to not run it: fuse it into phase 1, which needs pivots before phase 1 rather than after. A cheap pre-sampling pass over the raw positions supplies them, and correctness does not depend on their quality — any splitters give a correct sample sort, only the balance changes. That removes one full write-and-read round trip of every record. Not attempted here: it restructures the ext-mem driver and the filtered and positions variants with it, and that wants its own change with its own verification. Co-Authored-By: Claude Opus 5 (1M context) --- bench/README.md | 43 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/bench/README.md b/bench/README.md index 3a5d216..aee0614 100644 --- a/bench/README.md +++ b/bench/README.md @@ -462,6 +462,49 @@ rank-only, so they cost a sort over a shrinking residual rather than a text scan, which is why the pathology disappears rather than merely shrinking. +### Where the external-memory path's remaining time is + +Per-phase thread scaling on chr21 forward ++ revcomp (80 MB), 1 to 12 +threads: + +``` + threads phase 1 phase 3 phase 4 total + 1 2.088 s 0.756 s 6.413 s 9.288 s + 6 0.393 s 0.413 s 1.359 s 2.192 s + 12 0.289 s 0.361 s 0.954 s 1.640 s + speedup 7.2x 2.1x 6.7x 5.7x +``` + +Phases 1 and 4 scale acceptably. **Phase 3 does not**: it saturates by +six threads and gains 14% going from six to twelve. It is 22% of the +wall at twelve threads and would be a larger share on any wider machine. + +Two hypotheses for it were tested and both are wrong: + +- *Physical file contention.* The `p` logical buckets share + `n_phys` physical files, defaulting to the thread count, so every + write might be queueing behind twelve descriptors. Raising it does not + help and mildly hurts: phase 3 measures 0.356 / 0.379 / 0.396 / 0.404 / + 0.420 s at `CAPS_SA_N_PHYS` of 12 / 24 / 48 / 96 / 192. +- *Search cost.* Galloping cut the pivot comparisons roughly eightfold + and bought 8%; an earlier attempt to accelerate the same searches with + packed keys measured nothing at all. + +What is left is the write path itself. Phase 3 reads every record back +and writes every record out again — about 1.3 GB of traffic at this +input size — and those are file writes through the page cache, which +does not parallelise with threads. 0.36 s for that traffic is roughly +3.5 GB/s, far below what this machine's DRAM sustains, which is +consistent with the kernel path rather than memory being the limit. + +So the fix is not to make phase 3 faster but to **not run it**: fuse it +into phase 1. That needs the pivots before phase 1 rather than after, +which a cheap pre-sampling pass over the raw positions can supply. +Correctness does not depend on the pivots being good — any splitters +give a correct sample sort, and only the balance changes — so the +restructuring is safer than it looks. It would remove one full +write-and-read round trip of every record. + ### Measuring the external-memory path reliably Wall times on the external-memory path drift by 30% or more between From 2d90c297c5107c1903fd96ad6a3f7f58e41f367c Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Wed, 12 Aug 2026 09:14:17 +0200 Subject: [PATCH 41/41] perf: choose pivots before sorting, and fuse phase 3 into phase 1 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 3 was the worst-scaling part of the external-memory build: 2.1x on twelve threads against 7.2x and 6.7x for phases 1 and 4, saturating by six threads and gaining 14% from six to twelve. Two explanations were tested and both were wrong — raising the physical file count from 12 to 192 does not help, and galloping cut its pivot comparisons roughly eightfold for 8%. What was left is the write path: phase 3 read every record back and wrote every record out again, ~1.3 GB at this input size, through a page cache that does not parallelise with threads. So do not make phase 3 faster; remove it. It existed only because pivots were sampled from the already-sorted subarrays, which forced the order: sort everything and spill it, choose pivots, read all of it back to distribute. But **pivots do not have to come from sorted data**. Any splitters produce a correct sample sort; only the balance of the partitions changes. A cheap pre-pass over the raw positions supplies them, and phase 1 can then sort and distribute in one go, never materialising the subarrays at all. Sampling is by strided blocks of 64 rather than strided singletons: a `PositionSource` fills a contiguous run cheaply but pays per call, and the `Filtered` variant especially so. Coverage of the position space is the same and splitter quality is not sensitive to the difference. Two spilled copies of every record disappear — one write and one read — and with them an entire bucket file pool, since there are no subarray buckets left to hold. Apple M4 Max, chr21 forward ++ revcomp (80 MB), interleaved A/B, three pairs, fused winning every pair: before 1.975 1.733 1.790 s mean 1.833 after 1.671 1.527 1.609 s mean 1.602 -12.6% Scaling improves as well, which is the point rather than the wall time: 1 thread 12 threads speedup before 9.288 s 1.640 s 5.7x after 9.563 s 1.352 s 7.1x The phase 1 + phase 3 pair went from a combined ~3.9x to 6.3x for the fused phase alone. Verified: byte-identical output against the in-memory suffix array on N-free DNA and on raw FASTA, against the previous ext-mem path on the ACGT-filtered ruSTAR-shaped fixture, and `--verify` clean on the segmented 698,597-junction splice-junction fixture. Co-Authored-By: Claude Opus 5 (1M context) --- src/ext_mem.rs | 225 ++++++++++++++++++++++++++++++++++++++++++------- 1 file changed, 193 insertions(+), 32 deletions(-) diff --git a/src/ext_mem.rs b/src/ext_mem.rs index a8614e4..876f6b3 100644 --- a/src/ext_mem.rs +++ b/src/ext_mem.rs @@ -533,7 +533,8 @@ where // local-disk wall time is neutral or marginally improved. See // `bench/README.md` for the empirical sizing. let n_phys = effective_physical_file_count(opts.physical_file_count); - let phase1_pool = BucketPool::new(n_phys, &work_dir)?; + // One pool now, not two: the fused phase 1 writes partition buckets + // directly, so there are no subarray buckets to hold. let phase3_pool = BucketPool::new(n_phys, &work_dir)?; profile_log(&format!( @@ -541,22 +542,29 @@ where std::mem::size_of::() * 8 )); - let sub_factory = |i: usize| phase1_pool.new_bucket::>(i); let part_factory = |j: usize| phase3_pool.new_bucket::>(j); let t = Instant::now(); - let (mut subarray_buckets, samples) = phase1_sort_sample_spill::( + let pivots = phase0_presample_pivots::(text, lp, &source, p, opts, cmp); + profile_log(&format!( + "phase0 (presample pivots) {:.3}s", + t.elapsed().as_secs_f64() + )); + + let t = Instant::now(); + let mut partition_buckets = phase1_sort_and_distribute::( text, lp, &source, + &pivots, p, opts, cmp, seed_params, - sub_factory, + part_factory, )?; profile_log(&format!( - "phase1 (sort+sample+spill) {:.3}s", + "phase1 (sort+distribute) {:.3}s", t.elapsed().as_secs_f64() )); @@ -565,35 +573,9 @@ where // the caller's `Vec` (e.g. ~47 GB on a human-scale // _for_positions build); for `PositionSource::Filtered` it // frees the bitmap + cumsum (~770 MB); for `Identity` it's a - // no-op. The text and the spilled `subarray_buckets` are all - // phase 2+ needs. + // no-op. Phase 4 needs only the text and the partition buckets. drop(source); - let t = Instant::now(); - let pivots = phase2_select_pivots::(text, lp, samples, p, opts.max_context, cmp); - profile_log(&format!( - "phase2 (select pivots) {:.3}s", - t.elapsed().as_secs_f64() - )); - - let t = Instant::now(); - let mut partition_buckets = phase3_distribute::( - text, - lp, - &mut subarray_buckets, - &pivots, - p, - opts, - cmp, - part_factory, - )?; - profile_log(&format!( - "phase3 (distribute) {:.3}s", - t.elapsed().as_secs_f64() - )); - - drop(subarray_buckets); - let t = Instant::now(); let result = phase4_merge_and_emit::( text, @@ -1471,6 +1453,185 @@ where .collect()) } +/// Phase 0: choose the `p - 1` pivots *before* sorting anything. +/// +/// The old flow sampled from the already-sorted subarrays, which forced an +/// ordering: sort everything and spill it, pick pivots, then read all of it +/// back to distribute. That round trip is the single worst-scaling part of the +/// build — phase 3 gained 14% going from six threads to twelve, because it is +/// bound by the page-cache write path rather than by anything threads help +/// with. +/// +/// Pivots do not have to come from sorted data. **Any** splitters produce a +/// correct sample sort; only the balance of the partitions changes. So a cheap +/// pre-pass over the raw positions can supply them, and phase 1 can then sort +/// and distribute in one go, never materialising the subarrays at all. +/// +/// Sampling is by strided blocks rather than strided singletons: a +/// [`PositionSource`] fills a contiguous run cheaply but pays per call, and +/// the `Filtered` variant especially so. Blocks give the same coverage of the +/// position space for a fraction of the calls, and splitter quality is not +/// sensitive to the difference. +fn phase0_presample_pivots( + text: &[S], + lp: &L, + source: &PositionSource<'_>, + p: usize, + opts: &ExtMemOpts, + cmp: Cmp<'_>, +) -> Vec +where + S: Symbol, + I: Index, + L: LimitProvider, +{ + let n = source.len(); + if p <= 1 || n == 0 { + return Vec::new(); + } + const BLOCK: usize = 64; + let target = sample_target_total(n, p).min(n); + let n_blocks = target.div_ceil(BLOCK).max(1); + let stride = (n / n_blocks).max(1); + + let mut sample: Vec = Vec::with_capacity(n_blocks * BLOCK); + let mut start = 0usize; + while start < n && sample.len() < target { + let len = BLOCK.min(n - start); + let base = sample.len(); + sample.resize(base + len, I::zero()); + source.fill_chunk(start, &mut sample[base..]); + start += stride; + } + if sample.is_empty() { + return Vec::new(); + } + + let m = sample.len(); + let mut sa_w = vec![I::zero(); m]; + let mut lcp = vec![I::zero(); m]; + let mut lcp_w = vec![I::zero(); m]; + sample_sort::merge_sort( + text, + lp, + &mut sample, + &mut sa_w, + &mut lcp, + &mut lcp_w, + 0, + opts.max_context, + cmp, + ); + (1..p).map(|i| sample[(i * m / p).min(m - 1)]).collect() +} + +/// Phase 1, fused with the old phase 3: sort each subarray and write its +/// pieces straight into the partition buckets. +/// +/// Because [`phase0_presample_pivots`] has already chosen the splitters, a +/// subarray never has to be spilled and read back. That removes one complete +/// write-and-read round trip of every record from the build. +#[allow(clippy::too_many_arguments)] +fn phase1_sort_and_distribute( + text: &[S], + lp: &L, + source: &PositionSource<'_>, + pivots: &[I], + p: usize, + opts: &ExtMemOpts, + cmp: Cmp<'_>, + seed_params: Option<&crate::radix::Packer>, + mk_bucket: MkB, +) -> io::Result> +where + S: Symbol, + I: Index, + L: LimitProvider, + SaLcp: BucketRecord, + B: SaLcpBucketStore + Send, + MkB: Fn(usize) -> B + Send + Sync, +{ + let n = source.len(); + let chunk_size = n.div_ceil(p); + let partition_buckets: Vec> = (0..p).map(|j| Mutex::new(mk_bucket(j))).collect(); + + (0..p).into_par_iter().try_for_each(|i| -> io::Result<()> { + let start = (i * chunk_size).min(n); + let end = ((i + 1) * chunk_size).min(n); + let len = end - start; + if len == 0 { + return Ok(()); + } + + let mut sa: Vec = vec![I::zero(); len]; + source.fill_chunk(start, &mut sa); + let mut sa_w = vec![I::zero(); len]; + let mut lcp_arr = vec![I::zero(); len]; + let mut lcp_w = vec![I::zero(); len]; + if !crate::radix::seed_subarray( + text, + lp, + seed_params, + &mut sa, + &mut lcp_arr, + &mut sa_w, + &mut lcp_w, + opts.max_context, + cmp, + ) { + sample_sort::merge_sort( + text, + lp, + &mut sa, + &mut sa_w, + &mut lcp_arr, + &mut lcp_w, + 0, + opts.max_context, + cmp, + ); + } + drop(sa_w); + drop(lcp_w); + + // Split the sorted subarray at the pivots and hand each piece to its + // partition. Splits are non-decreasing, so each search gallops from + // the previous one. + let records: Vec> = sa + .iter() + .zip(lcp_arr.iter()) + .map(|(&pos, &lcp)| SaLcp { pos, lcp }) + .collect(); + drop(sa); + drop(lcp_arr); + + let mut splits = Vec::with_capacity(p + 1); + splits.push(0usize); + let mut from = 0usize; + for &pivot in pivots { + from = upper_bound_from(&records, from, pivot, text, lp, opts.max_context, cmp); + splits.push(from); + } + splits.push(records.len()); + + for j in 0..p { + let (lo, hi) = (splits[j], splits[j + 1]); + if lo >= hi { + continue; + } + let mut bucket = partition_buckets[j].lock().unwrap(); + bucket.add_slice_reset_first_lcp(&records[lo..hi])?; + bucket.mark_boundary(); + } + Ok(()) + })?; + + Ok(partition_buckets + .into_iter() + .map(|m| m.into_inner().expect("partition mutex poisoned")) + .collect()) +} + /// Upper bound of `pivot` in `records`, searched forward from `from`. /// /// Returns the first index at or after `from` whose suffix is strictly