2/4 perf: MSD counting-sort seed, and cover subsets + in-memory sample sort

README.md

@@ -16,7 +16,53 @@ 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.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
+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 +151,67 @@ 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<u32> = 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.
+
+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
 
 The in-memory kernel is a parallel merge-sort whose two-way merge uses

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

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

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<usize>,
+    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<usize> = 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 <input> <output> [--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<I: caps_sa::Index>(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<u32> = 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<u64> = 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()