From ffe3680ac471daafef43eb8c95171ce08a2df031 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 21:33:42 +0200 Subject: [PATCH 01/20] 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 | 47 ++++++ 4 files changed, 465 insertions(+) create mode 100644 src/radix.rs diff --git a/src/lib.rs b/src/lib.rs index 06fc661..e9ef2bc 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -22,6 +22,7 @@ mod ext_mem; mod lcp; mod lcp_memo; mod limits; +mod radix; mod sample_sort; pub use ext_mem::{ diff --git a/src/limits.rs b/src/limits.rs index e4e5ead..3a1d9e8 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 b3343aa..9eabf86 100644 --- a/src/sample_sort.rs +++ b/src/sample_sort.rs @@ -121,11 +121,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. From 148f266f0b0de4b8a2e92591c385b9d06ad153a4 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 21:33:42 +0200 Subject: [PATCH 02/20] 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 17945d8..53f5de3 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() @@ -174,6 +198,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 e9ef2bc..407a822 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -43,6 +43,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 9a39247020075108f7a39c03db5d7d827a8700f1 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 21:33:42 +0200 Subject: [PATCH 03/20] 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 2a4ebbb..c4fe1ab 100644 --- a/src/ext_mem.rs +++ b/src/ext_mem.rs @@ -43,7 +43,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 9369034c6ac16e2eedff1c03e5c84d1ff53b081d Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 21:33:42 +0200 Subject: [PATCH 04/20] 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 cda07cb1034eb451287778bb0f3ff22ada8deeda Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 21:33:42 +0200 Subject: [PATCH 05/20] 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 | 86 ++++++++++++++++++++++++++++++++++++++++++++++++- bench/README.md | 86 +++++++++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 169 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 5d9a382..1a2b29e 100644 --- a/README.md +++ b/README.md @@ -16,7 +16,48 @@ streams the SA out as positions are emitted. Both the in-memory and external-memory paths are implemented, tested on Linux, macOS, and Windows, and differentially verified against direct suffix -comparison on small, random, segmented, filtered, and finite-context inputs. +comparison on small, random, segmented, filtered, and finite-context 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 + +The external-memory and sample-sort paths still use the merge kernel, so +they retain the scan cost described above on repeat-heavy input. On the complete ruSTAR-shaped GENCODE Human v50 input (6.56 billion text symbols, 6.18 billion retained suffixes, and 1.40 million segments), caps-sa @@ -105,6 +146,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 f5ef8e24d818cfa33d50428290a377a26870483e Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 21:59:35 +0200 Subject: [PATCH 06/20] 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 fe411d81abb201524b92b41bbca0260966febd77 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 21:59:35 +0200 Subject: [PATCH 07/20] 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 c4fe1ab..0087e5a 100644 --- a/src/ext_mem.rs +++ b/src/ext_mem.rs @@ -739,6 +739,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 9eabf86..353fa28 100644 --- a/src/sample_sort.rs +++ b/src/sample_sort.rs @@ -35,6 +35,7 @@ use crate::lcp::{LcpDispatch, Symbol}; use crate::lcp_memo::GeometricMemo; use crate::limits::{LimitProvider, PlainText}; 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 @@ -173,6 +174,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. @@ -223,6 +283,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(); @@ -901,6 +965,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 14f3107a31c4dedd1fa2d143e1ba2146e19e2c43 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 21:59:35 +0200 Subject: [PATCH 08/20] 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 1a2b29e..a137913 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 @@ -185,9 +190,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 ef6a78b609107c7d24868e9e76264c5edd348276 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 22:53:06 +0200 Subject: [PATCH 09/20] 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/radix.rs | 316 +++++++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 269 insertions(+), 47 deletions(-) diff --git a/src/radix.rs b/src/radix.rs index 7eab7db..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. @@ -56,42 +59,234 @@ 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; -/// 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 + } + + /// 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))) + } +} + +/// 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 { + if size_of::() != 1 { + return None; } - let end = (p + k).min(text.len()); - let mut key: u64 = 0; - for &s in &text[p..end] { - key = (key << bits) | s as u64; + // 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(Packer::new(bytes)) +} + +/// 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, + packer: Option<&Packer>, + sa: &mut [I], + lcp: &mut [I], + sa_w: &mut [I], + lcp_w: &mut [I], + max_ctx: usize, + cmp: Cmp<'_>, +) -> bool { + 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; + } + 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(); + (packer.key_at(bytes, p), 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; } - // Shift the packed prefix up so the missing trailing fields read as zero. - key << (bits as usize * (k - (end - p))) + 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 @@ -119,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(); @@ -140,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 }) @@ -178,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 @@ -241,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. ---- // @@ -265,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() @@ -501,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 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 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_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 f9c42c940dc0aec105e723a5e9bfd8a6553506bf Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 22:55:59 +0200 Subject: [PATCH 10/20] 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 a137913..9348403 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 83f721eea39e8813f8f36235fe072d50703fc732 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 23:03:39 +0200 Subject: [PATCH 11/20] 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 cd28bfb7600465ef39f8681a40182ce3d35583b6 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 23:09:02 +0200 Subject: [PATCH 12/20] 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 9348403..1f7b5c8 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 54b6c5796f5001745c5ab874815326794266df8b Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 23:16:11 +0200 Subject: [PATCH 13/20] 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 3400b1c9255009624f3da69cb873cce547419a25 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 23:16:28 +0200 Subject: [PATCH 14/20] 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 1f7b5c8..7f57881 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 51016012b21146a8406de0ab523ba33cf93ad258 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Thu, 13 Aug 2026 19:45:00 +0200 Subject: [PATCH 15/20] refactor: drop the external-memory seed from the doubling module `radix.rs` carried a phase-1 subarray seed alongside the in-memory doubling path, and that seed took its comparator as a `runs::Cmp`, the run-skipping wrapper closed as #10. It is also unreachable on 0.7.0, whose phase 1 fuses sorting and distribution and never called it. The seed is worth having, but as a re-derivation against the phase 1 that exists, which is #15. Removing it here leaves this module with the in-memory doubling path alone, which is what this PR is about, and drops the last dependency on the closed work. Co-Authored-By: Claude Opus 5 (1M context) --- src/radix.rs | 135 +-------------------------------------------------- 1 file changed, 2 insertions(+), 133 deletions(-) diff --git a/src/radix.rs b/src/radix.rs index 96fc9dc..4dcf5a7 100644 --- a/src/radix.rs +++ b/src/radix.rs @@ -59,10 +59,6 @@ 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; @@ -156,6 +152,8 @@ impl Packer { } } + /// Bits per packed field. Used by the key-geometry tests. + #[cfg(test)] #[inline] pub(crate) fn bits(&self) -> u32 { self.bits @@ -239,135 +237,6 @@ impl Packer { } } -/// 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 { - 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(Packer::new(bytes)) -} - -/// 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, - packer: Option<&Packer>, - sa: &mut [I], - lcp: &mut [I], - sa_w: &mut [I], - lcp_w: &mut [I], - max_ctx: usize, - cmp: Cmp<'_>, -) -> bool { - 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; - } - 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(); - (packer.key_at(bytes, p), 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) -} - /// 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; From 94921f6512d493c9daa469eac2906f37260cfc15 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Thu, 13 Aug 2026 20:36:03 +0200 Subject: [PATCH 16/20] perf: stop materialising the seed key array, and narrow the group bounds The doubling path's peak was not the suffix array or the ranks, it was the seed's key array: eight bytes per position against the four `sa` costs, still resident when the doubling rounds allocated theirs. Three changes take it out. The key array goes entirely. Pass 2 computes each key to pick its bucket and then discards it, and pass 3 recomputes the keys of one bucket at a time into a local buffer. That is one extra key per position, spread across workers, against the largest allocation in the build. Group starts move into a bit per slot, set while each bucket's keys are still alive. Grouping never needed the keys themselves, only where one tied group ends and the next begins, and a bucket boundary is always a group boundary because two buckets differ in the key by construction. Two buckets can share a word, so the bits are set with relaxed fetch-or. The per-round rank scratch is sized to the tied population instead of to the text, with each group taking a prefix-sum window, and group bounds are stored in the index type rather than in `usize`. Only slots inside a group are ever rewritten, and even the first round has far fewer of those than the text has positions. chr21 forward ++ revcomp, coded `0..=3`, 80 MB, 12 threads, output verified: peak RSS 1.83 GB -> 0.90 GB (merge kernel: 1.27 GB) wall 0.49 s -> 0.73 s (merge kernel: 4.35 s) So the doubling path now costs less memory than the kernel it replaces, rather than 44% more, and stays 5.9x faster. Recomputing the keys is what the wall time pays for that. Raw FASTA is unchanged in character: 1.46 GB against the kernel's 0.75 GB, 1.08 s against 23 s. Its wider alphabet forces 8-bit fields, so the seed resolves 8 symbols instead of 32 and leaves very large tied groups inside the `N` runs, whose per-group buffers set the peak. Texts that reach caps-sa through a parser rather than as raw FASTA do not have this shape. Co-Authored-By: Claude Opus 5 (1M context) --- src/radix.rs | 204 +++++++++++++++++++++++++++++++++++---------------- 1 file changed, 141 insertions(+), 63 deletions(-) diff --git a/src/radix.rs b/src/radix.rs index 4dcf5a7..3b35d45 100644 --- a/src/radix.rs +++ b/src/radix.rs @@ -60,6 +60,7 @@ use crate::Index; use crate::ext_mem::profile_log; use rayon::prelude::*; +use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Instant; /// An order-preserving remap of the bytes that actually occur in a text onto @@ -251,7 +252,13 @@ 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. +/// sorted positions and a bit per slot marking where a tied group starts. +/// +/// The key array does not come back. Grouping only needs to know where one +/// group ends and the next begins, which is one bit per slot rather than the +/// eight bytes a key costs, and dropping the keys here rather than after the +/// grouping pass is what keeps the doubling path's peak below the merge +/// kernel's. See [`build_sa`] for the accounting. /// /// This is an MSD counting sort rather than a comparison sort, for three /// reasons that all matter at genome scale: @@ -268,7 +275,7 @@ fn seed_sort( text: &[u8], packer: &Packer, visible_len: &(dyn Fn(usize) -> usize + Sync), -) -> (Vec, Vec) { +) -> (Vec, Vec) { let n = text.len(); let bucket_of = |key: u64| -> usize { (key >> (64 - RADIX_BITS)) as usize }; @@ -310,12 +317,18 @@ fn seed_sort( 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]; + // Pass 2: scatter positions into their buckets. Each chunk owns a + // disjoint slice of every bucket, so the writes never collide even though + // they are not contiguous. + // + // The key is computed here to pick the bucket and then thrown away. An + // `n`-entry key array would be the largest allocation in the whole build, + // 8 bytes per position against the 4 that `sa` costs, and it would still + // be resident when the doubling rounds need their own arrays. Pass 3 + // recomputes the keys of one bucket at a time instead, which is one extra + // key per position spread across the workers. 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() @@ -324,50 +337,70 @@ fn seed_sort( let mut cursor: Vec = offsets[c * RADIX_BUCKETS..(c + 1) * RADIX_BUCKETS].to_vec(); for p in start..end { - let key = packer.key_at(text, p); - let slot = &mut cursor[bucket_of(key)]; + let slot = &mut cursor[bucket_of(packer.key_at(text, p))]; // 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)); - } + unsafe { 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; + // Pass 3: order within each bucket, and record where tied groups start. + // + // 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, and two positions in different buckets differ in the + // key by construction, so a bucket's first slot always starts a group and + // the group-start bits can be filled in here, per bucket, while that + // bucket's keys exist. + // Atomic words: two buckets can share the word their boundary bits live + // in, so the bits are set with relaxed fetch-or. Reads are relaxed loads + // and cost nothing once the fill is done. + let starts: Vec = (0..n.div_ceil(64)).map(|_| AtomicU64::new(0)).collect(); let mut rest_sa: &mut [I] = &mut sa; - let mut slices: Vec<(&mut [u64], &mut [I])> = Vec::with_capacity(RADIX_BUCKETS); + let mut slices: Vec<(usize, &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; + slices.push((bucket_start[b], sb)); 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()))) + { + let start_bits = &starts; + let set_bit = |h: usize| { + start_bits[h / 64].fetch_or(1 << (h % 64), Ordering::Relaxed); + }; + slices.into_par_iter().for_each(|(base, sb)| { + if sb.is_empty() { + return; + } + set_bit(base); + if sb.len() > 1 { + let mut pairs: Vec<(u64, I)> = sb + .iter() + .map(|&p| (packer.key_at(text, p.to_usize()), p)) + .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, &(_, pos)) in pairs.iter().enumerate() { + sb[i] = pos; + } + for i in 1..pairs.len() { + let (ka, pa) = pairs[i - 1]; + let (kb, pb) = pairs[i]; + if ka != kb || visible_len(pa.to_usize()) != visible_len(pb.to_usize()) { + set_bit(base + i); + } + } + } }); - for (i, &(key, pos)) in pairs.iter().enumerate() { - kb[i] = key; - sb[i] = pos; - } - }); + } - (keys, sa) + (starts, sa) } /// Build the standard lexicographic suffix array of `text` by radix-seeded @@ -411,17 +444,17 @@ pub(crate) fn build_sa(text: &[u8]) -> Vec { t0.elapsed().as_secs_f64() )); let t1 = Instant::now(); - let (keys, mut sa) = seed_sort::(text, &packer, &visible_len); + let (starts, mut sa) = seed_sort::(text, &packer, &visible_len); profile_log(&format!( "radix seed sort {:.3}s", t1.elapsed().as_secs_f64() )); let t2 = Instant::now(); - // 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()) - }; + // Slot `h` begins a tied group exactly when the seed marked it, so two + // adjacent slots tie exactly when the later one is not a start. + let starts_group = + |h: usize| -> bool { starts[h / 64].load(Ordering::Relaxed) >> (h % 64) & 1 == 1 }; // `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, @@ -435,14 +468,14 @@ pub(crate) fn build_sa(text: &[u8]) -> Vec { // `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) + let groups: Vec<(I, I)> = (0..n) .into_par_iter() .filter_map(|h| { - if h > 0 && seed_eq(h - 1, h) { + if !starts_group(h) { return None; } let mut e = h + 1; - while e < n && seed_eq(e, h) { + while e < n && !starts_group(e) { e += 1; } let g = I::from_usize(h); @@ -452,11 +485,16 @@ pub(crate) fn build_sa(text: &[u8]) -> Vec { // no other thread writes. unsafe { ranks.set(entry.to_usize(), g) }; } - (e - h > 1).then_some((h, e)) + (e - h > 1).then_some((I::from_usize(h), I::from_usize(e))) }) .collect(); let mut groups = groups; - drop(keys); + profile_log(&format!( + "radix groups {} groups, {} MB", + groups.len(), + groups.len() * std::mem::size_of::<(I, I)>() / (1 << 20) + )); + drop(starts); profile_log(&format!( "radix grouping {:.3}s", t2.elapsed().as_secs_f64() @@ -465,9 +503,17 @@ pub(crate) fn build_sa(text: &[u8]) -> Vec { // ---- 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]; + // Scratch for the new rank of each *tied* `sa` slot, so a round's reads + // of `rank` never observe that same round's writes. + // + // Sized to the tied population rather than to `n`. Only slots inside a + // group are rewritten, and even the first round has far fewer of those + // than the text has positions (18% of them on chr21), so a full second + // rank array would be mostly untouched pages. Each group gets a + // contiguous window at its prefix-sum offset, and the buffer is reused + // across rounds, which shrink monotonically. + let mut next_rank: Vec = Vec::new(); + let mut offsets: Vec = Vec::new(); while !groups.is_empty() { let round_t = Instant::now(); @@ -484,17 +530,38 @@ pub(crate) fn build_sa(text: &[u8]) -> Vec { // 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. + // Window each group takes in the tied-slot buffer. + offsets.clear(); + offsets.reserve(groups.len() + 1); + let mut tied = 0usize; + for &(start, end) in &groups { + offsets.push(tied); + tied += end.to_usize() - start.to_usize(); + } + offsets.push(tied); + if next_rank.len() < tied { + next_rank.resize(tied, I::zero()); + } + 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 + let offsets_ref = &offsets; + let sub: Vec<(I, I)> = groups .par_iter() - .flat_map_iter(|&(start, end)| { + .enumerate() + .flat_map_iter(|(gi, &(start, end))| { + let (start, end) = (start.to_usize(), end.to_usize()); 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)) }; + // group is the sole owner of `start..end` in `sa` and of its + // own prefix-sum window in the tied-slot buffer. + let (sa_g, nr_g) = unsafe { + ( + sa_cell.slice_mut(start, len), + nr_cell.slice_mut(offsets_ref[gi], len), + ) + }; let succ = |p: usize| -> u64 { // End-of-text sorts first: the shorter suffix is smaller. match p.checked_add(depth) { @@ -517,7 +584,7 @@ pub(crate) fn build_sa(text: &[u8]) -> Vec { }; keyed.sort_unstable(); - let mut fresh = Vec::new(); + let mut fresh: Vec<(I, I)> = Vec::new(); let mut i = 0; while i < len { let key = keyed[i].0; @@ -530,7 +597,7 @@ pub(crate) fn build_sa(text: &[u8]) -> Vec { *slot = g; } if j - i > 1 { - fresh.push((start + i, start + j)); + fresh.push((I::from_usize(start + i), I::from_usize(start + j))); } i = j; } @@ -545,18 +612,29 @@ pub(crate) fn build_sa(text: &[u8]) -> Vec { // 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 { - // 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]) }; - } - }); + groups + .par_iter() + .enumerate() + .for_each(|(gi, &(start, end))| { + let base = offsets[gi]; + for (i, slot) in (start.to_usize()..end.to_usize()).enumerate() { + // SAFETY: `sa[start..end]` are distinct positions owned + // solely by this group, and the groups partition their + // index range. + unsafe { ranks.set(sa[slot].to_usize(), next_rank[base + i]) }; + } + }); - let before: usize = groups.iter().map(|&(s, e)| e - s).sum(); + let before: usize = groups + .iter() + .map(|&(s, e)| e.to_usize() - s.to_usize()) + .sum(); let n_groups = groups.len(); groups = sub; - let after: usize = groups.iter().map(|&(s, e)| e - s).sum(); + let after: usize = groups + .iter() + .map(|&(s, e)| e.to_usize() - s.to_usize()) + .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 From b53d0975c763356a79875fadeb2e8c47f39638af Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Wed, 12 Aug 2026 07:46:31 +0200 Subject: [PATCH 17/20] bench: add an ACGT-start filter to the benchmark CLI `--filter-acgt` routes the external-memory path through `build_ext_mem_for_filter`, sorting only suffixes that start below code 4. That is the shape a STAR-style genome index uses, where A/C/G/T participate and N and spacers do not, and it makes the CLI able to reproduce that workload's cost without a segmented fixture. Co-Authored-By: Claude Opus 5 (1M context) --- examples/caps_sa.rs | 28 ++++++++++++++++++++++++++-- src/radix.rs | 15 +++++++++++++++ 2 files changed, 41 insertions(+), 2 deletions(-) diff --git a/examples/caps_sa.rs b/examples/caps_sa.rs index 53f5de3..da6de17 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, } } @@ -171,7 +184,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/radix.rs b/src/radix.rs index 3b35d45..6577911 100644 --- a/src/radix.rs +++ b/src/radix.rs @@ -731,6 +731,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""); From f823bee86e4639e37716592fc2d84efe8f0eb54b Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Wed, 12 Aug 2026 08:29:02 +0200 Subject: [PATCH 18/20] 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 353fa28..beceada 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 @@ -1004,6 +1034,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 04ae2e222c32a78a5529bbc7a99b02694bec3af6 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Wed, 12 Aug 2026 08:37:25 +0200 Subject: [PATCH 19/20] 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 407a822..e406e6a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -111,6 +111,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 6577911..8821d3f 100644 --- a/src/radix.rs +++ b/src/radix.rs @@ -238,6 +238,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 +} + /// 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; @@ -746,6 +788,46 @@ mod tests { assert!(seed_params(&bytes).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 252b0aab8c781b21b5d093325d441ed8a4e9c83f Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Thu, 13 Aug 2026 19:48:37 +0200 Subject: [PATCH 20/20] fix: keep the signed-symbol guard covered on the surviving path The `i8` test asserted on the external-memory seed's eligibility check, which no longer exists here. The guard it protects does: the in-memory fast path demands exactly `u8`, because `Symbol` covers `i8` too and a packed key orders its fields as unsigned, so `-1` (byte `0xFF`) would sort above `1`. The test now goes through `build_in_memory` and asserts an `i8` text comes back in signed order. Also fills in `subset_full_sa_budget` in a test's `Opts` literal, which the budget commit added to the struct. Co-Authored-By: Claude Opus 5 (1M context) --- src/radix.rs | 15 --------------- src/sample_sort.rs | 16 ++++++++++++++++ 2 files changed, 16 insertions(+), 15 deletions(-) diff --git a/src/radix.rs b/src/radix.rs index 8821d3f..4f3c1eb 100644 --- a/src/radix.rs +++ b/src/radix.rs @@ -773,21 +773,6 @@ 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()); - } - /// 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] diff --git a/src/sample_sort.rs b/src/sample_sort.rs index beceada..f73b1f0 100644 --- a/src/sample_sort.rs +++ b/src/sample_sort.rs @@ -906,6 +906,21 @@ mod tests { } } + /// `Symbol` is implemented for `i8`, and a one-byte-wide check alone would + /// let a signed text through a packer that orders its fields as unsigned: + /// `-1` has byte `0xFF`, so it would sort above `1`, inverting the text's + /// real order. The fast-path guard demands exactly `u8`, so a signed text + /// stays on the merge kernel and keeps signed order. + #[test] + fn signed_symbol_texts_keep_signed_order() { + let text: Vec = vec![-1, 0, -1, 1, -2, 0, 1, -1, 0, -2, 1, 0]; + let got: Vec = build_in_memory(&text); + let dispatch = LcpDispatch::detect(); + let mut want: Vec = (0..text.len() as u32).collect(); + want.sort_by(|&a, &b| dispatch.suffix_cmp(&text, a as usize, b as usize, usize::MAX)); + assert_eq!(got, want, "an i8 text must sort by signed order"); + } + #[test] fn suffix_array_respects_finite_max_context() { use rand::{RngExt, SeedableRng}; @@ -917,6 +932,7 @@ mod tests { let text: Vec = (0..n).map(|_| rng.random_range(0..3u8)).collect(); let opts = Opts { max_context: max_ctx, + ..Opts::default() }; let got: Vec = build_in_memory_with_opts(&text, &opts); let mut want: Vec = (0..n as u32).collect();