From 610245feee94a7f674b582fefa7f7733e442f68d Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 7 Jul 2026 16:17:39 -0300 Subject: [PATCH] Parallelize the two dominant bitwise-histogram sources internally MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- prover/src/tables/trace_builder.rs | 59 ++++++++++++++++++------------ 1 file changed, 36 insertions(+), 23 deletions(-) diff --git a/prover/src/tables/trace_builder.rs b/prover/src/tables/trace_builder.rs index 2d6e426e8..73b7675a8 100644 --- a/prover/src/tables/trace_builder.rs +++ b/prover/src/tables/trace_builder.rs @@ -3055,6 +3055,8 @@ fn build_traces( type Collector<'a> = Box; let mul_chunk = max_rows.mul; let dvrm_chunk = max_rows.dvrm; + // Every source except the two dominant ones (the in-walk lookups and MEMW_R, which are + // split into row-ranges in the parallel path below) stays a single whole-source collector. let mut collectors: Vec = vec![ Box::new(|h| h.add_ops(&collect_bitwise_from_lt(<_ops))), Box::new(|h| h.add_ops(&collect_bitwise_from_mul(&mul_ops, mul_chunk))), @@ -3077,7 +3079,6 @@ fn build_traces( } }), Box::new(|h| h.add_ops(&collect_bitwise_from_memw_aligned(&memw_aligned_ops))), - Box::new(|h| memw_register::collect_bitwise_from_memw_register(&memw_register_rows, h)), Box::new(|h| h.add_ops(&collect_bitwise_from_commit(&commit_ops))), Box::new(|h| h.add_ops(&collect_bitwise_from_keccak(&keccak_ops))), Box::new(|h| h.add_ops(&collect_bitwise_from_ecsm(&ecsm_ops))), @@ -3092,33 +3093,42 @@ fn build_traces( })); } - // Fold the in-walk lookups (from the serial p2a walk + CPU32) into the base histogram, - // then free that Vec: it is no longer needed once counted. let mut base = bitwise::BitwiseHistogram::new(); - base.add_ops(&bitwise_ops); - drop(bitwise_ops); #[cfg(feature = "parallel")] { use rayon::prelude::*; - // To keep peak memory from scaling with core count, cap the number of - // concurrent histograms: split the collectors into at most `max_par` - // chunks (one histogram each), run the chunks in parallel, and process - // the collectors within a chunk sequentially (so at most one transient - // Vec is live per chunk — the heavy collectors bump their chunk's - // histogram directly and have none). `reduce_with` pairs the chunk - // histograms directly; `reduce(identity, ..)` would instead allocate a - // zeroed 80 MiB identity histogram per leaf, so prefer `reduce_with`. - // Tree-reduce is valid because bump/add_ops/merge form a commutative - // monoid, so the sum is independent of chunk order and matches the - // non-`parallel` fallback below. - let max_par = rayon::current_num_threads().clamp(1, 8); - let chunk_size = collectors.len().div_ceil(max_par).max(1); - if let Some(reduced) = collectors - .par_chunks(chunk_size) - .map(|chunk| { + // Cap concurrent 80 MiB histograms at `cap` to bound peak memory. The two dominant + // sources — the in-walk lookups and MEMW_R (each tens of millions of items) — are + // split into ~`cap` row-range slices so they parallelize INTERNALLY instead of each + // pinning one core while the rest idle. Every unit (whole collectors + the heavy + // slices) is round-robined into exactly `cap` buckets, one histogram each, so the + // split heavy work is spread across buckets rather than piled into one. + // add_ops/bump/merge form a commutative monoid, so any partition yields + // byte-identical multiplicities (same as the serial fallback below). + let cap = rayon::current_num_threads().clamp(1, 8); + let mut units: Vec = Vec::with_capacity(collectors.len() + 2 * cap); + let iw_chunk = bitwise_ops.len().div_ceil(cap).max(1); + for slice in bitwise_ops.chunks(iw_chunk) { + units.push(Box::new(move |h| h.add_ops(slice))); + } + let reg_chunk = memw_register_rows.len().div_ceil(cap).max(1); + for slice in memw_register_rows.chunks(reg_chunk) { + units.push(Box::new(move |h| { + memw_register::collect_bitwise_from_memw_register(slice, h) + })); + } + units.extend(collectors); + + let mut buckets: Vec> = (0..cap).map(|_| Vec::new()).collect(); + for (i, unit) in units.into_iter().enumerate() { + buckets[i % cap].push(unit); + } + if let Some(reduced) = buckets + .par_iter() + .map(|bucket| { let mut h = bitwise::BitwiseHistogram::new(); - for f in chunk { + for f in bucket { f(&mut h); } h @@ -3133,12 +3143,15 @@ fn build_traces( } #[cfg(not(feature = "parallel"))] { + base.add_ops(&bitwise_ops); + memw_register::collect_bitwise_from_memw_register(&memw_register_rows, &mut base); for f in &collectors { f(&mut base); } } let bitwise_histogram = base; - drop(collectors); + // The in-walk lookup Vec has been counted into the histogram; free it now. + drop(bitwise_ops); #[cfg(feature = "instruments")] drop(__sp);