From ffe3680ac471daafef43eb8c95171ce08a2df031 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Tue, 11 Aug 2026 21:33:42 +0200 Subject: [PATCH 1/5] 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 2/5] 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 3/5] 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 4/5] 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 5/5] 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