perf(collect): amortize per-pivot tokenize via a rolling forward key cache - #65
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (3)
WalkthroughAdds a rolling forward 32-mer cache in ChangesKeyed forward search via rolling key cache
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant collect_smems_unsorted_into
participant zz_step1
participant zz_right_emit
participant mem_search_keyed
collect_smems_unsorted_into->>collect_smems_unsorted_into: fwd_keys = rolling_forward_keys_into(read)
collect_smems_unsorted_into->>zz_step1: fwd_keys
zz_step1->>zz_right_emit: pivot, fwd_keys
zz_right_emit->>mem_search_keyed: query, optional query_key
mem_search_keyed-->>zz_right_emit: MemMatch
Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
3d8fb77 to
5491c6d
Compare
1a42aea to
0baae2c
Compare
0baae2c to
3091fe8
Compare
|
@coderabbitai review |
✅ Action performedReview finished.
|
|
@coderabbitai review |
✅ Action performedReview finished.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
prmi/src/index/spectrum.rs (1)
2058-2108: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winByte-identity is correct; guard the caller-key contract in-function.
The
qm_keyreuse is sound:match_len >= KMER_LEN⇒query.len() >= 32, soquery_key == tokenize_32mer(query, 32) == tokenize_32mer(qm, 32). Correct.The residual footgun:
query_key: Option<u64>is used as the actual compare key (not a seed likehint), so a wrongSome(k)from a future caller yields a silently non-byte-identical result in release. The in-tree call sitesdebug_assertthis, but thesepubprimitives do not. Mirror the assertion inside so misuse is caught regardless of caller:debug_assert!( query_key.is_none() || query_key == Some(tokenize_32mer(query, query.len().min(KMER_LEN))), "mem_search_keyed: query_key != fresh tokenize" );Same applies to
mem_search_warmstart_keyed(Line 2215) andforward_truncate_below_maximal_keyed(Line 3171).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@prmi/src/index/spectrum.rs` around lines 2058 - 2108, The keyed search functions trust an optional caller-provided key too much, which can silently produce incorrect non-byte-identical results if a future caller passes the wrong value. Add an internal debug_assert in mem_search_keyed to verify query_key matches a fresh tokenize_32mer over query, and mirror the same caller-key contract check in mem_search_warmstart_keyed and forward_truncate_below_maximal_keyed so misuse is caught even when in-tree call-site assertions are absent.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@prmi/src/encoding.rs`:
- Around line 94-128: The hot path in `CollectScratch::fwd_keys` is still
allocating because it replaces the buffer with `rolling_forward_keys(read)`
instead of reusing the existing `Vec<u64>`. Add or switch to an in-place fill
API (for example an `*_into` variant in `crate::encoding`) and have
`CollectScratch::fwd_keys` write directly into the existing scratch buffer
rather than assigning a new vector. Use the `rolling_forward_keys` logic as the
source of truth, but preserve and refill the existing `fwd_keys` storage.
In `@prmi/src/index/collect.rs`:
- Around line 587-597: The `collect.rs` read-processing path is still
reallocating `fwd_keys` for every read by replacing the scratch buffer with a
new `Vec<u64>`. Update the `rolling_forward_keys` flow so `collect` reuses the
existing `fwd_keys` storage in place, ideally by adding a
`rolling_forward_keys_into` helper and calling it from the loop where `scratch`,
`fill_next_n`, and `fwd_keys` are handled. Make sure the buffer is cleared or
resized without allocation so the per-read path stays allocation-free once
warmed up.
---
Outside diff comments:
In `@prmi/src/index/spectrum.rs`:
- Around line 2058-2108: The keyed search functions trust an optional
caller-provided key too much, which can silently produce incorrect
non-byte-identical results if a future caller passes the wrong value. Add an
internal debug_assert in mem_search_keyed to verify query_key matches a fresh
tokenize_32mer over query, and mirror the same caller-key contract check in
mem_search_warmstart_keyed and forward_truncate_below_maximal_keyed so misuse is
caught even when in-tree call-site assertions are absent.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: ab319c7b-a2b7-4d81-9fb6-18e08492229c
📒 Files selected for processing (3)
prmi/src/encoding.rsprmi/src/index/collect.rsprmi/src/index/spectrum.rs
…cache Build the per-read forward 32-mer key cache once (encoding::rolling_forward_keys, one O(read.len()) right-to-left pass) into CollectScratch.fwd_keys, and thread it through the pass-1/pass-2 forward searches so a pivot reads fwd_keys[pivot] instead of re-tokenizing read[pivot..]. Consulted only under a query.len() >= 32 guard: an N clamps the query below 32, so the cache is read only on N-free full 32-mer windows, where the rolling key is byte-identical to a fresh tokenize_32mer. Adds keyed entry points that accept the precomputed key -- mem_search_keyed, mem_search_warmstart_keyed, forward_truncate_below_maximal_keyed (each the existing function taking an Option<u64> key; None re-tokenizes internally, byte-identical). They also reuse the full-32 query key as the maximal match's qm key when match_len >= 32 (qm is a prefix, same first 32 bases), dropping a second tokenize per hit. This is a byte-identical work reduction, NOT a wall-clock speedup: the removed compute is hidden behind the DRAM-miss latency that bounds cold seeding, so the cold wall is neutral (measured ~-4..-8% retired instructions/read on Graviton4, gate-subtracted). Gated by the rolling_forward_keys_* proptests (key == tokenize at every offset; N-safe and identical on N-free windows) plus per-pivot debug_asserts (cache key == fresh tokenize) that run across the existing collect_smems / zz_* byte-identity corpus.
3091fe8 to
d8dbd50
Compare
|
@coderabbitai review |
✅ Action performedReview finished.
|
|
@coderabbitai review |
✅ Action performedReview finished.
|
|
@coderabbitai review |
✅ Action performedReview finished.
|
1 participant
Third of a three-PR stack. Stacked on #64 (
perf/boundary-search) — review #63 then #64 first; this PR's diff is against #64.What
Amortize the per-pivot
tokenize_32merin SMEM collection. Build the per-read forward 32-mer key cache once (encoding::rolling_forward_keys, oneO(read.len())right-to-left pass) intoCollectScratch.fwd_keys, and thread it through the pass-1/pass-2 forward searches so a pivot readsfwd_keys[pivot]instead of re-tokenizingread[pivot..]. Consulted only under aquery.len() >= 32guard: anNclamps the query below 32, so the cache is read only onN-free full 32-mer windows, where the rolling key is byte-identical to a freshtokenize_32mer.Adds keyed entry points that accept the precomputed key —
mem_search_keyed,mem_search_warmstart_keyed,forward_truncate_below_maximal_keyed(each the existing function taking anOption<u64>key;Nonere-tokenizes internally, byte-identical). They also reuse the full-32 query key as the maximal match'sqmkey whenmatch_len >= 32(qmis a prefix, so the same first 32 bases), dropping a second tokenize per hit.Wall / instructions
This is a byte-identical work reduction, not a wall-clock speedup: the removed compute is hidden behind the DRAM-miss latency that bounds cold seeding, so the cold wall is neutral. Measured ≈ −4..−8% retired instructions/read on Graviton4 (gate-subtracted, built without instrumentation).
Correctness
Byte-identical.
rolling_forward_keysisN-safe (masks each base, never panics) and gated byrolling_forward_keys_*proptests (key ==tokenize_32merat every offset;N-safe and identical onN-free full windows). Per-pivotdebug_asserts (cached key == fresh tokenize) run at every guarded call site across the existingcollect_smems/zz_*byte-identity corpus.Summary by CodeRabbit
Performance Improvements
Bug Fixes
Tests