From c44d624b6b27553202dd303603d00a2d992a9fac Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Mon, 6 Jul 2026 11:47:32 -0300 Subject: [PATCH 01/17] opt --- prover/src/tables/bitwise.rs | 173 ++++- prover/src/tables/memw.rs | 35 +- prover/src/tables/memw_aligned.rs | 4 +- prover/src/tables/memw_register.rs | 8 +- prover/src/tables/trace_builder.rs | 648 +++++++++++++++--- .../tests/bitwise_histogram_parity_tests.rs | 130 ++++ .../memw_register_direct_parity_tests.rs | 166 +++++ prover/src/tests/mod.rs | 10 + 8 files changed, 1045 insertions(+), 129 deletions(-) create mode 100644 prover/src/tests/bitwise_histogram_parity_tests.rs create mode 100644 prover/src/tests/memw_register_direct_parity_tests.rs diff --git a/prover/src/tables/bitwise.rs b/prover/src/tables/bitwise.rs index c4871765f..388ba0bd4 100644 --- a/prover/src/tables/bitwise.rs +++ b/prover/src/tables/bitwise.rs @@ -412,18 +412,7 @@ pub fn update_multiplicities( ) { for op in ops { let row = row_index(op.x, op.y, op.z); - let mu_col = match op.lookup_type { - BitwiseOperationType::Msb8 => cols::MU_MSB8, - BitwiseOperationType::Msb16 => cols::MU_MSB16, - BitwiseOperationType::Zero => cols::MU_ZERO, - BitwiseOperationType::AreBytes => cols::MU_ARE_BYTES, - BitwiseOperationType::IsHalf => cols::MU_IS_HALF, - BitwiseOperationType::IsB20 => cols::MU_IS_B20, - BitwiseOperationType::Hwsl => cols::MU_HWSL, - BitwiseOperationType::ByteAluAnd => cols::MU_BYTE_ALU_AND, - BitwiseOperationType::ByteAluOr => cols::MU_BYTE_ALU_OR, - BitwiseOperationType::ByteAluXor => cols::MU_BYTE_ALU_XOR, - }; + let mu_col = mu_column(op.lookup_type); // Increment multiplicity let current = trace.main_table.get_row(row)[mu_col]; @@ -431,6 +420,166 @@ pub fn update_multiplicities( } } +/// Number of distinct BITWISE lookup types (one multiplicity column each). +pub const NUM_LOOKUP_TYPES: usize = 10; + +/// Dense index in `[0, NUM_LOOKUP_TYPES)` for a lookup type. Ordering is an +/// internal detail of the histogram; only [`type_mu_column`] (its inverse for +/// the fill) needs to agree with it. +#[inline] +pub const fn lookup_type_index(t: BitwiseOperationType) -> usize { + match t { + BitwiseOperationType::Msb8 => 0, + BitwiseOperationType::Msb16 => 1, + BitwiseOperationType::Zero => 2, + BitwiseOperationType::AreBytes => 3, + BitwiseOperationType::IsHalf => 4, + BitwiseOperationType::IsB20 => 5, + BitwiseOperationType::Hwsl => 6, + BitwiseOperationType::ByteAluAnd => 7, + BitwiseOperationType::ByteAluOr => 8, + BitwiseOperationType::ByteAluXor => 9, + } +} + +/// Multiplicity column for a lookup type (used by the legacy per-op path). +#[inline] +pub const fn mu_column(t: BitwiseOperationType) -> usize { + match t { + BitwiseOperationType::Msb8 => cols::MU_MSB8, + BitwiseOperationType::Msb16 => cols::MU_MSB16, + BitwiseOperationType::Zero => cols::MU_ZERO, + BitwiseOperationType::AreBytes => cols::MU_ARE_BYTES, + BitwiseOperationType::IsHalf => cols::MU_IS_HALF, + BitwiseOperationType::IsB20 => cols::MU_IS_B20, + BitwiseOperationType::Hwsl => cols::MU_HWSL, + BitwiseOperationType::ByteAluAnd => cols::MU_BYTE_ALU_AND, + BitwiseOperationType::ByteAluOr => cols::MU_BYTE_ALU_OR, + BitwiseOperationType::ByteAluXor => cols::MU_BYTE_ALU_XOR, + } +} + +/// Multiplicity column for the histogram lane at dense index `type_idx` +/// (inverse of [`lookup_type_index`]). Used by [`BitwiseHistogram::fill_multiplicities`]. +#[inline] +const fn type_mu_column(type_idx: usize) -> usize { + // Columns 11..=20, one per lookup type, in `lookup_type_index` order. + cols::MU_MSB8 + type_idx +} + +/// Guards the fragile assumption that the histogram's arithmetic column map +/// (`type_mu_column ∘ lookup_type_index`) agrees with the authoritative per-type +/// `mu_column` match. If anyone reorders `lookup_type_index` or renumbers the `MU_*` +/// columns so they are no longer contiguous 11..=20 in that order, this fails loudly in a +/// plain `cargo test` (not just the #[ignore]d parity gate) — a wrong MU column would +/// silently unbalance the BITWISE bus and produce an unsound proof. +#[cfg(test)] +mod histogram_column_map_guard { + use super::*; + + #[test] + fn type_mu_column_matches_mu_column_for_all_types() { + use BitwiseOperationType::*; + const ALL: [BitwiseOperationType; NUM_LOOKUP_TYPES] = [ + Msb8, Msb16, Zero, AreBytes, IsHalf, IsB20, Hwsl, ByteAluAnd, ByteAluOr, ByteAluXor, + ]; + for t in ALL { + assert_eq!( + type_mu_column(lookup_type_index(t)), + mu_column(t), + "histogram column map disagrees with mu_column for {t:?}" + ); + } + // Also confirm the dense indices are a bijection over 0..NUM_LOOKUP_TYPES. + let mut seen = [false; NUM_LOOKUP_TYPES]; + for t in ALL { + seen[lookup_type_index(t)] = true; + } + assert!(seen.iter().all(|&s| s), "lookup_type_index is not a bijection"); + } +} + +/// "Histogram-on-the-fly" accumulator for BITWISE lookup multiplicities. +/// +/// Replaces materializing the giant `Vec` (whose only consumer +/// is the multiplicity count) with a dense counter array. Each lookup increments +/// `counters[type_idx * NUM_ROWS + row_index(x, y, z)]`. +/// +/// The histogram is a commutative monoid: increments and [`merge`](Self::merge) +/// are order-independent, so per-thread histograms can be tree-reduced and the +/// resulting multiplicities are byte-identical to the serial per-op count that +/// [`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. +pub struct BitwiseHistogram { + counters: Box<[u64]>, +} + +impl BitwiseHistogram { + /// Allocate a zeroed histogram (80 MiB). + pub fn new() -> Self { + Self { + counters: vec![0u64; NUM_ROWS * NUM_LOOKUP_TYPES].into_boxed_slice(), + } + } + + /// Increment the counter for one lookup. + #[inline] + pub fn bump(&mut self, op: BitwiseOperation) { + let idx = lookup_type_index(op.lookup_type) * NUM_ROWS + row_index(op.x, op.y, op.z); + // `debug_assert` guards the invariant; row_index already bounds-checks z<16 + // in debug and (x,y) are u8 so idx is always in range. + self.counters[idx] += 1; + } + + /// Fold a slice of lookups into the histogram. + #[inline] + pub fn add_ops(&mut self, ops: &[BitwiseOperation]) { + for &op in ops { + self.bump(op); + } + } + + /// Merge another histogram into this one (commutative, order-independent). + pub fn merge(&mut self, other: &BitwiseHistogram) { + for (a, b) in self.counters.iter_mut().zip(other.counters.iter()) { + *a += *b; + } + } + + /// Total number of lookups counted (for instrumentation only). + pub fn total(&self) -> u64 { + self.counters.iter().sum() + } + + /// Write the accumulated multiplicities into the BITWISE trace's MU columns. + /// + /// Produces exactly the same MU columns as calling [`update_multiplicities`] + /// with the full op vector, because both sum one increment per lookup into + /// the `(row, mu_col)` cell. + pub fn fill_multiplicities( + &self, + trace: &mut TraceTable, + ) { + for type_idx in 0..NUM_LOOKUP_TYPES { + let mu_col = type_mu_column(type_idx); + let base = type_idx * NUM_ROWS; + for row in 0..NUM_ROWS { + let count = self.counters[base + row]; + if count != 0 { + trace.main_table.set_fe(row, mu_col, FE::from(count)); + } + } + } + } +} + +impl Default for BitwiseHistogram { + fn default() -> Self { + Self::new() + } +} + /// Types of lookups the BITWISE table provides. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum BitwiseOperationType { diff --git a/prover/src/tables/memw.rs b/prover/src/tables/memw.rs index 2b240747c..dceb5b1a7 100644 --- a/prover/src/tables/memw.rs +++ b/prover/src/tables/memw.rs @@ -108,17 +108,22 @@ pub struct MemwOperation { pub is_register: bool, /// Base address (64-bit) pub base_address: u64, - /// Values to write (8 bytes) - pub value: [u64; 8], + /// Values to write. Each element is one memory byte (0-255) or, for register + /// accesses, a 32-bit half of the register word — both fit in u32, so this is + /// `[u32; 8]` rather than `[u64; 8]` to halve the struct's footprint (the walk + /// materializes tens of millions of these; it is memory-bandwidth-bound). + pub value: [u32; 8], /// Timestamp of this access pub timestamp: u64, /// Access width: 1, 2, 4, or 8 bytes pub width: u8, /// Whether this is a read (true) or write (false) pub is_read: bool, - /// Previous values at the addresses (filled by memory model) - pub old: [u64; 8], - /// Previous timestamps at the addresses (filled by memory model) + /// Previous values at the addresses (filled by memory model). Same element + /// domain as `value` (byte or 32-bit register half) → `[u32; 8]`. + pub old: [u32; 8], + /// Previous timestamps at the addresses (filled by memory model). Timestamps + /// can reach u64::MAX (HALT), so these stay `[u64; 8]`. pub old_timestamp: [u64; 8], } @@ -135,7 +140,16 @@ impl MemwOperation { Self { is_register, base_address, - value, + // Callers build a [u64; 8] transiently on the stack; we store the u32 + // domain (byte / 32-bit register half) so the persisted struct is half + // the size. Values never exceed u32 (every element is a byte or a + // pack_register_value 32-bit limb). The debug_assert fails loudly if a + // future caller ever passes an out-of-domain value instead of silently + // truncating it (which would be a soundness bug). + value: value.map(|v| { + debug_assert!(v <= u32::MAX as u64, "MemwOperation value element exceeds u32: {v}"); + v as u32 + }), timestamp, width, is_read, @@ -146,7 +160,10 @@ impl MemwOperation { /// Set the old values (from memory model). pub fn with_old(mut self, old: [u64; 8], old_timestamp: [u64; 8]) -> Self { - self.old = old; + self.old = old.map(|v| { + debug_assert!(v <= u32::MAX as u64, "MemwOperation old element exceeds u32: {v}"); + v as u32 + }); self.old_timestamp = old_timestamp; self } @@ -192,7 +209,7 @@ pub fn generate_memw_trace( // value[8] for i in 0..8 { - table.set_u64(row_idx, cols::VALUE[i], op.value[i]); + table.set_u64(row_idx, cols::VALUE[i], op.value[i] as u64); } // timestamp as DWordWL (2 words) @@ -206,7 +223,7 @@ pub fn generate_memw_trace( // Output: old[8] for i in 0..8 { - table.set_u64(row_idx, cols::OLD[i], op.old[i]); + table.set_u64(row_idx, cols::OLD[i], op.old[i] as u64); } // Auxiliary: carry[7] diff --git a/prover/src/tables/memw_aligned.rs b/prover/src/tables/memw_aligned.rs index 8042d9052..eec30ff3d 100644 --- a/prover/src/tables/memw_aligned.rs +++ b/prover/src/tables/memw_aligned.rs @@ -107,7 +107,7 @@ pub fn generate_memw_aligned_trace( table.set_dword_whh(row_idx, cols::BASE_ADDRESS[0], op.base_address); for i in 0..8 { - table.set_u64(row_idx, cols::VALUE[i], op.value[i]); + table.set_u64(row_idx, cols::VALUE[i], op.value[i] as u64); } table.set_dword_wl(row_idx, cols::TIMESTAMP_0, op.timestamp); @@ -118,7 +118,7 @@ pub fn generate_memw_aligned_trace( table.set_bool(row_idx, cols::WRITE8, w8); for i in 0..8 { - table.set_u64(row_idx, cols::OLD[i], op.old[i]); + table.set_u64(row_idx, cols::OLD[i], op.old[i] as u64); } // Single old_timestamp (from old_timestamp[0], verified equal for all bytes) diff --git a/prover/src/tables/memw_register.rs b/prover/src/tables/memw_register.rs index 14a696cb9..49f7bbff2 100644 --- a/prover/src/tables/memw_register.rs +++ b/prover/src/tables/memw_register.rs @@ -125,12 +125,12 @@ pub fn generate_memw_register_trace( table.set_dword_wl(row_idx, cols::TIMESTAMP_0, op.timestamp); // Value: registers are DWordWL = 2 words - table.set_u64(row_idx, cols::VAL_0, op.value[0]); - table.set_u64(row_idx, cols::VAL_1, op.value[1]); + table.set_u64(row_idx, cols::VAL_0, op.value[0] as u64); + table.set_u64(row_idx, cols::VAL_1, op.value[1] as u64); // Old value - table.set_u64(row_idx, cols::OLD_0, op.old[0]); - table.set_u64(row_idx, cols::OLD_1, op.old[1]); + table.set_u64(row_idx, cols::OLD_0, op.old[0] as u64); + table.set_u64(row_idx, cols::OLD_1, op.old[1] as u64); // Old timestamp low (upper limb shared with TIMESTAMP_1) table.set_u64( diff --git a/prover/src/tables/trace_builder.rs b/prover/src/tables/trace_builder.rs index 93f3ba563..c61b64262 100644 --- a/prover/src/tables/trace_builder.rs +++ b/prover/src/tables/trace_builder.rs @@ -365,12 +365,307 @@ fn collect_cpu_ops( /// Returns: (memw_ops, load_ops, lt_ops, shift_ops, bitwise_ops, commit_ops, keccak_ops, /// cpu32_ops, ecsm_ops, ec_scalar_ops, ecdas_ops) #[allow(clippy::type_complexity)] +/// Whether to use the legacy MEMW_R trace-gen path (materialize `Vec` +/// for register accesses + fill columns from it), selected via the +/// `LAMBDA_VM_LEGACY_TRACEGEN` environment variable (set to any non-empty value). +/// +/// The default (unset) is the new direct-to-column register fill. The legacy path +/// exists for A/B measurement and the byte-parity gate test. +fn legacy_tracegen() -> bool { + std::env::var("LAMBDA_VM_LEGACY_TRACEGEN") + .map(|v| !v.is_empty()) + .unwrap_or(false) +} + +/// Compact, already-decomposed record for one MEMW_R (register fast-path) access. +/// +/// This is the "direct-to-column" carrier: it holds exactly the fields the MEMW_R +/// column fill (`generate_memw_register_trace_direct`) and its IS_HALFWORD bitwise +/// collector (`collect_bitwise_from_memw_register_direct`) need, and nothing else. +/// It replaces the full `MemwOperation` (216→~152 B via the E6 shrink, but still 8 +/// `[u32;8]`/`[u64;8]` arrays) for register accesses — the largest table by rows — +/// so the walk never materializes a `MemwOperation` for the register fast path. +/// +/// Field domains mirror `MemwOperation`'s so the produced table is byte-identical: +/// - `address` = `base_address / 2` (the register index 0..=255; ADDRESS column, +/// and `2*ADDRESS` on the memory/MEMW buses) +/// - `val0/val1` = `value[0]`/`value[1]` (the 32-bit register halves) +/// - `old0/old1` = `old[0]`/`old[1]` +/// - `old_ts_lo` = `old_timestamp[0] & 0xFFFF_FFFF` (the two words share old_timestamp, +/// enforced by `is_register_op`; the upper limb is TIMESTAMP_1 = timestamp>>32) +#[derive(Debug, Clone, Copy)] +struct RegRow { + address: u64, + timestamp: u64, + val0: u32, + val1: u32, + old0: u32, + old1: u32, + old_ts_lo: u32, + is_read: bool, +} + +impl RegRow { + /// Build a `RegRow` from a fully-formed register `MemwOperation`. Used on the + /// precompile / commit / keccak / halt paths, which construct a `MemwOperation` + /// first (their computation is unchanged) and only convert to the compact row + /// once the op is known to route to MEMW_R. + /// + /// Only valid for ops for which `is_register_op` is true (width==2, atomic + /// old_timestamp). The debug asserts mirror `generate_memw_register_trace`. + #[inline] + fn from_memw(op: &MemwOperation) -> Self { + debug_assert_eq!(op.base_address % 2, 0); + debug_assert_eq!(op.old_timestamp[0], op.old_timestamp[1]); + RegRow { + address: op.base_address / 2, + timestamp: op.timestamp, + val0: op.value[0], + val1: op.value[1], + old0: op.old[0], + old1: op.old[1], + old_ts_lo: (op.old_timestamp[0] & 0xFFFF_FFFF) as u32, + is_read: op.is_read, + } + } +} + +/// Direct-to-column MEMW_R trace fill from compact [`RegRow`]s. +/// +/// This is the fused replacement for `generate_memw_register_trace` (which fills from +/// `&[MemwOperation]`): it produces a BYTE-IDENTICAL table. Every column write below +/// mirrors `memw_register::generate_memw_register_trace` cell-for-cell — the only +/// difference is the source carrier (`RegRow` vs `MemwOperation`). Keep the two in sync. +fn generate_memw_register_trace_direct( + rows: &[RegRow], +) -> TraceTable { + use super::types::{FE, VmTable}; + + let num_rows = rows.len().next_power_of_two().max(4); + let mut trace = TraceTable::new_main( + vec![FE::zero(); num_rows * memw_register::cols::NUM_COLUMNS], + memw_register::cols::NUM_COLUMNS, + 1, + ); + let table = &mut trace.main_table; + + use memw_register::cols; + for (row_idx, r) in rows.iter().enumerate() { + // ADDRESS = base_address / 2 (already divided in RegRow). + table.set_u64(row_idx, cols::ADDRESS, r.address); + // Timestamp split into lo/hi 32-bit words. + table.set_dword_wl(row_idx, cols::TIMESTAMP_0, r.timestamp); + // Value: registers are DWordWL = 2 words. + table.set_u64(row_idx, cols::VAL_0, r.val0 as u64); + table.set_u64(row_idx, cols::VAL_1, r.val1 as u64); + // Old value. + table.set_u64(row_idx, cols::OLD_0, r.old0 as u64); + table.set_u64(row_idx, cols::OLD_1, r.old1 as u64); + // Old timestamp low (upper limb shared with TIMESTAMP_1). + table.set_u64(row_idx, cols::OLD_TIMESTAMP_LO, r.old_ts_lo as u64); + // Multiplicity. + table.set_bool(row_idx, cols::MU_READ, r.is_read); + table.set_bool(row_idx, cols::MU_WRITE, !r.is_read); + } + + trace +} + +/// The single IS_HALFWORD lookup a MEMW_R access sends: proves the timestamp delta +/// `ts_lo - old_ts_lo` is in [1, 2^16] by decomposing `ts_lo - old_ts_lo - 1` into two bytes. +/// Shared by BOTH the direct (`RegRow`) and `MemwOperation` collectors so they can never drift +/// (a divergence would be a silent soundness bug — see the E7 review). +#[inline] +fn memw_register_is_half_lookup(ts_lo: u32, old_ts_lo: u32) -> BitwiseOperation { + debug_assert!( + ts_lo > old_ts_lo, + "ts_lo must exceed old_ts_lo (enforced by reg_ts_delta_in_range)" + ); + let diff_minus_1 = (ts_lo - old_ts_lo - 1) as u16; + BitwiseOperation::halfword( + BitwiseOperationType::IsHalf, + (diff_minus_1 & 0xFF) as u8, + (diff_minus_1 >> 8) as u8, + ) +} + +/// IS_HALFWORD bitwise lookups for MEMW_R, computed directly from [`RegRow`]s. +/// Byte-identical to `collect_bitwise_from_memw_register` (both call the shared helper). +fn collect_bitwise_from_memw_register_direct(rows: &[RegRow]) -> Vec { + rows.iter() + .map(|r| memw_register_is_half_lookup((r.timestamp & 0xFFFF_FFFF) as u32, r.old_ts_lo)) + .collect() +} + +/// Routes each `MemwOperation` into its destination table bucket at CREATION time +/// (register fast-path / aligned / general), so the walk fills the three buckets directly +/// and no separate partition pass is needed downstream. Classification order matches the old +/// two-stage partition (register first, then aligned), and push order within each bucket is +/// preserved, so the buckets are byte-identical to `partition`-ing one combined vec. +/// +/// ## Direct-to-column register fill +/// +/// For the register fast path we do NOT materialize a `Vec`. Ops that route +/// to MEMW_R are stored as compact [`RegRow`]s (`register_rows`) and later filled directly +/// into the MEMW_R columns. Under the `LAMBDA_VM_LEGACY_TRACEGEN` flag we instead keep the +/// old `Vec` (`register_ops`) and the old column fill, for A/B + byte-parity +/// testing. The `aligned` / `general` buckets are unchanged in both modes — an op that FAILS +/// `is_register_op` is always routed as a `MemwOperation` exactly as before. +#[derive(Default)] +struct MemwBuckets { + /// New path: compact register rows (empty in legacy mode). + register_rows: Vec, + /// Legacy path: full register `MemwOperation`s (empty in new mode). + register_ops: Vec, + aligned: Vec, + general: Vec, + /// When true, keep the old `Vec` register materialization + fill. + legacy: bool, +} + +impl MemwBuckets { + fn with_register_capacity(n: usize, legacy: bool) -> Self { + Self { + register_rows: if legacy { Vec::new() } else { Vec::with_capacity(n) }, + register_ops: if legacy { Vec::with_capacity(n) } else { Vec::new() }, + aligned: Vec::new(), + general: Vec::new(), + legacy, + } + } + + /// Store an op already known to route to MEMW_R. + #[inline] + fn push_register(&mut self, op: MemwOperation) { + if self.legacy { + self.register_ops.push(op); + } else { + self.register_rows.push(RegRow::from_memw(&op)); + } + } + + /// Store a compact register row directly (new path only). In legacy mode this + /// reconstructs the equivalent `MemwOperation` so the two modes are comparable. + #[inline] + fn push_register_row(&mut self, row: RegRow, op_if_legacy: impl FnOnce() -> MemwOperation) { + if self.legacy { + self.register_ops.push(op_if_legacy()); + } else { + self.register_rows.push(row); + } + } + + #[inline] + fn push(&mut self, op: MemwOperation) { + if is_register_op(&op) { + self.push_register(op); + } else if is_aligned_op(&op) { + self.aligned.push(op); + } else { + self.general.push(op); + } + } + fn extend_ops(&mut self, ops: impl IntoIterator) { + for op in ops { + self.push(op); + } + } +} + +/// Sink for `MemwOperation`s so `collect_register_ops_from_cpu` can feed either a plain +/// `Vec` (tests/scratch paths) or the classifying [`MemwBuckets`] (the walk). +trait MemwSink { + fn push_op(&mut self, op: MemwOperation); + + /// Fast path for a 2-word register access (M1/M3/M5 and precompile register I/O). + /// + /// The caller passes the compact, pre-decomposed fields. The sink decides routing + /// (via the same predicate as `is_register_op`): if the timestamp delta admits the + /// op into MEMW_R it fills a compact [`RegRow`] DIRECTLY — no `MemwOperation` is + /// built. Only on the (rare) fallback (delta out of IS_HALF range, or upper-limb + /// mismatch) does it call `build_fallback` to materialize the `MemwOperation` and + /// route it to the aligned/general bucket exactly as before. + /// + /// `reg_addr` is `2 * reg_index`; `[val0,val1]`/`[old0,old1]` are the 32-bit halves; + /// `old_ts` is the (shared) old_timestamp of both words. + #[inline] + #[allow(clippy::too_many_arguments)] + fn push_reg_access( + &mut self, + reg_addr: u64, + val0: u32, + val1: u32, + old0: u32, + old1: u32, + timestamp: u64, + old_ts: u64, + is_read: bool, + build_fallback: impl FnOnce() -> MemwOperation, + ) { + // Default impl (plain Vec): register accesses are still ordinary MemwOperations. + let _ = (reg_addr, val0, val1, old0, old1, timestamp, old_ts, is_read); + self.push_op(build_fallback()); + } +} +impl MemwSink for Vec { + #[inline] + fn push_op(&mut self, op: MemwOperation) { + self.push(op); + } +} +impl MemwSink for MemwBuckets { + #[inline] + fn push_op(&mut self, op: MemwOperation) { + self.push(op); + } + + #[inline] + #[allow(clippy::too_many_arguments)] + fn push_reg_access( + &mut self, + reg_addr: u64, + val0: u32, + val1: u32, + old0: u32, + old1: u32, + timestamp: u64, + old_ts: u64, + is_read: bool, + build_fallback: impl FnOnce() -> MemwOperation, + ) { + // Mirror `is_register_op` for a width-2 register access whose two words share + // `old_ts` (always true here by construction). If it passes, fill a RegRow + // directly; otherwise fall back to the general/aligned MemwOperation path. + if reg_ts_delta_in_range(timestamp, old_ts) { + let row = RegRow { + address: reg_addr / 2, + timestamp, + val0, + val1, + old0, + old1, + old_ts_lo: (old_ts & 0xFFFF_FFFF) as u32, + is_read, + }; + self.push_register_row(row, build_fallback); + } else { + let op = build_fallback(); + debug_assert!(!is_register_op(&op), "reg fallback must not be MEMW_R"); + if is_aligned_op(&op) { + self.aligned.push(op); + } else { + self.general.push(op); + } + } + } +} + fn collect_ops_from_cpu( cpu_ops: &[CpuOperation], memory_state: &mut MemoryState, register_state: &mut RegisterState, ) -> ( - Vec, + MemwBuckets, Vec, Vec, Vec, @@ -382,7 +677,7 @@ fn collect_ops_from_cpu( Vec, Vec, ) { - let mut memw_ops = Vec::with_capacity(cpu_ops.len() * 3); + let mut memw = MemwBuckets::with_register_capacity(cpu_ops.len() * 3, legacy_tracegen()); let mut load_ops = Vec::with_capacity(cpu_ops.len() / 8 + 1); let mut lt_ops = Vec::with_capacity(cpu_ops.len() / 10 + 1); let mut shift_ops = Vec::with_capacity(cpu_ops.len() / 10 + 1); @@ -414,16 +709,16 @@ fn collect_ops_from_cpu( // Collect memory operations for Load/Store instructions if op.decode.fields.is_load() { let (memw_op, load_op, lookups) = collect_load_op_from_cpu(op, memory_state); - memw_ops.push(memw_op); + memw.push(memw_op); load_ops.push(load_op); bitwise_ops.extend(lookups); } else if op.decode.fields.is_store() { let memw_op = collect_store_op_from_cpu(op, memory_state); - memw_ops.push(memw_op); + memw.push(memw_op); } // Collect register operations (M1, M3, M5) - collect_register_ops_from_cpu(op, register_state, &mut memw_ops); + collect_register_ops_from_cpu(op, register_state, &mut memw); // Collect COMMIT ECALL memory operations (register reads/writes + byte reads) if op.ecall_commit { @@ -433,7 +728,7 @@ fn collect_ops_from_cpu( current_commit_index as u64, )); let reg_commit_ops = collect_commit_memw_ops(op, register_state, memory_state); - memw_ops.extend(reg_commit_ops); + memw.extend_ops(reg_commit_ops); let count = u32::try_from(op.commit_count).expect("commit_count exceeds u32 range"); current_commit_index = current_commit_index .checked_add(count) @@ -469,7 +764,7 @@ fn collect_ops_from_cpu( // collect_keccak_memw_ops handles memory_state + register_state updates let keccak_memw_ops = collect_keccak_memw_ops(op, &input, &output, memory_state, register_state); - memw_ops.extend(keccak_memw_ops); + memw.extend_ops(keccak_memw_ops); keccak_ops.push(KeccakOperation { timestamp: op.timestamp, state_addr, @@ -482,7 +777,7 @@ fn collect_ops_from_cpu( if op.ecall_ecsm { let (ecsm_memw, ecsm_op, ec_scalar_rows, ecdas_rows) = collect_ecsm_ops(op, memory_state, register_state); - memw_ops.extend(ecsm_memw); + memw.extend_ops(ecsm_memw); ecsm_ops.push(ecsm_op); ec_scalar_ops.extend(ec_scalar_rows); ecdas_ops.extend(ecdas_rows); @@ -518,7 +813,9 @@ fn collect_ops_from_cpu( } } - // Collect CPU range-check bitwise lookups (ARE_BYTES + IS_HALF). + // Collect CPU range-check bitwise lookups (ARE_BYTES + IS_HALF). Kept serial here: + // it's only ~110 ms (a serial `.extend` into one growing Vec), and moving it to a + // rayon `flat_map`-collect over 6.8 M per-op Vecs regressed p4 ~4× (alloc + merge). bitwise_ops.extend(op.collect_bitwise_ops()); } @@ -531,7 +828,7 @@ fn collect_ops_from_cpu( ); ( - memw_ops, + memw, load_ops, lt_ops, shift_ops, @@ -773,10 +1070,10 @@ fn collect_ecsm_ops( /// Collects register read/write operations (M1, M3, M5) from CpuOperation, /// pushing them into `memw_ops`. -fn collect_register_ops_from_cpu( +fn collect_register_ops_from_cpu( op: &CpuOperation, register_state: &mut RegisterState, - memw_ops: &mut Vec, + memw_ops: &mut S, ) { let d = &op.decode.fields; // These register accesses happen for every real instruction. For non-word @@ -795,12 +1092,24 @@ fn collect_register_ops_from_cpu( } else { register_state.read(d.rs1) }; - // old_timestamps array is 8 elements but only first 2 are used for registers - let old_timestamps = [old_ts, old_ts, 0, 0, 0, 0, 0, 0]; - - let memw_op = MemwOperation::new(true, reg_addr, reg_value, op.timestamp, 2, true) - .with_old(reg_value, old_timestamps); - memw_ops.push(memw_op); + let ts = op.timestamp; + // Direct fast path: fill a RegRow when routing to MEMW_R; the closure rebuilds + // the identical MemwOperation only on the (rare) general/aligned fallback. + memw_ops.push_reg_access( + reg_addr, + reg_value[0] as u32, + reg_value[1] as u32, + reg_value[0] as u32, + reg_value[1] as u32, + ts, + old_ts, + true, + || { + let old_timestamps = [old_ts, old_ts, 0, 0, 0, 0, 0, 0]; + MemwOperation::new(true, reg_addr, reg_value, ts, 2, true) + .with_old(reg_value, old_timestamps) + }, + ); if d.rs1 == 255 { register_state.write_pc(op.rv1, op.timestamp); } else { @@ -813,12 +1122,22 @@ fn collect_register_ops_from_cpu( let reg_value = pack_register_value(op.rv2); let reg_addr = 2 * d.rs2 as u64; let (_old_val, old_ts) = register_state.read(d.rs2); - // old_timestamps array is 8 elements but only first 2 are used for registers - let old_timestamps = [old_ts, old_ts, 0, 0, 0, 0, 0, 0]; - - let memw_op = MemwOperation::new(true, reg_addr, reg_value, op.timestamp + 1, 2, true) - .with_old(reg_value, old_timestamps); - memw_ops.push(memw_op); + let ts = op.timestamp + 1; + memw_ops.push_reg_access( + reg_addr, + reg_value[0] as u32, + reg_value[1] as u32, + reg_value[0] as u32, + reg_value[1] as u32, + ts, + old_ts, + true, + || { + let old_timestamps = [old_ts, old_ts, 0, 0, 0, 0, 0, 0]; + MemwOperation::new(true, reg_addr, reg_value, ts, 2, true) + .with_old(reg_value, old_timestamps) + }, + ); register_state.write(d.rs2, op.rv2, op.timestamp + 1); } @@ -828,12 +1147,22 @@ fn collect_register_ops_from_cpu( let reg_addr = 2 * d.rd as u64; let (old_val, old_ts) = register_state.read(d.rd); let old_value = pack_register_value(old_val); - // old_timestamps array is 8 elements but only first 2 are used for registers - let old_timestamps = [old_ts, old_ts, 0, 0, 0, 0, 0, 0]; - - let memw_op = MemwOperation::new(true, reg_addr, reg_value, op.timestamp + 2, 2, false) - .with_old(old_value, old_timestamps); - memw_ops.push(memw_op); + let ts = op.timestamp + 2; + memw_ops.push_reg_access( + reg_addr, + reg_value[0] as u32, + reg_value[1] as u32, + old_value[0] as u32, + old_value[1] as u32, + ts, + old_ts, + false, + || { + let old_timestamps = [old_ts, old_ts, 0, 0, 0, 0, 0, 0]; + MemwOperation::new(true, reg_addr, reg_value, ts, 2, false) + .with_old(old_value, old_timestamps) + }, + ); register_state.write(d.rd, op.rvd, op.timestamp + 2); } @@ -1377,11 +1706,23 @@ pub(crate) fn is_register_op(op: &MemwOperation) -> bool { if op.old_timestamp[0] != op.old_timestamp[1] { return false; } - let ts = op.timestamp; - let old_ts = op.old_timestamp[0]; - let ts_lo = ts & 0xFFFF_FFFF; + reg_ts_delta_in_range(op.timestamp, op.old_timestamp[0]) +} + +/// The timestamp-delta admission test for MEMW_R (conditions 3-5 of `is_register_op`), +/// factored out so the direct fast path (`push_reg_access`) and the `MemwOperation` +/// classifier (`is_register_op`) share EXACTLY the same routing logic: +/// - `ts_hi == old_ts_hi` (upper limbs match) +/// - `ts_lo > old_ts_lo` (lower-limb ordering) +/// - `ts_lo - old_ts_lo <= 2^16` (delta fits the IS_HALF range [1, 2^16]) +/// +/// The fast path only calls this for width-2 register accesses whose two words share +/// `old_ts` by construction, so conditions 1-2 of `is_register_op` always hold there. +#[inline] +fn reg_ts_delta_in_range(timestamp: u64, old_ts: u64) -> bool { + let ts_lo = timestamp & 0xFFFF_FFFF; let old_ts_lo = old_ts & 0xFFFF_FFFF; - let ts_hi = ts >> 32; + let ts_hi = timestamp >> 32; let old_ts_hi = old_ts >> 32; ts_hi == old_ts_hi && ts_lo > old_ts_lo && (ts_lo - old_ts_lo) <= 0x10000 } @@ -1393,17 +1734,9 @@ pub(crate) fn is_register_op(op: &MemwOperation) -> bool { fn collect_bitwise_from_memw_register(ops: &[MemwOperation]) -> Vec { ops.iter() .map(|op| { - let ts_lo = op.timestamp & 0xFFFF_FFFF; - let old_ts_lo = op.old_timestamp[0] & 0xFFFF_FFFF; - debug_assert!( - ts_lo > old_ts_lo, - "ts_lo must exceed old_ts_lo (enforced by is_register_op)" - ); - let diff_minus_1 = (ts_lo - old_ts_lo - 1) as u16; - BitwiseOperation::halfword( - BitwiseOperationType::IsHalf, - (diff_minus_1 & 0xFF) as u8, - (diff_minus_1 >> 8) as u8, + memw_register_is_half_lookup( + (op.timestamp & 0xFFFF_FFFF) as u32, + (op.old_timestamp[0] & 0xFFFF_FFFF) as u32, ) }) .collect() @@ -2584,7 +2917,10 @@ struct CollectedOps { cpu_ops: Vec, memw_ops: Vec, memw_aligned_ops: Vec, + /// Legacy MEMW_R ops (non-empty only under `LAMBDA_VM_LEGACY_TRACEGEN`). memw_register_ops: Vec, + /// New direct-fill MEMW_R rows (non-empty in the default path). + memw_register_rows: Vec, load_ops: Vec, lt_ops: Vec, shift_ops: Vec, @@ -2647,7 +2983,7 @@ fn chunk_and_generate( #[allow(clippy::too_many_arguments)] fn collect_all_ops( cpu_ops: Vec, - mut memw_ops: Vec, + mut memw: MemwBuckets, load_ops: Vec, mut lt_ops: Vec, mut shift_ops: Vec, @@ -2666,16 +3002,22 @@ fn collect_all_ops( // Only the final epoch terminates; intermediate epochs keep their boundary // register state (no zeroizing) so it can seed the next epoch. if is_final { - let halt_memw_ops = collect_halt_ops(register_state); - memw_ops.extend(halt_memw_ops); + // Route halt ops through the same classifier; they append to the end of their + // buckets, matching the old "append then partition" order. + memw.extend_ops(collect_halt_ops(register_state)); } - // Route MEMW_R (register fast-path) first, then MEMW_A (aligned), rest → MEMW. - // Order matters: register ops would also pass is_aligned_op, so check first. - let (memw_register_ops, memw_ops): (Vec<_>, Vec<_>) = - memw_ops.into_iter().partition(is_register_op); - let (memw_aligned_ops, memw_ops): (Vec<_>, Vec<_>) = - memw_ops.into_iter().partition(is_aligned_op); + // The walk (`collect_ops_from_cpu`) already routed every MemwOperation into its bucket at + // creation via `MemwBuckets`, so there is no separate partition pass here — the old + // two-`partition` sweep (which moved millions of structs a second time, a bandwidth-bound + // cost) is gone. Order within each bucket matches the old stable partitions → byte-identical. + let MemwBuckets { + register_rows: memw_register_rows, + register_ops: memw_register_ops, + aligned: memw_aligned_ops, + general: memw_ops, + legacy: _, + } = memw; // Collect BRANCH operations from CPU ops where branch_cond = true let branch_ops: Vec = cpu_ops @@ -2781,6 +3123,7 @@ fn collect_all_ops( memw_ops, memw_aligned_ops, memw_register_ops, + memw_register_rows, load_ops, lt_ops, shift_ops, @@ -2825,6 +3168,7 @@ fn build_traces( memw_ops, memw_aligned_ops, memw_register_ops, + memw_register_rows, load_ops, mut lt_ops, shift_ops, @@ -2854,51 +3198,12 @@ fn build_traces( // ===================================================================== #[cfg(feature = "instruments")] let __sp = stark::instruments::span("p4_bitwise_collect"); - bitwise_ops.extend(collect_bitwise_from_lt(<_ops)); - // MUL/DVRM dedup their per-unique bit-gated lookups PER CHIP INSTANCE, so pass - // the same chunk size used to split them into instances (see chunk_and_generate - // below) so the BITWISE multiplicity matches the per-instance sends. - bitwise_ops.extend(collect_bitwise_from_mul(&mul_ops, max_rows.mul)); - bitwise_ops.extend(collect_bitwise_from_dvrm(&dvrm_ops, max_rows.dvrm)); - bitwise_ops.extend(collect_bitwise_from_branch(&branch_ops)); - bitwise_ops.extend(shift::collect_bitwise_from_shift(&shift_ops)); - // Auxiliary chips: BYTEWISE sends 8× BYTE_ALU/op; EQ sends 4× IS_HALF + ZERO. - for op in &bytewise_ops { - bitwise_ops.extend(op.collect_bitwise_ops()); - } - for op in &eq_ops { - bitwise_ops.extend(op.collect_bitwise_ops()); - } - for op in &store_ops { - bitwise_ops.extend(op.collect_bitwise_ops()); - } - bitwise_ops.extend(collect_bitwise_from_memw_aligned(&memw_aligned_ops)); - // MEMW_R sends IS_HALFWORD[timestamp_0 - old_timestamp_lo - 1] - bitwise_ops.extend(collect_bitwise_from_memw_register(&memw_register_ops)); - // PAGE tables do a batched ARE_BYTES[init, fini] lookup per row (C1+C2). - // Continuation epochs (l2g_memory_bookend) skip PAGE entirely (see the - // generate_page_tables call below), so they skip its AreBytes lookups too. - if let Some(image) = initial_image - && !l2g_memory_bookend - { - bitwise_ops.extend(collect_bitwise_from_page( - image, - memory_state, - l2g_memory_bookend, - )); - } let public_output_bytes: Vec = commit_ops .iter() .filter(|op| !op.end) .map(|op| op.value) .collect(); - // COMMIT table sends AreBytes and IsHalfword lookups - bitwise_ops.extend(collect_bitwise_from_commit(&commit_ops)); - // KECCAK_RND sends XOR/AND/ARE_BYTES/HWSL; KECCAK core sends IS_HALF - bitwise_ops.extend(collect_bitwise_from_keccak(&keccak_ops)); - bitwise_ops.extend(collect_bitwise_from_ecsm(&ecsm_ops)); - bitwise_ops.extend(collect_bitwise_from_ecdas(&ecdas_ops)); // CPU padding rows send ARE_BYTES with all-zero values. // Add corresponding ops so the bitwise table multiplicities balance. @@ -2906,7 +3211,124 @@ fn build_traces( .chunks(max_rows.cpu) .map(|chunk| chunk.len().next_power_of_two().max(4) - chunk.len()) .sum(); - bitwise_ops.extend(collect_byte_check_ops_for_padding(num_padding_rows)); + + // The per-source bitwise collectors are all pure functions of their inputs and the + // BITWISE multiplicities are order-independent (they ride a permutation-invariant bus), + // so we collect every source in parallel and concatenate. The result is byte-identical + // to the previous serial `.extend()` chain regardless of order. + // + // MUL/DVRM dedup their per-unique bit-gated lookups PER CHIP INSTANCE, so pass the same + // chunk size used to split them into instances so multiplicities match the per-instance + // sends. MEMW_R sends IS_HALFWORD[timestamp_0 - old_timestamp_lo - 1]. PAGE does a + // batched ARE_BYTES[init, fini] per row (skipped in continuation epochs, which the L2G + // table owns). COMMIT sends AreBytes+IsHalfword; KECCAK_RND sends XOR/AND/ARE_BYTES/HWSL. + // Every source's bitwise lookups are collected by a pure `collect_*` function. + // In the DEFAULT ("histogram-on-the-fly") path we never concatenate them into one + // giant `Vec` (~140 M ops / ~560 MB at 10-tx whose only consumer + // is the multiplicity count). Instead each collector's transient per-source Vec is + // folded into a `BitwiseHistogram` and dropped, and the histograms are tree-reduced. + // The histogram is a commutative monoid, so the summed multiplicities are + // byte-identical to the serial per-op `update_multiplicities` count regardless of + // accumulation order (they ride a permutation-invariant bus). + // + // Under `LAMBDA_VM_LEGACY_TRACEGEN` we keep the original materialize-then-count + // path (build one `bitwise_ops` Vec, then `update_multiplicities`) for A/B and the + // byte-parity gate. + let use_histogram = !legacy_tracegen(); + // Holds the accumulated multiplicities in the default path; `None` under the legacy flag. + let mut bitwise_histogram: Option = None; + + // Shared collector list: each is a pure `Fn() -> Vec` of its source. + // The order of the list does not affect the histogram (commutative) nor the legacy + // concatenation's multiplicities (order-independent bus), matching the prior design. + type Collector<'a> = Box Vec + Sync + 'a>; + let mul_chunk = max_rows.mul; + let dvrm_chunk = max_rows.dvrm; + let mut collectors: Vec = vec![ + Box::new(|| collect_bitwise_from_lt(<_ops)), + Box::new(|| collect_bitwise_from_mul(&mul_ops, mul_chunk)), + Box::new(|| collect_bitwise_from_dvrm(&dvrm_ops, dvrm_chunk)), + Box::new(|| collect_bitwise_from_branch(&branch_ops)), + Box::new(|| shift::collect_bitwise_from_shift(&shift_ops)), + Box::new(|| bytewise_ops.iter().flat_map(|op| op.collect_bitwise_ops()).collect()), + Box::new(|| eq_ops.iter().flat_map(|op| op.collect_bitwise_ops()).collect()), + Box::new(|| store_ops.iter().flat_map(|op| op.collect_bitwise_ops()).collect()), + Box::new(|| collect_bitwise_from_memw_aligned(&memw_aligned_ops)), + Box::new(|| { + // MEMW_R IS_HALFWORD lookups from whichever carrier is populated + // (direct RegRows by default; legacy MemwOperations under the flag). + // The two produce byte-identical lookups; exactly one is non-empty. + let mut v = collect_bitwise_from_memw_register(&memw_register_ops); + v.extend(collect_bitwise_from_memw_register_direct(&memw_register_rows)); + v + }), + Box::new(|| collect_bitwise_from_commit(&commit_ops)), + Box::new(|| collect_bitwise_from_keccak(&keccak_ops)), + Box::new(|| collect_bitwise_from_ecsm(&ecsm_ops)), + Box::new(|| collect_bitwise_from_ecdas(&ecdas_ops)), + Box::new(|| collect_byte_check_ops_for_padding(num_padding_rows)), + ]; + if let Some(image) = initial_image + && !l2g_memory_bookend + { + collectors.push(Box::new(move || { + collect_bitwise_from_page(image, memory_state, l2g_memory_bookend) + })); + } + + if use_histogram { + // --- Histogram-on-the-fly (default) --- + // 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); + bitwise_ops = Vec::new(); + + #[cfg(feature = "parallel")] + { + use rayon::prelude::*; + // Each collector produces its transient Vec, which is folded into a per-WORKER + // histogram and dropped. Using `fold` (not `map`) means one 80 MiB histogram per + // worker thread reused across all collectors it handles — not one allocated+zeroed + // per collector. Tree-reduce the per-worker histograms (commutative monoid) so we + // never hold the concatenated op stream. + let reduced = collectors + .par_iter() + .fold(bitwise::BitwiseHistogram::new, |mut h, f| { + h.add_ops(&f()); + h + }) + .reduce(bitwise::BitwiseHistogram::new, |mut a, b| { + a.merge(&b); + a + }); + base.merge(&reduced); + } + #[cfg(not(feature = "parallel"))] + { + for f in &collectors { + base.add_ops(&f()); + } + } + bitwise_histogram = Some(base); + } else { + // --- Legacy materialize-then-count (LAMBDA_VM_LEGACY_TRACEGEN) --- + #[cfg(feature = "parallel")] + { + use rayon::prelude::*; + let parts: Vec> = collectors.par_iter().map(|f| f()).collect(); + for part in parts { + bitwise_ops.extend(part); + } + } + #[cfg(not(feature = "parallel"))] + { + for f in &collectors { + bitwise_ops.extend(f()); + } + } + } + drop(collectors); #[cfg(feature = "instruments")] drop(__sp); @@ -2973,13 +3395,29 @@ fn build_traces( ) }; let gen_memw_registers = || { - chunk_and_generate( - &memw_register_ops, - max_rows.memw_register, - memw_register::generate_memw_register_trace, - #[cfg(feature = "disk-spill")] - storage_mode, - ) + // Default: direct-to-column fill from compact RegRows (no MemwOperation + // materialization). Legacy (`LAMBDA_VM_LEGACY_TRACEGEN`): fill from the + // MemwOperation carrier. Dispatch on the SAME flag that chose the carrier at + // walk time — not on carrier-emptiness, which would silently pick the direct + // path for zero-register-op workloads even under the legacy flag (masking the + // A/B parity gate). Both paths chunk identically and produce byte-identical tables. + if !legacy_tracegen() { + chunk_and_generate( + &memw_register_rows, + max_rows.memw_register, + generate_memw_register_trace_direct, + #[cfg(feature = "disk-spill")] + storage_mode, + ) + } else { + chunk_and_generate( + &memw_register_ops, + max_rows.memw_register, + memw_register::generate_memw_register_trace, + #[cfg(feature = "disk-spill")] + storage_mode, + ) + } }; let gen_loads = || { chunk_and_generate( @@ -3076,7 +3514,13 @@ fn build_traces( }; let gen_bitwise = || { let mut bitwise = bitwise::generate_bitwise_trace(); - bitwise::update_multiplicities(&mut bitwise, &bitwise_ops); + // Default: fill MU columns from the accumulated histogram (no giant op Vec). + // Legacy (`LAMBDA_VM_LEGACY_TRACEGEN`): scatter-add over the materialized Vec. + // Both write the same summed multiplicities into columns 11..=20. + match &bitwise_histogram { + Some(h) => h.fill_multiplicities(&mut bitwise), + None => bitwise::update_multiplicities(&mut bitwise, &bitwise_ops), + } bitwise }; // Each CPU operation looks up the DECODE table once; padding rows look up diff --git a/prover/src/tests/bitwise_histogram_parity_tests.rs b/prover/src/tests/bitwise_histogram_parity_tests.rs new file mode 100644 index 000000000..e7bb40291 --- /dev/null +++ b/prover/src/tests/bitwise_histogram_parity_tests.rs @@ -0,0 +1,130 @@ +//! Byte-parity gate for the "histogram-on-the-fly" BITWISE multiplicity fill. +//! +//! The default trace-gen path accumulates BITWISE lookup multiplicities directly into a +//! dense `BitwiseHistogram` (per-thread, tree-reduced) instead of materializing the giant +//! `Vec` (~140 M ops at 10-tx) whose only consumer was the multiplicity +//! count. Under `LAMBDA_VM_LEGACY_TRACEGEN` the old path (materialize the Vec, then +//! `update_multiplicities`) is used. +//! +//! SOUNDNESS-CRITICAL: the two paths MUST produce a byte-identical BITWISE table — every +//! one of the 21 columns (11 preprocessed + 10 multiplicity) over all 2^20 rows. Because +//! the multiplicities are summed (a commutative monoid) the equality holds regardless of +//! accumulation order. These tests assert exact equality on real ethrex workloads. +//! +//! `#[ignore]`d (they execute + trace-build ethrex twice, which is slow); run explicitly: +//! `cargo test -p lambda-vm-prover --release bitwise_histogram_parity -- --ignored --nocapture` + +use std::path::PathBuf; + +use executor::elf::Elf; +use executor::vm::execution::Executor; + +use crate::tables::MaxRowsConfig; +use crate::tables::trace_builder::Traces; + +/// The `LAMBDA_VM_LEGACY_TRACEGEN` env var is process-global, so ALL tests that toggle it — +/// across every file — must share ONE mutex, otherwise a concurrent file's build can read the +/// flag mid-flip. Use the crate-wide lock, not a per-file one. +use crate::tests::TRACEGEN_ENV_LOCK as ENV_LOCK; + +fn workspace_root() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .parent() + .expect("workspace root") + .to_path_buf() +} + +fn build_traces(elf: &Elf, logs: &[executor::vm::logs::Log], private_input: &[u8]) -> Traces { + Traces::from_elf_and_logs( + elf, + logs, + &MaxRowsConfig::default(), + private_input, + #[cfg(feature = "disk-spill")] + stark::storage_mode::StorageMode::Ram, + ) + .expect("trace build") +} + +/// Build once with the legacy materialize-then-count path and once with the histogram +/// path, holding the env lock so the flag is stable across each build. +fn build_both( + elf: &Elf, + logs: &[executor::vm::logs::Log], + private_input: &[u8], +) -> (Traces, Traces) { + let _guard = ENV_LOCK.lock().unwrap(); + + // SAFETY: single-threaded within the lock; no other test reads the var concurrently. + unsafe { std::env::set_var("LAMBDA_VM_LEGACY_TRACEGEN", "1") }; + let legacy = build_traces(elf, logs, private_input); + + unsafe { std::env::remove_var("LAMBDA_VM_LEGACY_TRACEGEN") }; + let histogram = build_traces(elf, logs, private_input); + + (legacy, histogram) +} + +fn assert_parity(name: &str, elf: &Elf, logs: &[executor::vm::logs::Log], private_input: &[u8]) { + let (legacy, histogram) = build_both(elf, logs, private_input); + + assert_eq!( + legacy.bitwise.num_rows(), + histogram.bitwise.num_rows(), + "[{name}] BITWISE row count differs" + ); + assert_eq!( + legacy.bitwise.num_cols(), + histogram.bitwise.num_cols(), + "[{name}] BITWISE col count differs" + ); + + // Full byte-parity: all 21 columns, all 2^20 rows, row-major. + assert_eq!( + legacy.bitwise.main_table.row_major_data(), + histogram.bitwise.main_table.row_major_data(), + "[{name}] BITWISE table data differs (NOT byte-identical)" + ); + + // Guard against a vacuous pass: at least one multiplicity column must be non-zero. + use math::field::element::FieldElement; + let any_mult = histogram + .bitwise + .main_table + .row_major_data() + .iter() + .any(|fe: &FieldElement| *fe.value() != 0); + assert!(any_mult, "[{name}] BITWISE table all-zero — test is vacuous"); + + eprintln!( + "[{name}] BITWISE byte-parity OK ({} rows x {} cols)", + legacy.bitwise.num_rows(), + legacy.bitwise.num_cols() + ); +} + +fn run_ethrex(bin: &str) -> (Elf, Vec, Vec) { + let root = workspace_root(); + let elf_bytes = std::fs::read(root.join("executor/program_artifacts/rust/ethrex.elf")) + .expect("need ethrex.elf"); + let input = std::fs::read(root.join("executor/tests").join(bin)) + .unwrap_or_else(|_| panic!("need {bin}")); + let elf = Elf::load(&elf_bytes).expect("ELF load"); + let executor = Executor::new(&elf, input.clone()).expect("executor"); + let result = executor.run().expect("execute"); + (elf, result.logs, input) +} + +#[test] +#[ignore = "slow: executes + trace-builds ethrex twice"] +fn bitwise_histogram_parity_ethrex_1tx() { + let (elf, logs, input) = run_ethrex("ethrex_simple_tx.bin"); + assert_parity("ethrex-1tx", &elf, &logs, &input); +} + +#[test] +#[ignore = "slow: executes + trace-builds ethrex twice"] +fn bitwise_histogram_parity_ethrex_10tx() { + let (elf, logs, input) = run_ethrex("ethrex_10_transfers.bin"); + assert_parity("ethrex-10tx", &elf, &logs, &input); +} diff --git a/prover/src/tests/memw_register_direct_parity_tests.rs b/prover/src/tests/memw_register_direct_parity_tests.rs new file mode 100644 index 000000000..9ebeb8f81 --- /dev/null +++ b/prover/src/tests/memw_register_direct_parity_tests.rs @@ -0,0 +1,166 @@ +//! Byte-parity gate for the direct-to-column MEMW_R trace fill. +//! +//! The default trace-gen path builds the MEMW_R (register fast-path) columns DIRECTLY +//! from compact `RegRow`s, without materializing an intermediate `Vec` +//! for register accesses. Under `LAMBDA_VM_LEGACY_TRACEGEN` the old path (materialize +//! `Vec` + fill from it) is used. +//! +//! SOUNDNESS-CRITICAL: the two paths MUST produce a byte-identical MEMW_R table (every +//! column, every chunk, padding included), and MUST leave the general/aligned MEMW +//! buckets unchanged (as a multiset — those tables ride a permutation-invariant bus). +//! These tests assert exactly that on real ethrex workloads (1-tx and 10-tx). +//! +//! These are `#[ignore]`d (they execute + trace-build ethrex twice, which is slow) and +//! run explicitly, e.g. `cargo test -p lambda-vm-prover --release +//! memw_register_direct_parity -- --ignored --nocapture`. + +use std::path::PathBuf; + +use executor::elf::Elf; +use executor::vm::execution::Executor; + +use crate::tables::MaxRowsConfig; +use crate::tables::trace_builder::Traces; + +/// Shared crate-wide lock: the `LAMBDA_VM_LEGACY_TRACEGEN` env var is process-global, so this +/// MUST be the same mutex the other parity-test file uses (a per-file lock would not serialize +/// them across files under a parallel `--ignored` run). +use crate::tests::TRACEGEN_ENV_LOCK as ENV_LOCK; + +fn workspace_root() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .parent() + .expect("workspace root") + .to_path_buf() +} + +fn build_traces(elf: &Elf, logs: &[executor::vm::logs::Log], private_input: &[u8]) -> Traces { + Traces::from_elf_and_logs( + elf, + logs, + &MaxRowsConfig::default(), + private_input, + #[cfg(feature = "disk-spill")] + stark::storage_mode::StorageMode::Ram, + ) + .expect("trace build") +} + +/// Build traces once with legacy MEMW_R fill and once with the direct fill, holding the +/// env lock so the flag is stable across each build. +fn build_both(elf: &Elf, logs: &[executor::vm::logs::Log], private_input: &[u8]) -> (Traces, Traces) { + let _guard = ENV_LOCK.lock().unwrap(); + + // SAFETY: single-threaded within the lock; no other test reads the var concurrently. + unsafe { std::env::set_var("LAMBDA_VM_LEGACY_TRACEGEN", "1") }; + let legacy = build_traces(elf, logs, private_input); + + unsafe { std::env::remove_var("LAMBDA_VM_LEGACY_TRACEGEN") }; + let direct = build_traces(elf, logs, private_input); + + (legacy, direct) +} + +/// Multiset of a table's rows (row-major field-element bytes per row), for +/// permutation-invariant (bus) comparison. +fn row_multiset( + tables: &[stark::trace::TraceTable< + crate::tables::types::GoldilocksField, + crate::tables::types::GoldilocksExtension, + >], +) -> std::collections::BTreeMap, usize> { + use math::field::element::FieldElement; + let mut ms = std::collections::BTreeMap::new(); + for t in tables { + let cols = t.num_cols(); + let data = t.main_table.row_major_data(); + for row in data.chunks(cols) { + let key: Vec = row + .iter() + .map(|fe: &FieldElement| *fe.value()) + .collect(); + *ms.entry(key).or_insert(0) += 1; + } + } + ms +} + +fn assert_parity(name: &str, elf: &Elf, logs: &[executor::vm::logs::Log], private_input: &[u8]) { + let (legacy, direct) = build_both(elf, logs, private_input); + + // ---- MEMW_R: byte-identical, chunk-for-chunk, including padding ---- + assert_eq!( + legacy.memw_registers.len(), + direct.memw_registers.len(), + "[{name}] MEMW_R chunk count differs" + ); + let mut total_rows = 0usize; + for (i, (l, d)) in legacy + .memw_registers + .iter() + .zip(direct.memw_registers.iter()) + .enumerate() + { + assert_eq!( + l.num_rows(), + d.num_rows(), + "[{name}] MEMW_R chunk {i} row count differs" + ); + assert_eq!( + l.num_cols(), + d.num_cols(), + "[{name}] MEMW_R chunk {i} col count differs" + ); + assert_eq!( + l.main_table.row_major_data(), + d.main_table.row_major_data(), + "[{name}] MEMW_R chunk {i} data differs (NOT byte-identical)" + ); + total_rows += l.num_rows(); + } + assert!(total_rows > 0, "[{name}] MEMW_R has no rows — test is vacuous"); + + // ---- General + aligned MEMW buckets: unchanged as a multiset ---- + assert_eq!( + row_multiset(&legacy.memws), + row_multiset(&direct.memws), + "[{name}] general MEMW bucket multiset changed" + ); + assert_eq!( + row_multiset(&legacy.memw_aligneds), + row_multiset(&direct.memw_aligneds), + "[{name}] aligned MEMW bucket multiset changed" + ); + + eprintln!( + "[{name}] MEMW_R byte-parity OK ({} chunks, {total_rows} rows); \ + general/aligned buckets multiset-identical", + legacy.memw_registers.len() + ); +} + +fn run_ethrex(bin: &str) -> (Elf, Vec, Vec) { + let root = workspace_root(); + let elf_bytes = std::fs::read(root.join("executor/program_artifacts/rust/ethrex.elf")) + .expect("need ethrex.elf"); + let input = std::fs::read(root.join("executor/tests").join(bin)) + .unwrap_or_else(|_| panic!("need {bin}")); + let elf = Elf::load(&elf_bytes).expect("ELF load"); + let executor = Executor::new(&elf, input.clone()).expect("executor"); + let result = executor.run().expect("execute"); + (elf, result.logs, input) +} + +#[test] +#[ignore = "slow: executes + trace-builds ethrex twice"] +fn memw_register_direct_parity_ethrex_1tx() { + let (elf, logs, input) = run_ethrex("ethrex_simple_tx.bin"); + assert_parity("ethrex-1tx", &elf, &logs, &input); +} + +#[test] +#[ignore = "slow: executes + trace-builds ethrex twice"] +fn memw_register_direct_parity_ethrex_10tx() { + let (elf, logs, input) = run_ethrex("ethrex_10_transfers.bin"); + assert_parity("ethrex-10tx", &elf, &logs, &input); +} diff --git a/prover/src/tests/mod.rs b/prover/src/tests/mod.rs index 9e650422f..2a0a0b047 100644 --- a/prover/src/tests/mod.rs +++ b/prover/src/tests/mod.rs @@ -1,3 +1,9 @@ +/// Serializes tests that toggle the process-global `LAMBDA_VM_LEGACY_TRACEGEN` env var. +/// MUST be shared by every such test (across files) — a per-file mutex would not serialize +/// them, letting one file's `remove_var` race a concurrent file's legacy build. +#[cfg(test)] +pub static TRACEGEN_ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + #[cfg(all(test, feature = "disk-spill"))] pub mod auto_storage_tests; #[cfg(test)] @@ -51,6 +57,10 @@ pub mod lt_tests; #[cfg(test)] pub mod memw_aligned_tests; #[cfg(test)] +pub mod bitwise_histogram_parity_tests; +#[cfg(test)] +pub mod memw_register_direct_parity_tests; +#[cfg(test)] pub mod memw_register_tests; #[cfg(test)] pub mod memw_tests; From 69685bf40e7d0e594219b290355828593c8806c3 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Mon, 6 Jul 2026 11:52:01 -0300 Subject: [PATCH 02/17] fmt --- prover/src/tables/bitwise.rs | 5 ++- prover/src/tables/memw.rs | 10 ++++- prover/src/tables/trace_builder.rs | 37 ++++++++++++++++--- .../tests/bitwise_histogram_parity_tests.rs | 5 ++- .../memw_register_direct_parity_tests.rs | 11 +++++- prover/src/tests/mod.rs | 4 +- 6 files changed, 58 insertions(+), 14 deletions(-) diff --git a/prover/src/tables/bitwise.rs b/prover/src/tables/bitwise.rs index 388ba0bd4..00132cf67 100644 --- a/prover/src/tables/bitwise.rs +++ b/prover/src/tables/bitwise.rs @@ -495,7 +495,10 @@ mod histogram_column_map_guard { for t in ALL { seen[lookup_type_index(t)] = true; } - assert!(seen.iter().all(|&s| s), "lookup_type_index is not a bijection"); + assert!( + seen.iter().all(|&s| s), + "lookup_type_index is not a bijection" + ); } } diff --git a/prover/src/tables/memw.rs b/prover/src/tables/memw.rs index dceb5b1a7..083f2d509 100644 --- a/prover/src/tables/memw.rs +++ b/prover/src/tables/memw.rs @@ -147,7 +147,10 @@ impl MemwOperation { // future caller ever passes an out-of-domain value instead of silently // truncating it (which would be a soundness bug). value: value.map(|v| { - debug_assert!(v <= u32::MAX as u64, "MemwOperation value element exceeds u32: {v}"); + debug_assert!( + v <= u32::MAX as u64, + "MemwOperation value element exceeds u32: {v}" + ); v as u32 }), timestamp, @@ -161,7 +164,10 @@ impl MemwOperation { /// Set the old values (from memory model). pub fn with_old(mut self, old: [u64; 8], old_timestamp: [u64; 8]) -> Self { self.old = old.map(|v| { - debug_assert!(v <= u32::MAX as u64, "MemwOperation old element exceeds u32: {v}"); + debug_assert!( + v <= u32::MAX as u64, + "MemwOperation old element exceeds u32: {v}" + ); v as u32 }); self.old_timestamp = old_timestamp; diff --git a/prover/src/tables/trace_builder.rs b/prover/src/tables/trace_builder.rs index c61b64262..dee512a38 100644 --- a/prover/src/tables/trace_builder.rs +++ b/prover/src/tables/trace_builder.rs @@ -526,8 +526,16 @@ struct MemwBuckets { impl MemwBuckets { fn with_register_capacity(n: usize, legacy: bool) -> Self { Self { - register_rows: if legacy { Vec::new() } else { Vec::with_capacity(n) }, - register_ops: if legacy { Vec::with_capacity(n) } else { Vec::new() }, + register_rows: if legacy { + Vec::new() + } else { + Vec::with_capacity(n) + }, + register_ops: if legacy { + Vec::with_capacity(n) + } else { + Vec::new() + }, aligned: Vec::new(), general: Vec::new(), legacy, @@ -3250,16 +3258,33 @@ fn build_traces( Box::new(|| collect_bitwise_from_dvrm(&dvrm_ops, dvrm_chunk)), Box::new(|| collect_bitwise_from_branch(&branch_ops)), Box::new(|| shift::collect_bitwise_from_shift(&shift_ops)), - Box::new(|| bytewise_ops.iter().flat_map(|op| op.collect_bitwise_ops()).collect()), - Box::new(|| eq_ops.iter().flat_map(|op| op.collect_bitwise_ops()).collect()), - Box::new(|| store_ops.iter().flat_map(|op| op.collect_bitwise_ops()).collect()), + Box::new(|| { + bytewise_ops + .iter() + .flat_map(|op| op.collect_bitwise_ops()) + .collect() + }), + Box::new(|| { + eq_ops + .iter() + .flat_map(|op| op.collect_bitwise_ops()) + .collect() + }), + Box::new(|| { + store_ops + .iter() + .flat_map(|op| op.collect_bitwise_ops()) + .collect() + }), Box::new(|| collect_bitwise_from_memw_aligned(&memw_aligned_ops)), Box::new(|| { // MEMW_R IS_HALFWORD lookups from whichever carrier is populated // (direct RegRows by default; legacy MemwOperations under the flag). // The two produce byte-identical lookups; exactly one is non-empty. let mut v = collect_bitwise_from_memw_register(&memw_register_ops); - v.extend(collect_bitwise_from_memw_register_direct(&memw_register_rows)); + v.extend(collect_bitwise_from_memw_register_direct( + &memw_register_rows, + )); v }), Box::new(|| collect_bitwise_from_commit(&commit_ops)), diff --git a/prover/src/tests/bitwise_histogram_parity_tests.rs b/prover/src/tests/bitwise_histogram_parity_tests.rs index e7bb40291..ac07942f3 100644 --- a/prover/src/tests/bitwise_histogram_parity_tests.rs +++ b/prover/src/tests/bitwise_histogram_parity_tests.rs @@ -94,7 +94,10 @@ fn assert_parity(name: &str, elf: &Elf, logs: &[executor::vm::logs::Log], privat .row_major_data() .iter() .any(|fe: &FieldElement| *fe.value() != 0); - assert!(any_mult, "[{name}] BITWISE table all-zero — test is vacuous"); + assert!( + any_mult, + "[{name}] BITWISE table all-zero — test is vacuous" + ); eprintln!( "[{name}] BITWISE byte-parity OK ({} rows x {} cols)", diff --git a/prover/src/tests/memw_register_direct_parity_tests.rs b/prover/src/tests/memw_register_direct_parity_tests.rs index 9ebeb8f81..9b0c2b1f1 100644 --- a/prover/src/tests/memw_register_direct_parity_tests.rs +++ b/prover/src/tests/memw_register_direct_parity_tests.rs @@ -48,7 +48,11 @@ fn build_traces(elf: &Elf, logs: &[executor::vm::logs::Log], private_input: &[u8 /// Build traces once with legacy MEMW_R fill and once with the direct fill, holding the /// env lock so the flag is stable across each build. -fn build_both(elf: &Elf, logs: &[executor::vm::logs::Log], private_input: &[u8]) -> (Traces, Traces) { +fn build_both( + elf: &Elf, + logs: &[executor::vm::logs::Log], + private_input: &[u8], +) -> (Traces, Traces) { let _guard = ENV_LOCK.lock().unwrap(); // SAFETY: single-threaded within the lock; no other test reads the var concurrently. @@ -118,7 +122,10 @@ fn assert_parity(name: &str, elf: &Elf, logs: &[executor::vm::logs::Log], privat ); total_rows += l.num_rows(); } - assert!(total_rows > 0, "[{name}] MEMW_R has no rows — test is vacuous"); + assert!( + total_rows > 0, + "[{name}] MEMW_R has no rows — test is vacuous" + ); // ---- General + aligned MEMW buckets: unchanged as a multiset ---- assert_eq!( diff --git a/prover/src/tests/mod.rs b/prover/src/tests/mod.rs index 2a0a0b047..d70241c60 100644 --- a/prover/src/tests/mod.rs +++ b/prover/src/tests/mod.rs @@ -9,6 +9,8 @@ pub mod auto_storage_tests; #[cfg(test)] pub mod bitwise_bus_tests; #[cfg(test)] +pub mod bitwise_histogram_parity_tests; +#[cfg(test)] pub mod bitwise_tests; #[cfg(test)] pub mod branch_bus_tests; @@ -57,8 +59,6 @@ pub mod lt_tests; #[cfg(test)] pub mod memw_aligned_tests; #[cfg(test)] -pub mod bitwise_histogram_parity_tests; -#[cfg(test)] pub mod memw_register_direct_parity_tests; #[cfg(test)] pub mod memw_register_tests; From f56a1a0e974bd1161100cd7dd732148469cf7155 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Mon, 6 Jul 2026 12:39:51 -0300 Subject: [PATCH 03/17] rm old code --- prover/src/tables/trace_builder.rs | 239 +++++------------- .../tests/bitwise_histogram_parity_tests.rs | 133 ---------- .../memw_register_direct_parity_tests.rs | 173 ------------- prover/src/tests/mod.rs | 10 - 4 files changed, 62 insertions(+), 493 deletions(-) delete mode 100644 prover/src/tests/bitwise_histogram_parity_tests.rs delete mode 100644 prover/src/tests/memw_register_direct_parity_tests.rs diff --git a/prover/src/tables/trace_builder.rs b/prover/src/tables/trace_builder.rs index dee512a38..7a41956c0 100644 --- a/prover/src/tables/trace_builder.rs +++ b/prover/src/tables/trace_builder.rs @@ -365,18 +365,6 @@ fn collect_cpu_ops( /// Returns: (memw_ops, load_ops, lt_ops, shift_ops, bitwise_ops, commit_ops, keccak_ops, /// cpu32_ops, ecsm_ops, ec_scalar_ops, ecdas_ops) #[allow(clippy::type_complexity)] -/// Whether to use the legacy MEMW_R trace-gen path (materialize `Vec` -/// for register accesses + fill columns from it), selected via the -/// `LAMBDA_VM_LEGACY_TRACEGEN` environment variable (set to any non-empty value). -/// -/// The default (unset) is the new direct-to-column register fill. The legacy path -/// exists for A/B measurement and the byte-parity gate test. -fn legacy_tracegen() -> bool { - std::env::var("LAMBDA_VM_LEGACY_TRACEGEN") - .map(|v| !v.is_empty()) - .unwrap_or(false) -} - /// Compact, already-decomposed record for one MEMW_R (register fast-path) access. /// /// This is the "direct-to-column" carrier: it holds exactly the fields the MEMW_R @@ -507,60 +495,35 @@ fn collect_bitwise_from_memw_register_direct(rows: &[RegRow]) -> Vec`. Ops that route /// to MEMW_R are stored as compact [`RegRow`]s (`register_rows`) and later filled directly -/// into the MEMW_R columns. Under the `LAMBDA_VM_LEGACY_TRACEGEN` flag we instead keep the -/// old `Vec` (`register_ops`) and the old column fill, for A/B + byte-parity -/// testing. The `aligned` / `general` buckets are unchanged in both modes — an op that FAILS -/// `is_register_op` is always routed as a `MemwOperation` exactly as before. +/// into the MEMW_R columns. The `aligned` / `general` buckets hold `MemwOperation`s — an op +/// that FAILS `is_register_op` is routed there exactly as the old two-stage partition did. #[derive(Default)] struct MemwBuckets { - /// New path: compact register rows (empty in legacy mode). + /// Compact register rows (filled directly into the MEMW_R columns). register_rows: Vec, - /// Legacy path: full register `MemwOperation`s (empty in new mode). - register_ops: Vec, aligned: Vec, general: Vec, - /// When true, keep the old `Vec` register materialization + fill. - legacy: bool, } impl MemwBuckets { - fn with_register_capacity(n: usize, legacy: bool) -> Self { + fn with_register_capacity(n: usize) -> Self { Self { - register_rows: if legacy { - Vec::new() - } else { - Vec::with_capacity(n) - }, - register_ops: if legacy { - Vec::with_capacity(n) - } else { - Vec::new() - }, + register_rows: Vec::with_capacity(n), aligned: Vec::new(), general: Vec::new(), - legacy, } } /// Store an op already known to route to MEMW_R. #[inline] fn push_register(&mut self, op: MemwOperation) { - if self.legacy { - self.register_ops.push(op); - } else { - self.register_rows.push(RegRow::from_memw(&op)); - } + self.register_rows.push(RegRow::from_memw(&op)); } - /// Store a compact register row directly (new path only). In legacy mode this - /// reconstructs the equivalent `MemwOperation` so the two modes are comparable. + /// Store a compact register row directly. #[inline] - fn push_register_row(&mut self, row: RegRow, op_if_legacy: impl FnOnce() -> MemwOperation) { - if self.legacy { - self.register_ops.push(op_if_legacy()); - } else { - self.register_rows.push(row); - } + fn push_register_row(&mut self, row: RegRow) { + self.register_rows.push(row); } #[inline] @@ -655,7 +618,7 @@ impl MemwSink for MemwBuckets { old_ts_lo: (old_ts & 0xFFFF_FFFF) as u32, is_read, }; - self.push_register_row(row, build_fallback); + self.push_register_row(row); } else { let op = build_fallback(); debug_assert!(!is_register_op(&op), "reg fallback must not be MEMW_R"); @@ -685,7 +648,7 @@ fn collect_ops_from_cpu( Vec, Vec, ) { - let mut memw = MemwBuckets::with_register_capacity(cpu_ops.len() * 3, legacy_tracegen()); + let mut memw = MemwBuckets::with_register_capacity(cpu_ops.len() * 3); let mut load_ops = Vec::with_capacity(cpu_ops.len() / 8 + 1); let mut lt_ops = Vec::with_capacity(cpu_ops.len() / 10 + 1); let mut shift_ops = Vec::with_capacity(cpu_ops.len() / 10 + 1); @@ -1735,21 +1698,6 @@ fn reg_ts_delta_in_range(timestamp: u64, old_ts: u64) -> bool { ts_hi == old_ts_hi && ts_lo > old_ts_lo && (ts_lo - old_ts_lo) <= 0x10000 } -/// Collects IS_HALFWORD bitwise lookups for MEMW_R operations. -/// -/// For each register op: checks that `timestamp[0] - old_timestamp_lo - 1` fits -/// in a halfword (proving the timestamp delta is in range [1, 2^16]). -fn collect_bitwise_from_memw_register(ops: &[MemwOperation]) -> Vec { - ops.iter() - .map(|op| { - memw_register_is_half_lookup( - (op.timestamp & 0xFFFF_FFFF) as u32, - (op.old_timestamp[0] & 0xFFFF_FFFF) as u32, - ) - }) - .collect() -} - // ============================================================================= // Phase 4: All → Bitwise lookups // ============================================================================= @@ -2925,9 +2873,7 @@ struct CollectedOps { cpu_ops: Vec, memw_ops: Vec, memw_aligned_ops: Vec, - /// Legacy MEMW_R ops (non-empty only under `LAMBDA_VM_LEGACY_TRACEGEN`). - memw_register_ops: Vec, - /// New direct-fill MEMW_R rows (non-empty in the default path). + /// Direct-fill MEMW_R rows (register fast path). memw_register_rows: Vec, load_ops: Vec, lt_ops: Vec, @@ -3021,10 +2967,8 @@ fn collect_all_ops( // cost) is gone. Order within each bucket matches the old stable partitions → byte-identical. let MemwBuckets { register_rows: memw_register_rows, - register_ops: memw_register_ops, aligned: memw_aligned_ops, general: memw_ops, - legacy: _, } = memw; // Collect BRANCH operations from CPU ops where branch_cond = true @@ -3130,7 +3074,6 @@ fn collect_all_ops( cpu_ops, memw_ops, memw_aligned_ops, - memw_register_ops, memw_register_rows, load_ops, lt_ops, @@ -3175,12 +3118,11 @@ fn build_traces( cpu_ops, memw_ops, memw_aligned_ops, - memw_register_ops, memw_register_rows, load_ops, mut lt_ops, shift_ops, - mut bitwise_ops, + bitwise_ops, branch_ops, mul_ops, dvrm_ops, @@ -3230,25 +3172,16 @@ fn build_traces( // sends. MEMW_R sends IS_HALFWORD[timestamp_0 - old_timestamp_lo - 1]. PAGE does a // batched ARE_BYTES[init, fini] per row (skipped in continuation epochs, which the L2G // table owns). COMMIT sends AreBytes+IsHalfword; KECCAK_RND sends XOR/AND/ARE_BYTES/HWSL. - // Every source's bitwise lookups are collected by a pure `collect_*` function. - // In the DEFAULT ("histogram-on-the-fly") path we never concatenate them into one - // giant `Vec` (~140 M ops / ~560 MB at 10-tx whose only consumer - // is the multiplicity count). Instead each collector's transient per-source Vec is - // folded into a `BitwiseHistogram` and dropped, and the histograms are tree-reduced. - // The histogram is a commutative monoid, so the summed multiplicities are - // byte-identical to the serial per-op `update_multiplicities` count regardless of - // accumulation order (they ride a permutation-invariant bus). - // - // Under `LAMBDA_VM_LEGACY_TRACEGEN` we keep the original materialize-then-count - // path (build one `bitwise_ops` Vec, then `update_multiplicities`) for A/B and the - // byte-parity gate. - let use_histogram = !legacy_tracegen(); - // Holds the accumulated multiplicities in the default path; `None` under the legacy flag. - let mut bitwise_histogram: Option = None; + // Every source's bitwise lookups are collected by a pure `collect_*` function. We never + // concatenate them into one giant `Vec` (~140 M ops / ~560 MB at 10-tx + // whose only consumer is the multiplicity count). Instead each collector's transient + // per-source Vec is folded into a `BitwiseHistogram` and dropped, and the per-worker + // histograms are tree-reduced. The histogram is a commutative monoid, so the summed + // multiplicities are independent of accumulation order (the lookups ride a + // permutation-invariant bus). // Shared collector list: each is a pure `Fn() -> Vec` of its source. - // The order of the list does not affect the histogram (commutative) nor the legacy - // concatenation's multiplicities (order-independent bus), matching the prior design. + // Order does not affect the histogram (commutative). type Collector<'a> = Box Vec + Sync + 'a>; let mul_chunk = max_rows.mul; let dvrm_chunk = max_rows.dvrm; @@ -3277,16 +3210,7 @@ fn build_traces( .collect() }), Box::new(|| collect_bitwise_from_memw_aligned(&memw_aligned_ops)), - Box::new(|| { - // MEMW_R IS_HALFWORD lookups from whichever carrier is populated - // (direct RegRows by default; legacy MemwOperations under the flag). - // The two produce byte-identical lookups; exactly one is non-empty. - let mut v = collect_bitwise_from_memw_register(&memw_register_ops); - v.extend(collect_bitwise_from_memw_register_direct( - &memw_register_rows, - )); - v - }), + Box::new(|| collect_bitwise_from_memw_register_direct(&memw_register_rows)), Box::new(|| collect_bitwise_from_commit(&commit_ops)), Box::new(|| collect_bitwise_from_keccak(&keccak_ops)), Box::new(|| collect_bitwise_from_ecsm(&ecsm_ops)), @@ -3301,58 +3225,38 @@ fn build_traces( })); } - if use_histogram { - // --- Histogram-on-the-fly (default) --- - // 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); - bitwise_ops = Vec::new(); + // 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::*; - // Each collector produces its transient Vec, which is folded into a per-WORKER - // histogram and dropped. Using `fold` (not `map`) means one 80 MiB histogram per - // worker thread reused across all collectors it handles — not one allocated+zeroed - // per collector. Tree-reduce the per-worker histograms (commutative monoid) so we - // never hold the concatenated op stream. - let reduced = collectors - .par_iter() - .fold(bitwise::BitwiseHistogram::new, |mut h, f| { - h.add_ops(&f()); - h - }) - .reduce(bitwise::BitwiseHistogram::new, |mut a, b| { - a.merge(&b); - a - }); - base.merge(&reduced); - } - #[cfg(not(feature = "parallel"))] - { - for f in &collectors { - base.add_ops(&f()); - } - } - bitwise_histogram = Some(base); - } else { - // --- Legacy materialize-then-count (LAMBDA_VM_LEGACY_TRACEGEN) --- - #[cfg(feature = "parallel")] - { - use rayon::prelude::*; - let parts: Vec> = collectors.par_iter().map(|f| f()).collect(); - for part in parts { - bitwise_ops.extend(part); - } - } - #[cfg(not(feature = "parallel"))] - { - for f in &collectors { - bitwise_ops.extend(f()); - } + #[cfg(feature = "parallel")] + { + use rayon::prelude::*; + // Each collector produces its transient Vec, which is folded into a per-worker + // histogram and dropped. `fold` (not `map`) reuses one histogram per work-split + // instead of allocating one per collector. Tree-reduce them (commutative monoid) so + // we never hold the concatenated op stream. + let reduced = collectors + .par_iter() + .fold(bitwise::BitwiseHistogram::new, |mut h, f| { + h.add_ops(&f()); + h + }) + .reduce(bitwise::BitwiseHistogram::new, |mut a, b| { + a.merge(&b); + a + }); + base.merge(&reduced); + } + #[cfg(not(feature = "parallel"))] + { + for f in &collectors { + base.add_ops(&f()); } } + let bitwise_histogram = base; drop(collectors); #[cfg(feature = "instruments")] drop(__sp); @@ -3420,29 +3324,15 @@ fn build_traces( ) }; let gen_memw_registers = || { - // Default: direct-to-column fill from compact RegRows (no MemwOperation - // materialization). Legacy (`LAMBDA_VM_LEGACY_TRACEGEN`): fill from the - // MemwOperation carrier. Dispatch on the SAME flag that chose the carrier at - // walk time — not on carrier-emptiness, which would silently pick the direct - // path for zero-register-op workloads even under the legacy flag (masking the - // A/B parity gate). Both paths chunk identically and produce byte-identical tables. - if !legacy_tracegen() { - chunk_and_generate( - &memw_register_rows, - max_rows.memw_register, - generate_memw_register_trace_direct, - #[cfg(feature = "disk-spill")] - storage_mode, - ) - } else { - chunk_and_generate( - &memw_register_ops, - max_rows.memw_register, - memw_register::generate_memw_register_trace, - #[cfg(feature = "disk-spill")] - storage_mode, - ) - } + // Direct-to-column fill from compact RegRows — the register fast path never + // materializes a `Vec`. + chunk_and_generate( + &memw_register_rows, + max_rows.memw_register, + generate_memw_register_trace_direct, + #[cfg(feature = "disk-spill")] + storage_mode, + ) }; let gen_loads = || { chunk_and_generate( @@ -3539,13 +3429,8 @@ fn build_traces( }; let gen_bitwise = || { let mut bitwise = bitwise::generate_bitwise_trace(); - // Default: fill MU columns from the accumulated histogram (no giant op Vec). - // Legacy (`LAMBDA_VM_LEGACY_TRACEGEN`): scatter-add over the materialized Vec. - // Both write the same summed multiplicities into columns 11..=20. - match &bitwise_histogram { - Some(h) => h.fill_multiplicities(&mut bitwise), - None => bitwise::update_multiplicities(&mut bitwise, &bitwise_ops), - } + // Fill the MU columns (11..=20) from the accumulated histogram. + bitwise_histogram.fill_multiplicities(&mut bitwise); bitwise }; // Each CPU operation looks up the DECODE table once; padding rows look up diff --git a/prover/src/tests/bitwise_histogram_parity_tests.rs b/prover/src/tests/bitwise_histogram_parity_tests.rs deleted file mode 100644 index ac07942f3..000000000 --- a/prover/src/tests/bitwise_histogram_parity_tests.rs +++ /dev/null @@ -1,133 +0,0 @@ -//! Byte-parity gate for the "histogram-on-the-fly" BITWISE multiplicity fill. -//! -//! The default trace-gen path accumulates BITWISE lookup multiplicities directly into a -//! dense `BitwiseHistogram` (per-thread, tree-reduced) instead of materializing the giant -//! `Vec` (~140 M ops at 10-tx) whose only consumer was the multiplicity -//! count. Under `LAMBDA_VM_LEGACY_TRACEGEN` the old path (materialize the Vec, then -//! `update_multiplicities`) is used. -//! -//! SOUNDNESS-CRITICAL: the two paths MUST produce a byte-identical BITWISE table — every -//! one of the 21 columns (11 preprocessed + 10 multiplicity) over all 2^20 rows. Because -//! the multiplicities are summed (a commutative monoid) the equality holds regardless of -//! accumulation order. These tests assert exact equality on real ethrex workloads. -//! -//! `#[ignore]`d (they execute + trace-build ethrex twice, which is slow); run explicitly: -//! `cargo test -p lambda-vm-prover --release bitwise_histogram_parity -- --ignored --nocapture` - -use std::path::PathBuf; - -use executor::elf::Elf; -use executor::vm::execution::Executor; - -use crate::tables::MaxRowsConfig; -use crate::tables::trace_builder::Traces; - -/// The `LAMBDA_VM_LEGACY_TRACEGEN` env var is process-global, so ALL tests that toggle it — -/// across every file — must share ONE mutex, otherwise a concurrent file's build can read the -/// flag mid-flip. Use the crate-wide lock, not a per-file one. -use crate::tests::TRACEGEN_ENV_LOCK as ENV_LOCK; - -fn workspace_root() -> PathBuf { - PathBuf::from(env!("CARGO_MANIFEST_DIR")) - .parent() - .expect("workspace root") - .to_path_buf() -} - -fn build_traces(elf: &Elf, logs: &[executor::vm::logs::Log], private_input: &[u8]) -> Traces { - Traces::from_elf_and_logs( - elf, - logs, - &MaxRowsConfig::default(), - private_input, - #[cfg(feature = "disk-spill")] - stark::storage_mode::StorageMode::Ram, - ) - .expect("trace build") -} - -/// Build once with the legacy materialize-then-count path and once with the histogram -/// path, holding the env lock so the flag is stable across each build. -fn build_both( - elf: &Elf, - logs: &[executor::vm::logs::Log], - private_input: &[u8], -) -> (Traces, Traces) { - let _guard = ENV_LOCK.lock().unwrap(); - - // SAFETY: single-threaded within the lock; no other test reads the var concurrently. - unsafe { std::env::set_var("LAMBDA_VM_LEGACY_TRACEGEN", "1") }; - let legacy = build_traces(elf, logs, private_input); - - unsafe { std::env::remove_var("LAMBDA_VM_LEGACY_TRACEGEN") }; - let histogram = build_traces(elf, logs, private_input); - - (legacy, histogram) -} - -fn assert_parity(name: &str, elf: &Elf, logs: &[executor::vm::logs::Log], private_input: &[u8]) { - let (legacy, histogram) = build_both(elf, logs, private_input); - - assert_eq!( - legacy.bitwise.num_rows(), - histogram.bitwise.num_rows(), - "[{name}] BITWISE row count differs" - ); - assert_eq!( - legacy.bitwise.num_cols(), - histogram.bitwise.num_cols(), - "[{name}] BITWISE col count differs" - ); - - // Full byte-parity: all 21 columns, all 2^20 rows, row-major. - assert_eq!( - legacy.bitwise.main_table.row_major_data(), - histogram.bitwise.main_table.row_major_data(), - "[{name}] BITWISE table data differs (NOT byte-identical)" - ); - - // Guard against a vacuous pass: at least one multiplicity column must be non-zero. - use math::field::element::FieldElement; - let any_mult = histogram - .bitwise - .main_table - .row_major_data() - .iter() - .any(|fe: &FieldElement| *fe.value() != 0); - assert!( - any_mult, - "[{name}] BITWISE table all-zero — test is vacuous" - ); - - eprintln!( - "[{name}] BITWISE byte-parity OK ({} rows x {} cols)", - legacy.bitwise.num_rows(), - legacy.bitwise.num_cols() - ); -} - -fn run_ethrex(bin: &str) -> (Elf, Vec, Vec) { - let root = workspace_root(); - let elf_bytes = std::fs::read(root.join("executor/program_artifacts/rust/ethrex.elf")) - .expect("need ethrex.elf"); - let input = std::fs::read(root.join("executor/tests").join(bin)) - .unwrap_or_else(|_| panic!("need {bin}")); - let elf = Elf::load(&elf_bytes).expect("ELF load"); - let executor = Executor::new(&elf, input.clone()).expect("executor"); - let result = executor.run().expect("execute"); - (elf, result.logs, input) -} - -#[test] -#[ignore = "slow: executes + trace-builds ethrex twice"] -fn bitwise_histogram_parity_ethrex_1tx() { - let (elf, logs, input) = run_ethrex("ethrex_simple_tx.bin"); - assert_parity("ethrex-1tx", &elf, &logs, &input); -} - -#[test] -#[ignore = "slow: executes + trace-builds ethrex twice"] -fn bitwise_histogram_parity_ethrex_10tx() { - let (elf, logs, input) = run_ethrex("ethrex_10_transfers.bin"); - assert_parity("ethrex-10tx", &elf, &logs, &input); -} diff --git a/prover/src/tests/memw_register_direct_parity_tests.rs b/prover/src/tests/memw_register_direct_parity_tests.rs deleted file mode 100644 index 9b0c2b1f1..000000000 --- a/prover/src/tests/memw_register_direct_parity_tests.rs +++ /dev/null @@ -1,173 +0,0 @@ -//! Byte-parity gate for the direct-to-column MEMW_R trace fill. -//! -//! The default trace-gen path builds the MEMW_R (register fast-path) columns DIRECTLY -//! from compact `RegRow`s, without materializing an intermediate `Vec` -//! for register accesses. Under `LAMBDA_VM_LEGACY_TRACEGEN` the old path (materialize -//! `Vec` + fill from it) is used. -//! -//! SOUNDNESS-CRITICAL: the two paths MUST produce a byte-identical MEMW_R table (every -//! column, every chunk, padding included), and MUST leave the general/aligned MEMW -//! buckets unchanged (as a multiset — those tables ride a permutation-invariant bus). -//! These tests assert exactly that on real ethrex workloads (1-tx and 10-tx). -//! -//! These are `#[ignore]`d (they execute + trace-build ethrex twice, which is slow) and -//! run explicitly, e.g. `cargo test -p lambda-vm-prover --release -//! memw_register_direct_parity -- --ignored --nocapture`. - -use std::path::PathBuf; - -use executor::elf::Elf; -use executor::vm::execution::Executor; - -use crate::tables::MaxRowsConfig; -use crate::tables::trace_builder::Traces; - -/// Shared crate-wide lock: the `LAMBDA_VM_LEGACY_TRACEGEN` env var is process-global, so this -/// MUST be the same mutex the other parity-test file uses (a per-file lock would not serialize -/// them across files under a parallel `--ignored` run). -use crate::tests::TRACEGEN_ENV_LOCK as ENV_LOCK; - -fn workspace_root() -> PathBuf { - PathBuf::from(env!("CARGO_MANIFEST_DIR")) - .parent() - .expect("workspace root") - .to_path_buf() -} - -fn build_traces(elf: &Elf, logs: &[executor::vm::logs::Log], private_input: &[u8]) -> Traces { - Traces::from_elf_and_logs( - elf, - logs, - &MaxRowsConfig::default(), - private_input, - #[cfg(feature = "disk-spill")] - stark::storage_mode::StorageMode::Ram, - ) - .expect("trace build") -} - -/// Build traces once with legacy MEMW_R fill and once with the direct fill, holding the -/// env lock so the flag is stable across each build. -fn build_both( - elf: &Elf, - logs: &[executor::vm::logs::Log], - private_input: &[u8], -) -> (Traces, Traces) { - let _guard = ENV_LOCK.lock().unwrap(); - - // SAFETY: single-threaded within the lock; no other test reads the var concurrently. - unsafe { std::env::set_var("LAMBDA_VM_LEGACY_TRACEGEN", "1") }; - let legacy = build_traces(elf, logs, private_input); - - unsafe { std::env::remove_var("LAMBDA_VM_LEGACY_TRACEGEN") }; - let direct = build_traces(elf, logs, private_input); - - (legacy, direct) -} - -/// Multiset of a table's rows (row-major field-element bytes per row), for -/// permutation-invariant (bus) comparison. -fn row_multiset( - tables: &[stark::trace::TraceTable< - crate::tables::types::GoldilocksField, - crate::tables::types::GoldilocksExtension, - >], -) -> std::collections::BTreeMap, usize> { - use math::field::element::FieldElement; - let mut ms = std::collections::BTreeMap::new(); - for t in tables { - let cols = t.num_cols(); - let data = t.main_table.row_major_data(); - for row in data.chunks(cols) { - let key: Vec = row - .iter() - .map(|fe: &FieldElement| *fe.value()) - .collect(); - *ms.entry(key).or_insert(0) += 1; - } - } - ms -} - -fn assert_parity(name: &str, elf: &Elf, logs: &[executor::vm::logs::Log], private_input: &[u8]) { - let (legacy, direct) = build_both(elf, logs, private_input); - - // ---- MEMW_R: byte-identical, chunk-for-chunk, including padding ---- - assert_eq!( - legacy.memw_registers.len(), - direct.memw_registers.len(), - "[{name}] MEMW_R chunk count differs" - ); - let mut total_rows = 0usize; - for (i, (l, d)) in legacy - .memw_registers - .iter() - .zip(direct.memw_registers.iter()) - .enumerate() - { - assert_eq!( - l.num_rows(), - d.num_rows(), - "[{name}] MEMW_R chunk {i} row count differs" - ); - assert_eq!( - l.num_cols(), - d.num_cols(), - "[{name}] MEMW_R chunk {i} col count differs" - ); - assert_eq!( - l.main_table.row_major_data(), - d.main_table.row_major_data(), - "[{name}] MEMW_R chunk {i} data differs (NOT byte-identical)" - ); - total_rows += l.num_rows(); - } - assert!( - total_rows > 0, - "[{name}] MEMW_R has no rows — test is vacuous" - ); - - // ---- General + aligned MEMW buckets: unchanged as a multiset ---- - assert_eq!( - row_multiset(&legacy.memws), - row_multiset(&direct.memws), - "[{name}] general MEMW bucket multiset changed" - ); - assert_eq!( - row_multiset(&legacy.memw_aligneds), - row_multiset(&direct.memw_aligneds), - "[{name}] aligned MEMW bucket multiset changed" - ); - - eprintln!( - "[{name}] MEMW_R byte-parity OK ({} chunks, {total_rows} rows); \ - general/aligned buckets multiset-identical", - legacy.memw_registers.len() - ); -} - -fn run_ethrex(bin: &str) -> (Elf, Vec, Vec) { - let root = workspace_root(); - let elf_bytes = std::fs::read(root.join("executor/program_artifacts/rust/ethrex.elf")) - .expect("need ethrex.elf"); - let input = std::fs::read(root.join("executor/tests").join(bin)) - .unwrap_or_else(|_| panic!("need {bin}")); - let elf = Elf::load(&elf_bytes).expect("ELF load"); - let executor = Executor::new(&elf, input.clone()).expect("executor"); - let result = executor.run().expect("execute"); - (elf, result.logs, input) -} - -#[test] -#[ignore = "slow: executes + trace-builds ethrex twice"] -fn memw_register_direct_parity_ethrex_1tx() { - let (elf, logs, input) = run_ethrex("ethrex_simple_tx.bin"); - assert_parity("ethrex-1tx", &elf, &logs, &input); -} - -#[test] -#[ignore = "slow: executes + trace-builds ethrex twice"] -fn memw_register_direct_parity_ethrex_10tx() { - let (elf, logs, input) = run_ethrex("ethrex_10_transfers.bin"); - assert_parity("ethrex-10tx", &elf, &logs, &input); -} diff --git a/prover/src/tests/mod.rs b/prover/src/tests/mod.rs index d70241c60..9e650422f 100644 --- a/prover/src/tests/mod.rs +++ b/prover/src/tests/mod.rs @@ -1,16 +1,8 @@ -/// Serializes tests that toggle the process-global `LAMBDA_VM_LEGACY_TRACEGEN` env var. -/// MUST be shared by every such test (across files) — a per-file mutex would not serialize -/// them, letting one file's `remove_var` race a concurrent file's legacy build. -#[cfg(test)] -pub static TRACEGEN_ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); - #[cfg(all(test, feature = "disk-spill"))] pub mod auto_storage_tests; #[cfg(test)] pub mod bitwise_bus_tests; #[cfg(test)] -pub mod bitwise_histogram_parity_tests; -#[cfg(test)] pub mod bitwise_tests; #[cfg(test)] pub mod branch_bus_tests; @@ -59,8 +51,6 @@ pub mod lt_tests; #[cfg(test)] pub mod memw_aligned_tests; #[cfg(test)] -pub mod memw_register_direct_parity_tests; -#[cfg(test)] pub mod memw_register_tests; #[cfg(test)] pub mod memw_tests; From 1b66905c43bb5b2a843f3b287ae3413cd136dde1 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Mon, 6 Jul 2026 16:43:01 -0300 Subject: [PATCH 04/17] fmt --- prover/src/tables/trace_builder.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/prover/src/tables/trace_builder.rs b/prover/src/tables/trace_builder.rs index 7b6f17b14..93b9d60b3 100644 --- a/prover/src/tables/trace_builder.rs +++ b/prover/src/tables/trace_builder.rs @@ -631,6 +631,7 @@ impl MemwSink for MemwBuckets { } } +#[allow(clippy::type_complexity)] fn collect_ops_from_cpu( cpu_ops: &[CpuOperation], memory_state: &mut MemoryState, From 3b6840e7341a64c623aa37bfcbb37ff5ac9ec255 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Mon, 6 Jul 2026 17:33:41 -0300 Subject: [PATCH 05/17] doc cleanup --- prover/src/tables/bitwise.rs | 9 +++------ prover/src/tables/memw.rs | 11 ++++++----- prover/src/tables/trace_builder.rs | 30 +++++++++++++++--------------- 3 files changed, 24 insertions(+), 26 deletions(-) diff --git a/prover/src/tables/bitwise.rs b/prover/src/tables/bitwise.rs index 00132cf67..875e5e37c 100644 --- a/prover/src/tables/bitwise.rs +++ b/prover/src/tables/bitwise.rs @@ -520,6 +520,9 @@ pub struct BitwiseHistogram { impl BitwiseHistogram { /// Allocate a zeroed histogram (80 MiB). + // No `Default` impl on purpose: `new()` allocates 80 MiB, so a stray + // `..Default::default()` / `#[derive(Default)]` must not silently do that. + #[allow(clippy::new_without_default)] pub fn new() -> Self { Self { counters: vec![0u64; NUM_ROWS * NUM_LOOKUP_TYPES].into_boxed_slice(), @@ -577,12 +580,6 @@ impl BitwiseHistogram { } } -impl Default for BitwiseHistogram { - fn default() -> Self { - Self::new() - } -} - /// Types of lookups the BITWISE table provides. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum BitwiseOperationType { diff --git a/prover/src/tables/memw.rs b/prover/src/tables/memw.rs index fa4449101..f0f8f2df7 100644 --- a/prover/src/tables/memw.rs +++ b/prover/src/tables/memw.rs @@ -142,11 +142,12 @@ impl MemwOperation { // Callers build a [u64; 8] transiently on the stack; we store the u32 // domain (byte / 32-bit register half) so the persisted struct is half // the size. Values never exceed u32 (every element is a byte or a - // pack_register_value 32-bit limb). The debug_assert fails loudly if a - // future caller ever passes an out-of-domain value instead of silently - // truncating it (which would be a soundness bug). + // pack_register_value 32-bit limb). This is an always-on `assert!` (not + // `debug_assert!`): a caller passing an out-of-domain value would + // otherwise silently truncate here and produce an unsound witness in + // release builds, so the invariant is enforced in every build profile. value: value.map(|v| { - debug_assert!( + assert!( v <= u32::MAX as u64, "MemwOperation value element exceeds u32: {v}" ); @@ -163,7 +164,7 @@ impl MemwOperation { /// Set the old values (from memory model). pub fn with_old(mut self, old: [u64; 8], old_timestamp: [u64; 8]) -> Self { self.old = old.map(|v| { - debug_assert!( + assert!( v <= u32::MAX as u64, "MemwOperation old element exceeds u32: {v}" ); diff --git a/prover/src/tables/trace_builder.rs b/prover/src/tables/trace_builder.rs index 93b9d60b3..c08da6bde 100644 --- a/prover/src/tables/trace_builder.rs +++ b/prover/src/tables/trace_builder.rs @@ -352,19 +352,6 @@ fn collect_cpu_ops( // Phase 2: CPU ops → MEMW, LOAD, LT, Bitwise // ============================================================================= -/// Collects all derived operations from CPU operations in a single pass. -/// -/// This includes: -/// - MEMW ops (register reads/writes M1/M3/M5, memory loads/stores M6/M7) -/// - LOAD ops (memory loads with sign/zero extension) -/// - LT ops (from SLT/BLT instructions) -/// - Bitwise lookups (from CPU operations) -/// -/// MEMW and LOAD collection requires sequential processing with state tracking. -/// -/// Returns: (memw_ops, load_ops, lt_ops, shift_ops, bitwise_ops, commit_ops, keccak_ops, -/// cpu32_ops, ecsm_ops, ec_scalar_ops, ecdas_ops) -#[allow(clippy::type_complexity)] /// Compact, already-decomposed record for one MEMW_R (register fast-path) access. /// /// This is the "direct-to-column" carrier: it holds exactly the fields the MEMW_R @@ -477,8 +464,9 @@ fn memw_register_is_half_lookup(ts_lo: u32, old_ts_lo: u32) -> BitwiseOperation ) } -/// IS_HALFWORD bitwise lookups for MEMW_R, computed directly from [`RegRow`]s. -/// Byte-identical to `collect_bitwise_from_memw_register` (both call the shared helper). +/// IS_HALFWORD bitwise lookups for MEMW_R, computed directly from [`RegRow`]s via +/// the shared [`memw_register_is_half_lookup`] helper (the same lookup the MEMW_R +/// trace fill uses), so the multiplicities stay consistent with that table. fn collect_bitwise_from_memw_register_direct(rows: &[RegRow]) -> Vec { rows.iter() .map(|r| memw_register_is_half_lookup((r.timestamp & 0xFFFF_FFFF) as u32, r.old_ts_lo)) @@ -631,6 +619,18 @@ impl MemwSink for MemwBuckets { } } +/// Collects all derived operations from CPU operations in a single pass. +/// +/// This includes: +/// - MEMW ops (register reads/writes M1/M3/M5, memory loads/stores M6/M7) +/// - LOAD ops (memory loads with sign/zero extension) +/// - LT ops (from SLT/BLT instructions) +/// - Bitwise lookups (from CPU operations) +/// +/// MEMW and LOAD collection requires sequential processing with state tracking. +/// +/// Returns: (memw_buckets, load_ops, lt_ops, shift_ops, bitwise_ops, commit_ops, +/// keccak_ops, cpu32_ops, ecsm_ops, ec_scalar_ops, ecdas_ops) #[allow(clippy::type_complexity)] fn collect_ops_from_cpu( cpu_ops: &[CpuOperation], From 3ba4aac74848a2ca39d44700cfc2690f7fb38850 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Mon, 6 Jul 2026 18:15:31 -0300 Subject: [PATCH 06/17] fix --- prover/src/tables/trace_builder.rs | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/prover/src/tables/trace_builder.rs b/prover/src/tables/trace_builder.rs index c08da6bde..eb5d19cf3 100644 --- a/prover/src/tables/trace_builder.rs +++ b/prover/src/tables/trace_builder.rs @@ -3228,14 +3228,23 @@ fn build_traces( #[cfg(feature = "parallel")] { use rayon::prelude::*; - // Each collector produces its transient Vec, which is folded into a per-worker - // histogram and dropped. `fold` (not `map`) reuses one histogram per work-split - // instead of allocating one per collector. Tree-reduce them (commutative monoid) so - // we never hold the concatenated op stream. + // Each collector produces a transient Vec of lookups that is folded into a + // dense 80 MiB histogram and dropped. 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). Tree-reduce the chunk histograms — + // add_ops/merge form a commutative monoid, so the result is order-independent + // and byte-identical to the sequential path. + let max_par = rayon::current_num_threads().clamp(1, 8); + let chunk_size = collectors.len().div_ceil(max_par).max(1); let reduced = collectors - .par_iter() - .fold(bitwise::BitwiseHistogram::new, |mut h, f| { - h.add_ops(&f()); + .par_chunks(chunk_size) + .map(|chunk| { + let mut h = bitwise::BitwiseHistogram::new(); + for f in chunk { + h.add_ops(&f()); + } h }) .reduce(bitwise::BitwiseHistogram::new, |mut a, b| { From ae7d5d465106da3fa6c0adf5bebb5198a89eaaae Mon Sep 17 00:00:00 2001 From: MauroFab Date: Mon, 6 Jul 2026 19:02:50 -0300 Subject: [PATCH 07/17] Fix stale trace-gen docs and drop dead BitwiseHistogram::total - 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 --- prover/src/tables/bitwise.rs | 19 ++++++++++--------- prover/src/tables/trace_builder.rs | 16 ++++++++++------ 2 files changed, 20 insertions(+), 15 deletions(-) diff --git a/prover/src/tables/bitwise.rs b/prover/src/tables/bitwise.rs index 875e5e37c..eed78f81a 100644 --- a/prover/src/tables/bitwise.rs +++ b/prover/src/tables/bitwise.rs @@ -442,7 +442,9 @@ pub const fn lookup_type_index(t: BitwiseOperationType) -> usize { } } -/// Multiplicity column for a lookup type (used by the legacy per-op path). +/// Multiplicity column for a lookup type. Used by the per-op path +/// ([`update_multiplicities`]), which is still live production code: continuation +/// epochs add their L2G lookups through it on top of the histogram-filled trace. #[inline] pub const fn mu_column(t: BitwiseOperationType) -> usize { match t { @@ -553,16 +555,15 @@ impl BitwiseHistogram { } } - /// Total number of lookups counted (for instrumentation only). - pub fn total(&self) -> u64 { - self.counters.iter().sum() - } - /// Write the accumulated multiplicities into the BITWISE trace's MU columns. /// - /// Produces exactly the same MU columns as calling [`update_multiplicities`] - /// with the full op vector, because both sum one increment per lookup into - /// the `(row, mu_col)` cell. + /// OVERWRITES each nonzero cell with its count (it does not add to what is + /// there), so it assumes the MU columns are still zero — true for a fresh + /// [`generate_bitwise_trace`] output, where it produces exactly the same MU + /// columns as calling [`update_multiplicities`] with the full op vector. + /// Callers that layer additional lookups on top (continuation epochs add + /// their L2G lookups via `update_multiplicities`, which increments) must do + /// so strictly AFTER this fill, never before. pub fn fill_multiplicities( &self, trace: &mut TraceTable, diff --git a/prover/src/tables/trace_builder.rs b/prover/src/tables/trace_builder.rs index eb5d19cf3..7bdff290f 100644 --- a/prover/src/tables/trace_builder.rs +++ b/prover/src/tables/trace_builder.rs @@ -357,9 +357,10 @@ fn collect_cpu_ops( /// This is the "direct-to-column" carrier: it holds exactly the fields the MEMW_R /// column fill (`generate_memw_register_trace_direct`) and its IS_HALFWORD bitwise /// collector (`collect_bitwise_from_memw_register_direct`) need, and nothing else. -/// It replaces the full `MemwOperation` (216→~152 B via the E6 shrink, but still 8 -/// `[u32;8]`/`[u64;8]` arrays) for register accesses — the largest table by rows — -/// so the walk never materializes a `MemwOperation` for the register fast path. +/// It replaces the full `MemwOperation` (~152 B after the `[u32; 8]` value/old +/// shrink, but still 8-element arrays) for register accesses — the largest table +/// by rows — so the walk never materializes a `MemwOperation` for the register +/// fast path. /// /// Field domains mirror `MemwOperation`'s so the produced table is byte-identical: /// - `address` = `base_address / 2` (the register index 0..=255; ADDRESS column, @@ -448,8 +449,10 @@ fn generate_memw_register_trace_direct( /// The single IS_HALFWORD lookup a MEMW_R access sends: proves the timestamp delta /// `ts_lo - old_ts_lo` is in [1, 2^16] by decomposing `ts_lo - old_ts_lo - 1` into two bytes. -/// Shared by BOTH the direct (`RegRow`) and `MemwOperation` collectors so they can never drift -/// (a divergence would be a silent soundness bug — see the E7 review). +/// +/// Must stay in lockstep with the IS_HALFWORD send in +/// `memw_register::bus_interactions()`: the lookup counted here has to be exactly +/// the lookup each MEMW_R row sends, or the BITWISE bus goes unbalanced. #[inline] fn memw_register_is_half_lookup(ts_lo: u32, old_ts_lo: u32) -> BitwiseOperation { debug_assert!( @@ -622,7 +625,8 @@ impl MemwSink for MemwBuckets { /// Collects all derived operations from CPU operations in a single pass. /// /// This includes: -/// - MEMW ops (register reads/writes M1/M3/M5, memory loads/stores M6/M7) +/// - MEMW ops (register reads/writes M1/M3/M5, memory loads/stores M6/M7), +/// already routed into their MEMW_R / MEMW_A / MEMW buckets (see [`MemwBuckets`]) /// - LOAD ops (memory loads with sign/zero extension) /// - LT ops (from SLT/BLT instructions) /// - Bitwise lookups (from CPU operations) From bf2aa3e03d0883b10ba1be8c4e2d171a12272223 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Mon, 6 Jul 2026 19:17:16 -0300 Subject: [PATCH 08/17] Single-source the MEMW_R fill and the BITWISE type-column map MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- prover/src/tables/bitwise.rs | 85 ++++++------ prover/src/tables/memw_register.rs | 191 +++++++++++++++++++------ prover/src/tables/trace_builder.rs | 214 +++++++---------------------- 3 files changed, 243 insertions(+), 247 deletions(-) diff --git a/prover/src/tables/bitwise.rs b/prover/src/tables/bitwise.rs index eed78f81a..fc4a14f8d 100644 --- a/prover/src/tables/bitwise.rs +++ b/prover/src/tables/bitwise.rs @@ -421,11 +421,13 @@ pub fn update_multiplicities( } /// Number of distinct BITWISE lookup types (one multiplicity column each). -pub const NUM_LOOKUP_TYPES: usize = 10; +/// Derived from [`BitwiseOperationType::ALL`], which the compile-time guard +/// below keeps in lockstep with [`lookup_type_index`]. +pub const NUM_LOOKUP_TYPES: usize = BitwiseOperationType::ALL.len(); /// Dense index in `[0, NUM_LOOKUP_TYPES)` for a lookup type. Ordering is an -/// internal detail of the histogram; only [`type_mu_column`] (its inverse for -/// the fill) needs to agree with it. +/// internal detail of the histogram; [`BitwiseOperationType::ALL`] is its +/// inverse, enforced at compile time. #[inline] pub const fn lookup_type_index(t: BitwiseOperationType) -> usize { match t { @@ -463,46 +465,27 @@ pub const fn mu_column(t: BitwiseOperationType) -> usize { /// Multiplicity column for the histogram lane at dense index `type_idx` /// (inverse of [`lookup_type_index`]). Used by [`BitwiseHistogram::fill_multiplicities`]. +/// +/// Defined as `mu_column ∘ ALL`, so it cannot disagree with the authoritative +/// per-type match — no assumption about MU column contiguity or ordering. #[inline] const fn type_mu_column(type_idx: usize) -> usize { - // Columns 11..=20, one per lookup type, in `lookup_type_index` order. - cols::MU_MSB8 + type_idx + mu_column(BitwiseOperationType::ALL[type_idx]) } -/// Guards the fragile assumption that the histogram's arithmetic column map -/// (`type_mu_column ∘ lookup_type_index`) agrees with the authoritative per-type -/// `mu_column` match. If anyone reorders `lookup_type_index` or renumbers the `MU_*` -/// columns so they are no longer contiguous 11..=20 in that order, this fails loudly in a -/// plain `cargo test` (not just the #[ignore]d parity gate) — a wrong MU column would -/// silently unbalance the BITWISE bus and produce an unsound proof. -#[cfg(test)] -mod histogram_column_map_guard { - use super::*; - - #[test] - fn type_mu_column_matches_mu_column_for_all_types() { - use BitwiseOperationType::*; - const ALL: [BitwiseOperationType; NUM_LOOKUP_TYPES] = [ - Msb8, Msb16, Zero, AreBytes, IsHalf, IsB20, Hwsl, ByteAluAnd, ByteAluOr, ByteAluXor, - ]; - for t in ALL { - assert_eq!( - type_mu_column(lookup_type_index(t)), - mu_column(t), - "histogram column map disagrees with mu_column for {t:?}" - ); - } - // Also confirm the dense indices are a bijection over 0..NUM_LOOKUP_TYPES. - let mut seen = [false; NUM_LOOKUP_TYPES]; - for t in ALL { - seen[lookup_type_index(t)] = true; - } - assert!( - seen.iter().all(|&s| s), - "lookup_type_index is not a bijection" - ); +// Compile-time guard: `ALL` must list every lookup type exactly once, in +// `lookup_type_index` order (i.e. it is the exact inverse of that mapping). +// Adding a variant forces the `lookup_type_index` match to be extended +// (exhaustiveness), and this assert then forces `ALL` — and with it +// `NUM_LOOKUP_TYPES` — to follow. A mismatch is a compile error, not a test +// failure: a wrong MU column would silently unbalance the BITWISE bus. +const _: () = { + let mut i = 0; + while i < NUM_LOOKUP_TYPES { + assert!(lookup_type_index(BitwiseOperationType::ALL[i]) == i); + i += 1; } -} +}; /// "Histogram-on-the-fly" accumulator for BITWISE lookup multiplicities. /// @@ -535,8 +518,12 @@ impl BitwiseHistogram { #[inline] pub fn bump(&mut self, op: BitwiseOperation) { let idx = lookup_type_index(op.lookup_type) * NUM_ROWS + row_index(op.x, op.y, op.z); - // `debug_assert` guards the invariant; row_index already bounds-checks z<16 - // in debug and (x,y) are u8 so idx is always in range. + // (x, y) are u8, and row_index debug-asserts z < 16, so in debug builds a + // corrupt op fails loudly here. In release an out-of-domain z would NOT + // panic: the flat index can land in another type's lane and silently + // mis-count both cells — the proof then fails verification instead of the + // prover crashing. What actually upholds the invariant is that every + // `BitwiseOperation` constructor masks or debug-asserts z < 16. self.counters[idx] += 1; } @@ -596,6 +583,24 @@ pub enum BitwiseOperationType { ByteAluXor, } +impl BitwiseOperationType { + /// Every lookup type exactly once, in [`lookup_type_index`] order (the + /// compile-time guard next to [`type_mu_column`] enforces this). The array + /// length is the single origin of [`NUM_LOOKUP_TYPES`]. + pub const ALL: [Self; 10] = [ + Self::Msb8, + Self::Msb16, + Self::Zero, + Self::AreBytes, + Self::IsHalf, + Self::IsB20, + Self::Hwsl, + Self::ByteAluAnd, + Self::ByteAluOr, + Self::ByteAluXor, + ]; +} + /// A lookup request to the BITWISE precomputed table. /// /// The BITWISE table has 2^20 rows indexed by `(x, y, z)`. diff --git a/prover/src/tables/memw_register.rs b/prover/src/tables/memw_register.rs index aa45cf6db..d47b9744a 100644 --- a/prover/src/tables/memw_register.rs +++ b/prover/src/tables/memw_register.rs @@ -43,6 +43,7 @@ use stark::trace::TraceTable; use stark::constraints::builder::{ConstraintBuilder, ConstraintSet}; +use super::bitwise::{BitwiseOperation, BitwiseOperationType}; use super::memw::MemwOperation; use super::types::{BusId, FE, GoldilocksExtension, GoldilocksField, VmTable}; use crate::constraints::templates::emit_is_bit; @@ -85,28 +86,81 @@ pub mod cols { // Trace generation // ========================================================================= -/// Generates the MEMW_R trace table from register operations. +/// Compact, already-decomposed record for one MEMW_R (register fast-path) access. /// -/// Reuses `MemwOperation` -- the trace generator divides `base_address` by 2 -/// to recover the register index (CPU sends `2 * register_index`). -pub fn generate_memw_register_trace( - operations: &[MemwOperation], -) -> TraceTable { - let num_rows = operations.len().next_power_of_two().max(4); - let mut trace = TraceTable::new_main( - vec![FE::zero(); num_rows * cols::NUM_COLUMNS], - cols::NUM_COLUMNS, - 1, - ); - let table = &mut trace.main_table; +/// This is the "direct-to-column" carrier: it holds exactly the fields the MEMW_R +/// column fill ([`generate_memw_register_trace_from_rows`]) and its IS_HALFWORD +/// bitwise collector ([`collect_bitwise_from_memw_register`]) need, and nothing +/// else. It replaces the full `MemwOperation` (~152 B after the `[u32; 8]` +/// value/old shrink, but still 8-element arrays) for register accesses — the +/// largest table by rows — so the walk never materializes a `MemwOperation` for +/// the register fast path. +/// +/// Field domains mirror `MemwOperation`'s: +/// - `address` = `base_address / 2` (the register index 0..=255; ADDRESS column, +/// and `2*ADDRESS` on the memory/MEMW buses) +/// - `val0/val1` = `value[0]`/`value[1]` (the 32-bit register halves) +/// - `old0/old1` = `old[0]`/`old[1]` +/// - `old_ts_lo` = `old_timestamp[0] & 0xFFFF_FFFF` (the two words share old_timestamp, +/// enforced by `is_register_op`; the upper limb is TIMESTAMP_1 = timestamp>>32) +#[derive(Debug, Clone, Copy)] +pub(crate) struct RegRow { + address: u64, + timestamp: u64, + val0: u32, + val1: u32, + old0: u32, + old1: u32, + old_ts_lo: u32, + is_read: bool, +} - for (row_idx, op) in operations.iter().enumerate() { +impl RegRow { + /// Build a `RegRow` from pre-decomposed register-access fields. + /// + /// `reg_addr` is `2 * reg_index` as sent by the CPU; `old_ts` is the (shared) + /// old_timestamp of both register words. This is the ONLY place the MEMW_R + /// row encoding (halved address, masked `old_ts_lo`) is defined — + /// [`Self::from_memw`] delegates here, so the walk fast path and the + /// `MemwOperation` paths cannot drift. + #[inline] + #[allow(clippy::too_many_arguments)] + pub(crate) fn new( + reg_addr: u64, + timestamp: u64, + val0: u32, + val1: u32, + old0: u32, + old1: u32, + old_ts: u64, + is_read: bool, + ) -> Self { debug_assert_eq!( - op.base_address % 2, + reg_addr % 2, 0, - "register base_address must be even (got {})", - op.base_address + "register base_address must be even (got {reg_addr})" ); + RegRow { + address: reg_addr / 2, + timestamp, + val0, + val1, + old0, + old1, + old_ts_lo: (old_ts & 0xFFFF_FFFF) as u32, + is_read, + } + } + + /// Build a `RegRow` from a fully-formed register `MemwOperation`. Used on the + /// precompile / commit / keccak / halt paths, which construct a `MemwOperation` + /// first and only convert to the compact row once the op is known to route to + /// MEMW_R. + /// + /// Only valid for ops for which `is_register_op` is true (width==2, atomic + /// old_timestamp). + #[inline] + pub(crate) fn from_memw(op: &MemwOperation) -> Self { // Both register words must have been last accessed at the same timestamp. // MEMW_R stores a single old_timestamp_lo and shares TIMESTAMP_1 as the // upper limb, so if the two words differ, the wrong token would be sent @@ -116,36 +170,95 @@ pub fn generate_memw_register_trace( "register words must share old_timestamp ({} != {})", op.old_timestamp[0], op.old_timestamp[1] ); + Self::new( + op.base_address, + op.timestamp, + op.value[0], + op.value[1], + op.old[0], + op.old[1], + op.old_timestamp[0], + op.is_read, + ) + } +} - // ADDRESS = base_address / 2 (CPU sends 2 * register_index) - table.set_u64(row_idx, cols::ADDRESS, op.base_address / 2); - - // Timestamp split into lo/hi 32-bit words - table.set_dword_wl(row_idx, cols::TIMESTAMP_0, op.timestamp); - - // Value: registers are DWordWL = 2 words - table.set_u64(row_idx, cols::VAL_0, op.value[0] as u64); - table.set_u64(row_idx, cols::VAL_1, op.value[1] as u64); - - // Old value - table.set_u64(row_idx, cols::OLD_0, op.old[0] as u64); - table.set_u64(row_idx, cols::OLD_1, op.old[1] as u64); +/// Generates the MEMW_R trace table from register operations. +/// +/// Thin wrapper over [`generate_memw_register_trace_from_rows`] (via +/// [`RegRow::from_memw`]) so there is exactly one MEMW_R column-write sequence. +pub fn generate_memw_register_trace( + operations: &[MemwOperation], +) -> TraceTable { + let rows: Vec = operations.iter().map(RegRow::from_memw).collect(); + generate_memw_register_trace_from_rows(&rows) +} - // Old timestamp low (upper limb shared with TIMESTAMP_1) - table.set_u64( - row_idx, - cols::OLD_TIMESTAMP_LO, - op.old_timestamp[0] & 0xFFFF_FFFF, - ); +/// The MEMW_R column fill from compact [`RegRow`]s. This is the single source of +/// truth for the MEMW_R trace layout; both the walk's direct fast path and the +/// `MemwOperation`-based [`generate_memw_register_trace`] land here. +pub(crate) fn generate_memw_register_trace_from_rows( + rows: &[RegRow], +) -> TraceTable { + let num_rows = rows.len().next_power_of_two().max(4); + let mut trace = TraceTable::new_main( + vec![FE::zero(); num_rows * cols::NUM_COLUMNS], + cols::NUM_COLUMNS, + 1, + ); + let table = &mut trace.main_table; - // Multiplicity - table.set_bool(row_idx, cols::MU_READ, op.is_read); - table.set_bool(row_idx, cols::MU_WRITE, !op.is_read); + for (row_idx, r) in rows.iter().enumerate() { + // ADDRESS = base_address / 2 (already divided in RegRow). + table.set_u64(row_idx, cols::ADDRESS, r.address); + // Timestamp split into lo/hi 32-bit words. + table.set_dword_wl(row_idx, cols::TIMESTAMP_0, r.timestamp); + // Value: registers are DWordWL = 2 words. + table.set_u64(row_idx, cols::VAL_0, r.val0 as u64); + table.set_u64(row_idx, cols::VAL_1, r.val1 as u64); + // Old value. + table.set_u64(row_idx, cols::OLD_0, r.old0 as u64); + table.set_u64(row_idx, cols::OLD_1, r.old1 as u64); + // Old timestamp low (upper limb shared with TIMESTAMP_1). + table.set_u64(row_idx, cols::OLD_TIMESTAMP_LO, r.old_ts_lo as u64); + // Multiplicity. + table.set_bool(row_idx, cols::MU_READ, r.is_read); + table.set_bool(row_idx, cols::MU_WRITE, !r.is_read); } trace } +/// The single IS_HALFWORD lookup a MEMW_R access sends: proves the timestamp delta +/// `ts_lo - old_ts_lo` is in [1, 2^16] by decomposing `ts_lo - old_ts_lo - 1` into +/// two bytes. +/// +/// Must stay in lockstep with the IS_HALFWORD send in [`bus_interactions`]: the +/// lookup counted here has to be exactly the lookup each MEMW_R row sends, or the +/// BITWISE bus goes unbalanced. +#[inline] +fn memw_register_is_half_lookup(ts_lo: u32, old_ts_lo: u32) -> BitwiseOperation { + debug_assert!( + ts_lo > old_ts_lo, + "ts_lo must exceed old_ts_lo (enforced by the MEMW_R routing predicate)" + ); + let diff_minus_1 = (ts_lo - old_ts_lo - 1) as u16; + BitwiseOperation::halfword( + BitwiseOperationType::IsHalf, + (diff_minus_1 & 0xFF) as u8, + (diff_minus_1 >> 8) as u8, + ) +} + +/// IS_HALFWORD bitwise lookups for MEMW_R, computed directly from [`RegRow`]s via +/// the shared [`memw_register_is_half_lookup`] helper (the same lookup the MEMW_R +/// trace fill uses), so the multiplicities stay consistent with that table. +pub(crate) fn collect_bitwise_from_memw_register(rows: &[RegRow]) -> Vec { + rows.iter() + .map(|r| memw_register_is_half_lookup((r.timestamp & 0xFFFF_FFFF) as u32, r.old_ts_lo)) + .collect() +} + // ========================================================================= // Bus interactions (7 total) // ========================================================================= diff --git a/prover/src/tables/trace_builder.rs b/prover/src/tables/trace_builder.rs index 7bdff290f..ad56ef757 100644 --- a/prover/src/tables/trace_builder.rs +++ b/prover/src/tables/trace_builder.rs @@ -60,7 +60,7 @@ use super::local_to_global; use super::lt::{self, LtOperation}; use super::memw::{self, MemwOperation}; use super::memw_aligned; -use super::memw_register; +use super::memw_register::{self, RegRow}; use super::mul::{self, MulOperation}; use super::page::{self, FinalByteState, FinalStateMap, PageConfig}; use super::register::{self, FinalRegisterStateMap, FinalRegisterWordState}; @@ -352,128 +352,30 @@ fn collect_cpu_ops( // Phase 2: CPU ops → MEMW, LOAD, LT, Bitwise // ============================================================================= -/// Compact, already-decomposed record for one MEMW_R (register fast-path) access. +/// Destination table for a `MemwOperation`. /// -/// This is the "direct-to-column" carrier: it holds exactly the fields the MEMW_R -/// column fill (`generate_memw_register_trace_direct`) and its IS_HALFWORD bitwise -/// collector (`collect_bitwise_from_memw_register_direct`) need, and nothing else. -/// It replaces the full `MemwOperation` (~152 B after the `[u32; 8]` value/old -/// shrink, but still 8-element arrays) for register accesses — the largest table -/// by rows — so the walk never materializes a `MemwOperation` for the register -/// fast path. -/// -/// Field domains mirror `MemwOperation`'s so the produced table is byte-identical: -/// - `address` = `base_address / 2` (the register index 0..=255; ADDRESS column, -/// and `2*ADDRESS` on the memory/MEMW buses) -/// - `val0/val1` = `value[0]`/`value[1]` (the 32-bit register halves) -/// - `old0/old1` = `old[0]`/`old[1]` -/// - `old_ts_lo` = `old_timestamp[0] & 0xFFFF_FFFF` (the two words share old_timestamp, -/// enforced by `is_register_op`; the upper limb is TIMESTAMP_1 = timestamp>>32) -#[derive(Debug, Clone, Copy)] -struct RegRow { - address: u64, - timestamp: u64, - val0: u32, - val1: u32, - old0: u32, - old1: u32, - old_ts_lo: u32, - is_read: bool, -} - -impl RegRow { - /// Build a `RegRow` from a fully-formed register `MemwOperation`. Used on the - /// precompile / commit / keccak / halt paths, which construct a `MemwOperation` - /// first (their computation is unchanged) and only convert to the compact row - /// once the op is known to route to MEMW_R. - /// - /// Only valid for ops for which `is_register_op` is true (width==2, atomic - /// old_timestamp). The debug asserts mirror `generate_memw_register_trace`. - #[inline] - fn from_memw(op: &MemwOperation) -> Self { - debug_assert_eq!(op.base_address % 2, 0); - debug_assert_eq!(op.old_timestamp[0], op.old_timestamp[1]); - RegRow { - address: op.base_address / 2, - timestamp: op.timestamp, - val0: op.value[0], - val1: op.value[1], - old0: op.old[0], - old1: op.old[1], - old_ts_lo: (op.old_timestamp[0] & 0xFFFF_FFFF) as u32, - is_read: op.is_read, - } - } -} - -/// Direct-to-column MEMW_R trace fill from compact [`RegRow`]s. -/// -/// This is the fused replacement for `generate_memw_register_trace` (which fills from -/// `&[MemwOperation]`): it produces a BYTE-IDENTICAL table. Every column write below -/// mirrors `memw_register::generate_memw_register_trace` cell-for-cell — the only -/// difference is the source carrier (`RegRow` vs `MemwOperation`). Keep the two in sync. -fn generate_memw_register_trace_direct( - rows: &[RegRow], -) -> TraceTable { - use super::types::{FE, VmTable}; - - let num_rows = rows.len().next_power_of_two().max(4); - let mut trace = TraceTable::new_main( - vec![FE::zero(); num_rows * memw_register::cols::NUM_COLUMNS], - memw_register::cols::NUM_COLUMNS, - 1, - ); - let table = &mut trace.main_table; - - use memw_register::cols; - for (row_idx, r) in rows.iter().enumerate() { - // ADDRESS = base_address / 2 (already divided in RegRow). - table.set_u64(row_idx, cols::ADDRESS, r.address); - // Timestamp split into lo/hi 32-bit words. - table.set_dword_wl(row_idx, cols::TIMESTAMP_0, r.timestamp); - // Value: registers are DWordWL = 2 words. - table.set_u64(row_idx, cols::VAL_0, r.val0 as u64); - table.set_u64(row_idx, cols::VAL_1, r.val1 as u64); - // Old value. - table.set_u64(row_idx, cols::OLD_0, r.old0 as u64); - table.set_u64(row_idx, cols::OLD_1, r.old1 as u64); - // Old timestamp low (upper limb shared with TIMESTAMP_1). - table.set_u64(row_idx, cols::OLD_TIMESTAMP_LO, r.old_ts_lo as u64); - // Multiplicity. - table.set_bool(row_idx, cols::MU_READ, r.is_read); - table.set_bool(row_idx, cols::MU_WRITE, !r.is_read); - } - - trace +/// The order of the checks matters and must never change: register ops would +/// also pass `is_aligned_op`, so MEMW_R is decided first, then MEMW_A, and the +/// rest goes to the general MEMW table. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum MemwRoute { + Register, + Aligned, + General, } -/// The single IS_HALFWORD lookup a MEMW_R access sends: proves the timestamp delta -/// `ts_lo - old_ts_lo` is in [1, 2^16] by decomposing `ts_lo - old_ts_lo - 1` into two bytes. -/// -/// Must stay in lockstep with the IS_HALFWORD send in -/// `memw_register::bus_interactions()`: the lookup counted here has to be exactly -/// the lookup each MEMW_R row sends, or the BITWISE bus goes unbalanced. +/// The single classification used everywhere a `MemwOperation` is routed to a +/// table — the walk's [`MemwBuckets`] and the sizing pass (`count_table_lengths`) +/// share it, so their routing cannot drift. #[inline] -fn memw_register_is_half_lookup(ts_lo: u32, old_ts_lo: u32) -> BitwiseOperation { - debug_assert!( - ts_lo > old_ts_lo, - "ts_lo must exceed old_ts_lo (enforced by reg_ts_delta_in_range)" - ); - let diff_minus_1 = (ts_lo - old_ts_lo - 1) as u16; - BitwiseOperation::halfword( - BitwiseOperationType::IsHalf, - (diff_minus_1 & 0xFF) as u8, - (diff_minus_1 >> 8) as u8, - ) -} - -/// IS_HALFWORD bitwise lookups for MEMW_R, computed directly from [`RegRow`]s via -/// the shared [`memw_register_is_half_lookup`] helper (the same lookup the MEMW_R -/// trace fill uses), so the multiplicities stay consistent with that table. -fn collect_bitwise_from_memw_register_direct(rows: &[RegRow]) -> Vec { - rows.iter() - .map(|r| memw_register_is_half_lookup((r.timestamp & 0xFFFF_FFFF) as u32, r.old_ts_lo)) - .collect() +fn classify_memw(op: &MemwOperation) -> MemwRoute { + if is_register_op(op) { + MemwRoute::Register + } else if is_aligned_op(op) { + MemwRoute::Aligned + } else { + MemwRoute::General + } } /// Routes each `MemwOperation` into its destination table bucket at CREATION time @@ -505,26 +407,12 @@ impl MemwBuckets { } } - /// Store an op already known to route to MEMW_R. - #[inline] - fn push_register(&mut self, op: MemwOperation) { - self.register_rows.push(RegRow::from_memw(&op)); - } - - /// Store a compact register row directly. - #[inline] - fn push_register_row(&mut self, row: RegRow) { - self.register_rows.push(row); - } - #[inline] fn push(&mut self, op: MemwOperation) { - if is_register_op(&op) { - self.push_register(op); - } else if is_aligned_op(&op) { - self.aligned.push(op); - } else { - self.general.push(op); + match classify_memw(&op) { + MemwRoute::Register => self.register_rows.push(RegRow::from_memw(&op)), + MemwRoute::Aligned => self.aligned.push(op), + MemwRoute::General => self.general.push(op), } } fn extend_ops(&mut self, ops: impl IntoIterator) { @@ -599,25 +487,13 @@ impl MemwSink for MemwBuckets { // `old_ts` (always true here by construction). If it passes, fill a RegRow // directly; otherwise fall back to the general/aligned MemwOperation path. if reg_ts_delta_in_range(timestamp, old_ts) { - let row = RegRow { - address: reg_addr / 2, - timestamp, - val0, - val1, - old0, - old1, - old_ts_lo: (old_ts & 0xFFFF_FFFF) as u32, - is_read, - }; - self.push_register_row(row); + self.register_rows.push(RegRow::new( + reg_addr, timestamp, val0, val1, old0, old1, old_ts, is_read, + )); } else { let op = build_fallback(); debug_assert!(!is_register_op(&op), "reg fallback must not be MEMW_R"); - if is_aligned_op(&op) { - self.aligned.push(op); - } else { - self.general.push(op); - } + self.push(op); } } } @@ -3208,7 +3084,7 @@ fn build_traces( .collect() }), Box::new(|| collect_bitwise_from_memw_aligned(&memw_aligned_ops)), - Box::new(|| collect_bitwise_from_memw_register_direct(&memw_register_rows)), + Box::new(|| memw_register::collect_bitwise_from_memw_register(&memw_register_rows)), Box::new(|| collect_bitwise_from_commit(&commit_ops)), Box::new(|| collect_bitwise_from_keccak(&keccak_ops)), Box::new(|| collect_bitwise_from_ecsm(&ecsm_ops)), @@ -3336,7 +3212,7 @@ fn build_traces( chunk_and_generate( &memw_register_rows, max_rows.memw_register, - generate_memw_register_trace_direct, + memw_register::generate_memw_register_trace_from_rows, #[cfg(feature = "disk-spill")] storage_mode, ) @@ -3756,19 +3632,21 @@ pub fn count_table_lengths( by_width: &mut [usize; 4], aligned: &mut usize, register: &mut usize| { - if is_register_op(op) { - *register += 1; - } else if is_aligned_op(op) { - *aligned += 1; - } else { - let idx = match op.width { - 1 => 0, - 2 => 1, - 4 => 2, - 8 => 3, - _ => return, - }; - by_width[idx] += 1; + // Same classifier as the walk's MemwBuckets, so the sizing pass counts + // exactly the rows trace generation will produce. + match classify_memw(op) { + MemwRoute::Register => *register += 1, + MemwRoute::Aligned => *aligned += 1, + MemwRoute::General => { + let idx = match op.width { + 1 => 0, + 2 => 1, + 4 => 2, + 8 => 3, + _ => return, + }; + by_width[idx] += 1; + } } }; From 1287123846cd9ac76d0d3719ace1ab2fc3e65c71 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 7 Jul 2026 10:58:09 -0300 Subject: [PATCH 09/17] Make the u32 value domain a compile-time fact for MemwOperation 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. --- prover/src/tables/memw.rs | 27 ++--------- prover/src/tables/trace_builder.rs | 77 +++++++++++++++++------------- 2 files changed, 47 insertions(+), 57 deletions(-) diff --git a/prover/src/tables/memw.rs b/prover/src/tables/memw.rs index f0f8f2df7..f4af949e1 100644 --- a/prover/src/tables/memw.rs +++ b/prover/src/tables/memw.rs @@ -131,7 +131,7 @@ impl MemwOperation { pub fn new( is_register: bool, base_address: u64, - value: [u64; 8], + value: [u32; 8], timestamp: u64, width: u8, is_read: bool, @@ -139,20 +139,7 @@ impl MemwOperation { Self { is_register, base_address, - // Callers build a [u64; 8] transiently on the stack; we store the u32 - // domain (byte / 32-bit register half) so the persisted struct is half - // the size. Values never exceed u32 (every element is a byte or a - // pack_register_value 32-bit limb). This is an always-on `assert!` (not - // `debug_assert!`): a caller passing an out-of-domain value would - // otherwise silently truncate here and produce an unsound witness in - // release builds, so the invariant is enforced in every build profile. - value: value.map(|v| { - assert!( - v <= u32::MAX as u64, - "MemwOperation value element exceeds u32: {v}" - ); - v as u32 - }), + value, timestamp, width, is_read, @@ -162,14 +149,8 @@ impl MemwOperation { } /// Set the old values (from memory model). - pub fn with_old(mut self, old: [u64; 8], old_timestamp: [u64; 8]) -> Self { - self.old = old.map(|v| { - assert!( - v <= u32::MAX as u64, - "MemwOperation old element exceeds u32: {v}" - ); - v as u32 - }); + pub fn with_old(mut self, old: [u32; 8], old_timestamp: [u64; 8]) -> Self { + self.old = old; self.old_timestamp = old_timestamp; self } diff --git a/prover/src/tables/trace_builder.rs b/prover/src/tables/trace_builder.rs index ad56ef757..4406e0c75 100644 --- a/prover/src/tables/trace_builder.rs +++ b/prover/src/tables/trace_builder.rs @@ -133,12 +133,12 @@ impl MemoryState { } /// Read multiple bytes. Returns arrays of values and timestamps. - fn read_bytes(&self, base_address: u64, count: usize) -> ([u64; 8], [u64; 8]) { - let mut values = [0u64; 8]; + fn read_bytes(&self, base_address: u64, count: usize) -> ([u32; 8], [u64; 8]) { + let mut values = [0u32; 8]; let mut timestamps = [0u64; 8]; for i in 0..count { let (val, ts) = self.read_byte(base_address.wrapping_add(i as u64)); - values[i] = val as u64; + values[i] = val as u32; timestamps[i] = ts; } (values, timestamps) @@ -313,8 +313,17 @@ fn cpu_op_to_bytes_and_signed(op: &CpuOperation) -> (usize, bool) { /// Pack a 64-bit register value into the MEMW value format. /// /// For register operations, values are packed as [lo32, hi32, 0, 0, 0, 0, 0, 0]. -fn pack_register_value(value: u64) -> [u64; 8] { - [value & 0xFFFF_FFFF, value >> 32, 0, 0, 0, 0, 0, 0] +fn pack_register_value(value: u64) -> [u32; 8] { + [ + (value & 0xFFFF_FFFF) as u32, + (value >> 32) as u32, + 0, + 0, + 0, + 0, + 0, + 0, + ] } // ============================================================================= @@ -711,9 +720,9 @@ fn collect_load_op_from_cpu( let (_old_values, old_timestamps) = memory_state.read_bytes(base_address, 8); // Extract individual bytes from loaded value - let mut value_bytes = [0u64; 8]; + let mut value_bytes = [0u32; 8]; for (j, byte) in value_bytes.iter_mut().take(byte_count).enumerate() { - *byte = (loaded_value >> (j * 8)) & 0xFF; + *byte = ((loaded_value >> (j * 8)) & 0xFF) as u32; } // Sign/zero extend the upper bytes @@ -744,7 +753,7 @@ fn collect_load_op_from_cpu( op.timestamp, byte_count as u8, signed, - res_bytes, + res_bytes.map(u64::from), ); // Collect MSB8 lookups for sign bit extraction @@ -774,9 +783,9 @@ fn collect_store_op_from_cpu(op: &CpuOperation, memory_state: &mut MemoryState) // of all 8 bytes. Must match CPU M7 which sends full rv2 as [lo32, hi32]. // Bus 16: only positions 0..byte_count participate (controlled by w2/w4/write8 // multiplicities), so extra bytes don't affect memory consistency. - let mut value_bytes = [0u64; 8]; + let mut value_bytes = [0u32; 8]; for (j, byte) in value_bytes.iter_mut().enumerate() { - *byte = (store_value >> (j * 8)) & 0xFF; + *byte = ((store_value >> (j * 8)) & 0xFF) as u32; } // The STORE chip now owns this MEMW write (the CPU sends MEMORY instead of @@ -849,10 +858,10 @@ fn collect_ecsm_ops( for (base, bytes) in [(addr_xg, &witness.x_g), (addr_k, &witness.k)] { for i in 0..4 { let addr = base.wrapping_add((8 * i) as u64); - let mut value = [0u64; 8]; + let mut value = [0u32; 8]; let mut dword = 0u64; for j in 0..8 { - value[j] = bytes[8 * i + j] as u64; + value[j] = bytes[8 * i + j] as u32; dword |= (bytes[8 * i + j] as u64) << (8 * j); } let (_old, old_ts) = memory_state.read_bytes(addr, 8); @@ -877,7 +886,7 @@ fn collect_ecsm_ops( for offset in 0..32u64 { let addr = addr_k.wrapping_add(offset); let byte = k[offset as usize]; - let value = [byte as u64, 0, 0, 0, 0, 0, 0, 0]; + let value = [byte as u32, 0, 0, 0, 0, 0, 0, 0]; let (_v, old_ts) = memory_state.read_byte(addr); memw_ops.push( MemwOperation::new(false, addr, value, t + 1, 1, true) @@ -889,10 +898,10 @@ fn collect_ecsm_ops( // xR writes at T + 2 (4 doublewords). for i in 0..4 { let addr = addr_xr.wrapping_add((8 * i) as u64); - let mut value = [0u64; 8]; + let mut value = [0u32; 8]; let mut dword = 0u64; for j in 0..8 { - value[j] = witness.x_r[8 * i + j] as u64; + value[j] = witness.x_r[8 * i + j] as u32; dword |= (witness.x_r[8 * i + j] as u64) << (8 * j); } let (old_vals, old_ts) = memory_state.read_bytes(addr, 8); @@ -949,10 +958,10 @@ fn collect_register_ops_from_cpu( // the identical MemwOperation only on the (rare) general/aligned fallback. memw_ops.push_reg_access( reg_addr, - reg_value[0] as u32, - reg_value[1] as u32, - reg_value[0] as u32, - reg_value[1] as u32, + reg_value[0], + reg_value[1], + reg_value[0], + reg_value[1], ts, old_ts, true, @@ -977,10 +986,10 @@ fn collect_register_ops_from_cpu( let ts = op.timestamp + 1; memw_ops.push_reg_access( reg_addr, - reg_value[0] as u32, - reg_value[1] as u32, - reg_value[0] as u32, - reg_value[1] as u32, + reg_value[0], + reg_value[1], + reg_value[0], + reg_value[1], ts, old_ts, true, @@ -1002,10 +1011,10 @@ fn collect_register_ops_from_cpu( let ts = op.timestamp + 2; memw_ops.push_reg_access( reg_addr, - reg_value[0] as u32, - reg_value[1] as u32, - old_value[0] as u32, - old_value[1] as u32, + reg_value[0], + reg_value[1], + old_value[0], + old_value[1], ts, old_ts, false, @@ -1249,8 +1258,8 @@ fn collect_commit_memw_ops( let new_index = old_index .checked_add(u32::try_from(count).expect("commit_count exceeds u32 range")) .expect("commit index exceeds u32 range"); - let old_value = [old_index as u64, 0, 0, 0, 0, 0, 0, 0]; - let new_value = [new_index as u64, 0, 0, 0, 0, 0, 0, 0]; + let old_value = [old_index, 0, 0, 0, 0, 0, 0, 0]; + let new_value = [new_index, 0, 0, 0, 0, 0, 0, 0]; let old_timestamps = [old_ts, 0, 0, 0, 0, 0, 0, 0]; let memw_op = MemwOperation::new( true, @@ -1269,7 +1278,7 @@ fn collect_commit_memw_ops( for i in 0..count { let addr = buf_addr.wrapping_add(i); let (byte_val, old_ts) = memory_state.read_byte(addr); - let value = [byte_val as u64, 0, 0, 0, 0, 0, 0, 0]; + let value = [byte_val as u32, 0, 0, 0, 0, 0, 0, 0]; let old_timestamps = [old_ts, 0, 0, 0, 0, 0, 0, 0]; let memw_op = MemwOperation::new(false, addr, value, ts, 1, true).with_old(value, old_timestamps); @@ -1377,10 +1386,10 @@ fn collect_keccak_memw_ops( .checked_add(lane_idx as u64 * 8) .expect("keccak state address range must be validated by the executor"); - let mut old_bytes = [0u64; 8]; + let mut old_bytes = [0u32; 8]; let mut old_timestamps = [0u64; 8]; for b in 0..8 { - old_bytes[b] = (in_lane >> (b * 8)) & 0xFF; + old_bytes[b] = ((in_lane >> (b * 8)) & 0xFF) as u32; let byte_addr = lane_addr .checked_add(b as u64) .expect("keccak state address range must be validated by the executor"); @@ -1388,9 +1397,9 @@ fn collect_keccak_memw_ops( old_timestamps[b] = old_ts; } - let mut value_bytes = [0u64; 8]; + let mut value_bytes = [0u32; 8]; for (b, byte) in value_bytes.iter_mut().enumerate() { - *byte = (out_lane >> (b * 8)) & 0xFF; + *byte = ((out_lane >> (b * 8)) & 0xFF) as u32; } let memw_op = MemwOperation::new(false, lane_addr, value_bytes, ts, 8, true) From 47948a26a6a8ba73ff01fb59e344082d4a7a80c7 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 7 Jul 2026 11:14:00 -0300 Subject: [PATCH 10/17] Feed bitwise collectors straight into the histogram; cut reduce allocations - Collectors now take &mut BitwiseHistogram instead of returning Vec. 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. --- prover/src/tables/bitwise.rs | 9 +- prover/src/tables/memw_register.rs | 32 ++++-- prover/src/tables/trace_builder.rs | 163 +++++++++++++---------------- 3 files changed, 104 insertions(+), 100 deletions(-) diff --git a/prover/src/tables/bitwise.rs b/prover/src/tables/bitwise.rs index fc4a14f8d..25d7c960d 100644 --- a/prover/src/tables/bitwise.rs +++ b/prover/src/tables/bitwise.rs @@ -517,6 +517,13 @@ impl BitwiseHistogram { /// Increment the counter for one lookup. #[inline] pub fn bump(&mut self, op: BitwiseOperation) { + self.bump_n(op, 1); + } + + /// Add `n` occurrences of one lookup in a single step (e.g. CPU padding rows, + /// which all send identical all-zero lookups). + #[inline] + pub fn bump_n(&mut self, op: BitwiseOperation, n: u64) { let idx = lookup_type_index(op.lookup_type) * NUM_ROWS + row_index(op.x, op.y, op.z); // (x, y) are u8, and row_index debug-asserts z < 16, so in debug builds a // corrupt op fails loudly here. In release an out-of-domain z would NOT @@ -524,7 +531,7 @@ impl BitwiseHistogram { // mis-count both cells — the proof then fails verification instead of the // prover crashing. What actually upholds the invariant is that every // `BitwiseOperation` constructor masks or debug-asserts z < 16. - self.counters[idx] += 1; + self.counters[idx] += n; } /// Fold a slice of lookups into the histogram. diff --git a/prover/src/tables/memw_register.rs b/prover/src/tables/memw_register.rs index d47b9744a..7e5ff7c02 100644 --- a/prover/src/tables/memw_register.rs +++ b/prover/src/tables/memw_register.rs @@ -43,7 +43,7 @@ use stark::trace::TraceTable; use stark::constraints::builder::{ConstraintBuilder, ConstraintSet}; -use super::bitwise::{BitwiseOperation, BitwiseOperationType}; +use super::bitwise::{BitwiseHistogram, BitwiseOperation, BitwiseOperationType}; use super::memw::MemwOperation; use super::types::{BusId, FE, GoldilocksExtension, GoldilocksField, VmTable}; use crate::constraints::templates::emit_is_bit; @@ -105,7 +105,9 @@ pub mod cols { /// enforced by `is_register_op`; the upper limb is TIMESTAMP_1 = timestamp>>32) #[derive(Debug, Clone, Copy)] pub(crate) struct RegRow { - address: u64, + /// Register index 0..=255 (`base_address / 2`); u16 keeps the struct at + /// 32 bytes — it is the largest persisted array of the walk. + address: u16, timestamp: u64, val0: u32, val1: u32, @@ -140,8 +142,12 @@ impl RegRow { 0, "register base_address must be even (got {reg_addr})" ); + debug_assert!( + reg_addr / 2 <= u16::MAX as u64, + "register index exceeds u16 (got base_address {reg_addr})" + ); RegRow { - address: reg_addr / 2, + address: (reg_addr / 2) as u16, timestamp, val0, val1, @@ -210,7 +216,7 @@ pub(crate) fn generate_memw_register_trace_from_rows( for (row_idx, r) in rows.iter().enumerate() { // ADDRESS = base_address / 2 (already divided in RegRow). - table.set_u64(row_idx, cols::ADDRESS, r.address); + table.set_u64(row_idx, cols::ADDRESS, r.address as u64); // Timestamp split into lo/hi 32-bit words. table.set_dword_wl(row_idx, cols::TIMESTAMP_0, r.timestamp); // Value: registers are DWordWL = 2 words. @@ -250,13 +256,17 @@ fn memw_register_is_half_lookup(ts_lo: u32, old_ts_lo: u32) -> BitwiseOperation ) } -/// IS_HALFWORD bitwise lookups for MEMW_R, computed directly from [`RegRow`]s via -/// the shared [`memw_register_is_half_lookup`] helper (the same lookup the MEMW_R -/// trace fill uses), so the multiplicities stay consistent with that table. -pub(crate) fn collect_bitwise_from_memw_register(rows: &[RegRow]) -> Vec { - rows.iter() - .map(|r| memw_register_is_half_lookup((r.timestamp & 0xFFFF_FFFF) as u32, r.old_ts_lo)) - .collect() +/// IS_HALFWORD bitwise lookups for MEMW_R, bumped straight into the histogram +/// via the shared [`memw_register_is_half_lookup`] helper (the same lookup the +/// MEMW_R trace fill uses), one per row. No intermediate op vector: register +/// rows number in the tens of millions and the histogram is the only consumer. +pub(crate) fn collect_bitwise_from_memw_register(rows: &[RegRow], hist: &mut BitwiseHistogram) { + for r in rows { + hist.bump(memw_register_is_half_lookup( + (r.timestamp & 0xFFFF_FFFF) as u32, + r.old_ts_lo, + )); + } } // ========================================================================= diff --git a/prover/src/tables/trace_builder.rs b/prover/src/tables/trace_builder.rs index 4406e0c75..5eb9cfd0f 100644 --- a/prover/src/tables/trace_builder.rs +++ b/prover/src/tables/trace_builder.rs @@ -1983,32 +1983,21 @@ fn collect_bitwise_from_branch(branch_ops: &[BranchOperation]) -> Vec Vec { - if num_padding_rows == 0 { - return Vec::new(); - } - - let mut ops = Vec::with_capacity(num_padding_rows * 14); - for _ in 0..num_padding_rows { - // The shrunk CPU sends, per row (incl. padding where all values are 0): - // 3× ARE_BYTES (rs1/rs2, rd/instruction_length, alu_flags/mem_flags) and - // 4× IS_HALF (the four `res` halves). - for _ in 0..3 { - ops.push(BitwiseOperation::byte_op( - BitwiseOperationType::AreBytes, - 0, - 0, - )); - } - for _ in 0..4 { - ops.push(BitwiseOperation::halfword( - BitwiseOperationType::IsHalf, - 0, - 0, - )); - } - } - ops +fn add_padding_byte_checks(hist: &mut bitwise::BitwiseHistogram, num_padding_rows: usize) { + // The shrunk CPU sends, per row (incl. padding where all values are 0): + // 3× ARE_BYTES (rs1/rs2, rd/instruction_length, alu_flags/mem_flags) and + // 4× IS_HALF (the four `res` halves). Padding rows all send the identical + // all-zero lookups, so this is two O(1) histogram bumps instead of an + // O(num_padding_rows) op vector. + let n = num_padding_rows as u64; + hist.bump_n( + BitwiseOperation::byte_op(BitwiseOperationType::AreBytes, 0, 0), + 3 * n, + ); + hist.bump_n( + BitwiseOperation::halfword(BitwiseOperationType::IsHalf, 0, 0), + 4 * n, + ); } /// Collects ARE_BYTES lookups from PAGE data (init and fini values). @@ -2128,11 +2117,11 @@ fn collect_bitwise_from_page( image: &I, memory_state: &MemoryState, exclude_touched: bool, -) -> Vec { + hist: &mut bitwise::BitwiseHistogram, +) { use std::collections::BTreeSet; let page_size = page::DEFAULT_PAGE_SIZE; - let mut bitwise_ops = Vec::new(); let init_page_data = build_init_page_data(image); @@ -2164,16 +2153,16 @@ fn collect_bitwise_from_page( // Get fini value (from final_state or init if never accessed) let fini = final_state.get(&addr).map_or(init, |state| state.value); - // C1+C2: ARE_BYTES[init, fini] — batched range check for both bytes - bitwise_ops.push(BitwiseOperation::byte_op( + // C1+C2: ARE_BYTES[init, fini] — batched range check for both bytes. + // Bumped straight into the histogram: this loop visits every byte of + // every touched page, and the histogram is the only consumer. + hist.bump(BitwiseOperation::byte_op( BitwiseOperationType::AreBytes, init, fini, )); } } - - bitwise_ops } // ============================================================================= @@ -3047,64 +3036,58 @@ fn build_traces( // The per-source bitwise collectors are all pure functions of their inputs and the // BITWISE multiplicities are order-independent (they ride a permutation-invariant bus), - // so we collect every source in parallel and concatenate. The result is byte-identical - // to the previous serial `.extend()` chain regardless of order. + // so we collect every source in parallel. The result is byte-identical to the previous + // serial `.extend()` chain regardless of order. // // MUL/DVRM dedup their per-unique bit-gated lookups PER CHIP INSTANCE, so pass the same // chunk size used to split them into instances so multiplicities match the per-instance // sends. MEMW_R sends IS_HALFWORD[timestamp_0 - old_timestamp_lo - 1]. PAGE does a // batched ARE_BYTES[init, fini] per row (skipped in continuation epochs, which the L2G // table owns). COMMIT sends AreBytes+IsHalfword; KECCAK_RND sends XOR/AND/ARE_BYTES/HWSL. - // Every source's bitwise lookups are collected by a pure `collect_*` function. We never - // concatenate them into one giant `Vec` (~140 M ops / ~560 MB at 10-tx - // whose only consumer is the multiplicity count). Instead each collector's transient - // per-source Vec is folded into a `BitwiseHistogram` and dropped, and the per-worker - // histograms are tree-reduced. The histogram is a commutative monoid, so the summed - // multiplicities are independent of accumulation order (the lookups ride a - // permutation-invariant bus). - - // Shared collector list: each is a pure `Fn() -> Vec` of its source. - // Order does not affect the histogram (commutative). - type Collector<'a> = Box Vec + Sync + 'a>; + // 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 + // 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; let mut collectors: Vec = vec![ - Box::new(|| collect_bitwise_from_lt(<_ops)), - Box::new(|| collect_bitwise_from_mul(&mul_ops, mul_chunk)), - Box::new(|| collect_bitwise_from_dvrm(&dvrm_ops, dvrm_chunk)), - Box::new(|| collect_bitwise_from_branch(&branch_ops)), - Box::new(|| shift::collect_bitwise_from_shift(&shift_ops)), - Box::new(|| { - bytewise_ops - .iter() - .flat_map(|op| op.collect_bitwise_ops()) - .collect() + 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))), + Box::new(|h| h.add_ops(&collect_bitwise_from_dvrm(&dvrm_ops, dvrm_chunk))), + Box::new(|h| h.add_ops(&collect_bitwise_from_branch(&branch_ops))), + Box::new(|h| h.add_ops(&shift::collect_bitwise_from_shift(&shift_ops))), + Box::new(|h| { + for op in &bytewise_ops { + h.add_ops(&op.collect_bitwise_ops()); + } }), - Box::new(|| { - eq_ops - .iter() - .flat_map(|op| op.collect_bitwise_ops()) - .collect() + Box::new(|h| { + for op in &eq_ops { + h.add_ops(&op.collect_bitwise_ops()); + } }), - Box::new(|| { - store_ops - .iter() - .flat_map(|op| op.collect_bitwise_ops()) - .collect() + Box::new(|h| { + for op in &store_ops { + h.add_ops(&op.collect_bitwise_ops()); + } }), - Box::new(|| collect_bitwise_from_memw_aligned(&memw_aligned_ops)), - Box::new(|| memw_register::collect_bitwise_from_memw_register(&memw_register_rows)), - Box::new(|| collect_bitwise_from_commit(&commit_ops)), - Box::new(|| collect_bitwise_from_keccak(&keccak_ops)), - Box::new(|| collect_bitwise_from_ecsm(&ecsm_ops)), - Box::new(|| collect_bitwise_from_ecdas(&ecdas_ops)), - Box::new(|| collect_byte_check_ops_for_padding(num_padding_rows)), + 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))), + Box::new(|h| h.add_ops(&collect_bitwise_from_ecdas(&ecdas_ops))), + Box::new(|h| add_padding_byte_checks(h, num_padding_rows)), ]; if let Some(image) = initial_image && !l2g_memory_bookend { - collectors.push(Box::new(move || { - collect_bitwise_from_page(image, memory_state, l2g_memory_bookend) + collectors.push(Box::new(move |h| { + collect_bitwise_from_page(image, memory_state, l2g_memory_bookend, h) })); } @@ -3117,35 +3100,39 @@ fn build_traces( #[cfg(feature = "parallel")] { use rayon::prelude::*; - // Each collector produces a transient Vec of lookups that is folded into a - // dense 80 MiB histogram and dropped. 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). Tree-reduce the chunk histograms — - // add_ops/merge form a commutative monoid, so the result is order-independent - // and byte-identical to the sequential path. + // 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; unlike `reduce(identity, ..)` it allocates NO + // per-leaf identity histograms (each a zeroed 80 MiB). Tree-reduce is + // valid because bump/add_ops/merge form a commutative monoid, so the + // result is order-independent and byte-identical to the sequential path. let max_par = rayon::current_num_threads().clamp(1, 8); let chunk_size = collectors.len().div_ceil(max_par).max(1); - let reduced = collectors + if let Some(reduced) = collectors .par_chunks(chunk_size) .map(|chunk| { let mut h = bitwise::BitwiseHistogram::new(); for f in chunk { - h.add_ops(&f()); + f(&mut h); } h }) - .reduce(bitwise::BitwiseHistogram::new, |mut a, b| { + .reduce_with(|mut a, b| { a.merge(&b); a - }); - base.merge(&reduced); + }) + { + base.merge(&reduced); + } } #[cfg(not(feature = "parallel"))] { for f in &collectors { - base.add_ops(&f()); + f(&mut base); } } let bitwise_histogram = base; From ee3fd85e476d0a47c3fe8da5470f37b0be9afb1c Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 7 Jul 2026 13:45:39 -0300 Subject: [PATCH 11/17] Rewrite trace-gen comments to stand alone for a fresh reader MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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). --- prover/src/tables/trace_builder.rs | 27 +++++++++++++-------------- 1 file changed, 13 insertions(+), 14 deletions(-) diff --git a/prover/src/tables/trace_builder.rs b/prover/src/tables/trace_builder.rs index 5eb9cfd0f..c77982489 100644 --- a/prover/src/tables/trace_builder.rs +++ b/prover/src/tables/trace_builder.rs @@ -1981,14 +1981,12 @@ fn collect_bitwise_from_branch(branch_ops: &[BranchOperation]) -> Vec( .map(|chunk| chunk.len().next_power_of_two().max(4) - chunk.len()) .sum(); - // The per-source bitwise collectors are all pure functions of their inputs and the + // The per-source bitwise collectors are all pure functions of their inputs, and the // BITWISE multiplicities are order-independent (they ride a permutation-invariant bus), - // so we collect every source in parallel. The result is byte-identical to the previous - // serial `.extend()` chain regardless of order. + // so every source can be collected in parallel and the per-worker histograms summed in + // any order. // // MUL/DVRM dedup their per-unique bit-gated lookups PER CHIP INSTANCE, so pass the same // chunk size used to split them into instances so multiplicities match the per-instance @@ -3106,10 +3104,11 @@ fn build_traces( // 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; unlike `reduce(identity, ..)` it allocates NO - // per-leaf identity histograms (each a zeroed 80 MiB). Tree-reduce is - // valid because bump/add_ops/merge form a commutative monoid, so the - // result is order-independent and byte-identical to the sequential path. + // 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 From ef13fd1c5fe696ab867f450d6f47e0ce86dae249 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 7 Jul 2026 13:54:59 -0300 Subject: [PATCH 12/17] Clean up stale trace-gen comments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- prover/src/tables/memw_register.rs | 4 ++-- prover/src/tables/trace_builder.rs | 17 +++++++++-------- 2 files changed, 11 insertions(+), 10 deletions(-) diff --git a/prover/src/tables/memw_register.rs b/prover/src/tables/memw_register.rs index 7e5ff7c02..e29c29287 100644 --- a/prover/src/tables/memw_register.rs +++ b/prover/src/tables/memw_register.rs @@ -17,7 +17,7 @@ //! //! ## Column layout (10 columns) //! -//! - `ADDRESS`: Byte (register index 0-31) +//! - `ADDRESS`: Byte (register index 0-255: x0-x31, plus x254/x255) //! - `TIMESTAMP_0`: Word (low 32 bits) //! - `TIMESTAMP_1`: Word (high 32 bits) //! - `VAL_0`: Word (low 32 bits of register value) @@ -53,7 +53,7 @@ use crate::constraints::templates::emit_is_bit; // ========================================================================= pub mod cols { - /// Register index (0-31). CPU sends base_address = 2*reg_index. + /// Register index (0-255: x0-x31, plus x254/x255). CPU sends base_address = 2*reg_index. pub const ADDRESS: usize = 0; /// Timestamp low 32 bits diff --git a/prover/src/tables/trace_builder.rs b/prover/src/tables/trace_builder.rs index c77982489..2e7bd4925 100644 --- a/prover/src/tables/trace_builder.rs +++ b/prover/src/tables/trace_builder.rs @@ -389,16 +389,17 @@ fn classify_memw(op: &MemwOperation) -> MemwRoute { /// Routes each `MemwOperation` into its destination table bucket at CREATION time /// (register fast-path / aligned / general), so the walk fills the three buckets directly -/// and no separate partition pass is needed downstream. Classification order matches the old -/// two-stage partition (register first, then aligned), and push order within each bucket is -/// preserved, so the buckets are byte-identical to `partition`-ing one combined vec. +/// and no separate routing pass is needed downstream. Classification order is register +/// first, then aligned (see [`classify_memw`]), and push order within each bucket is the +/// walk's insertion order — the buckets are fully deterministic, which the per-cell +/// multiplicity counts rely on. /// /// ## Direct-to-column register fill /// /// For the register fast path we do NOT materialize a `Vec`. Ops that route /// to MEMW_R are stored as compact [`RegRow`]s (`register_rows`) and later filled directly /// into the MEMW_R columns. The `aligned` / `general` buckets hold `MemwOperation`s — an op -/// that FAILS `is_register_op` is routed there exactly as the old two-stage partition did. +/// that FAILS `is_register_op` is routed there (aligned if `is_aligned_op`, else general). #[derive(Default)] struct MemwBuckets { /// Compact register rows (filled directly into the MEMW_R columns). @@ -2827,14 +2828,14 @@ fn collect_all_ops( // register state (no zeroizing) so it can seed the next epoch. if is_final { // Route halt ops through the same classifier; they append to the end of their - // buckets, matching the old "append then partition" order. + // buckets. memw.extend_ops(collect_halt_ops(register_state)); } // The walk (`collect_ops_from_cpu`) already routed every MemwOperation into its bucket at - // creation via `MemwBuckets`, so there is no separate partition pass here — the old - // two-`partition` sweep (which moved millions of structs a second time, a bandwidth-bound - // cost) is gone. Order within each bucket matches the old stable partitions → byte-identical. + // creation via `MemwBuckets`, so there is no separate routing pass here: the ops are not + // moved a second time. Order within each bucket is the walk's insertion order, which the + // multiplicity counts depend on being deterministic. let MemwBuckets { register_rows: memw_register_rows, aligned: memw_aligned_ops, From 4799da31824e8338252ef5c88547be8a13650683 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Tue, 7 Jul 2026 14:08:52 -0300 Subject: [PATCH 13/17] Fix stale CPU M7 reference in store MEMW comment 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. --- prover/src/tables/trace_builder.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/prover/src/tables/trace_builder.rs b/prover/src/tables/trace_builder.rs index 2e7bd4925..2d6e426e8 100644 --- a/prover/src/tables/trace_builder.rs +++ b/prover/src/tables/trace_builder.rs @@ -780,8 +780,10 @@ fn collect_store_op_from_cpu(op: &CpuOperation, memory_state: &mut MemoryState) let (old_values, old_timestamps) = memory_state.read_bytes(base_address, 8); // Pack ALL 8 bytes of store_value into value_bytes. - // Bus 14: MEMW Memory Write receiver reconstructs lo32/hi32 via linear combination - // of all 8 bytes. Must match CPU M7 which sends full rv2 as [lo32, hi32]. + // Bus 14: the MEMW Memory Write receiver reconstructs lo32/hi32 via a linear + // combination of all 8 bytes, so it must match the store value the CPU sends + // as [lo32, hi32] on the MEMORY bus (MEMOP) and that this STORE chip forwards + // as the MEMW write (the CPU no longer emits an inline store MEMW — see below). // Bus 16: only positions 0..byte_count participate (controlled by w2/w4/write8 // multiplicities), so extra bytes don't affect memory consistency. let mut value_bytes = [0u32; 8]; From baa9e736c9a640dfdac73454a4874e58fa8fd293 Mon Sep 17 00:00:00 2001 From: Mauro Toscano <12560266+MauroToscano@users.noreply.github.com> Date: Wed, 8 Jul 2026 12:17:52 -0300 Subject: [PATCH 14/17] Parallelize the two dominant bitwise-histogram sources internally (#794) 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); From 9a004b78a74b67994e27093a354bf59774ba309b Mon Sep 17 00:00:00 2001 From: Mauro Toscano <12560266+MauroToscano@users.noreply.github.com> Date: Wed, 8 Jul 2026 12:18:26 -0300 Subject: [PATCH 15/17] Trace-gen cleanup: single MU-column source, simpler push_reg_access, 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). --- prover/src/tables/bitwise.rs | 88 ++++++++++++++++++----------- prover/src/tables/memw_register.rs | 8 ++- prover/src/tables/trace_builder.rs | 91 ++++++++++++++---------------- 3 files changed, 102 insertions(+), 85 deletions(-) diff --git a/prover/src/tables/bitwise.rs b/prover/src/tables/bitwise.rs index 25d7c960d..d280cbdbd 100644 --- a/prover/src/tables/bitwise.rs +++ b/prover/src/tables/bitwise.rs @@ -383,8 +383,9 @@ pub fn generate_bitwise_trace() -> TraceTable usize { +pub(crate) const fn lookup_type_index(t: BitwiseOperationType) -> usize { match t { BitwiseOperationType::Msb8 => 0, BitwiseOperationType::Msb16 => 1, @@ -444,45 +445,64 @@ pub const fn lookup_type_index(t: BitwiseOperationType) -> usize { } } +/// The MU_* multiplicity column for each lookup type, in [`lookup_type_index`] +/// order. This is the single source of truth for the type→column mapping: both +/// the per-op path ([`mu_column`]) and the histogram fill ([`type_mu_column`]) +/// index into this one array. The compile-time block below checks the entries +/// are pairwise distinct, so a duplicate column is a build error rather than a +/// silent overwrite in [`BitwiseHistogram::fill_multiplicities`]. +const MU_COLUMNS: [usize; NUM_LOOKUP_TYPES] = [ + cols::MU_MSB8, // Msb8 + cols::MU_MSB16, // Msb16 + cols::MU_ZERO, // Zero + cols::MU_ARE_BYTES, // AreBytes + cols::MU_IS_HALF, // IsHalf + cols::MU_IS_B20, // IsB20 + cols::MU_HWSL, // Hwsl + cols::MU_BYTE_ALU_AND, // ByteAluAnd + cols::MU_BYTE_ALU_OR, // ByteAluOr + cols::MU_BYTE_ALU_XOR, // ByteAluXor +]; + /// Multiplicity column for a lookup type. Used by the per-op path /// ([`update_multiplicities`]), which is still live production code: continuation /// epochs add their L2G lookups through it on top of the histogram-filled trace. #[inline] -pub const fn mu_column(t: BitwiseOperationType) -> usize { - match t { - BitwiseOperationType::Msb8 => cols::MU_MSB8, - BitwiseOperationType::Msb16 => cols::MU_MSB16, - BitwiseOperationType::Zero => cols::MU_ZERO, - BitwiseOperationType::AreBytes => cols::MU_ARE_BYTES, - BitwiseOperationType::IsHalf => cols::MU_IS_HALF, - BitwiseOperationType::IsB20 => cols::MU_IS_B20, - BitwiseOperationType::Hwsl => cols::MU_HWSL, - BitwiseOperationType::ByteAluAnd => cols::MU_BYTE_ALU_AND, - BitwiseOperationType::ByteAluOr => cols::MU_BYTE_ALU_OR, - BitwiseOperationType::ByteAluXor => cols::MU_BYTE_ALU_XOR, - } +pub(crate) const fn mu_column(t: BitwiseOperationType) -> usize { + MU_COLUMNS[lookup_type_index(t)] } /// Multiplicity column for the histogram lane at dense index `type_idx` /// (inverse of [`lookup_type_index`]). Used by [`BitwiseHistogram::fill_multiplicities`]. /// -/// Defined as `mu_column ∘ ALL`, so it cannot disagree with the authoritative -/// per-type match — no assumption about MU column contiguity or ordering. +/// Reads directly from [`MU_COLUMNS`], the single type→column source of truth. #[inline] const fn type_mu_column(type_idx: usize) -> usize { - mu_column(BitwiseOperationType::ALL[type_idx]) + MU_COLUMNS[type_idx] } -// Compile-time guard: `ALL` must list every lookup type exactly once, in -// `lookup_type_index` order (i.e. it is the exact inverse of that mapping). -// Adding a variant forces the `lookup_type_index` match to be extended -// (exhaustiveness), and this assert then forces `ALL` — and with it -// `NUM_LOOKUP_TYPES` — to follow. A mismatch is a compile error, not a test -// failure: a wrong MU column would silently unbalance the BITWISE bus. +// Compile-time guards on the type↔column bookkeeping. +// +// 1. `ALL` must list every lookup type exactly once, in `lookup_type_index` +// order (i.e. it is the exact inverse of that mapping). Adding a variant +// forces the `lookup_type_index` match to be extended (exhaustiveness), and +// this assert then forces `ALL` — and with it `NUM_LOOKUP_TYPES` — to follow. +// 2. The type→column map is now derived from the single `MU_COLUMNS` array, and +// its entries are checked pairwise distinct (injective). A wrong or duplicated +// MU column would silently unbalance the BITWISE bus, so both are compile +// errors, not test failures. const _: () = { let mut i = 0; while i < NUM_LOOKUP_TYPES { assert!(lookup_type_index(BitwiseOperationType::ALL[i]) == i); + let mut j = i + 1; + while j < NUM_LOOKUP_TYPES { + assert!( + MU_COLUMNS[i] != MU_COLUMNS[j], + "MU_COLUMNS entries must map distinct lookup types to distinct columns" + ); + j += 1; + } i += 1; } }; @@ -499,7 +519,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. -pub struct BitwiseHistogram { +pub(crate) struct BitwiseHistogram { counters: Box<[u64]>, } @@ -508,7 +528,7 @@ impl BitwiseHistogram { // No `Default` impl on purpose: `new()` allocates 80 MiB, so a stray // `..Default::default()` / `#[derive(Default)]` must not silently do that. #[allow(clippy::new_without_default)] - pub fn new() -> Self { + pub(crate) fn new() -> Self { Self { counters: vec![0u64; NUM_ROWS * NUM_LOOKUP_TYPES].into_boxed_slice(), } @@ -516,14 +536,14 @@ impl BitwiseHistogram { /// Increment the counter for one lookup. #[inline] - pub fn bump(&mut self, op: BitwiseOperation) { + pub(crate) fn bump(&mut self, op: BitwiseOperation) { self.bump_n(op, 1); } /// Add `n` occurrences of one lookup in a single step (e.g. CPU padding rows, /// which all send identical all-zero lookups). #[inline] - pub fn bump_n(&mut self, op: BitwiseOperation, n: u64) { + pub(crate) fn bump_n(&mut self, op: BitwiseOperation, n: u64) { let idx = lookup_type_index(op.lookup_type) * NUM_ROWS + row_index(op.x, op.y, op.z); // (x, y) are u8, and row_index debug-asserts z < 16, so in debug builds a // corrupt op fails loudly here. In release an out-of-domain z would NOT @@ -536,14 +556,14 @@ impl BitwiseHistogram { /// Fold a slice of lookups into the histogram. #[inline] - pub fn add_ops(&mut self, ops: &[BitwiseOperation]) { + pub(crate) fn add_ops(&mut self, ops: &[BitwiseOperation]) { for &op in ops { self.bump(op); } } /// Merge another histogram into this one (commutative, order-independent). - pub fn merge(&mut self, other: &BitwiseHistogram) { + pub(crate) fn merge(&mut self, other: &BitwiseHistogram) { for (a, b) in self.counters.iter_mut().zip(other.counters.iter()) { *a += *b; } @@ -558,7 +578,7 @@ impl BitwiseHistogram { /// Callers that layer additional lookups on top (continuation epochs add /// their L2G lookups via `update_multiplicities`, which increments) must do /// so strictly AFTER this fill, never before. - pub fn fill_multiplicities( + pub(crate) fn fill_multiplicities( &self, trace: &mut TraceTable, ) { @@ -594,7 +614,7 @@ impl BitwiseOperationType { /// Every lookup type exactly once, in [`lookup_type_index`] order (the /// compile-time guard next to [`type_mu_column`] enforces this). The array /// length is the single origin of [`NUM_LOOKUP_TYPES`]. - pub const ALL: [Self; 10] = [ + pub(crate) const ALL: [Self; 10] = [ Self::Msb8, Self::Msb16, Self::Zero, diff --git a/prover/src/tables/memw_register.rs b/prover/src/tables/memw_register.rs index e29c29287..fdad08f5e 100644 --- a/prover/src/tables/memw_register.rs +++ b/prover/src/tables/memw_register.rs @@ -193,7 +193,11 @@ impl RegRow { /// /// Thin wrapper over [`generate_memw_register_trace_from_rows`] (via /// [`RegRow::from_memw`]) so there is exactly one MEMW_R column-write sequence. -pub fn generate_memw_register_trace( +/// +/// Test-only: production code fills MEMW_R directly from [`RegRow`]s, so the walk +/// never routes through this `MemwOperation`-based entry point. +#[cfg(test)] +pub(crate) fn generate_memw_register_trace( operations: &[MemwOperation], ) -> TraceTable { let rows: Vec = operations.iter().map(RegRow::from_memw).collect(); @@ -202,7 +206,7 @@ pub fn generate_memw_register_trace( /// The MEMW_R column fill from compact [`RegRow`]s. This is the single source of /// truth for the MEMW_R trace layout; both the walk's direct fast path and the -/// `MemwOperation`-based [`generate_memw_register_trace`] land here. +/// `MemwOperation`-based `generate_memw_register_trace` test wrapper land here. pub(crate) fn generate_memw_register_trace_from_rows( rows: &[RegRow], ) -> TraceTable { diff --git a/prover/src/tables/trace_builder.rs b/prover/src/tables/trace_builder.rs index 73b7675a8..36fd107c3 100644 --- a/prover/src/tables/trace_builder.rs +++ b/prover/src/tables/trace_builder.rs @@ -433,7 +433,8 @@ impl MemwBuckets { } /// Sink for `MemwOperation`s so `collect_register_ops_from_cpu` can feed either a plain -/// `Vec` (tests/scratch paths) or the classifying [`MemwBuckets`] (the walk). +/// `Vec` (the `count_table_lengths` trace-sizing pass) or the classifying +/// [`MemwBuckets`] (the walk). trait MemwSink { fn push_op(&mut self, op: MemwOperation); @@ -443,28 +444,25 @@ trait MemwSink { /// (via the same predicate as `is_register_op`): if the timestamp delta admits the /// op into MEMW_R it fills a compact [`RegRow`] DIRECTLY — no `MemwOperation` is /// built. Only on the (rare) fallback (delta out of IS_HALF range, or upper-limb - /// mismatch) does it call `build_fallback` to materialize the `MemwOperation` and + /// mismatch) does it build the `MemwOperation` (via [`build_reg_fallback`]) and /// route it to the aligned/general bucket exactly as before. /// - /// `reg_addr` is `2 * reg_index`; `[val0,val1]`/`[old0,old1]` are the 32-bit halves; - /// `old_ts` is the (shared) old_timestamp of both words. + /// `reg_addr` is `2 * reg_index`; `val`/`old` are the two 32-bit halves of the new + /// and previous register words; `old_ts` is the (shared) old_timestamp of both words. #[inline] - #[allow(clippy::too_many_arguments)] fn push_reg_access( &mut self, reg_addr: u64, - val0: u32, - val1: u32, - old0: u32, - old1: u32, + val: [u32; 2], + old: [u32; 2], timestamp: u64, old_ts: u64, is_read: bool, - build_fallback: impl FnOnce() -> MemwOperation, ) { // Default impl (plain Vec): register accesses are still ordinary MemwOperations. - let _ = (reg_addr, val0, val1, old0, old1, timestamp, old_ts, is_read); - self.push_op(build_fallback()); + self.push_op(build_reg_fallback( + reg_addr, val, old, timestamp, old_ts, is_read, + )); } } impl MemwSink for Vec { @@ -480,34 +478,49 @@ impl MemwSink for MemwBuckets { } #[inline] - #[allow(clippy::too_many_arguments)] fn push_reg_access( &mut self, reg_addr: u64, - val0: u32, - val1: u32, - old0: u32, - old1: u32, + val: [u32; 2], + old: [u32; 2], timestamp: u64, old_ts: u64, is_read: bool, - build_fallback: impl FnOnce() -> MemwOperation, ) { // Mirror `is_register_op` for a width-2 register access whose two words share // `old_ts` (always true here by construction). If it passes, fill a RegRow // directly; otherwise fall back to the general/aligned MemwOperation path. if reg_ts_delta_in_range(timestamp, old_ts) { self.register_rows.push(RegRow::new( - reg_addr, timestamp, val0, val1, old0, old1, old_ts, is_read, + reg_addr, timestamp, val[0], val[1], old[0], old[1], old_ts, is_read, )); } else { - let op = build_fallback(); + let op = build_reg_fallback(reg_addr, val, old, timestamp, old_ts, is_read); debug_assert!(!is_register_op(&op), "reg fallback must not be MEMW_R"); self.push(op); } } } +/// Materialize the aligned/general `MemwOperation` for a register access that does +/// NOT fit the MEMW_R fast path. Register values pack as `[lo, hi, 0, …]` (see +/// [`pack_register_value`]) and both words share `old_ts`, so this rebuilds exactly +/// the op the fast-path callers would otherwise have routed to the buckets. +fn build_reg_fallback( + reg_addr: u64, + val: [u32; 2], + old: [u32; 2], + timestamp: u64, + old_ts: u64, + is_read: bool, +) -> MemwOperation { + let value = [val[0], val[1], 0, 0, 0, 0, 0, 0]; + let old_value = [old[0], old[1], 0, 0, 0, 0, 0, 0]; + let old_timestamps = [old_ts, old_ts, 0, 0, 0, 0, 0, 0]; + MemwOperation::new(true, reg_addr, value, timestamp, 2, is_read) + .with_old(old_value, old_timestamps) +} + /// Collects all derived operations from CPU operations in a single pass. /// /// This includes: @@ -957,22 +970,16 @@ fn collect_register_ops_from_cpu( register_state.read(d.rs1) }; let ts = op.timestamp; - // Direct fast path: fill a RegRow when routing to MEMW_R; the closure rebuilds - // the identical MemwOperation only on the (rare) general/aligned fallback. + // Direct fast path: fill a RegRow when routing to MEMW_R; push_reg_access + // rebuilds the identical MemwOperation only on the (rare) general/aligned + // fallback. Reads leave the value unchanged, so old == new here. memw_ops.push_reg_access( reg_addr, - reg_value[0], - reg_value[1], - reg_value[0], - reg_value[1], + [reg_value[0], reg_value[1]], + [reg_value[0], reg_value[1]], ts, old_ts, true, - || { - let old_timestamps = [old_ts, old_ts, 0, 0, 0, 0, 0, 0]; - MemwOperation::new(true, reg_addr, reg_value, ts, 2, true) - .with_old(reg_value, old_timestamps) - }, ); if d.rs1 == 255 { register_state.write_pc(op.rv1, op.timestamp); @@ -989,18 +996,11 @@ fn collect_register_ops_from_cpu( let ts = op.timestamp + 1; memw_ops.push_reg_access( reg_addr, - reg_value[0], - reg_value[1], - reg_value[0], - reg_value[1], + [reg_value[0], reg_value[1]], + [reg_value[0], reg_value[1]], ts, old_ts, true, - || { - let old_timestamps = [old_ts, old_ts, 0, 0, 0, 0, 0, 0]; - MemwOperation::new(true, reg_addr, reg_value, ts, 2, true) - .with_old(reg_value, old_timestamps) - }, ); register_state.write(d.rs2, op.rv2, op.timestamp + 1); } @@ -1014,18 +1014,11 @@ fn collect_register_ops_from_cpu( let ts = op.timestamp + 2; memw_ops.push_reg_access( reg_addr, - reg_value[0], - reg_value[1], - old_value[0], - old_value[1], + [reg_value[0], reg_value[1]], + [old_value[0], old_value[1]], ts, old_ts, false, - || { - let old_timestamps = [old_ts, old_ts, 0, 0, 0, 0, 0, 0]; - MemwOperation::new(true, reg_addr, reg_value, ts, 2, false) - .with_old(old_value, old_timestamps) - }, ); register_state.write(d.rd, op.rvd, op.timestamp + 2); } From af9751c9ad52e577e4bb092dc8f22ef95312653d Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Wed, 8 Jul 2026 12:29:54 -0300 Subject: [PATCH 16/17] opt --- prover/src/tables/bitwise.rs | 2 +- prover/src/tables/branch.rs | 4 ++-- prover/src/tables/bytewise.rs | 4 ++-- prover/src/tables/commit.rs | 2 +- prover/src/tables/cpu.rs | 4 ++-- prover/src/tables/cpu32.rs | 2 +- prover/src/tables/decode.rs | 4 ++-- prover/src/tables/dvrm.rs | 4 ++-- prover/src/tables/ec_scalar.rs | 2 +- prover/src/tables/ecdas.rs | 2 +- prover/src/tables/ecsm.rs | 2 +- prover/src/tables/eq.rs | 4 ++-- prover/src/tables/global_memory.rs | 2 +- prover/src/tables/halt.rs | 8 +++++-- prover/src/tables/keccak.rs | 2 +- prover/src/tables/keccak_rc.rs | 2 +- prover/src/tables/keccak_rnd.rs | 2 +- prover/src/tables/load.rs | 2 +- prover/src/tables/local_to_global.rs | 2 +- prover/src/tables/lt.rs | 4 ++-- prover/src/tables/memw.rs | 4 ++-- prover/src/tables/memw_aligned.rs | 4 ++-- prover/src/tables/memw_register.rs | 4 ++-- prover/src/tables/mul.rs | 4 ++-- prover/src/tables/page.rs | 6 ++--- prover/src/tables/register.rs | 12 +++++----- prover/src/tables/shift.rs | 4 ++-- prover/src/tables/store.rs | 2 +- prover/src/tables/types.rs | 34 ++++++++++++++++++++++++++++ 29 files changed, 86 insertions(+), 48 deletions(-) diff --git a/prover/src/tables/bitwise.rs b/prover/src/tables/bitwise.rs index d280cbdbd..45bddb636 100644 --- a/prover/src/tables/bitwise.rs +++ b/prover/src/tables/bitwise.rs @@ -342,7 +342,7 @@ pub fn preprocessed_commitment(options: &ProofOptions) -> Commitment { /// to zero and will be updated when other tables send lookups. pub fn generate_bitwise_trace() -> TraceTable { let mut trace = TraceTable::new_main( - vec![FE::zero(); NUM_ROWS * cols::NUM_COLUMNS], + crate::tables::types::zeroed_fe_vec(NUM_ROWS * cols::NUM_COLUMNS), cols::NUM_COLUMNS, 1, ); diff --git a/prover/src/tables/branch.rs b/prover/src/tables/branch.rs index d4baf10c5..b5bfe83b6 100644 --- a/prover/src/tables/branch.rs +++ b/prover/src/tables/branch.rs @@ -32,7 +32,7 @@ use stark::trace::TraceTable; use std::collections::HashMap; -use super::types::{BusId, FE, GoldilocksExtension, GoldilocksField, SHIFT_16, VmTable, alu_op}; +use super::types::{BusId, GoldilocksExtension, GoldilocksField, SHIFT_16, VmTable, alu_op}; // ========================================================================= // Column indices for BRANCH table @@ -166,7 +166,7 @@ pub fn generate_branch_trace( let unique_ops: Vec<_> = op_map.into_iter().collect(); let num_rows = unique_ops.len().next_power_of_two().max(4); let mut trace = TraceTable::new_main( - vec![FE::zero(); num_rows * cols::NUM_COLUMNS], + crate::tables::types::zeroed_fe_vec(num_rows * cols::NUM_COLUMNS), cols::NUM_COLUMNS, 1, ); diff --git a/prover/src/tables/bytewise.rs b/prover/src/tables/bytewise.rs index 82d7c8772..2808365c6 100644 --- a/prover/src/tables/bytewise.rs +++ b/prover/src/tables/bytewise.rs @@ -19,7 +19,7 @@ use stark::lookup::{BusInteraction, BusValue, Multiplicity, Packing}; use stark::trace::TraceTable; -use super::types::{BusId, FE, GoldilocksExtension, GoldilocksField, VmTable, alu_op}; +use super::types::{BusId, GoldilocksExtension, GoldilocksField, VmTable, alu_op}; // ========================================================================= // Column indices for BYTEWISE table @@ -107,7 +107,7 @@ pub fn generate_bytewise_trace( let unique_ops: Vec<_> = op_map.into_iter().collect(); let num_rows = unique_ops.len().next_power_of_two().max(4); let mut trace = TraceTable::new_main( - vec![FE::zero(); num_rows * cols::NUM_COLUMNS], + crate::tables::types::zeroed_fe_vec(num_rows * cols::NUM_COLUMNS), cols::NUM_COLUMNS, 1, ); diff --git a/prover/src/tables/commit.rs b/prover/src/tables/commit.rs index cd1ca264b..4660c7fb0 100644 --- a/prover/src/tables/commit.rs +++ b/prover/src/tables/commit.rs @@ -163,7 +163,7 @@ pub fn generate_commit_trace( let n = ops.len(); let num_rows = n.next_power_of_two().max(4); let mut trace = TraceTable::new_main( - vec![FE::zero(); num_rows * cols::NUM_COLUMNS], + crate::tables::types::zeroed_fe_vec(num_rows * cols::NUM_COLUMNS), cols::NUM_COLUMNS, 1, ); diff --git a/prover/src/tables/cpu.rs b/prover/src/tables/cpu.rs index 42d197942..781bb02b0 100644 --- a/prover/src/tables/cpu.rs +++ b/prover/src/tables/cpu.rs @@ -24,7 +24,7 @@ //! JALR bit (the memory-width bits are 0), so `mem_flags ∈ {0,1} = JALR` and the //! `mem_flags` column is used directly as `JALR` wherever it is gated by `BRANCH`. -use super::types::{BusId, DecodeEntry, FE, GoldilocksExtension, GoldilocksField, VmTable, alu_op}; +use super::types::{BusId, DecodeEntry, GoldilocksExtension, GoldilocksField, VmTable, alu_op}; use crate::Error; use executor::vm::{ instruction::{decoding::Instruction, execution::SyscallNumbers}, @@ -440,7 +440,7 @@ pub fn generate_cpu_trace( let n = operations.len(); let num_rows = n.next_power_of_two().max(4); let mut trace = TraceTable::new_main( - vec![FE::zero(); num_rows * cols::NUM_COLUMNS], + crate::tables::types::zeroed_fe_vec(num_rows * cols::NUM_COLUMNS), cols::NUM_COLUMNS, 1, ); diff --git a/prover/src/tables/cpu32.rs b/prover/src/tables/cpu32.rs index e0931ff3a..8b1bf86d6 100644 --- a/prover/src/tables/cpu32.rs +++ b/prover/src/tables/cpu32.rs @@ -197,7 +197,7 @@ pub fn generate_cpu32_trace( ) -> TraceTable { let num_rows = operations.len().next_power_of_two().max(4); let mut trace = TraceTable::new_main( - vec![FE::zero(); num_rows * cols::NUM_COLUMNS], + crate::tables::types::zeroed_fe_vec(num_rows * cols::NUM_COLUMNS), cols::NUM_COLUMNS, 1, ); diff --git a/prover/src/tables/decode.rs b/prover/src/tables/decode.rs index 509f86991..bfd1ddb90 100644 --- a/prover/src/tables/decode.rs +++ b/prover/src/tables/decode.rs @@ -128,7 +128,7 @@ pub fn generate_decode_trace( let num_entries = entries.len() + 1; let num_rows = num_entries.next_power_of_two().max(2); let mut trace = TraceTable::new_main( - vec![FE::zero(); num_rows * cols::NUM_COLUMNS], + crate::tables::types::zeroed_fe_vec(num_rows * cols::NUM_COLUMNS), cols::NUM_COLUMNS, 1, ); @@ -393,7 +393,7 @@ fn build_decode_table( let num_entries = entries.len() + 1; let num_rows = num_entries.next_power_of_two().max(2); let mut trace = TraceTable::new_main( - vec![FE::zero(); num_rows * cols::NUM_COLUMNS], + crate::tables::types::zeroed_fe_vec(num_rows * cols::NUM_COLUMNS), cols::NUM_COLUMNS, 1, ); diff --git a/prover/src/tables/dvrm.rs b/prover/src/tables/dvrm.rs index 032963c82..9f979742b 100644 --- a/prover/src/tables/dvrm.rs +++ b/prover/src/tables/dvrm.rs @@ -35,7 +35,7 @@ use stark::trace::TraceTable; use std::collections::HashMap; use super::types::{ - BusId, FE, GoldilocksExtension, GoldilocksField, NEG_INV_2_16, NEG_INV_2_32, NEG_INV_2_48, + BusId, GoldilocksExtension, GoldilocksField, NEG_INV_2_16, NEG_INV_2_32, NEG_INV_2_48, NEG_INV_2_64, SHIFT_16, VmTable, alu_op, }; @@ -298,7 +298,7 @@ pub fn generate_dvrm_trace( let unique_ops: Vec<_> = op_map.into_iter().collect(); let num_rows = unique_ops.len().next_power_of_two().max(4); let mut trace = TraceTable::new_main( - vec![FE::zero(); num_rows * cols::NUM_COLUMNS], + crate::tables::types::zeroed_fe_vec(num_rows * cols::NUM_COLUMNS), cols::NUM_COLUMNS, 1, ); diff --git a/prover/src/tables/ec_scalar.rs b/prover/src/tables/ec_scalar.rs index 5589a14e5..08f797f47 100644 --- a/prover/src/tables/ec_scalar.rs +++ b/prover/src/tables/ec_scalar.rs @@ -83,7 +83,7 @@ pub fn generate_ec_scalar_trace( let n = ops.len(); let num_rows = n.next_power_of_two().max(4); let mut trace = TraceTable::new_main( - vec![FE::zero(); num_rows * cols::NUM_COLUMNS], + crate::tables::types::zeroed_fe_vec(num_rows * cols::NUM_COLUMNS), cols::NUM_COLUMNS, 1, ); diff --git a/prover/src/tables/ecdas.rs b/prover/src/tables/ecdas.rs index 26bbe44e4..2ffde164d 100644 --- a/prover/src/tables/ecdas.rs +++ b/prover/src/tables/ecdas.rs @@ -94,7 +94,7 @@ pub fn generate_ecdas_trace( let n = ops.len(); let num_rows = n.next_power_of_two().max(4); let mut trace = TraceTable::new_main( - vec![FE::zero(); num_rows * cols::NUM_COLUMNS], + crate::tables::types::zeroed_fe_vec(num_rows * cols::NUM_COLUMNS), cols::NUM_COLUMNS, 1, ); diff --git a/prover/src/tables/ecsm.rs b/prover/src/tables/ecsm.rs index 0dba13910..846c31812 100644 --- a/prover/src/tables/ecsm.rs +++ b/prover/src/tables/ecsm.rs @@ -148,7 +148,7 @@ pub fn generate_ecsm_trace( let n = ops.len(); let num_rows = n.next_power_of_two().max(4); let mut trace = TraceTable::new_main( - vec![FE::zero(); num_rows * cols::NUM_COLUMNS], + crate::tables::types::zeroed_fe_vec(num_rows * cols::NUM_COLUMNS), cols::NUM_COLUMNS, 1, ); diff --git a/prover/src/tables/eq.rs b/prover/src/tables/eq.rs index 117d8426b..0f20ca695 100644 --- a/prover/src/tables/eq.rs +++ b/prover/src/tables/eq.rs @@ -26,7 +26,7 @@ use stark::trace::TraceTable; use stark::constraints::builder::{ConstraintBuilder, ConstraintSet}; -use super::types::{BusId, FE, GoldilocksExtension, GoldilocksField, VmTable, alu_op}; +use super::types::{BusId, GoldilocksExtension, GoldilocksField, VmTable, alu_op}; use crate::constraints::templates::{AddOperand, emit_add_pair, emit_is_bit}; // ========================================================================= @@ -128,7 +128,7 @@ pub fn generate_eq_trace( let unique_ops: Vec<_> = op_map.into_iter().collect(); let num_rows = unique_ops.len().next_power_of_two().max(4); let mut trace = TraceTable::new_main( - vec![FE::zero(); num_rows * cols::NUM_COLUMNS], + crate::tables::types::zeroed_fe_vec(num_rows * cols::NUM_COLUMNS), cols::NUM_COLUMNS, 1, ); diff --git a/prover/src/tables/global_memory.rs b/prover/src/tables/global_memory.rs index de6d95d7d..d7bc2ebe9 100644 --- a/prover/src/tables/global_memory.rs +++ b/prover/src/tables/global_memory.rs @@ -123,7 +123,7 @@ pub fn generate_global_trace( ); let num_rows = page_size; // One row per byte in the page - let mut data = vec![FE::zero(); num_rows * cols::NUM_COLUMNS]; + let mut data = crate::tables::types::zeroed_fe_vec(num_rows * cols::NUM_COLUMNS); for offset in 0..page_size { let byte_addr = page_base + (offset as u64); diff --git a/prover/src/tables/halt.rs b/prover/src/tables/halt.rs index 44bbf26cb..675e10170 100644 --- a/prover/src/tables/halt.rs +++ b/prover/src/tables/halt.rs @@ -30,7 +30,7 @@ use stark::lookup::{BusInteraction, BusValue, LinearTerm, Multiplicity, Packing}; use stark::trace::TraceTable; -use super::types::{BusId, FE, GoldilocksExtension, GoldilocksField, VmTable}; +use super::types::{BusId, GoldilocksExtension, GoldilocksField, VmTable}; // ========================================================================= // Column indices for HALT table @@ -72,7 +72,11 @@ pub fn generate_halt_trace( timestamp <= u32::MAX as u64, "HALT timestamp {timestamp} exceeds u32 range" ); - let mut trace = TraceTable::new_main(vec![FE::zero(); cols::NUM_COLUMNS], cols::NUM_COLUMNS, 1); + let mut trace = TraceTable::new_main( + crate::tables::types::zeroed_fe_vec(cols::NUM_COLUMNS), + cols::NUM_COLUMNS, + 1, + ); let table = &mut trace.main_table; table.set_dword_wl(0, cols::TIMESTAMP_0, timestamp); diff --git a/prover/src/tables/keccak.rs b/prover/src/tables/keccak.rs index 832869012..9626b7e3b 100644 --- a/prover/src/tables/keccak.rs +++ b/prover/src/tables/keccak.rs @@ -98,7 +98,7 @@ pub fn generate_keccak_trace( let n = ops.len(); let num_rows = n.next_power_of_two().max(4); let mut trace = TraceTable::new_main( - vec![FE::zero(); num_rows * cols::NUM_COLUMNS], + crate::tables::types::zeroed_fe_vec(num_rows * cols::NUM_COLUMNS), cols::NUM_COLUMNS, 1, ); diff --git a/prover/src/tables/keccak_rc.rs b/prover/src/tables/keccak_rc.rs index f9f0d1cc4..142b5bdde 100644 --- a/prover/src/tables/keccak_rc.rs +++ b/prover/src/tables/keccak_rc.rs @@ -190,7 +190,7 @@ pub fn preprocessed_commitment(options: &ProofOptions) -> Commitment { /// updated via `update_multiplicities` after all round-chip lookups are known. pub fn generate_keccak_rc_trace() -> TraceTable { let mut trace = TraceTable::new_main( - vec![FE::zero(); NUM_ROWS * cols::NUM_COLUMNS], + crate::tables::types::zeroed_fe_vec(NUM_ROWS * cols::NUM_COLUMNS), cols::NUM_COLUMNS, 1, ); diff --git a/prover/src/tables/keccak_rnd.rs b/prover/src/tables/keccak_rnd.rs index 30a50e0b2..afc5dee3a 100644 --- a/prover/src/tables/keccak_rnd.rs +++ b/prover/src/tables/keccak_rnd.rs @@ -244,7 +244,7 @@ pub fn generate_keccak_rnd_trace( ) -> TraceTable { let n_rows = (ops.len() * 24).next_power_of_two().max(4); let mut trace = TraceTable::new_main( - vec![FE::zero(); n_rows * cols::NUM_COLUMNS], + crate::tables::types::zeroed_fe_vec(n_rows * cols::NUM_COLUMNS), cols::NUM_COLUMNS, 1, ); diff --git a/prover/src/tables/load.rs b/prover/src/tables/load.rs index c2bf389dc..1da3cf564 100644 --- a/prover/src/tables/load.rs +++ b/prover/src/tables/load.rs @@ -180,7 +180,7 @@ pub fn generate_load_trace( ) -> TraceTable { let num_rows = operations.len().next_power_of_two().max(4); let mut trace = TraceTable::new_main( - vec![FE::zero(); num_rows * cols::NUM_COLUMNS], + crate::tables::types::zeroed_fe_vec(num_rows * cols::NUM_COLUMNS), cols::NUM_COLUMNS, 1, ); diff --git a/prover/src/tables/local_to_global.rs b/prover/src/tables/local_to_global.rs index 668dc1353..3de72f62a 100644 --- a/prover/src/tables/local_to_global.rs +++ b/prover/src/tables/local_to_global.rs @@ -267,7 +267,7 @@ pub fn generate_local_to_global_trace( boundaries: &[CellBoundary], ) -> TraceTable { let num_rows = boundaries.len().next_power_of_two().max(1); - let mut data = vec![FE::zero(); num_rows * cols::NUM_COLUMNS]; + let mut data = crate::tables::types::zeroed_fe_vec(num_rows * cols::NUM_COLUMNS); for (row, b) in boundaries.iter().enumerate() { let base = row * cols::NUM_COLUMNS; diff --git a/prover/src/tables/lt.rs b/prover/src/tables/lt.rs index a68191f37..86e88a9f7 100644 --- a/prover/src/tables/lt.rs +++ b/prover/src/tables/lt.rs @@ -31,7 +31,7 @@ use stark::trace::TraceTable; use std::collections::HashMap; -use super::types::{BusId, FE, GoldilocksExtension, GoldilocksField, SHIFT_16, VmTable, alu_op}; +use super::types::{BusId, GoldilocksExtension, GoldilocksField, SHIFT_16, VmTable, alu_op}; // ========================================================================= // Column indices for LT table @@ -168,7 +168,7 @@ pub fn generate_lt_trace( let unique_ops: Vec<_> = op_map.into_iter().collect(); let num_rows = unique_ops.len().next_power_of_two().max(4); let mut trace = TraceTable::new_main( - vec![FE::zero(); num_rows * cols::NUM_COLUMNS], + crate::tables::types::zeroed_fe_vec(num_rows * cols::NUM_COLUMNS), cols::NUM_COLUMNS, 1, ); diff --git a/prover/src/tables/memw.rs b/prover/src/tables/memw.rs index f4af949e1..338b85467 100644 --- a/prover/src/tables/memw.rs +++ b/prover/src/tables/memw.rs @@ -35,7 +35,7 @@ use stark::trace::TraceTable; use stark::constraints::builder::{ConstraintBuilder, ConstraintSet}; -use super::types::{BusId, FE, GoldilocksExtension, GoldilocksField, VmTable, alu_op}; +use super::types::{BusId, GoldilocksExtension, GoldilocksField, VmTable, alu_op}; use crate::constraints::templates::emit_is_bit; /// Maximum number of rows per MEMW table chunk. @@ -180,7 +180,7 @@ pub fn generate_memw_trace( ) -> TraceTable { let num_rows = operations.len().next_power_of_two().max(4); let mut trace = TraceTable::new_main( - vec![FE::zero(); num_rows * cols::NUM_COLUMNS], + crate::tables::types::zeroed_fe_vec(num_rows * cols::NUM_COLUMNS), cols::NUM_COLUMNS, 1, ); diff --git a/prover/src/tables/memw_aligned.rs b/prover/src/tables/memw_aligned.rs index 0e262efb7..0853bf5ff 100644 --- a/prover/src/tables/memw_aligned.rs +++ b/prover/src/tables/memw_aligned.rs @@ -42,7 +42,7 @@ use stark::trace::TraceTable; use stark::constraints::builder::{ConstraintBuilder, ConstraintSet}; use super::memw::MemwOperation; -use super::types::{BusId, FE, GoldilocksExtension, GoldilocksField, VmTable, alu_op}; +use super::types::{BusId, GoldilocksExtension, GoldilocksField, VmTable, alu_op}; use crate::constraints::templates::emit_is_bit; /// Maximum number of rows per MEMW_A table chunk. @@ -95,7 +95,7 @@ pub fn generate_memw_aligned_trace( ) -> TraceTable { let num_rows = operations.len().next_power_of_two().max(4); let mut trace = TraceTable::new_main( - vec![FE::zero(); num_rows * cols::NUM_COLUMNS], + crate::tables::types::zeroed_fe_vec(num_rows * cols::NUM_COLUMNS), cols::NUM_COLUMNS, 1, ); diff --git a/prover/src/tables/memw_register.rs b/prover/src/tables/memw_register.rs index fdad08f5e..590a55100 100644 --- a/prover/src/tables/memw_register.rs +++ b/prover/src/tables/memw_register.rs @@ -45,7 +45,7 @@ use stark::constraints::builder::{ConstraintBuilder, ConstraintSet}; use super::bitwise::{BitwiseHistogram, BitwiseOperation, BitwiseOperationType}; use super::memw::MemwOperation; -use super::types::{BusId, FE, GoldilocksExtension, GoldilocksField, VmTable}; +use super::types::{BusId, GoldilocksExtension, GoldilocksField, VmTable}; use crate::constraints::templates::emit_is_bit; // ========================================================================= @@ -212,7 +212,7 @@ pub(crate) fn generate_memw_register_trace_from_rows( ) -> TraceTable { let num_rows = rows.len().next_power_of_two().max(4); let mut trace = TraceTable::new_main( - vec![FE::zero(); num_rows * cols::NUM_COLUMNS], + crate::tables::types::zeroed_fe_vec(num_rows * cols::NUM_COLUMNS), cols::NUM_COLUMNS, 1, ); diff --git a/prover/src/tables/mul.rs b/prover/src/tables/mul.rs index 2f0fa1d0e..181fba514 100644 --- a/prover/src/tables/mul.rs +++ b/prover/src/tables/mul.rs @@ -36,7 +36,7 @@ use stark::trace::TraceTable; use std::collections::HashMap; use super::types::{ - BusId, FE, GoldilocksExtension, GoldilocksField, INV_2_32, INV_2_64, INV_2_96, INV_2_128, + BusId, GoldilocksExtension, GoldilocksField, INV_2_32, INV_2_64, INV_2_96, INV_2_128, NEG_INV_2_16, NEG_INV_2_32, NEG_INV_2_48, NEG_INV_2_64, NEG_INV_2_80, NEG_INV_2_96, NEG_INV_2_112, NEG_INV_2_128, SHIFT_16, VmTable, alu_op, }; @@ -306,7 +306,7 @@ pub fn generate_mul_trace( let unique_ops: Vec<_> = op_map.into_iter().collect(); let num_rows = unique_ops.len().next_power_of_two().max(4); let mut trace = TraceTable::new_main( - vec![FE::zero(); num_rows * cols::NUM_COLUMNS], + crate::tables::types::zeroed_fe_vec(num_rows * cols::NUM_COLUMNS), cols::NUM_COLUMNS, 1, ); diff --git a/prover/src/tables/page.rs b/prover/src/tables/page.rs index 5cba30435..18ce6b52b 100644 --- a/prover/src/tables/page.rs +++ b/prover/src/tables/page.rs @@ -249,7 +249,7 @@ pub fn generate_page_trace( let num_rows = page_size; // One row per byte in the page let mut trace = TraceTable::new_main( - vec![FE::zero(); num_rows * cols::NUM_COLUMNS], + crate::tables::types::zeroed_fe_vec(num_rows * cols::NUM_COLUMNS), cols::NUM_COLUMNS, 1, ); @@ -361,8 +361,8 @@ pub fn compute_precomputed_commitment(config: &PageConfig, options: &ProofOption // bytes loaded from the binary. Either way the column is fully determined // before execution, so the verifier can check it against a preprocessed // commitment instead of including it in the main trace. - let mut offset_col = vec![FE::zero(); num_rows]; - let mut init_col = vec![FE::zero(); num_rows]; + let mut offset_col = crate::tables::types::zeroed_fe_vec(num_rows); + let mut init_col = crate::tables::types::zeroed_fe_vec(num_rows); for i in 0..page_size { offset_col[i] = FE::from(i as u64); diff --git a/prover/src/tables/register.rs b/prover/src/tables/register.rs index 46c675b65..4da1f5efe 100644 --- a/prover/src/tables/register.rs +++ b/prover/src/tables/register.rs @@ -220,7 +220,7 @@ pub fn generate_register_trace( ) -> TraceTable { let num_rows = NUM_REGISTER_ADDRESSES.next_power_of_two(); let mut trace = TraceTable::new_main( - vec![FE::zero(); num_rows * cols::NUM_COLUMNS], + crate::tables::types::zeroed_fe_vec(num_rows * cols::NUM_COLUMNS), cols::NUM_COLUMNS, 1, ); @@ -281,8 +281,8 @@ pub fn compute_precomputed_commitment(options: &ProofOptions, init: &[u32]) -> C let num_rows = NUM_REGISTER_ADDRESSES.next_power_of_two(); let addr_list = register_word_address_list(); - let mut offset_col = vec![FE::zero(); num_rows]; - let mut init_col = vec![FE::zero(); num_rows]; + let mut offset_col = crate::tables::types::zeroed_fe_vec(num_rows); + let mut init_col = crate::tables::types::zeroed_fe_vec(num_rows); for i in 0..NUM_REGISTER_ADDRESSES { offset_col[i] = FE::from(addr_list[i]); @@ -308,9 +308,9 @@ pub fn compute_precomputed_commitment_with_fini( let num_rows = NUM_REGISTER_ADDRESSES.next_power_of_two(); let addr_list = register_word_address_list(); - let mut offset_col = vec![FE::zero(); num_rows]; - let mut init_col = vec![FE::zero(); num_rows]; - let mut fini_col = vec![FE::zero(); num_rows]; + let mut offset_col = crate::tables::types::zeroed_fe_vec(num_rows); + let mut init_col = crate::tables::types::zeroed_fe_vec(num_rows); + let mut fini_col = crate::tables::types::zeroed_fe_vec(num_rows); for i in 0..NUM_REGISTER_ADDRESSES { offset_col[i] = FE::from(addr_list[i]); diff --git a/prover/src/tables/shift.rs b/prover/src/tables/shift.rs index 77a8ae32a..5ac5a393f 100644 --- a/prover/src/tables/shift.rs +++ b/prover/src/tables/shift.rs @@ -21,7 +21,7 @@ use stark::lookup::{BusInteraction, BusValue, LinearTerm, Multiplicity, Packing}; use stark::trace::TraceTable; -use super::types::{BusId, FE, GoldilocksExtension, GoldilocksField, SHIFT_16, VmTable, alu_op}; +use super::types::{BusId, GoldilocksExtension, GoldilocksField, SHIFT_16, VmTable, alu_op}; // ========================================================================= // Column indices @@ -357,7 +357,7 @@ pub fn generate_shift_trace( // Spec declares μ: Bit. let num_rows = operations.len().next_power_of_two().max(4); let mut trace = TraceTable::new_main( - vec![FE::zero(); num_rows * cols::NUM_COLUMNS], + crate::tables::types::zeroed_fe_vec(num_rows * cols::NUM_COLUMNS), cols::NUM_COLUMNS, 1, ); diff --git a/prover/src/tables/store.rs b/prover/src/tables/store.rs index c1dfc937a..ac30832e1 100644 --- a/prover/src/tables/store.rs +++ b/prover/src/tables/store.rs @@ -97,7 +97,7 @@ pub fn generate_store_trace( ) -> TraceTable { let num_rows = operations.len().next_power_of_two().max(4); let mut trace = TraceTable::new_main( - vec![FE::zero(); num_rows * cols::NUM_COLUMNS], + crate::tables::types::zeroed_fe_vec(num_rows * cols::NUM_COLUMNS), cols::NUM_COLUMNS, 1, ); diff --git a/prover/src/tables/types.rs b/prover/src/tables/types.rs index 71c83284a..3581dc8ec 100644 --- a/prover/src/tables/types.rs +++ b/prover/src/tables/types.rs @@ -43,6 +43,40 @@ pub fn dword_wl(x: u64) -> [FE; 2] { [FE::from(x & 0xFFFF_FFFF), FE::from(x >> 32)] } +/// Allocate a zeroed `Vec` via the allocator's calloc path (demand-zeroed OS +/// pages, touched lazily on first write) instead of the element-wise clone loop +/// that `vec![FE::zero(); n]` runs. Trace fill is memory-bandwidth-bound, so +/// skipping the eager zeroing sweep of the (multi-GB) trace is a real win. +/// +/// Sound because `FE` is `#[repr(transparent)]` over its `u64` base value +/// (Goldilocks has no Montgomery form), so `FE`'s canonical zero is the all-zero +/// bit pattern — identical to `0u64`. The `zeroed_fe_vec_matches_fe_zero` test +/// guards this invariant. Technique borrowed from SP1's `zeroed_f_vec`. +#[inline] +pub fn zeroed_fe_vec(len: usize) -> Vec { + const _: () = assert!(core::mem::size_of::() == core::mem::size_of::()); + const _: () = assert!(core::mem::align_of::() == core::mem::align_of::()); + let zeros: Vec = vec![0u64; len]; + // SAFETY: `FE` is `#[repr(transparent)]` over `u64` with identical size and + // alignment (asserted above), and `0u64` is exactly `FE::zero()`'s bit + // pattern (no Montgomery form), so reinterpreting the buffer is valid. + unsafe { core::mem::transmute::, Vec>(zeros) } +} + +#[cfg(test)] +mod zeroed_fe_vec_tests { + use super::*; + + /// Guards the `zeroed_fe_vec` invariant: a calloc'd all-zero buffer + /// reinterpreted as `Vec` must equal an element-wise `FE::zero()` fill. + #[test] + fn zeroed_fe_vec_matches_fe_zero() { + for len in [0usize, 1, 7, 64, 1024] { + assert_eq!(zeroed_fe_vec(len), vec![FE::zero(); len], "len={len}"); + } + } +} + /// Decompose a `u64` into its four little-endian 16-bit limbs as field elements: /// `[x[0..16], x[16..32], x[32..48], x[48..64]]` (the `DWordHL` column encoding). #[inline] From 5fa2f76766d0dc60ce9f34d9e956e79c0fa5c563 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Wed, 8 Jul 2026 17:02:56 -0300 Subject: [PATCH 17/17] zeroed_fe_vec: reinterpret via from_raw_parts, not Vec transmute MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit zeroed_fe_vec built its Vec with mem::transmute::, Vec>. 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. --- prover/src/tables/types.rs | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/prover/src/tables/types.rs b/prover/src/tables/types.rs index 3581dc8ec..85eaca17a 100644 --- a/prover/src/tables/types.rs +++ b/prover/src/tables/types.rs @@ -57,10 +57,22 @@ pub fn zeroed_fe_vec(len: usize) -> Vec { const _: () = assert!(core::mem::size_of::() == core::mem::size_of::()); const _: () = assert!(core::mem::align_of::() == core::mem::align_of::()); let zeros: Vec = vec![0u64; len]; + // Reinterpret the buffer as `Vec` via its raw parts rather than + // `mem::transmute::, Vec>`. `Vec`'s field layout is unspecified + // and may depend on its element type, so transmuting one `Vec` to another + // relies on that unspecified layout (the std `mem::transmute` docs call this + // out and recommend `from_raw_parts`). Rebuilding from `(ptr, len, cap)` + // reuses the same allocation and carries no `Vec`-layout assumption. + let mut zeros = core::mem::ManuallyDrop::new(zeros); // SAFETY: `FE` is `#[repr(transparent)]` over `u64` with identical size and // alignment (asserted above), and `0u64` is exactly `FE::zero()`'s bit - // pattern (no Montgomery form), so reinterpreting the buffer is valid. - unsafe { core::mem::transmute::, Vec>(zeros) } + // pattern (Goldilocks has no Montgomery form), so the zeroed `u64` buffer is + // a valid `[FE]` of all `FE::zero()`. `len`/`capacity` are element counts and + // the element sizes are equal, so they carry over unchanged; the eventual + // dealloc uses the same `size * capacity` and alignment as the original + // allocation. `ManuallyDrop` stops the source `Vec` from freeing the buffer + // that the returned `Vec` now owns. + unsafe { Vec::from_raw_parts(zeros.as_mut_ptr() as *mut FE, zeros.len(), zeros.capacity()) } } #[cfg(test)]