diff --git a/prover/src/tables/bitwise.rs b/prover/src/tables/bitwise.rs index c73e1e341..8af7fa594 100644 --- a/prover/src/tables/bitwise.rs +++ b/prover/src/tables/bitwise.rs @@ -523,6 +523,7 @@ const _: () = { /// [`update_multiplicities`] produces (both just sum the same lookups per cell). /// /// Memory: `NUM_ROWS * NUM_LOOKUP_TYPES * 8` bytes = 2^20 * 10 * 8 = 80 MiB. +#[cfg_attr(test, derive(PartialEq, Eq))] pub(crate) struct BitwiseHistogram { counters: Box<[u64]>, } diff --git a/prover/src/tables/trace_builder.rs b/prover/src/tables/trace_builder.rs index c3b695a80..1d769b78d 100644 --- a/prover/src/tables/trace_builder.rs +++ b/prover/src/tables/trace_builder.rs @@ -2437,22 +2437,23 @@ pub(crate) fn collect_bitwise_from_ecdas(ops: &[ecdas::EcdasOperation]) -> Vec Vec { +pub(crate) fn for_each_keccak_bitwise_lookup( + keccak_ops: &[KeccakOperation], + mut emit: impl FnMut(BitwiseOperation), +) { use executor::vm::instruction::execution::{KECCAK_RC, KECCAK_RHO}; - let mut ops = Vec::new(); - for kop in keccak_ops { let state_addr = kop.state_addr; - ops.push(BitwiseOperation::byte_op( + emit(BitwiseOperation::byte_op( BitwiseOperationType::ByteAluAnd, (state_addr & 0xFF) as u8, 7, @@ -2465,7 +2466,7 @@ pub(crate) fn collect_bitwise_from_keccak(keccak_ops: &[KeccakOperation]) -> Vec for i in 0..4 { let lo = ((state_addr >> (2 * i * 8)) & 0xFF) as u8; let hi = ((state_addr >> ((2 * i + 1) * 8)) & 0xFF) as u8; - ops.push(BitwiseOperation::byte_op( + emit(BitwiseOperation::byte_op( BitwiseOperationType::AreBytes, lo, hi, @@ -2479,7 +2480,7 @@ pub(crate) fn collect_bitwise_from_keccak(keccak_ops: &[KeccakOperation]) -> Vec .expect("keccak state address range must be validated by the executor"); for shift in [0, 16, 32, 48] { let half = ((ptr >> shift) & 0xFFFF) as u16; - ops.push(BitwiseOperation::halfword( + emit(BitwiseOperation::halfword( BitwiseOperationType::IsHalf, (half & 0xFF) as u8, ((half >> 8) & 0xFF) as u8, @@ -2497,7 +2498,7 @@ pub(crate) fn collect_bitwise_from_keccak(keccak_ops: &[KeccakOperation]) -> Vec let v0 = ((state[x] >> (b * 8)) & 0xFF) as u8; let v1 = ((state[x + 5] >> (b * 8)) & 0xFF) as u8; cxz[x][0][b] = v0 ^ v1; - ops.push(BitwiseOperation::byte_op( + emit(BitwiseOperation::byte_op( BitwiseOperationType::ByteAluXor, v0, v1, @@ -2509,7 +2510,7 @@ pub(crate) fn collect_bitwise_from_keccak(keccak_ops: &[KeccakOperation]) -> Vec let prev = cxz[x][stage - 1][b]; let sv = ((state[x + 5 * y] >> (b * 8)) & 0xFF) as u8; cxz[x][stage][b] = prev ^ sv; - ops.push(BitwiseOperation::byte_op( + emit(BitwiseOperation::byte_op( BitwiseOperationType::ByteAluXor, prev, sv, @@ -2530,7 +2531,7 @@ pub(crate) fn collect_bitwise_from_keccak(keccak_ops: &[KeccakOperation]) -> Vec let shifted = halfword << 1; // u16 wraps // ARE_BYTES for cxz_left bytes: paired (low, high) of the halfword, // matching `(cxz_left[x][2i], cxz_left[x][2i+1])` sender pairing. - ops.push(BitwiseOperation::byte_op( + emit(BitwiseOperation::byte_op( BitwiseOperationType::AreBytes, (shifted & 0xFF) as u8, ((shifted >> 8) & 0xFF) as u8, @@ -2563,7 +2564,7 @@ pub(crate) fn collect_bitwise_from_keccak(keccak_ops: &[KeccakOperation]) -> Vec let a = cxz[(x + 4) % 5][3][b]; let rb = rotated_c[(x + 1) % 5][b]; d_bytes[x][b] = a ^ rb; - ops.push(BitwiseOperation::byte_op( + emit(BitwiseOperation::byte_op( BitwiseOperationType::ByteAluXor, a, rb, @@ -2583,7 +2584,7 @@ pub(crate) fn collect_bitwise_from_keccak(keccak_ops: &[KeccakOperation]) -> Vec theta_lanes[x + 5 * y] = lane ^ d_lane; for b in 0..8 { let s = ((lane >> (b * 8)) & 0xFF) as u8; - ops.push(BitwiseOperation::byte_op( + emit(BitwiseOperation::byte_op( BitwiseOperationType::ByteAluXor, s, d_bytes[x][b], @@ -2609,12 +2610,12 @@ pub(crate) fn collect_bitwise_from_keccak(keccak_ops: &[KeccakOperation]) -> Vec // ARE_BYTES paired as (rot_left[b], rot_right[b]) for // each byte of the halfword, matching the sender pairing // in keccak_rnd::bus_interactions. - ops.push(BitwiseOperation::byte_op( + emit(BitwiseOperation::byte_op( BitwiseOperationType::AreBytes, (shifted & 0xFF) as u8, (carry & 0xFF) as u8, )); - ops.push(BitwiseOperation::byte_op( + emit(BitwiseOperation::byte_op( BitwiseOperationType::AreBytes, ((shifted >> 8) & 0xFF) as u8, ((carry >> 8) & 0xFF) as u8, @@ -2645,14 +2646,14 @@ pub(crate) fn collect_bitwise_from_keccak(keccak_ops: &[KeccakOperation]) -> Vec for b in 0..8 { let not_byte = ((not_next >> (b * 8)) & 0xFF) as u8; let n2_byte = ((next2 >> (b * 8)) & 0xFF) as u8; - ops.push(BitwiseOperation::byte_op( + emit(BitwiseOperation::byte_op( BitwiseOperationType::ByteAluAnd, not_byte, n2_byte, )); let pi_byte = ((pi_lanes[x + 5 * y] >> (b * 8)) & 0xFF) as u8; let and_byte = ((and_val >> (b * 8)) & 0xFF) as u8; - ops.push(BitwiseOperation::byte_op( + emit(BitwiseOperation::byte_op( BitwiseOperationType::ByteAluXor, pi_byte, and_byte, @@ -2666,7 +2667,7 @@ pub(crate) fn collect_bitwise_from_keccak(keccak_ops: &[KeccakOperation]) -> Vec for b in 0..8 { let chi_byte = ((chi_lanes[0] >> (b * 8)) & 0xFF) as u8; let rc_byte = ((rc_val >> (b * 8)) & 0xFF) as u8; - ops.push(BitwiseOperation::byte_op( + emit(BitwiseOperation::byte_op( BitwiseOperationType::ByteAluXor, chi_byte, rc_byte, @@ -2678,8 +2679,15 @@ pub(crate) fn collect_bitwise_from_keccak(keccak_ops: &[KeccakOperation]) -> Vec state = chi_lanes; } } +} - ops +/// Count Keccak lookups directly, without materializing 24,777 four-byte +/// records per permutation. Shared by CPU and CUDA trace construction. +pub(crate) fn collect_bitwise_from_keccak( + keccak_ops: &[KeccakOperation], + histogram: &mut bitwise::BitwiseHistogram, +) { + for_each_keccak_bitwise_lookup(keccak_ops, |op| histogram.bump(op)); } /// every address accessed during execution (ELF init + runtime stores/loads). @@ -3290,15 +3298,15 @@ fn build_traces( // We never concatenate the lookups into one giant `Vec` (~140 M ops / // ~560 MB at 10-tx whose only consumer is the multiplicity count). Each collector bumps // the `BitwiseHistogram` it is handed: the heavy sources (MEMW_R one-per-row, PAGE - // one-per-byte, padding) count directly with no per-source Vec at all, and the small + // one-per-byte, Keccak, padding) count directly with no per-source Vec, and the small // sources fold their transient `collect_*` Vec in and drop it. The histogram is a // commutative monoid, so per-worker histograms tree-reduce to multiplicities that are // independent of accumulation order. 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. + // In-walk, MEMW_R and Keccak lookups are split into ranges in the parallel + // path below. Other sources remain whole-source collectors. 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))), @@ -3322,7 +3330,6 @@ fn build_traces( }), Box::new(|h| h.add_ops(&collect_bitwise_from_memw_aligned(&memw_aligned_ops))), 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))), Box::new(|h| h.add_ops(&collect_bitwise_from_ecdas(&ecdas_ops))), Box::new(|h| h.add_ops(&collect_bitwise_from_hint(&hint_ops))), @@ -3341,8 +3348,8 @@ fn build_traces( #[cfg(feature = "parallel")] { use rayon::prelude::*; - // 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 + // Cap concurrent 80 MiB worker histograms at `cap` to bound peak memory. Heavy + // sources — in-walk lookups, MEMW_R, and Keccak (24,777 lookups per permutation) — 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 @@ -3350,7 +3357,7 @@ fn build_traces( // 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 mut units: Vec = Vec::with_capacity(collectors.len() + 3 * 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))); @@ -3361,6 +3368,13 @@ fn build_traces( memw_register::collect_bitwise_from_memw_register(slice, h) })); } + // Split by complete permutations: round state is local to each operation. + // Reuse the existing capped histogram buckets rather than spawning an + // additional set of 80 MiB accumulators. + let keccak_chunk = keccak_ops.len().div_ceil(cap).max(1); + for slice in keccak_ops.chunks(keccak_chunk) { + units.push(Box::new(move |h| collect_bitwise_from_keccak(slice, h))); + } units.extend(collectors); let mut buckets: Vec> = (0..cap).map(|_| Vec::new()).collect(); @@ -3387,6 +3401,7 @@ fn build_traces( #[cfg(not(feature = "parallel"))] { base.add_ops(&bitwise_ops); + collect_bitwise_from_keccak(&keccak_ops, &mut base); memw_register::collect_bitwise_from_memw_register(&memw_register_rows, &mut base); for f in &collectors { f(&mut base); diff --git a/prover/src/tests/keccak_lookup_tests.rs b/prover/src/tests/keccak_lookup_tests.rs new file mode 100644 index 000000000..626fa9828 --- /dev/null +++ b/prover/src/tests/keccak_lookup_tests.rs @@ -0,0 +1,193 @@ +//! Streaming and chunking must preserve every Keccak lookup multiplicity. +use crate::tables::{ + bitwise::{BitwiseHistogram, BitwiseOperation, BitwiseOperationType}, + keccak::KeccakOperation, + trace_builder::{collect_bitwise_from_keccak, for_each_keccak_bitwise_lookup}, +}; +use executor::vm::instruction::execution::keccak_f1600; + +fn operations(n: usize) -> Vec { + let mut seed = 17u64; + (0..n) + .map(|i| { + let input = core::array::from_fn(|lane| match i % 4 { + 0 => 0, + 1 => u64::MAX, + 2 => 0xaaaa_5555_aaaa_5555u64.rotate_left(lane as u32), + _ => { + seed ^= seed << 13; + seed ^= seed >> 7; + seed ^= seed << 17; + seed + } + }); + let mut output = input; + keccak_f1600(&mut output); + KeccakOperation { + timestamp: i as u64, + // Exercise carries in the 25 aligned state pointers, and repeat + // some inputs/addresses to test accumulation of duplicate lookups. + state_addr: [0x1000, 0xfff8, 0xffff_fff8, 0x1000][i % 4], + input, + output, + } + }) + .collect() +} + +fn buffered(ops: &[KeccakOperation], histogram: &mut BitwiseHistogram) { + // Reproduce the previous materialize-then-count consumer. The operation + // generator is unchanged by this refactor; existing round/bus tests check + // its arithmetic, while this reference checks how its output is counted. + let mut records = Vec::new(); + for_each_keccak_bitwise_lookup(ops, |op| records.push(op)); + histogram.add_ops(&records); +} + +#[test] +fn keccak_streaming_preserves_buffered_multiplicities() { + let ops = operations(9); + let mut expected = BitwiseHistogram::new(); + let mut actual = BitwiseHistogram::new(); + let sentinel = BitwiseOperation::byte_op(BitwiseOperationType::AreBytes, 17, 29); + expected.bump(sentinel); + actual.bump(sentinel); + collect_bitwise_from_keccak(&[], &mut actual); + assert!(actual == expected, "empty input must preserve prior counts"); + for end in [1, 3, 9] { + buffered(&ops[..end], &mut expected); + collect_bitwise_from_keccak(&ops[..end], &mut actual); + assert!(actual == expected, "lookup multiplicities differ at {end}"); + } +} + +#[cfg(feature = "parallel")] +#[test] +fn keccak_parallel_chunks_preserve_all_multiplicities() { + use rayon::prelude::*; + let ops = operations(11); + let mut expected = BitwiseHistogram::new(); + buffered(&ops, &mut expected); + // One worker, an uneven split and more workers than operations. This + // exercises the same permutation boundaries as the trace-build scheduler. + for (workers, input) in [(1, &ops[..]), (3, &ops[..]), (4, &ops[..2]), (4, &ops[..0])] { + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(workers) + .build() + .unwrap(); + let chunk = input.len().div_ceil(workers).max(1); + let merged = pool + .install(|| { + input + .par_chunks(chunk) + .map(|slice| { + let mut h = BitwiseHistogram::new(); + collect_bitwise_from_keccak(slice, &mut h); + h + }) + .reduce_with(|mut a, b| { + a.merge(&b); + a + }) + }) + .unwrap_or_else(BitwiseHistogram::new); + if input.len() == ops.len() { + assert!( + merged == expected, + "parallel counts differ with {workers} workers" + ); + } else { + let mut reference = BitwiseHistogram::new(); + buffered(input, &mut reference); + assert!(merged == reference, "short/empty parallel input differs"); + } + } +} + +#[test] +#[ignore = "collector microbenchmark; run with --release --ignored --nocapture"] +fn bench_keccak_lookup_collectors() { + use std::{hint::black_box, time::Instant}; + let n = std::env::var("KECCAK_LOOKUP_BENCH_N") + .ok() + .map(|s| s.parse().unwrap()) + .unwrap_or(1024); + let ops = operations(n); + println!( + "permutations={n} removed_vector_payload_bytes={}", + n * 24_777 * std::mem::size_of::() + ); + // Alternate order across runs to reduce warm-cache/order bias. Include + // allocation and first-touch of the same 80 MiB histogram in both paths. + for run in 0..6 { + let order = if run % 2 == 0 { + [false, true] + } else { + [true, false] + }; + for streaming in order { + let start = Instant::now(); + let mut h = BitwiseHistogram::new(); + if streaming { + collect_bitwise_from_keccak(black_box(&ops), &mut h); + } else { + buffered(black_box(&ops), &mut h); + } + black_box(&h); + let elapsed = start.elapsed(); + println!( + "run={run} streaming={streaming} elapsed_ms={:.3}", + elapsed.as_secs_f64() * 1000. + ); + } + } + #[cfg(feature = "parallel")] + { + use rayon::prelude::*; + let workers = 4; + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(workers) + .build() + .unwrap(); + // Both schedules use the same number of histogram buckets and the + // same reduction. The old schedule assigns all Keccak work to one; + // the new schedule distributes whole permutations among the buckets. + for run in 0..6 { + let order = if run % 2 == 0 { + [false, true] + } else { + [true, false] + }; + for split in order { + let start = Instant::now(); + let merged = pool.install(|| { + (0..workers) + .into_par_iter() + .map(|i| { + let mut h = BitwiseHistogram::new(); + if split { + let chunk = ops.len().div_ceil(workers).max(1); + let lo = (i * chunk).min(ops.len()); + let hi = (lo + chunk).min(ops.len()); + collect_bitwise_from_keccak(&ops[lo..hi], &mut h); + } else if i == 0 { + buffered(&ops, &mut h); + } + h + }) + .reduce_with(|mut a, b| { + a.merge(&b); + a + }) + .unwrap() + }); + black_box(&merged); + let elapsed = start.elapsed(); + println!( + "run={run} workers={workers} split_streaming={split} elapsed_ms={:.3}", + elapsed.as_secs_f64() * 1000. + ); + } + } + } +} diff --git a/prover/src/tests/mod.rs b/prover/src/tests/mod.rs index 73ff6ee45..44d9e272b 100644 --- a/prover/src/tests/mod.rs +++ b/prover/src/tests/mod.rs @@ -51,6 +51,8 @@ pub mod hint_tests; #[cfg(test)] pub mod ir_stats_dump; #[cfg(test)] +mod keccak_lookup_tests; +#[cfg(test)] pub mod keccak_rnd_tests; #[cfg(test)] pub mod load_tests; diff --git a/prover/src/tests/trace_builder_tests.rs b/prover/src/tests/trace_builder_tests.rs index 428fd4700..9c1558fd9 100644 --- a/prover/src/tests/trace_builder_tests.rs +++ b/prover/src/tests/trace_builder_tests.rs @@ -598,7 +598,8 @@ mod keccak_tests { #[test] fn test_keccak_bitwise_ops_count() { let (kop, _) = make_keccak_ops(); - let ops = collect_bitwise_from_keccak(&[kop]); + let mut ops = Vec::new(); + for_each_keccak_bitwise_lookup(&[kop], |op| ops.push(op)); let xor = ops .iter()