perf: query hot-path optimization stack (x86-validated) by nh13 · Pull Request #61 · fg-labs/prmi

nh13 · 2026-06-27T00:42:33Z

Summary

Query hot-path performance pass over the SMEM search, model-locate, and bloom dispatch gate. Every commit is byte-identical to the prior behaviour except one deliberately format-revising bloom commit (.blm body bytes change; no false negatives — details below). The 9-commit stack is based on main (#57's bloom dispatch gate is now merged); this PR targets main.

All wins were validated on x86 (the deployment target) on ephemeral c7i (Sapphire Rapids, AVX-512) and c6a (Zen3, AVX2) instances, because the aarch64 dev box materially under-reports LTO/SIMD/division-removal effects. Each commit went through a two-stage review (spec compliance + adversarial code/byte-identity review) plus a final whole-branch review.

Measured results (x86, base → HEAD)

Change	x86 AVX-512	x86 AVX2	Notes
`next_n` memoization (kernel)	2.84×	3.94×	replaces O(rlen²) per-read `fwd_qlen` rescans with one O(rlen) precompute
Lemire `bit_at` reduction (kernel)	2.29×	1.18×	removes a true 64-bit `div` per bloom probe; larger Intel win (slow div)
`unpack_packed_forward` chunks_exact (kernel)	~neutral	1.56×	amortizes per-step bounds checks on the >32 bp decode tail
thin-LTO + `codegen-units=1` (whole search)	+4.5–6.9% spectrum, +2.6–3.4% backward mem_search	+1.3–1.9%	end-to-end; `target-cpu` left operator-driven (artifacts stay portable)

Kernel ratios are isolated-primitive before/after; the LTO row is the end-to-end whole-search improvement.

Reading order (commit by commit)

perf(build) — thin-LTO + codegen-units=1 release profile + native bench harness. Re-baselines everything; the whole-search win above is mostly this.
test(sa) — gate a debug_assert should_panic test behind cfg(debug_assertions) so cargo test --release is green.
perf(collect) — memoize next-ambiguous-base distances per read in CollectScratch (the next_n kernel; O(rlen²)→O(1) across the three SMEM passes).
perf(spectrum) — hoist the keyed-compare mask out of the boundary-probe loop (byte-identical; mostly LTO-subsumed, kept as a clean invariant).
perf(batch) — zip the backward lockstep batch arm, drop dead windowed-locate guards, shrink the lockstep key buffers.
perf(spectrum) — chunks_exact restructure of unpack_packed_forward's word loop (the unpack kernel).
perf(collect) — drop the dead stage-1 sort in sort_within_read (the (m,n)-tied entries are provably identical in all fields, so stage 1 only reordered byte-identical duplicates; kept regression-guard tests).
(Tried and dropped — not in this stack.) A bloom any-window roll was prototyped but the x86 batch showed it regresses (≈10% AVX-512, ≈26% AVX2 — LLVM auto-vectorizes the simple per-window scan better than the hand-rolled state machine), so it was excluded. present_anchor_bloom is unchanged from the base.
perf(bloom) — division-free Lemire reduction in bit_at (the Lemire kernel). Format-revising: changes which bits a key sets, so a .blm body built by a prior binary differs. Header layout and the shared FORMAT_VERSION are unchanged (bumping it would invalidate every sidecar). Writer and reader share bit_at, so a self-consistent .blm has no false negatives. A stale cross-binary .blm is rejected on open via a new BLOOM_BODY_VERSION field (the previously-reserved header byte 20) — its bits would otherwise be stale, harmless on the mem_search-confirmed any-window gate but surfaced as mis-routed reads on the unconfirmed bloom_first gate. The shared FORMAT_VERSION is unchanged (a bloom-only reduction change must not invalidate other sidecars). A .blm is a build artifact regenerated per index build.
chore(perf) — final whole-branch-review touch-ups: backward batch-arm parity debug_assert, a non-vacuous-guard assert on the tie census, and #[allow]s so clippy --all-targets -D warnings is clean (production code has zero warnings).

Byte-identity

Every commit preserves SMEM output bit-for-bit except the Lemire bit_at commit (the bloom .blm body, confined to the non-default bloom dispatch gate, no false negatives). New proptests/oracles guard the non-obvious cases: next_n vs fwd_qlen, the (m,n)-tie census, unpack vs scalar, and the bloom no-false-negatives round-trip. present_anchor_bloom is unchanged from the base (the prototyped roll was dropped, not included).

Scope notes

target-cpu is intentionally not pinned; shipped artifacts stay portable (documented x86-64-v3 / Graviton floor in prmi/benches/README.md, with scripts/bench-native.sh for host-tuned measurement).
Cut after measurement: a leftward-walk vectorization (a spike showed cold-SA probes dominate at ~14× the per-base decode cost, so ≤30%/<5% ceiling — not worth the fwd/RC-seam byte-identity risk). The bloom any-window roll (reverted, above).
The backward search path was measured to be near probe-minimal already (gallop + bsearch + ISA seeding); no probe-reduction headroom.

Testing

Full workspace suite green; clippy --all-targets --all-features -- -D warnings clean. x86 micro + whole-search benches recorded per the table above.

coderabbitai · 2026-06-27T00:42:39Z

Walkthrough

Release-profile and native benchmark settings were added, three primitive Criterion benches and a new baseline output were recorded, cached next_n was threaded through SMEM collection, masked keyed comparisons were threaded through spectrum search paths, and lockstep/sys writes plus bloom/test auxiliary changes were updated.

Changes

Performance and hot-path updates

Layer / File(s)	Summary
Release profile and native bench wrapper `Cargo.toml`, `prmi/benches/README.md`, `scripts/bench-native.sh`	Release-profile settings, benchmark guidance, and the native runner script are added together.
Primitive microbenches and baselines `prmi/benches/primitives_bench.rs`, `prmi/benches/baselines/pr0.txt`	Three Criterion benches are added and the recorded benchmark output is extended for the updated suite runs.
Read memoization and sort order `prmi/src/index/collect.rs`	`CollectScratch` caches `next_n`, collection passes consume the cache instead of rescanning, within-read sorting becomes a single unstable `(m, n)` sort, and tests validate `fill_next_n` against `fwd_qlen`.
Masked compare helpers and forward search `prmi/src/index/spectrum.rs`	The keyed compare helpers accept explicit mask parameters, `unpack_packed_forward` switches to chunked iteration, and the forward `mem_search` and `forward_spectrum_from_hint` paths pass masks into boundary recovery.
Backward search and LCP probes `prmi/src/index/spectrum.rs`	Backward boundary recovery, RC span walking, truncation, and maximal-length paths use the masked helpers and update boundary and LCP bookkeeping.
Lockstep buffers and sys step writes `prmi/src/index/spectrum.rs`, `prmi-sys/src/lib.rs`	Lockstep batchers store raw keys until `advance`, and the sys step writers copy `SmemStep` slices with `copy_from_slice` and zipped output iteration.

Bloom bit reduction

Layer / File(s)	Summary
Lemire bit reduction `prmi/src/sidecar/bloom_file.rs`	`bit_at` switches from modulo to multiply-shift reduction, the crate docs describe the new `.blm` behavior, and the false-positive-rate test logs its measurement.

Debug-only panic test

Layer / File(s)	Summary
Debug-gated panic test `prmi/src/sa.rs`	The doubled-text out-of-range-base panic test is compiled only when debug assertions are enabled.

Estimated review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

fg-labs/prmi#46: Changes CollectScratch and the same collection plumbing in prmi/src/index/collect.rs.
fg-labs/prmi#58: Touches forward_maximal_len_seeded and adjacent spectrum boundary handling in prmi/src/index/spectrum.rs.
fg-labs/prmi#11: Also updates prmi-sys/src/lib.rs, including the same ABI-facing step-writing path.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately reflects a broad performance-focused hot-path optimization pass and its x86 validation.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.

✨ Finishing Touches

📝 Generate docstrings

Create stacked PR
Commit on current branch

🧪 Generate unit tests (beta)

Create PR with unit tests
Commit unit tests in branch perf/query-hot-path

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.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands.}

