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);