perf: low-hanging allocation/compute sweep (model, lookup, spectrum, sa-file, train, revcomp, FFI fill) - #18
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (16)
✅ Files skipped from review due to trivial changes (2)
🚧 Files skipped from review as they are similar to previous changes (11)
📝 WalkthroughWalkthroughThis PR refactors spectrum generation to write directly into caller-provided buffers via a new StepSink abstraction, eliminating intermediate allocations. It also adds bit-packed N-position handling, refactors reverse-complement and encoding tests, switches verification histogram to a single-pass parallel fold, and adds sidecar I/O buffering and related tests. ChangesAllocation-free spectrum and integrated optimizations
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
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 |
5b273aa to
fa2e837
Compare
0a6d128 to
941ffc6
Compare
fa2e837 to
dcb4286
Compare
941ffc6 to
69c6f72
Compare
dcb4286 to
43bb1a3
Compare
69c6f72 to
98f7b95
Compare
98f7b95 to
4b7c1ad
Compare
|
@coderabbitai review |
✅ Action performedReview finished.
|
|
@coderabbitai review |
✅ Action performedReview finished.
|
|
@coderabbitai review |
✅ Action performedReview finished.
|
|
@coderabbitai review |
✅ Action performedReview finished.
|
|
@coderabbitai review |
✅ Action performedReview finished.
|
|
@coderabbitai review |
✅ Action performedReview finished.
|
|
@coderabbitai review |
✅ Action performedReview finished.
|
|
@coderabbitai review |
✅ Action performedReview finished.
|
|
@coderabbitai review |
✅ Action performedReview finished.
|
|
@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 (2)
prmi-sys/src/lib.rs (1)
1298-1345:⚠️ Potential issue | 🟠 Major | ⚡ Quick winDon't narrow the overlap contract without a compatibility path.
CHANGELOG.mdLines 170-177 still documentprmi_reverse_complement_2bitas handling overlap, and the previous temp-buffer implementation did too. This two-ended loop is only correct forout == in_or disjoint buffers, so partially overlapping callers now get silently corrupted output instead of the old behavior. That makes the change behavioral, not byte-identical.If you want the allocation-free fast path, keep it for disjoint/full-alias cases and fall back to a scratch buffer when the ranges partially overlap.
Compatibility-preserving approach
let n = len as usize; if n == 0 { return 0; } + let in_start = in_ as usize; + let out_start = out as usize; + let in_end = in_start.saturating_add(n); + let out_end = out_start.saturating_add(n); + let partial_overlap = + in_start != out_start && in_start < out_end && out_start < in_end; + + if partial_overlap { + let mut tmp = Vec::with_capacity(n); + for i in 0..n { + unsafe { tmp.push((*in_.add(n - 1 - i) & 0x3) ^ 0x3) }; + } + unsafe { std::ptr::copy_nonoverlapping(tmp.as_ptr(), out, n) }; + return 0; + } + unsafe { let mut lo = 0usize; let mut hi = n - 1; while lo < hi {🤖 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-sys/src/lib.rs` around lines 1298 - 1345, The new two-ended in-place algorithm in prmi_reverse_complement_2bit only works for full alias (out == in_) or fully disjoint buffers and will corrupt output for partial-overlap callers; change prmi_reverse_complement_2bit to detect partial overlap between the input range (in_.add(0) .. in_.add(n)) and the output range (out.add(0) .. out.add(n)) and if partial overlap is detected fall back to the old safe behavior: allocate a temporary buffer (or reuse an internal scratch buffer), perform the reverse-complement into it, then copy to out; otherwise keep the fast two-ended loop for the disjoint or out==in_ cases. Ensure the detection logic uses raw pointer comparisons on in_, out and n so it works in unsafe context and keep the same return/err handling in prmi_reverse_complement_2bit.prmi/src/train/verify.rs (1)
13-22:⚠️ Potential issue | 🟠 Major
HIST_DENSE_CAPstill permits ~2 GiB dense histograms per Rayonfoldaccumulator.In
compute_error_distribution, the fold includese <= cap(withcap = HIST_DENSE_CAP), and whene == capit grows the per-folddensevector tocap_len = cap + 1=268,435,457u64counters (~2,147,483,656 bytes ≈ 2.0 GiB). Because Rayonfoldkeeps independent accumulators per worker/chunk beforereduce, a badly-fit model can still drive multi-GiB peak memory on this verification path.🤖 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/train/verify.rs` around lines 13 - 22, HIST_DENSE_CAP is too large and can cause per-Rayon-fold allocs of ~2GiB; reduce the cap to a safe smaller value (e.g. 16 << 20 or another site-approved constant) and ensure vector sizing in compute_error_distribution uses that reduced cap as a usize when computing cap_len = (cap as usize) + 1; additionally, guard the per-fold growth so you only reserve/grow the per-worker dense Vec when cap_len is below a reasonable threshold (avoid blindly calling dense.resize(cap_len, 0) for the original huge cap) by checking the cap before resize and using min(cap_len, some_safe_max) or lazily spill to the overflow list for larger err values. Ensure references to HIST_DENSE_CAP and the per-fold `dense` vector growth in compute_error_distribution are updated accordingly.
🧹 Nitpick comments (1)
prmi/src/index/spectrum.rs (1)
1255-1265: ⚡ Quick winAvoid double-validating the packed pac on the auto path.
forward_spectrum_intoandforward_spectrum_tabled_intoalready perform this guard, soforward_spectrum_auto_intore-checks the same packed buffer on every packed query. Sincemem_searchand the FFI now funnel through this method, that adds avoidable work back into a hot path in a perf PR.Suggested simplification
pub(crate) fn forward_spectrum_auto_into<S: StepSink + ?Sized>( &self, query: &[u8], pac: &[u8], enc: PacEncoding, sink: &mut S, ) { - // Guard the packed pac before any walk (covers the `.kmt`-tabled and - // full paths, and the one-shot `mem_search` sink): a truncated buffer - // would otherwise be misread as a sentinel and yield a wrong interval. - if let PacEncoding::Packed { num_bases } = enc { - if validate_packed_pac(pac, num_bases, "forward_spectrum").is_err() { - return; - } - } match &self.kmt { Some(table) => self.forward_spectrum_tabled_into(query, pac, enc, table, sink), None => self.forward_spectrum_into(query, pac, enc, sink), } }🤖 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 1255 - 1265, In forward_spectrum_auto_into, remove the pre-match Packed validation block that calls validate_packed_pac for PacEncoding::Packed { num_bases } — forward_spectrum_into and forward_spectrum_tabled_into already perform this guard, so delete that duplicate validation (the if let PacEncoding::Packed { num_bases } = enc { ... validate_packed_pac(...) ... }) and let the downstream functions (forward_spectrum_into / forward_spectrum_tabled_into) handle packed PAC validation.
🤖 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-sys/src/lib.rs`:
- Around line 476-493: Guard the u64 max_steps before casting to usize to avoid
truncation on narrow-usize targets: in prmi_forward_spectrum (where
out_steps_as_smemstep(out_steps, max_steps as usize) is used) check if max_steps
> usize::MAX (or otherwise convert safely) and return an error code before
creating the slice; apply the same safe-guarded conversion in
prmi_backward_spectrum for its out_steps_as_smemstep call. Also reconcile the
docs for prmi_reverse_complement_2bit/CHANGELOG.md with the implementation:
either update the CHANGELOG/docs to state that partial overlap is not supported
(out must equal in_ or be fully disjoint) or alter the implementation to
correctly support partial overlap so the documented contract matches behavior.
In `@prmi/src/train/mod.rs`:
- Around line 239-250: Delay allocating n_positions_2x until you actually have
n_positions: instead of always calling NBitmap::zeros(text.len()), only create
the doubled-size bitmap inside the if let Some(ref np) { ... } block (use
NBitmap::zeros(text.len()) there), and when n_positions is None use a
lightweight empty bitmap (e.g. NBitmap::zeros(0) or a shared static empty
NBitmap) so masked_training_set(&sa, &text_bases, &n_positions_2x, &mask,
&config.prior) still gets a valid reference but you avoid the large allocation
for non-virtualized BuildSource::Pac builds; update the local variable
n_positions_2x and its scope accordingly.
---
Outside diff comments:
In `@prmi-sys/src/lib.rs`:
- Around line 1298-1345: The new two-ended in-place algorithm in
prmi_reverse_complement_2bit only works for full alias (out == in_) or fully
disjoint buffers and will corrupt output for partial-overlap callers; change
prmi_reverse_complement_2bit to detect partial overlap between the input range
(in_.add(0) .. in_.add(n)) and the output range (out.add(0) .. out.add(n)) and
if partial overlap is detected fall back to the old safe behavior: allocate a
temporary buffer (or reuse an internal scratch buffer), perform the
reverse-complement into it, then copy to out; otherwise keep the fast two-ended
loop for the disjoint or out==in_ cases. Ensure the detection logic uses raw
pointer comparisons on in_, out and n so it works in unsafe context and keep the
same return/err handling in prmi_reverse_complement_2bit.
In `@prmi/src/train/verify.rs`:
- Around line 13-22: HIST_DENSE_CAP is too large and can cause per-Rayon-fold
allocs of ~2GiB; reduce the cap to a safe smaller value (e.g. 16 << 20 or
another site-approved constant) and ensure vector sizing in
compute_error_distribution uses that reduced cap as a usize when computing
cap_len = (cap as usize) + 1; additionally, guard the per-fold growth so you
only reserve/grow the per-worker dense Vec when cap_len is below a reasonable
threshold (avoid blindly calling dense.resize(cap_len, 0) for the original huge
cap) by checking the cap before resize and using min(cap_len, some_safe_max) or
lazily spill to the overflow list for larger err values. Ensure references to
HIST_DENSE_CAP and the per-fold `dense` vector growth in
compute_error_distribution are updated accordingly.
---
Nitpick comments:
In `@prmi/src/index/spectrum.rs`:
- Around line 1255-1265: In forward_spectrum_auto_into, remove the pre-match
Packed validation block that calls validate_packed_pac for PacEncoding::Packed {
num_bases } — forward_spectrum_into and forward_spectrum_tabled_into already
perform this guard, so delete that duplicate validation (the if let
PacEncoding::Packed { num_bases } = enc { ... validate_packed_pac(...) ... })
and let the downstream functions (forward_spectrum_into /
forward_spectrum_tabled_into) handle packed PAC validation.
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: a5de2e62-ea47-4640-aecc-ac2e624d80b1
📒 Files selected for processing (16)
prmi-sys/src/lib.rsprmi-sys/tests/ffi_revcomp.rsprmi/proptest-regressions/encoding.txtprmi/src/encoding.rsprmi/src/index/lookup.rsprmi/src/index/spectrum.rsprmi/src/sidecar/model_file.rsprmi/src/sidecar/sa_file.rsprmi/src/train/mask.rsprmi/src/train/mod.rsprmi/src/train/training_set.rsprmi/src/train/verify.rsprmi/tests/index_lookup.rsprmi/tests/mask.rsprmi/tests/sidecar_model_file.rsprmi/tests/sidecar_sa_file.rs
4b7c1ad to
cc59733
Compare
|
@coderabbitai review |
✅ Action performedReview finished.
|
There was a problem hiding this comment.
🧹 Nitpick comments (1)
prmi/src/train/verify.rs (1)
210-230: 💤 Low valueDead variable
tscan be removed.
ts(line 213) is created but never used — onlyts2is used in the assertion. Thelet _ = ts;on line 227 just suppresses the warning for this unused binding.♻️ Proposed cleanup
fn error_distribution_empty_training_set() { // An empty training set should return all zeros without panicking. - let ts = make_ts(vec![], vec![]); let keys: Vec<u64> = vec![]; let sa_indices: Vec<u64> = vec![]; let ts2 = TrainingSet { keys: Keys::Materialized(Arc::new(keys)), sa_indices: SaIndices::Materialized(Arc::new(sa_indices)), sa_num: 0, weights: None, }; // We need a model; build one from a minimal non-empty ts, then test // compute_error_distribution against the empty ts. let ts_small = make_ts(vec![1u64 << 60], vec![0]); let config = TrainerConfig::default(); let model = train_with_config(&ts_small, 16, &config).unwrap(); - let _ = ts; // suppress unused warning assert_eq!(compute_error_distribution(&model, &ts2), (0, 0, 0, 0)); assert_eq!(compute_max_error_bound(&model, &ts2), 0); }🤖 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/train/verify.rs` around lines 210 - 230, Remove the unused temporary training set binding `ts` and its unused-suppression line; specifically delete the `let ts = make_ts(vec![], vec![]);` declaration and the `let _ = ts;` line in the `error_distribution_empty_training_set` test so the test uses only `ts2`, `ts_small`, `train_with_config`, `compute_error_distribution`, and `compute_max_error_bound` without dead variables.
🤖 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.
Nitpick comments:
In `@prmi/src/train/verify.rs`:
- Around line 210-230: Remove the unused temporary training set binding `ts` and
its unused-suppression line; specifically delete the `let ts = make_ts(vec![],
vec![]);` declaration and the `let _ = ts;` line in the
`error_distribution_empty_training_set` test so the test uses only `ts2`,
`ts_small`, `train_with_config`, `compute_error_distribution`, and
`compute_max_error_bound` without dead variables.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 430c5734-aa48-47dc-a1fe-60de89881b59
📒 Files selected for processing (16)
prmi-sys/src/lib.rsprmi-sys/tests/ffi_revcomp.rsprmi/proptest-regressions/encoding.txtprmi/src/encoding.rsprmi/src/index/lookup.rsprmi/src/index/spectrum.rsprmi/src/sidecar/model_file.rsprmi/src/sidecar/sa_file.rsprmi/src/train/mask.rsprmi/src/train/mod.rsprmi/src/train/training_set.rsprmi/src/train/verify.rsprmi/tests/index_lookup.rsprmi/tests/mask.rsprmi/tests/sidecar_model_file.rsprmi/tests/sidecar_sa_file.rs
✅ Files skipped from review due to trivial changes (1)
- prmi/proptest-regressions/encoding.txt
🚧 Files skipped from review as they are similar to previous changes (12)
- prmi/tests/sidecar_model_file.rs
- prmi/tests/index_lookup.rs
- prmi/src/sidecar/model_file.rs
- prmi/tests/sidecar_sa_file.rs
- prmi/src/train/mask.rs
- prmi/src/index/lookup.rs
- prmi/tests/mask.rs
- prmi/src/sidecar/sa_file.rs
- prmi/src/train/mod.rs
- prmi/src/encoding.rs
- prmi/src/train/training_set.rs
- prmi/src/index/spectrum.rs
…sa-file, train, revcomp, FFI fill)
cc59733 to
697e2c7
Compare
|
@coderabbitai review |
✅ Action performedReview finished.
|
1 participant
PR #5 of the v0.2 stack (carries commit
d1c447e). Base isfeat/v0.2-mem-search(#17) — stacked on #4. Tracked inV0.2_PR_STACK.mdrow #5.What this does
A low-hanging allocation/compute sweep across the hot paths — no behavior change, byte-identical results:
spectrum.rs— refactors the forward/backward spectra into sink-based cores (forward_spectrum_auto_into,forward_spectrum_into,backward_spectrum_inner_into) so theVec-returning entry points, the allocation-freemem_search(LastStepSink), and the FFI slice-fill all share one walk with no intermediateVec.train— bit-packedNBitmapfor the doubled-text N mask (replacesVec<bool>);.pacbuilds carryn_positions: None(no ~3.2 GB all-false alloc + scan).prmi-sys— the backward batch serial path fills each task's arena region directly viabackward_spectrum_fill(no per-task/whole-batchVec);+ ffi_revcompin-place test.encodingproptest regression seed.Carry-forward (the conflict resolution)
d1c447epredates the merged review fixes, so it reverted three and restructured the code they lived in. Re-applied each onto #5's new structure (compiler + full oracle suite as the safety net):NBitmap+Option<n_positions>materialized path (set(i)for the forward half,set(l_pac + (l_pac-1-i))for the RC half).forward_spectrum_into(the full-search core) andforward_spectrum_auto_into(covers the.kmt-tabled path and themem_searchsink).backward_spectrum_inner_intoand the reference oracle (thereturn steps→returnadaptation for the now-()-returning sink cores).-2(pac_num_basesoverflow) note with feat(prmi): sidecar file format — meta TOML, .sa, .l1/.l2, .skc #5's improved-4wording. The u64→usize guards sit in the non-conflicted validation loops and were preserved.Green-pass
cargo build --workspace✅cargo clippy --workspace --all-targets --all-features -- -D warnings✅cargo +nightly fmt --all -- --check✅cargo test --workspace✅ — 152 lib tests + all integration/FFI; the spectrum oracle proptests,mem_search ≡ maximal step,bwd_lockstep ≡ serial,*_fill_matches_vec, andstreamed_keys_match_materializedall confirm byte-identity; only Plan-3 deferralsignored.Summary by CodeRabbit
New Features
Performance
Bug Fixes
Tests