nh13 · 2026-06-27T03:51:56Z

@coderabbitai review

coderabbitai · 2026-06-27T03:52:01Z

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

…nch harness Enables cross-crate inlining of the query hot path. target-cpu stays operator-driven via RUSTFLAGS (documented x86-64-v3 / Graviton floor); shipped artifacts remain portable. Records the post-profile bench baseline that subsequent perf PRs are measured against.

…ions) doubled_text_rejects_out_of_range_base asserts via debug_assert!, compiled out at opt-level=3, so it spuriously failed under `cargo test --release` (now run as a gate by the release-profile change). Gate the test on cfg(debug_assertions); the assertion semantics are unchanged in debug builds. Pre-existing latent issue, surfaced by adding the release-test gate.

fwd_qlen rescanned the read forward on every zigzag/reseed/pass-3 step (O(rlen^2) per read). Precompute next_n once per read in CollectScratch and read it O(1) from all three passes. Byte-identical (proptest: next_n[p] == fwd_qlen(read, p)).

ref_less/shares_prefix/lcp_at and the windowed model-locate closure recomputed keyed_compare_mask(query.len()) on every SA probe; the query slice is invariant for the whole find_boundary search, so compute (nbases, mask) once per search and pass it to compare_query_vs_suffix_2x_keyed_with_mask. Byte-identical (the _with_mask variant is proptest-equivalent to the recompute variant).

