4/4 perf: alphabet ranking, SWAR key packing, and doubling-round overhead - #6
Open
BenjaminDEMAILLE wants to merge 16 commits into
Open
4/4 perf: alphabet ranking, SWAR key packing, and doubling-round overhead#6BenjaminDEMAILLE wants to merge 16 commits into
BenjaminDEMAILLE wants to merge 16 commits into
Conversation
Open
4 tasks
Contributor
Author
|
Part of the stack tracked in #7. |
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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
`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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
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) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
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<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) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This was referenced Aug 13, 2026
Contributor
Author
|
Status against 0.7.0: this does not rebase as it stands. Its The part of it that matters for the ruSTAR workload, the segment-aligned packed Leaving this open but stale rather than force-pushing something unreviewable. |
`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 COMBINE-lab#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 COMBINE-lab#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) <noreply@anthropic.com>
BenjaminDEMAILLE
force-pushed
the
stack/4-key-packing
branch
from
August 13, 2026 17:52
3ef11f2 to
5101601
Compare
Contributor
Author
|
Updated: this now rebases onto 0.7.0 and is reviewable. The blocker was that Fresh numbers, gate checks, and the commit-by-commit disposition are in the |
Contributor
Author
|
One correction to the numbers below, which I would rather state than have you
Prefix doubling needs a rank for every position, so this is inherent to the That is a trade for you to accept or refuse explicitly, not one to make |
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) <noreply@anthropic.com>
Contributor
Author
|
Peak RSS follow-up, superseding my note above for #6 and #8. The regression was real and it is now gone, on the coded-DNA input. The chr21 forward ++ revcomp, coded
So the path costs 29% less memory than the kernel it replaces rather than 44% Raw FASTA is not fixed and I am not claiming it is: 1.46 GB against the #3 and #4 still carry the old peak. They predate the seed structure this fixes, Segmented gate re-checked after the change, through the harness in #14 on the |
Contributor
Author
|
Follow-up on the raw-FASTA peak, which the previous note left open. I did not Where it is. On chr21 raw FASTA the seed alone reaches 0.98 GB against a live Why the obvious cause is not the cause. The input degenerates the radix pass: What does move it is the parallel granularity, with the loop body empty:
An empty closure over the bucket list reproduces the growth on its own, so I am not shipping a The coded-DNA result in the previous note stands and is unaffected: 0.90 GB |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Three changes to how keys are built and how doubling rounds spend their time. Both were found by measuring rather than reasoning, and in both cases my first guess was wrong.
Rank the alphabet before packing
The field width was chosen from the largest byte value rather than the number of distinct symbols, and for the most natural input format those differ sharply. 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 four symbols allow.Ranking the occurring bytes onto a dense range fixes it. The map is monotone by construction, so keys stay order-preserving and the zero-padding argument carries over: code 0 remains the minimum.
Both inputs benchmarked until now hid this: one is pre-coded to
0..3and already dense, the other is dominated byNruns where key depth is irrelevant. A FASTA-derived ACGT text is the common case neither covered.Cross-check: the suffix array of the ASCII text is identical to that of the
0..3-coded text, as it must be — the encodings are order-isomorphic.SWAR key packing
Building a key was a chain of
kdependent shift-or-lookup steps. The alphabet map is now applied to the whole text once (removing the dependent table load), and fields are gathered by a binary-tree SWAR shuffle: eight symbols cost three shift-or-mask pairs instead of eight dependent steps. Seed sort 0.281 s → 0.177 s.The ranked copy is only materialised when the identity map does not already rank the text, so pre-coded input pays nothing.
A bug worth noting: the group concatenation shifts by
8 * bits, which is 64 for 8-bit fields — not a legalu64shift. Release masks it to zero and happens to give the right answer; debug panics. Caught by running the tests in debug, which is why the CI workflow from stack 1 does.Doubling rounds
Instrumenting the rounds showed the cost was not where I assumed. The first round has 3.5 million tied groups averaging four elements:
At that size neither the sort nor the rank probes dominate — the bookkeeping does: one heap allocation per group, and a sequential
split_at_mutchain over all 3.5 M groups.Scatteralready encodes "disjoint ranges, one owner each", exactly the groups' property, so each takes its sub-slices directly with no prepass; groups up to 32 elements build their keys in a stack buffer.Doubling 0.502 s → 0.234 s, first round 0.124 s → 0.053 s.
The per-round log stays behind
CAPS_SA_PROFILE— it is what located this, and group-size distribution is the first thing to look at for any further work.Final state
0..3, 80 MBAll six outputs verified byte-identical to the original merge kernel;
verify_sapasses; 87 tests green in debug and release.Beyond this the remaining step is a linear-time algorithm (SA-IS / DC3), which replaces CaPS-SA rather than optimising it.
🤖 Generated with Claude Code