perf(prover): faster CPU trace generation#786
Conversation
|
/bench |
|
/ai-review |
Codex Code ReviewFindings
Notes I did not build or run tests, per the static-review constraints. |
Review: perf(prover) faster CPU trace generationReviewed against the diff. This is a solid, well-documented optimization. Routing at creation time (
Minor notes (non-blocking)
No blocking issues. |
AI ReviewPR #786 · 5 changed files Findings
Status column reflects the verdict from the verifier: deepseek-verifier (openrouter/deepseek/deepseek-v4-pro). AI-001: Soundness risk: silent truncation in MemwOperation::new value/old in release builds
Claim MemwOperation::new converts [u64; 8] to [u32; 8] via Evidence The PR adds Suggested fix Replace AI-002: Parallel BitwiseHistogram fold+reduce allocates many 80 MiB histograms concurrently
Claim The new Evidence Each Suggested fix Either reuse a single histogram across folds (e.g., pass AI-003: Significant peak-memory regression from per-worker BitwiseHistogram allocation
Claim The new Evidence trace_builder.rs lines 3222-3254: Suggested fix Either (a) split the large per-source Vecs (especially the page collector) into a streaming histogram-bump loop without materializing the Vec, or (b) pre-size the worker histograms smaller and overflow to a hashmap/dense rep, or (c) keep the legacy single-Vec extend path on memory-constrained configurations and only parallelize when peak memory budget allows. AI-004: BitwiseHistogram counters are u64; a single cell overflow would silently corrupt multiplicity
Claim
Evidence Single-cell lookups reaching 2^64 are unrealistic in practice, but Suggested fix Either use a wider counter (e.g., u128 or a packed overflow sentinel) or assert at fill time that no cell exceeds the Goldilocks prime. Document the explicit assumption. AI-005: fill_multiplicities is sequential and could be parallelized
Claim fill_multiplicities iterates NUM_LOOKUP_TYPES * NUM_ROWS = 10 * 2^20 = ~10.5M cells in a single thread, with no parallelism despite each cell being independent (no shared writes). Given the rest of the BitwiseHistogram path is parallelized for the bump phase, leaving the fill phase serial is inconsistent and a missed ~worker-count speedup for a step that runs unconditionally after every proof. Evidence bitwise.rs lines 563-577: nested Suggested fix Parallelize with rayon: AI-006: Dead code: BitwiseHistogram::default() allocates 80 MiB and is never called
Claim
Evidence grep over the crate for Suggested fix Remove the AI-007: MemwOperation::new and with_old silently truncate u64->u32 in release builds
Claim The new conversion Evidence
Suggested fix Change the constructor signature to accept AI-008: Stale doc comment referencing deleted collect_bitwise_from_memw_register
Claim The comment on Evidence grep for Suggested fix Rewrite the comment to refer to the actual shared helper, AI-009: MemwSink default push_reg_access silently ignores the decomposed fields
Claim The default Evidence For Suggested fix Make the default impl require the fields to be moved out (e.g., return a AI-010: MemwBuckets over-allocates register_rows capacity upfront
Claim
Evidence
Suggested fix Use a more conservative initial capacity (e.g., M1+M3+M5 upper bound minus a margin for fallback), or grow lazily. Also consider using Reviewer Lanes
Verification Lanes
Native Codex and Claude reviews run separately and post their own comments. They are not included in this structured provenance report. Raw lane outputs, candidates, final issues, and model metrics are uploaded as workflow artifacts. |
- Reattach the collect_ops_from_cpu doc block to the function (it was landing on struct RegRow together with a stray #[allow]) and update its Returns list to the MemwBuckets shape - Replace references to deleted code and internal review artifacts: the 'Shared by BOTH collectors' rationale (the MemwOperation collector is gone), the E6/E7 notes, and mu_column's 'legacy path' label (update_multiplicities is live for continuation epochs) - Document fill_multiplicities' overwrite semantics: it assumes zeroed MU columns and layered lookups (continuation L2G) must be added strictly after - Remove BitwiseHistogram::total, a pub method with no callers
MEMW_R: - Move RegRow and the direct column fill into memw_register.rs; generate_memw_register_trace is now a thin wrapper (ops -> RegRow::from_memw -> shared fill), so the unit-tested generator and the production fast path run the same code and cannot drift - RegRow::new is the single definition of the row encoding (halved address, masked old_ts_lo); the walk fast path and from_memw both delegate to it - Move the IS_HALFWORD collector next to bus_interactions() so the lookup a row sends and the lookup counted derive from the same module Routing: - Add MemwRoute + classify_memw, shared by MemwBuckets::push, the push_reg_access fallback (which now delegates to push), and count_table_lengths' partition_memw — one classification, three consumers - Drop the single-use push_register/push_register_row helpers BITWISE: - BitwiseOperationType::ALL is the single origin of NUM_LOOKUP_TYPES; type_mu_column is defined as mu_column(ALL[i]), removing the arithmetic column-contiguity assumption instead of guarding it - Replace the runtime histogram_column_map_guard test with a compile-time bijection assert (a mismatch is now a compile error, not a test failure) - Correct bump()'s comment: in release an out-of-domain z would mis-count in another lane rather than panic; constructors' masking upholds the invariant
MemwOperation::new and with_old take [u32; 8] directly and pack_register_value returns u32 limbs, so the u64->u32 narrowing debug_asserts (and the release-mode silent-truncation hazard behind them) are deleted rather than guarded: out-of-domain values are now unrepresentable. The register fast path's bare `as u32` casts, which skipped even the debug_assert, disappear for the same reason.
…ations - Collectors now take &mut BitwiseHistogram instead of returning Vec<BitwiseOperation>. The heavy sources count with no per-source vector at all: MEMW_R bumps one IS_HALFWORD per row (tens of millions of rows), PAGE bumps one ARE_BYTES per byte of every touched page, and CPU padding collapses to two bump_n calls instead of an O(padding_rows) vector of identical zero ops. - reduce_with replaces reduce(identity, ..) in the rayon tree-reduce, eliminating one zeroed 80 MiB identity histogram per reduce leaf. - RegRow.address shrinks to u16 (register index 0..=255), taking the largest persisted array of the walk from 40 to 32 bytes per row.
Docs and dead-code cleanup for tracegen optimizations
…of-truth Single-source the MEMW_R fill and the BITWISE type↔column map
The bitwise-histogram comments were phrased against the pre-refactor code
("byte-identical to the previous serial .extend() chain", "instead of an
O(n) op vector") — narration that is meaningless once this merges and the
old path is gone. Restate them as standalone rationale: the multiplicities
are order-independent (permutation-invariant bus) so parallel collection is
sound, and the reduce_with note now explains why to prefer it over
reduce(identity) rather than referencing a removed path. Also fix the stale
add_padding_byte_checks doc (it described the pre-shrink 14-op CPU layout).
Repo-wide sweep of the trace-generation code for comments that narrate a
change the next reader can't see, or that contradict the code:
- MemwBuckets doc and collect_all_ops no longer reference "the old two-stage
partition" / "byte-identical to partition-ing one combined vec" — that
partition sweep no longer exists anywhere (grep '.partition(' → 0 hits).
Restated as the current invariant (register-first-then-aligned, deterministic
insertion order the multiplicity counts rely on).
- Fix the MEMW_R ADDRESS column doc: register index range is 0-255 (x0-x31 plus
the x254 commit index and x255 PC), not 0-31.
The Bus-14 comment in collect_store_op_from_cpu said the packed value "must match CPU M7 which sends full rv2 as [lo32, hi32]", but M7 (the CPU's old inline store MEMW at timestamp+1) was removed in the ALU-bus migration. The store value now flows CPU -> MEMORY/MEMOP (rv2 as [lo32, hi32]) -> STORE chip -> MEMW write. Reword to describe that path; the byte-layout requirement it gestured at is unchanged.
Feed bitwise collectors straight into the histogram; cut reduce allocations
Make the u32 value domain a compile-time fact for MemwOperation
The bitwise multiplicity phase ran each source as one atomic collector across a capped par_chunks, so the two dominant sources — the in-walk lookups and MEMW_R (each tens of millions of items) — each pinned a single core while the rest idled, and the in-walk fold ran serially before the parallel region. Split those two sources into ~cap row-range slices and round-robin every unit (whole collectors + the heavy slices) into exactly cap buckets, one 80 MiB histogram each. Peak memory is unchanged (still cap concurrent histograms), but the heavy work is now spread across buckets instead of piled onto one core, and the in-walk fold moved into the parallel region. add_ops/bump/merge form a commutative monoid, so multiplicities are byte-identical (prove+verify pass). PAGE is left whole for now; if it becomes the bottleneck it can be sliced the same way in a follow-up.
…tighter visibility, stale docs (#795) - BITWISE MU columns now derive from one `MU_COLUMNS` array with a compile-time distinctness assert, so a duplicate column is a build error rather than a silent overwrite in fill_multiplicities. - push_reg_access takes [u32;2] val/old pairs and builds its own fallback op instead of an 8-scalar signature plus a repacking closure. - Narrow BitwiseHistogram + the MU-column/lookup-type helpers to pub(crate); scope the test-only generate_memw_register_trace wrapper. - Fix stale docs (MemwSink Vec impl is the sizing pass, not tests; generate_bitwise_trace MU columns are filled by fill_multiplicities).
|
/bench |
|
/bench |
|
/bench-gpu |
GPU Benchmark (ABBA) —
|
zeroed_fe_vec built its Vec<FE> with mem::transmute::<Vec<u64>, Vec<FE>>. That relies on Vec's field layout being identical across element types, which the language does not guarantee (Vec is #[repr(Rust)], and the std transmute docs flag exactly this Vec-to-Vec pattern as a Bad Idea). It happens to work on the pinned toolchain, but a layout change (e.g. -Zrandomize-layout) could swap the ptr/cap words and make it UB. Rebuild the Vec from its raw parts instead: same calloc'd allocation, same element size/align (already const-asserted), so the dealloc Layout matches and there is no dependence on Vec's internal layout. Behavior is unchanged — the zeroed_fe_vec_matches_fe_zero test and full prove+verify still pass.
…w-parts zeroed_fe_vec: reinterpret via from_raw_parts, not Vec transmute
2 participants
What & why
Trace generation was taking much longer than it should. The core problem was that trace-gen is memory-bandwidth-bound: it materialized millions of intermediate memory-operation structs and then re-read them to fill the trace. This PR removes that intermediate work and writes data straight into the trace columns, roughly halving trace-generation time on memory-heavy workloads.
Optimizations
Skip the intermediate copy of memory operations. Register accesses are now written directly into their trace columns (via a compact 48-byte
RegRow), instead of first building aMemwOperationand re-reading it. This is the biggest win — it eliminates a whole materialize + re-read pass over the hottest data.Sort memory ops into buckets up front. Operations are routed into their register / aligned / general buckets when created, instead of collecting everything into one list and classifying it afterwards.
Smaller memory-op structs. The
value/oldfields shran]` (the values always fit in 32 bits), halving the struct size andcutting cache/memory traffic.Faster bitwise multiplicity counting. Replaced per-operation multiplicity updates with a dense histogram that each thread fills locally and then merges with a tree-reduce — parallel-friendly and allocation-free.
Removed the old trace-gen code paths. Deleted the legacy con path.
Results
Measured on 2× AMD EPYC 9J14 (96-core) + RTX 5090, ethrex blocks,
Trace generation:
End-to-end proving time (GPU):
The end-to-end gain grows with workload size, since trace-gen is a bigger blocks.
Correctness
check.