nh13 · 2026-06-27T04:52:02Z

@coderabbitai review

coderabbitai · 2026-06-27T04:52:07Z

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

coderabbitai

Actionable comments posted: 2

🤖 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 961-966: The `copy_from_slice` path in
`out_steps_as_smemstep`/`SmemStep` relies on C ABI field order, but the current
size/alignment checks do not guard against field reordering. Add explicit
`offset_of!` assertions for `sa_start`, `occ_count`, and `match_len` in the
`SmemStep`/`prmi_smem_step_t` layout checks so the ABI contract is pinned even
if the struct fields are rearranged.

In `@prmi/src/sidecar/bloom_file.rs`:
- Around line 191-204: Update the stale comments around `bit_at` so their safety
invariant matches the new multiply-shift reduction. In `bit_at`, the bound comes
from Lemire reduction rather than `% num_bits`, so revise the SAFETY note near
the `contains` check to say `bit < num_bits` is guaranteed by the multiply-shift
mapping. Also adjust the writer guard comment in the related `contains`/write
path to remove the claim that `num_bits == 0` would panic in `bit_at`, and
instead describe the current behavior accurately.

🪄 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: 9bbd4f49-eac1-43b6-b5ac-e735a978597f

📥 Commits

Reviewing files that changed from the base of the PR and between 11b5125 and 37338f5.

📒 Files selected for processing (10)

Cargo.toml
prmi-sys/src/lib.rs
prmi/benches/README.md
prmi/benches/baselines/pr0.txt
prmi/benches/primitives_bench.rs
prmi/src/index/collect.rs
prmi/src/index/spectrum.rs
prmi/src/sa.rs
prmi/src/sidecar/bloom_file.rs
scripts/bench-native.sh

…, shrink lockstep key buffers A2: backward_spectrum_batch_impl mirrors the forward arm's zip (removes per-task bounds checks on the all_steps/tasks_s/out_ns triple). A3: write_fwd_task_steps and write_bwd_task_steps use copy_from_slice via the existing out_steps_as_smemstep helper instead of a field-by-field loop. A5: forward_boundary_windowed drops the always-true is_none_or guards on each bisection step (bisection monotonicity guarantees the most-recent probe is always the extreme; unconditional assignment is correct). A4: lockstep keyv Vec<Option<u64>> -> Vec<u64> + keys_present flag in all three lockstep loops (backward_spectrum_lockstep, forward_spectrum_lockstep, mem_search_lockstep); reduces per-element footprint from 16 B to 8 B. Byte-identical (lockstep/windowed equivalence oracles green, 253 tests passed).

