Skip to content

perf: stream Keccak lookups across histogram workers #1001

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
diegokingston wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

perf: stream Keccak lookups across histogram workers #1001

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions prover/src/tables/bitwise.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]>,
}
Expand Down
65 changes: 40 additions & 25 deletions prover/src/tables/trace_builder.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2437,22 +2437,23 @@ pub(crate) fn collect_bitwise_from_ecdas(ops: &[ecdas::EcdasOperation]) -> Vec<B
out
}

/// Collect BITWISE lookups generated by the keccak chips.
/// Emit BITWISE lookups generated by the Keccak chips.
///
/// The keccak round chip sends BYTE_ALU and ARE_BYTES interactions (the θ/ρ
/// halfword shifts are enforced by inline constraints, not HWSL lookups); the
/// keccak core chip sends IS_HALF interactions.
/// All of these must be registered so the BITWISE table's multiplicities are correct.
#[allow(clippy::needless_range_loop)]
pub(crate) fn collect_bitwise_from_keccak(keccak_ops: &[KeccakOperation]) -> Vec<BitwiseOperation> {
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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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],
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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).
Expand Down Expand Up @@ -3290,15 +3298,15 @@ fn build_traces<I: ImageSource + Sync>(
// We never concatenate the lookups into one giant `Vec<BitwiseOperation>` (~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<dyn Fn(&mut bitwise::BitwiseHistogram) + Sync + 'a>;
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<Collector> = vec![
Box::new(|h| h.add_ops(&collect_bitwise_from_lt(&lt_ops))),
Box::new(|h| h.add_ops(&collect_bitwise_from_mul(&mul_ops, mul_chunk))),
Expand All @@ -3322,7 +3330,6 @@ fn build_traces<I: ImageSource + Sync>(
}),
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))),
Expand All @@ -3341,16 +3348,16 @@ fn build_traces<I: ImageSource + Sync>(
#[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
// 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<Collector> = Vec::with_capacity(collectors.len() + 2 * cap);
let mut units: Vec<Collector> = 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)));
Expand All @@ -3361,6 +3368,13 @@ fn build_traces<I: ImageSource + Sync>(
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<Vec<Collector>> = (0..cap).map(|_| Vec::new()).collect();
Expand All @@ -3387,6 +3401,7 @@ fn build_traces<I: ImageSource + Sync>(
#[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);
Expand Down
Loading
Loading