…d loop The 32-base middle loop paid a try_into().unwrap() panic-branch and eight per-step out[] bounds checks (the i+4*j arithmetic defeats LLVM's proof from the i+32<=n guard). Drive it with chunks_exact_mut(32).zip(chunks_exact(8)) so bounds are checked once at the slice split with a known stride. Byte-identical (packed_word_decode + unaligned_word_step oracles green); no new unsafe. Add an inline before/after bench in primitives_bench.rs (unpack_packed_middle) that isolates the middle loop over a 1024-base / 256-corpus corpus. Measured ~35% wall-time reduction on aarch64 (19.5 µs → 12.5 µs, 13.4 → 20.9 Gelem/s); the real validation target is x86 (deferred EC2 batch).

The within-read sort ran two unstable sorts: stage 1 (m ASC, n DESC) then stage 2 (m ASC, n ASC). Stage 2 alone orders SMEMs with distinct (m, n); stage 1 could only reorder SMEMs sharing both m and n. Such a reorder is unobservable. Every emitted SMEM's (k, s) is mem_search(read[m..=n]), a deterministic function of the span content, so a fixed (m, n) determines a unique (k, s); with rid constant per read and l == 0 always, two entries sharing (m, n) are equal in all six fields. The multi-pass collection does emit (m, n) duplicates (a reseed or pass-3 round re-finding a pass-1 span), but each is a byte-identical copy, so permuting a tied group is the identity on the output bytes. Dropping stage 1 is therefore byte-identical to the C++ two-stage composition, including at the consumer box-gate (prmi's output bytes are unchanged). Add two regression guards: unsorted_mn_ties_are_byte_identical (a 4096-case proptest asserting every (m, n) tie in the unsorted set is byte-identical) and unsorted_mn_tie_census (a deterministic repeat-heavy corpus that finds 1440 tie-groups / 1764 duplicate entries across ~88k reads, all byte-identical).

nh13 · 2026-06-27T06:07:46Z

@coderabbitai review

coderabbitai · 2026-06-27T06:07:51Z

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

nh13 · 2026-06-27T07:08:02Z

@coderabbitai review

coderabbitai · 2026-06-27T07:08:08Z

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

nh13 · 2026-06-27T08:08:13Z

@coderabbitai review

coderabbitai · 2026-06-27T08:08:19Z

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

coderabbitai

Actionable comments posted: 1

🤖 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/sidecar/bloom_file.rs`:
- Around line 59-72: Add a `.blm`-specific version check in the bloom sidecar
loading path so stale bloom files are rejected on open instead of being accepted
by the shared sidecar bind check. Update the logic around the `BloomFile`/`.blm`
reader and its bind validation (`sa_num`, `ref_digest`, `keyset_digest`) to also
verify a bloom-only version discriminator or equivalent metadata before using
`bloom_first`. Keep the check scoped to `.blm` so other sidecars are unaffected,
and fail fast when the bloom payload was built by an incompatible binary.

🪄 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: 57d78cdf-14a2-4dee-b57d-75a428f6f3e6

📥 Commits

Reviewing files that changed from the base of the PR and between 37338f5 and 5828321.

📒 Files selected for processing (5)

prmi-sys/src/lib.rs
prmi/benches/primitives_bench.rs
prmi/src/index/collect.rs
prmi/src/index/spectrum.rs
prmi/src/sidecar/bloom_file.rs

Replace the per-probe `combined % num_bits` (true 64-bit div) with Lemire multiply-shift ((combined*num_bits)>>64), uniform into [0,num_bits) with no division; num_bits stays a multiple of 64 so sizing/header are unchanged. Writer and reader share bit_at so a self-consistent .blm has no false negatives (bloom_has_no_false_negatives green); FP rate stays near target (1.02%). Bench (1024 probes, num_bits=479296): before_mod: ~1.555 µs (~658 Melem/s) after_lemire: ~944 ns (~1084 Melem/s) ~1.65x throughput FORMAT NOTE: this changes which bits a key sets, so the .blm BODY bytes differ from a prior build (header layout unchanged; shared FORMAT_VERSION intentionally NOT bumped — it gates all sidecars). A .blm must be regenerated with a matching binary; a stale .blm yields benign bloom false negatives that the consumer's present-read fallback re-seeds (accelerator, not a correctness oracle). .blm body now diverges from #57's % reduction — flag for the #57 author at merge.

…cuous census guard, clippy-clean test targets - backward_spectrum_batch_impl: debug_assert_eq!(all_steps.len(), tasks_s.len()) for parity with the forward arm (zip would otherwise truncate silently on a length mismatch). - unsorted_mn_tie_census: assert tie_groups > 0 so the (m,n)-tie byte-identity guard cannot pass vacuously. - #[allow] the deliberately production-mirroring loops in the new proptests/bench so clippy --all-targets -D warnings is clean; production code already has zero warnings.

nh13 force-pushed the perf/query-hot-path branch from 221b609 to 5c85b65 Compare June 27, 2026 00:52

Base automatically changed from feat/z-bloom-gates to main June 27, 2026 03:46

nh13 force-pushed the perf/query-hot-path branch 2 times, most recently from db9c356 to 9c45f4e Compare June 27, 2026 03:50

nh13 added 4 commits June 26, 2026 20:54

nh13 force-pushed the perf/query-hot-path branch from 9c45f4e to 37338f5 Compare June 27, 2026 03:54

coderabbitai Bot reviewed Jun 27, 2026

View reviewed changes

Comment thread prmi-sys/src/lib.rs

Comment thread prmi/src/sidecar/bloom_file.rs

nh13 added 3 commits June 26, 2026 22:16

nh13 force-pushed the perf/query-hot-path branch from 37338f5 to 5828321 Compare June 27, 2026 05:17

coderabbitai Bot reviewed Jun 27, 2026

View reviewed changes

Comment thread prmi/src/sidecar/bloom_file.rs Outdated

nh13 added 2 commits June 27, 2026 01:40

nh13 force-pushed the perf/query-hot-path branch from 5828321 to 22119c6 Compare June 27, 2026 08:40

nh13 merged commit 3b9e2c6 into main Jun 27, 2026
4 checks passed

nh13 deleted the perf/query-hot-path branch June 27, 2026 08:46

Uh oh!

Conversation

nh13 commented Jun 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Measured results (x86, base → HEAD)

Reading order (commit by commit)

Byte-identity

Scope notes

Testing

Uh oh!

coderabbitai Bot commented Jun 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Estimated review effort

Possibly related PRs

Uh oh!

nh13 commented Jun 27, 2026

Uh oh!

coderabbitai Bot commented Jun 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

nh13 commented Jun 27, 2026

Uh oh!

coderabbitai Bot commented Jun 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

nh13 commented Jun 27, 2026

Uh oh!

coderabbitai Bot commented Jun 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

nh13 commented Jun 27, 2026

Uh oh!

coderabbitai Bot commented Jun 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

nh13 commented Jun 27, 2026

Uh oh!

coderabbitai Bot commented Jun 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

nh13 commented Jun 27, 2026 •

edited

Loading

coderabbitai Bot commented Jun 27, 2026 •

edited

Loading

coderabbitai Bot commented Jun 27, 2026 •

edited

Loading

coderabbitai Bot commented Jun 27, 2026 •

edited

Loading

coderabbitai Bot commented Jun 27, 2026 •

edited

Loading

coderabbitai Bot commented Jun 27, 2026 •

edited

Loading

coderabbitai Bot commented Jun 27, 2026 •

edited

Loading