From c44d624b6b27553202dd303603d00a2d992a9fac Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Mon, 6 Jul 2026 11:47:32 -0300 Subject: [PATCH 01/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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/21] 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 b68c0b31435c88120207f7a3ef0570afd164e899 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Mon, 13 Jul 2026 13:41:25 -0300 Subject: [PATCH 17/21] move tables to GPU --- Cargo.lock | 1 + crypto/math-cuda/build.rs | 1 + crypto/math-cuda/kernels/trace_cpu.cu | 503 +++++++++++++++++++++++ crypto/math-cuda/src/device.rs | 18 + crypto/math-cuda/src/inverse.rs | 24 +- crypto/math-cuda/src/lde.rs | 34 ++ crypto/math-cuda/src/lib.rs | 1 + crypto/math-cuda/src/logup.rs | 16 +- crypto/math-cuda/src/trace_cpu.rs | 391 ++++++++++++++++++ crypto/math-cuda/tests/lde_dev_parity.rs | 86 ++++ crypto/stark/src/gpu_lde.rs | 64 +++ crypto/stark/src/instruments.rs | 32 ++ crypto/stark/src/logup_gpu.rs | 18 +- crypto/stark/src/prover.rs | 52 +++ crypto/stark/src/trace.rs | 65 +++ prover/Cargo.toml | 3 +- prover/src/lib.rs | 10 + prover/src/tables/gpu_trace.rs | 482 ++++++++++++++++++++++ prover/src/tables/memw.rs | 2 +- prover/src/tables/memw_register.rs | 21 +- prover/src/tables/mod.rs | 2 + prover/src/tables/trace_builder.rs | 501 ++++++++++++++++++++++ 22 files changed, 2300 insertions(+), 27 deletions(-) create mode 100644 crypto/math-cuda/kernels/trace_cpu.cu create mode 100644 crypto/math-cuda/src/trace_cpu.rs create mode 100644 crypto/math-cuda/tests/lde_dev_parity.rs create mode 100644 prover/src/tables/gpu_trace.rs diff --git a/Cargo.lock b/Cargo.lock index d2f699b80..5bfab63d4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1687,6 +1687,7 @@ dependencies = [ "executor", "log", "math", + "math-cuda", "postcard", "rayon", "serde", diff --git a/crypto/math-cuda/build.rs b/crypto/math-cuda/build.rs index 8fc568e2b..1f696c257 100644 --- a/crypto/math-cuda/build.rs +++ b/crypto/math-cuda/build.rs @@ -121,4 +121,5 @@ fn main() { compile_ptx("fri.cu", "fri.ptx", have_nvcc); compile_ptx("inverse.cu", "inverse.ptx", have_nvcc); compile_ptx("logup.cu", "logup.ptx", have_nvcc); + compile_ptx("trace_cpu.cu", "trace_cpu.ptx", have_nvcc); } diff --git a/crypto/math-cuda/kernels/trace_cpu.cu b/crypto/math-cuda/kernels/trace_cpu.cu new file mode 100644 index 000000000..0f33e2f1d --- /dev/null +++ b/crypto/math-cuda/kernels/trace_cpu.cu @@ -0,0 +1,503 @@ +// On-GPU CPU trace-table fill: one thread per row, row-major output +// `out[row*NCOLS + col]` (the layout the device-input LDE seam consumes). +// +// Mirrors `prover/src/tables/cpu.rs::generate_cpu_trace`. The `CpuOperation` +// already carries the resolved values (rv1/rv2/arg2/res/next_pc/branch_cond), +// so the fill is pure bit-slicing — every column limb is < 2^32 (DWordWL lo/hi, +// DWordHL 16-bit halves, bytes, bools), so NO Goldilocks reduction is needed. +// +// Packed input, stride STRIDE u64 per op: +// [0]=timestamp [1]=pc [2]=imm [3]=next_pc [4]=rvd [5]=rv1 [6]=rv2 [7]=arg2 +// [8]=res [9]=flags [10]=bytes +// flags bits: 0 word_instr, 1 read_register1(raw), 2 read_register2(raw), +// 3 write_register(raw), 4 alu, 5 add, 6 sub, 7 memory, 8 branch, 9 ecall, +// 10 branch_cond. +// bytes: b0 rs1, b1 rs2, b2 rd, b3 half_instruction_length, b4 alu_flags, +// b5 mem_flags. +// +// The output buffer MUST be pre-zeroed (alloc_zeros); the kernel writes only the +// non-zero cells of real rows, and (TIMESTAMP, PC_0=1, NEXT_PC_0=1) for padding. + +#include + +#define NCOLS 38u +#define STRIDE 11u + +extern "C" __global__ void trace_cpu_fill(const uint64_t *ops, // n * STRIDE + uint64_t n, // real op count + uint64_t num_rows, // padded pow2 + uint64_t last_ts, // last real ts (0 if n==0) + uint64_t *out) // num_rows * NCOLS, zeroed +{ + uint64_t r = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (r >= num_rows) return; + uint64_t base = r * NCOLS; + + // Padding rows: continue the +4 timestamp cadence; PC_0 = NEXT_PC_0 = 1 + // (odd, unreachable); every other column stays zero. + if (r >= n) { + uint64_t j = r - n + 1; + out[base + 0] = last_ts + 4 * j; // TIMESTAMP + out[base + 1] = 1; // PC_0 + out[base + 21] = 1; // NEXT_PC_0 + return; + } + + const uint64_t *op = ops + r * STRIDE; + uint64_t ts = op[0]; + uint64_t pc = op[1]; + uint64_t imm = op[2]; + uint64_t npc = op[3]; + uint64_t rvd = op[4]; + uint64_t rv1 = op[5]; + uint64_t rv2 = op[6]; + uint64_t arg2 = op[7]; + uint64_t res = op[8]; + uint64_t fl = op[9]; + uint64_t by = op[10]; + + uint64_t word = fl & 1u; + uint64_t rr1 = (fl >> 1) & 1u; + uint64_t rr2 = (fl >> 2) & 1u; + uint64_t wr = (fl >> 3) & 1u; + uint64_t alu = (fl >> 4) & 1u; + uint64_t add = (fl >> 5) & 1u; + uint64_t sub = (fl >> 6) & 1u; + uint64_t mem = (fl >> 7) & 1u; + uint64_t br = (fl >> 8) & 1u; + uint64_t ec = (fl >> 9) & 1u; + uint64_t bcond = (fl >> 10) & 1u; + + uint64_t rs1 = by & 0xFFu; + uint64_t rs2 = (by >> 8) & 0xFFu; + uint64_t rd = (by >> 16) & 0xFFu; + uint64_t hil = (by >> 24) & 0xFFu; + uint64_t aluf = (by >> 32) & 0xFFu; + uint64_t memf = (by >> 40) & 0xFFu; + + // `effective(flag) = !word && flag`. Word-delegate rows suppress operational + // data (CPU32 owns it) — only the PC-advancing columns are set. + uint64_t nw = 1u - word; // !word + uint64_t rs1c = word ? 0u : rs1; + uint64_t rs2c = word ? 0u : rs2; + uint64_t rdc = word ? 0u : rd; + uint64_t immc = word ? 0u : imm; + uint64_t rvdc = word ? 0u : rvd; + uint64_t rv1c = word ? 0u : rv1; + uint64_t rv2c = word ? 0u : rv2; + uint64_t arg2c = word ? 0u : arg2; + uint64_t resc = word ? 0u : res; + + out[base + 0] = ts; // TIMESTAMP + out[base + 1] = pc & 0xFFFFFFFFu; // PC_0 + out[base + 2] = pc >> 32; // PC_1 + out[base + 3] = rs1c; // RS1 + out[base + 4] = rs2c; // RS2 + out[base + 5] = rdc; // RD + out[base + 6] = nw & rr1 & (rs1 != 0u); // READ_REGISTER1 = eff(rr1 && rs1!=0) + out[base + 7] = nw & rr2 & (rs2 != 0u); // READ_REGISTER2 + out[base + 8] = nw & wr & (rd != 0u); // WRITE_REGISTER + out[base + 9] = immc & 0xFFFFFFFFu; // IMM_0 + out[base + 10] = immc >> 32; // IMM_1 + out[base + 11] = hil; // HALF_INSTRUCTION_LENGTH (unmasked) + out[base + 12] = word; // WORD_INSTR + out[base + 13] = nw & alu; // ALU + out[base + 14] = word ? 0u : aluf; // ALU_FLAGS + out[base + 15] = nw & add; // ADD + out[base + 16] = nw & sub; // SUB + out[base + 17] = nw & mem; // MEMORY + out[base + 18] = word ? 0u : memf; // MEM_FLAGS + out[base + 19] = nw & br; // BRANCH + out[base + 20] = nw & ec; // ECALL + out[base + 21] = npc & 0xFFFFFFFFu; // NEXT_PC_0 + out[base + 22] = npc >> 32; // NEXT_PC_1 + out[base + 23] = rvdc & 0xFFFFFFFFu; // RVD_0 + out[base + 24] = rvdc >> 32; // RVD_1 + + // Inline-PC coordination columns. + uint64_t pcdr = nw & rr1 & (rs1 == 255u); // PC_DOUBLE_READ + uint64_t ts_lo = ts & 0xFFFFFFFFu; + uint64_t borrow = (pcdr == 0u && ts_lo < 3u) ? 1u : 0u; // PREV_PC_TIMESTAMP_BORROW + out[base + 25] = borrow; + out[base + 26] = pcdr; + out[base + 27] = rv1c & 0xFFFFFFFFu; // RV1_0 + out[base + 28] = rv1c >> 32; // RV1_1 + out[base + 29] = rv2c & 0xFFFFFFFFu; // RV2_0 + out[base + 30] = rv2c >> 32; // RV2_1 + out[base + 31] = arg2c & 0xFFFFFFFFu; // ARG2_0 + out[base + 32] = arg2c >> 32; // ARG2_1 + out[base + 33] = resc & 0xFFFFu; // RES_0 + out[base + 34] = (resc >> 16) & 0xFFFFu; // RES_1 + out[base + 35] = (resc >> 32) & 0xFFFFu; // RES_2 + out[base + 36] = (resc >> 48) & 0xFFFFu; // RES_3 + out[base + 37] = bcond; // BRANCH_COND (unmasked) +} + +// On-GPU MEMW_A (aligned memory) trace fill: one thread per row, row-major +// `out[row*MEMW_A_NCOLS + col]`. Mirrors +// `prover/src/tables/memw_aligned.rs::generate_memw_aligned_trace`. The op is +// already walked (old_value/old_timestamp filled by the memory-model walk), so +// this is pure bit-slicing — every limb is < 2^32 (halves/words/bytes/bools), no +// Goldilocks reduction. Padding rows (r >= n) stay all-zero (pre-zeroed buffer). +// +// Packed input, stride MEMW_A_STRIDE u64 per op: +// [0]=flags (bit0 is_register, bit1 is_read, bits8..16 width) +// [1]=base_address [2]=timestamp [3]=old_timestamp[0] +// [4..8]=value[0..8] packed 2×u32/u64 (value[2i] | value[2i+1]<<32) +// [8..12]=old[0..8] packed 2×u32/u64 +#define MEMW_A_NCOLS 29u +#define MEMW_A_STRIDE 12u + +extern "C" __global__ void memw_aligned_fill(const uint64_t *ops, uint64_t n, + uint64_t num_rows, uint64_t *out) { + uint64_t r = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (r >= num_rows) + return; + uint64_t base = r * MEMW_A_NCOLS; + if (r >= n) + return; // padding rows all-zero + + const uint64_t *op = ops + r * MEMW_A_STRIDE; + uint64_t fl = op[0]; + uint64_t is_register = fl & 1u; + uint64_t is_read = (fl >> 1) & 1u; + uint64_t width = (fl >> 8) & 0xFFu; + uint64_t addr = op[1]; + uint64_t ts = op[2]; + uint64_t old_ts = op[3]; + + out[base + 0] = is_register; // IS_REGISTER + out[base + 1] = addr & 0xFFFFu; // BASE_ADDRESS[0] (low half) + out[base + 2] = (addr >> 16) & 0xFFFFu; // BASE_ADDRESS[1] (mid half) + out[base + 3] = (addr >> 32) & 0xFFFFFFFFu; // BASE_ADDRESS[2] (high word) + // VALUE[0..8] at cols 4..11. + for (int i = 0; i < 4; ++i) { + uint64_t p = op[4 + i]; + out[base + 4 + 2 * i] = p & 0xFFFFFFFFu; + out[base + 4 + 2 * i + 1] = p >> 32; + } + out[base + 12] = ts & 0xFFFFFFFFu; // TIMESTAMP_0 + out[base + 13] = ts >> 32; // TIMESTAMP_1 + out[base + 14] = (width == 2u) ? 1u : 0u; // WRITE2 + out[base + 15] = (width == 4u) ? 1u : 0u; // WRITE4 + out[base + 16] = (width == 8u) ? 1u : 0u; // WRITE8 + // OLD[0..8] at cols 17..24. + for (int i = 0; i < 4; ++i) { + uint64_t p = op[8 + i]; + out[base + 17 + 2 * i] = p & 0xFFFFFFFFu; + out[base + 17 + 2 * i + 1] = p >> 32; + } + out[base + 25] = old_ts & 0xFFFFFFFFu; // OLD_TIMESTAMP_0 + out[base + 26] = old_ts >> 32; // OLD_TIMESTAMP_1 + out[base + 27] = is_read; // MU_READ + out[base + 28] = 1u - is_read; // MU_WRITE +} + +// On-GPU LOAD trace fill (18 cols). Mirrors +// `prover/src/tables/load.rs::generate_load_trace`. Packed stride LOAD_STRIDE: +// [0]=flags (bit0 signed, bits8..16 width) [1]=base_address [2]=timestamp +// [3..7]=res[0..8] packed 2×u32/u64. Padding rows (r>=n) all-zero. +#define LOAD_NCOLS 18u +#define LOAD_STRIDE 7u + +extern "C" __global__ void load_fill(const uint64_t *ops, uint64_t n, + uint64_t num_rows, uint64_t *out) { + uint64_t r = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (r >= num_rows) + return; + uint64_t base = r * LOAD_NCOLS; + if (r >= n) + return; + + const uint64_t *op = ops + r * LOAD_STRIDE; + uint64_t fl = op[0]; + uint64_t is_signed = fl & 1u; + uint64_t width = (fl >> 8) & 0xFFu; + uint64_t addr = op[1]; + uint64_t ts = op[2]; + + out[base + 0] = addr & 0xFFFFFFFFu; // BASE_ADDRESS_0 + out[base + 1] = addr >> 32; // BASE_ADDRESS_1 + out[base + 2] = ts & 0xFFFFFFFFu; // TIMESTAMP_0 + out[base + 3] = ts >> 32; // TIMESTAMP_1 + out[base + 4] = (width == 2u) ? 1u : 0u; // READ2 + out[base + 5] = (width == 4u) ? 1u : 0u; // READ4 + out[base + 6] = (width == 8u) ? 1u : 0u; // READ8 + out[base + 7] = is_signed; // SIGNED + + uint64_t res[8]; + for (int i = 0; i < 4; ++i) { + uint64_t p = op[3 + i]; + res[2 * i] = p & 0xFFFFFFFFu; + res[2 * i + 1] = p >> 32; + } + for (int i = 0; i < 8; ++i) + out[base + 8 + i] = res[i]; // RES[0..8] + + int bidx = (width == 8u) ? 7 : (width == 4u) ? 3 : (width == 2u) ? 1 : 0; + out[base + 16] = (res[bidx] >> 7) & 1u; // SIGN_BIT + out[base + 17] = 1u; // MU (active row) +} + +// On-GPU STORE trace fill (16 cols). Mirrors +// `prover/src/tables/store.rs::generate_store_trace`. Packed stride STORE_STRIDE: +// [0]=flags (bit0 write2, bit1 write4, bit2 write8) [1]=base_address +// [2]=timestamp [3]=value. Padding rows (r>=n) all-zero. +#define STORE_NCOLS 16u +#define STORE_STRIDE 4u + +extern "C" __global__ void store_fill(const uint64_t *ops, uint64_t n, + uint64_t num_rows, uint64_t *out) { + uint64_t r = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (r >= num_rows) + return; + uint64_t base = r * STORE_NCOLS; + if (r >= n) + return; + + const uint64_t *op = ops + r * STORE_STRIDE; + uint64_t fl = op[0]; + uint64_t addr = op[1]; + uint64_t ts = op[2]; + uint64_t val = op[3]; + + out[base + 0] = addr & 0xFFFFFFFFu; // BASE_ADDRESS_0 + out[base + 1] = addr >> 32; // BASE_ADDRESS_1 + out[base + 2] = ts & 0xFFFFFFFFu; // TIMESTAMP_0 + out[base + 3] = ts >> 32; // TIMESTAMP_1 + out[base + 4] = fl & 1u; // WRITE2 + out[base + 5] = (fl >> 1) & 1u; // WRITE4 + out[base + 6] = (fl >> 2) & 1u; // WRITE8 + for (int i = 0; i < 8; ++i) + out[base + 7 + i] = (val >> (i * 8)) & 0xFFu; // VALUE[0..8] (DWordBL) + out[base + 15] = 1u; // MU +} + +// On-GPU SHIFT trace fill (29 cols). Mirrors +// `prover/src/tables/shift.rs::{generate_shift_trace, compute_aux}` — the fill +// RECOMPUTES the aux (bit_shift, zbs, x/y half decomposition, limb_shift one-hot, +// shifted DWordHL) from the compact input, so only 3 u64/op are uploaded. No +// dedup: μ = 1 per op. Padding rows (r >= n) set ZBS = 1 (all else 0). +// +// Packed stride SHIFT_STRIDE: [0]=value (4×u16 in_halves) [1]=shift_amount +// [2]=flags (bit0 direction(=right), bit1 signed, bit2 word_instr). + +// HWSL: (halfword << z) & 0xFFFF (z in [0,15]). +__device__ __forceinline__ uint16_t hwsl_(uint16_t h, uint32_t z) { + return z == 0u ? h : (uint16_t)((uint32_t)h << z); +} +// HWSL carry: halfword >> (16 - z). +__device__ __forceinline__ uint16_t hwslc_(uint16_t h, uint32_t z) { + return z == 0u ? (uint16_t)0 : (uint16_t)(h >> (16u - z)); +} + +#define SHIFT_NCOLS 29u +#define SHIFT_STRIDE 3u + +extern "C" __global__ void shift_fill(const uint64_t *ops, uint64_t n, + uint64_t num_rows, uint64_t *out) { + uint64_t r = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (r >= num_rows) + return; + uint64_t base = r * SHIFT_NCOLS; + if (r >= n) { + out[base + 12] = 1u; // ZBS = 1 on padding rows + return; + } + + const uint64_t *op = ops + r * SHIFT_STRIDE; + uint64_t value = op[0]; + uint64_t shift_amount = op[1]; + uint64_t fl = op[2]; + uint32_t direction = (uint32_t)(fl & 1u); // right + uint32_t is_signed = (uint32_t)((fl >> 1) & 1u); + uint32_t word_instr = (uint32_t)((fl >> 2) & 1u); + + uint16_t in_h[4]; + in_h[0] = (uint16_t)(value & 0xFFFFu); + in_h[1] = (uint16_t)((value >> 16) & 0xFFFFu); + in_h[2] = (uint16_t)((value >> 32) & 0xFFFFu); + in_h[3] = (uint16_t)((value >> 48) & 0xFFFFu); + uint8_t shift = (uint8_t)(shift_amount & 0xFFu); + uint32_t left = 1u - direction; + uint32_t right = direction; + + uint32_t is_negative = (is_signed && ((in_h[3] >> 15) & 1u)) ? 1u : 0u; + uint16_t extension = is_negative ? (uint16_t)0xFFFF : (uint16_t)0; + + uint8_t bit_shift; + if (left) + bit_shift = (uint8_t)(shift & 15u); + else + bit_shift = (uint8_t)((256u - (uint32_t)shift) & 15u); + uint32_t zbs = (bit_shift == 0u) ? 1u : 0u; + + uint16_t x[5] = {0, 0, 0, 0, 0}; + uint16_t y[4] = {0, 0, 0, 0}; + if (zbs) { + for (int i = 0; i < 4; ++i) { + if (left) + x[i] = in_h[i]; + else + y[i] = in_h[i]; + } + x[4] = 0; + } else { + for (int i = 0; i < 4; ++i) { + x[i] = hwsl_(in_h[i], bit_shift); + y[i] = hwslc_(in_h[i], bit_shift); + } + x[4] = hwsl_(extension, bit_shift); + } + + // limb_shift: one-hot of (shift >> 4) & mask. + uint32_t limb_idx = word_instr ? (uint32_t)((shift >> 4) & 1u) + : (uint32_t)((shift >> 4) & 3u); + uint32_t ls[4] = {0, 0, 0, 0}; + ls[limb_idx] = 1u; + + // shifted[i] (DWordHL). intra_left(k)=x[0] if k==0 else x[k]+y[k-1]; + // intra_right(k)=y[k]+x[k+1]. + uint16_t shifted[4]; + for (int i = 0; i < 4; ++i) { + uint16_t v = 0; + if (left) { + for (int j = 0; j <= i; ++j) + if (ls[j]) { + int k = i - j; + v = (uint16_t)(v + (k == 0 ? x[0] : (uint16_t)(x[k] + y[k - 1]))); + } + } + if (right) { + for (int j = 0; j <= 3 - i; ++j) + if (ls[j]) { + int k = i + j; + v = (uint16_t)(v + (uint16_t)(y[k] + x[k + 1])); + } + for (int j = 4 - i; j < 4; ++j) + if (ls[j]) + v = (uint16_t)(v + extension); + } + shifted[i] = v; + } + uint32_t out0 = (uint32_t)shifted[0] | ((uint32_t)shifted[1] << 16); + uint32_t out1 = (uint32_t)shifted[2] | ((uint32_t)shifted[3] << 16); + + out[base + 0] = in_h[0]; // IN_0 + out[base + 1] = in_h[1]; + out[base + 2] = in_h[2]; + out[base + 3] = in_h[3]; + out[base + 4] = shift; // SHIFT_AMOUNT + out[base + 5] = direction; // DIRECTION + out[base + 6] = is_signed; // SIGNED + out[base + 7] = word_instr; // WORD_INSTR + out[base + 8] = out0; // OUT_0 + out[base + 9] = out1; // OUT_1 + out[base + 10] = is_negative; // IS_NEGATIVE + out[base + 11] = bit_shift; // BIT_SHIFT + out[base + 12] = zbs; // ZBS + out[base + 13] = x[0]; // X_0..X_4 + out[base + 14] = x[1]; + out[base + 15] = x[2]; + out[base + 16] = x[3]; + out[base + 17] = x[4]; + out[base + 18] = y[0]; // Y_0..Y_3 + out[base + 19] = y[1]; + out[base + 20] = y[2]; + out[base + 21] = y[3]; + out[base + 22] = ls[0]; // LIMB_SHIFT_RAW_0..2 + out[base + 23] = ls[1]; + out[base + 24] = ls[2]; + out[base + 25] = 1u; // MU + out[base + 26] = (shift_amount >> 8) & 0xFFu; // SHIFT_B1 + out[base + 27] = (shift_amount >> 16) & 0xFFFFu; // SHIFT_H1 + out[base + 28] = (shift_amount >> 32) & 0xFFFFFFFFu; // SHIFT_HIGH +} + +// On-GPU LT trace fill (17 cols). Mirrors the per-row compute of +// `prover/src/tables/lt.rs::generate_lt_trace`. Dedup is done on the HOST (the +// same per-chunk HashMap the CPU uses); this kernel receives already-unique ops +// with their summed multiplicity, one per row, and recomputes lt/out/sub/msbs. +// Order-independent (LogUp ALU bus) → validated by multiset/prove, not byte order. +// Padding rows (r >= n) stay all-zero. +// +// Packed stride LT_STRIDE: [0]=lhs [1]=rhs [2]=flags (bit0 signed, bit1 invert) +// [3]=multiplicity. +#define LT_NCOLS 17u +#define LT_STRIDE 4u + +extern "C" __global__ void lt_fill(const uint64_t *ops, uint64_t n, + uint64_t num_rows, uint64_t *out) { + uint64_t r = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (r >= num_rows) + return; + uint64_t base = r * LT_NCOLS; + if (r >= n) + return; + + const uint64_t *op = ops + r * LT_STRIDE; + uint64_t lhs = op[0]; + uint64_t rhs = op[1]; + uint64_t fl = op[2]; + uint64_t mult = op[3]; + uint64_t is_signed = fl & 1u; + uint64_t invert = (fl >> 1) & 1u; + + uint64_t lt = is_signed ? ((long long)lhs < (long long)rhs ? 1u : 0u) + : (lhs < rhs ? 1u : 0u); + uint64_t out_bit = lt ^ invert; + uint64_t sub = lhs - rhs; // wrapping + + out[base + 0] = lhs & 0xFFFFFFFFu; // LHS_0 (word) + out[base + 1] = (lhs >> 32) & 0xFFFFu; // LHS_1 (half) + out[base + 2] = (lhs >> 48) & 0xFFFFu; // LHS_2 (half) + out[base + 3] = rhs & 0xFFFFFFFFu; // RHS_0 + out[base + 4] = (rhs >> 32) & 0xFFFFu; // RHS_1 + out[base + 5] = (rhs >> 48) & 0xFFFFu; // RHS_2 + out[base + 6] = is_signed; // SIGNED + out[base + 7] = lt; // LT + out[base + 8] = sub & 0xFFFFu; // LHS_SUB_RHS_0 (DWordHL halves) + out[base + 9] = (sub >> 16) & 0xFFFFu; + out[base + 10] = (sub >> 32) & 0xFFFFu; + out[base + 11] = (sub >> 48) & 0xFFFFu; + out[base + 12] = (lhs >> 63) & 1u; // LHS_MSB + out[base + 13] = (rhs >> 63) & 1u; // RHS_MSB + out[base + 14] = invert; // INVERT + out[base + 15] = out_bit; // OUT + out[base + 16] = mult; // MU +} + +// On-GPU MEMW_R (register fast-path) fill: write the 10 MEMW_R columns ROW-MAJOR +// from the host-walked rows. One thread per row (row_index is the identity for +// the fill-from-walked-rows path). Columns mirror +// `prover/src/tables/memw_register.rs::generate_memw_register_trace_from_rows`: +// 0 ADDRESS=reg_addr/2, 1 TS0=ts&0xffffffff, 2 TS1=ts>>32, 3 VAL0, 4 VAL1, +// 5 OLD0, 6 OLD1, 7 OLD_TS_LO=old_ts&0xffffffff, 8 MU_READ=is_read, 9 MU_WRITE=!is_read. +// All limbs < 2^32 (no Goldilocks reduction). Padding rows are pre-zeroed. +extern "C" __global__ void memw_register_fill( + uint64_t n_acc, const uint32_t *reg_addr, const uint64_t *ts, + const uint64_t *value, const uint8_t *is_read, const long long *row_index, + const uint64_t *old_value, const uint64_t *old_ts, uint32_t ncols, + uint64_t *buf) { + uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n_acc) + return; + long long row = row_index[i]; + if (row < 0) + return; + uint64_t base = (uint64_t)row * ncols; + uint64_t v = value[i]; + uint64_t ov = old_value[i]; + uint64_t t = ts[i]; + uint64_t ot = old_ts[i]; + buf[base + 0] = (uint64_t)(reg_addr[i] / 2u); + buf[base + 1] = t & 0xFFFFFFFFull; + buf[base + 2] = t >> 32; + buf[base + 3] = v & 0xFFFFFFFFull; + buf[base + 4] = v >> 32; + buf[base + 5] = ov & 0xFFFFFFFFull; + buf[base + 6] = ov >> 32; + buf[base + 7] = ot & 0xFFFFFFFFull; + buf[base + 8] = is_read[i] ? 1ull : 0ull; + buf[base + 9] = is_read[i] ? 0ull : 1ull; +} diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index b1d6ca489..7a4ab90af 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -98,6 +98,7 @@ const DEEP_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/deep.ptx")); const FRI_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/fri.ptx")); const INVERSE_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/inverse.ptx")); const LOGUP_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/logup.ptx")); +const TRACE_CPU_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/trace_cpu.ptx")); /// Number of CUDA streams in the pool. Larger pools let many rayon-parallel /// callers overlap on the GPU without serializing on stream ownership. The @@ -188,6 +189,15 @@ pub struct Backend { pub logup_finalize_accum_ext3: CudaFunction, pub logup_assemble_aux_ext3: CudaFunction, + // On-GPU trace generation. + pub trace_cpu_fill: CudaFunction, + pub memw_aligned_fill: CudaFunction, + pub load_fill: CudaFunction, + pub store_fill: CudaFunction, + pub shift_fill: CudaFunction, + pub lt_fill: CudaFunction, + pub memw_register_fill: CudaFunction, + // Twiddle caches keyed by log_n. fwd_twiddles: Mutex>>>>, inv_twiddles: Mutex>>>>, @@ -291,6 +301,7 @@ impl Backend { let fri = ctx.load_module(Ptx::from_src(FRI_PTX))?; let inverse = ctx.load_module(Ptx::from_src(INVERSE_PTX))?; let logup = ctx.load_module(Ptx::from_src(LOGUP_PTX))?; + let trace_cpu = ctx.load_module(Ptx::from_src(TRACE_CPU_PTX))?; let mut streams = Vec::with_capacity(STREAM_POOL_SIZE); for _ in 0..STREAM_POOL_SIZE { @@ -378,6 +389,13 @@ impl Backend { logup_apply_offsets_add_ext3: logup.load_function("logup_apply_offsets_add_ext3")?, logup_finalize_accum_ext3: logup.load_function("logup_finalize_accum_ext3")?, logup_assemble_aux_ext3: logup.load_function("logup_assemble_aux_ext3")?, + trace_cpu_fill: trace_cpu.load_function("trace_cpu_fill")?, + memw_aligned_fill: trace_cpu.load_function("memw_aligned_fill")?, + load_fill: trace_cpu.load_function("load_fill")?, + store_fill: trace_cpu.load_function("store_fill")?, + shift_fill: trace_cpu.load_function("shift_fill")?, + lt_fill: trace_cpu.load_function("lt_fill")?, + memw_register_fill: trace_cpu.load_function("memw_register_fill")?, fwd_twiddles: Mutex::new(vec![None; max_log]), inv_twiddles: Mutex::new(vec![None; max_log]), ctx, diff --git a/crypto/math-cuda/src/inverse.rs b/crypto/math-cuda/src/inverse.rs index 485e005f8..36762b478 100644 --- a/crypto/math-cuda/src/inverse.rs +++ b/crypto/math-cuda/src/inverse.rs @@ -84,12 +84,15 @@ pub fn batch_inverse_ext3_dev( stream: &Arc, ) -> Result> { assert!(n >= 1, "batch_inverse_ext3_dev requires n >= 1"); - // Runtime guard (not debug_assert): a u32 grid_dim is truncated past - // u32::MAX / BLOCK_SIZE, which would silently launch too few blocks - // and leave a tail uninverted. Reachable on LDE size 2^23+ × multi- - // eval-point R4. Returning Err lets the dispatcher's Err(_) => None - // route the caller to the CPU `inplace_batch_inverse` fallback. - if n > u32::MAX as usize / BLOCK_SIZE as usize { + // Runtime guard (not debug_assert): `n` must fit in u32 so neither `n as u32` + // below nor the kernel's flat index truncates. The block count + // `n.div_ceil(BLOCK_SIZE)` is well within CUDA's grid.x bound (2^31-1) for any + // such `n` — grid.x is NOT the 65535 cap that applies to grid.y/z, so the old + // `u32::MAX / BLOCK_SIZE` bound was 256× too strict and spuriously declined + // legitimate launches (e.g. logup batch-invert of 18 interactions × 2^20 rows + // = ~18.9M), silently routing to a wrong fallback. Returning Err lets the + // dispatcher's Err(_) => None route the caller to the CPU fallback. + if n > u32::MAX as usize { return Err(cudarc::driver::DriverError( cudarc::driver::sys::CUresult::CUDA_ERROR_INVALID_VALUE, )); @@ -182,10 +185,11 @@ pub fn compute_and_invert_denoms_ext3_dev( let total = k_scalars .checked_mul(n) .expect("compute_and_invert_denoms_ext3_dev: k_scalars * n overflow"); - // See `batch_inverse_ext3_dev` for the rationale: runtime Err, not - // debug_assert, so release builds also route past the silent-truncation - // hazard via the caller's CPU fallback. - if total > u32::MAX as usize / BLOCK_SIZE as usize { + // See `batch_inverse_ext3_dev` for the rationale: guard on `u32::MAX` (the + // cast/index bound), NOT `u32::MAX / BLOCK_SIZE` (which is the grid.y/z cap, not + // grid.x, and 256× too strict). Runtime Err, not debug_assert, so release + // builds route past the truncation hazard via the caller's CPU fallback. + if total > u32::MAX as usize { return Err(cudarc::driver::DriverError( cudarc::driver::sys::CUresult::CUDA_ERROR_INVALID_VALUE, )); diff --git a/crypto/math-cuda/src/lde.rs b/crypto/math-cuda/src/lde.rs index 30e524c2c..fd5c529f2 100644 --- a/crypto/math-cuda/src/lde.rs +++ b/crypto/math-cuda/src/lde.rs @@ -583,6 +583,40 @@ pub fn coset_lde_row_major_with_merkle_tree_keep( Ok((handle, lde_out)) } +/// Like [`coset_lde_row_major_with_merkle_tree_keep`] but the input is an +/// already-resident device buffer (`n * m` base-field elements, row-major +/// `[row*m + col]`). No PCIe upload: the buffer is copied device-to-device into +/// the LDE scratch. Used by the on-GPU trace generator so a device-born main +/// trace feeds the LDE without a host round-trip. Retains the trace-domain +/// column-major snapshot exactly like the host keep path, so the resident LogUp +/// aux fingerprint (#762) still works unchanged. +pub fn coset_lde_row_major_with_merkle_tree_keep_dev( + input_dev: &CudaSlice, + n: usize, + m: usize, + blowup_factor: usize, + weights: &[u64], +) -> Result<(GpuLdeBase, Vec)> { + let (tree, col_major_dev, lde_out, trace_col_major) = coset_lde_row_major_inner( + InnerInput::Dev(input_dev), + n, + m, + blowup_factor, + weights, + "coset_lde_row_major_dev lde_size", + true, + )?; + let handle = GpuLdeBase { + buf: Arc::new(col_major_dev), + m, + lde_size: n * blowup_factor, + tree: Some(tree), + trace_dev: trace_col_major.map(Arc::new), + trace_rows: n, + }; + Ok((handle, lde_out)) +} + /// Row-major ext3 LDE + Keccak + Merkle, all on-device. /// /// `Fp3` is `[u64; 3]` in memory, so row-major ext3 with `m` ext3 columns is diff --git a/crypto/math-cuda/src/lib.rs b/crypto/math-cuda/src/lib.rs index abc487750..070b2edfd 100644 --- a/crypto/math-cuda/src/lib.rs +++ b/crypto/math-cuda/src/lib.rs @@ -14,6 +14,7 @@ pub mod lde; pub mod logup; pub mod merkle; pub mod ntt; +pub mod trace_cpu; // Re-exported for downstream crates so they can refer to CUDA primitive // types without depending on cudarc directly. diff --git a/crypto/math-cuda/src/logup.rs b/crypto/math-cuda/src/logup.rs index e9d120ee1..a022728bb 100644 --- a/crypto/math-cuda/src/logup.rs +++ b/crypto/math-cuda/src/logup.rs @@ -44,12 +44,16 @@ pub struct LogupDescriptor<'a> { } fn cfg(total: usize) -> Result { - // See `batch_inverse_ext3_dev` for the rationale: a u32 grid_dim is - // truncated past u32::MAX / BLOCK_SIZE, which would silently launch too - // few blocks and leave a tail of the (uninitialized) output unwritten. - // Runtime Err, not debug_assert, so release builds also route to the - // caller's CPU fallback. - if total > u32::MAX as usize / BLOCK_SIZE as usize { + // One thread per element: `total` must fit in u32 so neither the `total as u32` + // cast below nor the kernel's `blockIdx.x*blockDim.x + threadIdx.x` index + // truncates. The block count `total.div_ceil(BLOCK_SIZE)` is well within CUDA's + // grid.x bound (2^31-1) for any such `total` — grid.x is NOT capped at the + // 65535 that applies to grid.y/z. (The previous bound `u32::MAX / BLOCK_SIZE` + // conflated the two and spuriously declined large single-chunk tables — e.g. + // SHIFT at 2^20 rows × 18 interactions = ~18.9M > u32::MAX/256 — silently + // falling back to a wrong, zeroed-host aux.) Runtime Err (not debug_assert) so + // release builds route to the caller's CPU fallback rather than truncate. + if total > u32::MAX as usize { return Err(cudarc::driver::DriverError( cudarc::driver::sys::CUresult::CUDA_ERROR_INVALID_VALUE, )); diff --git a/crypto/math-cuda/src/trace_cpu.rs b/crypto/math-cuda/src/trace_cpu.rs new file mode 100644 index 000000000..b469efb3e --- /dev/null +++ b/crypto/math-cuda/src/trace_cpu.rs @@ -0,0 +1,391 @@ +//! On-GPU CPU trace-table generation. Uploads packed `CpuOperation` fields and +//! fills the CPU table row-major on device (see `kernels/trace_cpu.cu`), leaving +//! the result resident so it feeds `coset_lde_row_major_with_merkle_tree_keep_dev` +//! with no host round-trip. + +use std::sync::Arc; + +use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg}; + +use crate::Result; +use crate::device::{Backend, backend}; + +/// CPU table width (`prover::tables::cpu::cols::NUM_COLUMNS`). +pub const CPU_NCOLS: usize = 38; +/// Packed input stride, u64 per op (must match `trace_cpu.cu`). +pub const CPU_OP_STRIDE: usize = 11; + +/// Build one CPU trace-table chunk on device from packed ops. +/// +/// `packed_ops` is `n * CPU_OP_STRIDE` u64s (see `trace_cpu.cu` for the field +/// layout). `num_rows` is the padded power-of-two row count and `last_ts` the +/// timestamp of the last real op (0 when `n == 0`). Returns the row-major device +/// buffer `[row*CPU_NCOLS + col]` (`num_rows * CPU_NCOLS` u64), ready for the +/// device-input LDE. +pub fn gpu_build_cpu_trace( + packed_ops: &[u64], + n: usize, + num_rows: usize, + last_ts: u64, +) -> Result> { + debug_assert_eq!(packed_ops.len(), n * CPU_OP_STRIDE); + let be = backend()?; + let stream = be.next_stream(); + + // `clone_htod` rejects empty slices; a 1-element dummy is never read because + // every row is a padding row when `n == 0`. + let ops_dev = if packed_ops.is_empty() { + stream.alloc_zeros::(CPU_OP_STRIDE)? + } else { + stream.clone_htod(packed_ops)? + }; + + // Zero-initialised: the kernel only writes non-zero cells. + let mut out = stream.alloc_zeros::(num_rows * CPU_NCOLS)?; + + unsafe { + stream + .launch_builder(&be.trace_cpu_fill) + .arg(&ops_dev) + .arg(&(n as u64)) + .arg(&(num_rows as u64)) + .arg(&last_ts) + .arg(&mut out) + .launch(LaunchConfig::for_num_elems(num_rows as u32))?; + } + stream.synchronize()?; + Ok(out) +} + +/// MEMW_A table width (`prover::tables::memw_aligned::cols::NUM_COLUMNS`). +pub const MEMW_ALIGNED_NCOLS: usize = 29; +/// Packed MEMW_A input stride, u64 per op (must match `memw_aligned_fill` in `trace_cpu.cu`). +pub const MEMW_ALIGNED_STRIDE: usize = 12; + +/// Build one MEMW_A (aligned memory) trace-table chunk on device from packed ops. +/// +/// `packed_ops` is `n * MEMW_ALIGNED_STRIDE` u64s (see `trace_cpu.cu` for the field +/// layout). `num_rows` is the padded power-of-two row count. Returns the row-major +/// device buffer `[row*MEMW_ALIGNED_NCOLS + col]` (`num_rows * MEMW_ALIGNED_NCOLS` +/// u64), left resident so it feeds the device-input LDE with no full-column upload. +pub fn gpu_build_memw_aligned_trace( + packed_ops: &[u64], + n: usize, + num_rows: usize, +) -> Result> { + debug_assert_eq!(packed_ops.len(), n * MEMW_ALIGNED_STRIDE); + let be = backend()?; + let stream = be.next_stream(); + + let mut out = stream.alloc_zeros::(num_rows * MEMW_ALIGNED_NCOLS)?; + + // `clone_htod` rejects empty slices; a 1-op dummy is never read (all padding). + let ops_dev = if packed_ops.is_empty() { + stream.alloc_zeros::(MEMW_ALIGNED_STRIDE)? + } else { + stream.clone_htod(packed_ops)? + }; + + unsafe { + stream + .launch_builder(&be.memw_aligned_fill) + .arg(&ops_dev) + .arg(&(n as u64)) + .arg(&(num_rows as u64)) + .arg(&mut out) + .launch(LaunchConfig::for_num_elems(num_rows as u32))?; + } + stream.synchronize()?; + Ok(out) +} + +// ----------------------------------------------------------------------------- +// LOAD / STORE (per-row-map memory tables — same shape as MEMW_A). A shared +// interleaved-fill launcher keeps the per-table code to a packing convention + +// column/stride constants. +// ----------------------------------------------------------------------------- + +/// LOAD table width / packed input stride (must match `load_fill` in `trace_cpu.cu`). +pub const LOAD_NCOLS: usize = 18; +pub const LOAD_STRIDE: usize = 7; +/// STORE table width / packed input stride (must match `store_fill`). +pub const STORE_NCOLS: usize = 16; +pub const STORE_STRIDE: usize = 4; +/// SHIFT table width / packed input stride (must match `shift_fill`). +pub const SHIFT_NCOLS: usize = 29; +pub const SHIFT_STRIDE: usize = 3; +/// LT table width / packed input stride (must match `lt_fill`). Input is already +/// host-deduplicated: one unique op + summed multiplicity per row. +pub const LT_NCOLS: usize = 17; +pub const LT_STRIDE: usize = 4; + +/// Shared interleaved fill on `stream`: upload `packed` (n × stride u64) and run +/// `kernel(ops, n, num_rows, out)` into a zeroed row-major buffer. Returns the +/// (unsynchronized) device buffer. +#[allow(clippy::too_many_arguments)] +fn build_interleaved_on( + stream: &Arc, + kernel: &CudaFunction, + packed_ops: &[u64], + n: usize, + num_rows: usize, + ncols: usize, + stride: usize, +) -> Result> { + debug_assert_eq!(packed_ops.len(), n * stride); + let mut out = stream.alloc_zeros::(num_rows * ncols)?; + let ops_dev = if packed_ops.is_empty() { + stream.alloc_zeros::(stride)? + } else { + stream.clone_htod(packed_ops)? + }; + unsafe { + stream + .launch_builder(kernel) + .arg(&ops_dev) + .arg(&(n as u64)) + .arg(&(num_rows as u64)) + .arg(&mut out) + .launch(LaunchConfig::for_num_elems(num_rows as u32))?; + } + Ok(out) +} + +fn load_kernel(be: &Backend) -> &CudaFunction { + &be.load_fill +} +fn store_kernel(be: &Backend) -> &CudaFunction { + &be.store_fill +} +fn shift_kernel(be: &Backend) -> &CudaFunction { + &be.shift_fill +} +fn lt_kernel(be: &Backend) -> &CudaFunction { + &be.lt_fill +} + +/// Build one LT trace-table chunk on device from HOST-DEDUPLICATED ops (one unique +/// op + summed multiplicity per row; see `lt_fill` in `trace_cpu.cu`). +pub fn gpu_build_lt_trace(packed_ops: &[u64], n: usize, num_rows: usize) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let out = build_interleaved_on( + &stream, lt_kernel(be), packed_ops, n, num_rows, LT_NCOLS, LT_STRIDE, + )?; + stream.synchronize()?; + Ok(out) +} + +/// Host-returning LT build for multiset-equality tests. +pub fn gpu_build_lt_trace_host(packed_ops: &[u64], n: usize, num_rows: usize) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let out = build_interleaved_on( + &stream, lt_kernel(be), packed_ops, n, num_rows, LT_NCOLS, LT_STRIDE, + )?; + let host = stream.clone_dtoh(&out)?; + stream.synchronize()?; + Ok(host) +} + +/// Build one SHIFT trace-table chunk on device (residency-ready row-major buffer). +/// The kernel recomputes the shift aux from the packed inputs (see `trace_cpu.cu`). +pub fn gpu_build_shift_trace(packed_ops: &[u64], n: usize, num_rows: usize) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let out = build_interleaved_on( + &stream, shift_kernel(be), packed_ops, n, num_rows, SHIFT_NCOLS, SHIFT_STRIDE, + )?; + stream.synchronize()?; + Ok(out) +} + +/// Host-returning SHIFT build for byte-parity tests. +pub fn gpu_build_shift_trace_host(packed_ops: &[u64], n: usize, num_rows: usize) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let out = build_interleaved_on( + &stream, shift_kernel(be), packed_ops, n, num_rows, SHIFT_NCOLS, SHIFT_STRIDE, + )?; + let host = stream.clone_dtoh(&out)?; + stream.synchronize()?; + Ok(host) +} + +/// Build one LOAD trace-table chunk on device (residency-ready row-major buffer). +pub fn gpu_build_load_trace(packed_ops: &[u64], n: usize, num_rows: usize) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let out = build_interleaved_on( + &stream, load_kernel(be), packed_ops, n, num_rows, LOAD_NCOLS, LOAD_STRIDE, + )?; + stream.synchronize()?; + Ok(out) +} + +/// Host-returning LOAD build for byte-parity tests. +pub fn gpu_build_load_trace_host(packed_ops: &[u64], n: usize, num_rows: usize) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let out = build_interleaved_on( + &stream, load_kernel(be), packed_ops, n, num_rows, LOAD_NCOLS, LOAD_STRIDE, + )?; + let host = stream.clone_dtoh(&out)?; + stream.synchronize()?; + Ok(host) +} + +/// Build one STORE trace-table chunk on device (residency-ready row-major buffer). +pub fn gpu_build_store_trace(packed_ops: &[u64], n: usize, num_rows: usize) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let out = build_interleaved_on( + &stream, store_kernel(be), packed_ops, n, num_rows, STORE_NCOLS, STORE_STRIDE, + )?; + stream.synchronize()?; + Ok(out) +} + +/// Host-returning STORE build for byte-parity tests. +pub fn gpu_build_store_trace_host(packed_ops: &[u64], n: usize, num_rows: usize) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let out = build_interleaved_on( + &stream, store_kernel(be), packed_ops, n, num_rows, STORE_NCOLS, STORE_STRIDE, + )?; + let host = stream.clone_dtoh(&out)?; + stream.synchronize()?; + Ok(host) +} + +/// Host-returning wrapper over [`gpu_build_memw_aligned_trace`] for byte-parity +/// tests: builds the row-major MEMW_A buffer and copies it back on the same stream. +pub fn gpu_build_memw_aligned_trace_host( + packed_ops: &[u64], + n: usize, + num_rows: usize, +) -> Result> { + debug_assert_eq!(packed_ops.len(), n * MEMW_ALIGNED_STRIDE); + let be = backend()?; + let stream = be.next_stream(); + + let mut out = stream.alloc_zeros::(num_rows * MEMW_ALIGNED_NCOLS)?; + let ops_dev = if packed_ops.is_empty() { + stream.alloc_zeros::(MEMW_ALIGNED_STRIDE)? + } else { + stream.clone_htod(packed_ops)? + }; + unsafe { + stream + .launch_builder(&be.memw_aligned_fill) + .arg(&ops_dev) + .arg(&(n as u64)) + .arg(&(num_rows as u64)) + .arg(&mut out) + .launch(LaunchConfig::for_num_elems(num_rows as u32))?; + } + let host = stream.clone_dtoh(&out)?; + stream.synchronize()?; + Ok(host) +} + +// ----------------------------------------------------------------------------- +// MEMW_R (register fast-path) fill. Registers are pre-walked on the host (the +// sequential memory model recovers old_value/old_ts); this fills the 10 MEMW_R +// columns row-major on device from the walked rows, leaving the matrix resident +// for the device-input LDE (removes MEMW_R's full-column H2D). See the +// `memw_register_fill` kernel in `trace_cpu.cu`. +// ----------------------------------------------------------------------------- + +/// MEMW_R table width (`prover::tables::memw_register::cols::NUM_COLUMNS`). +pub const MEMW_REGISTER_NCOLS: usize = 10; + +/// Build the MEMW_R trace table on device from already-walked rows (the +/// `old_value`/`old_ts` recovered by the host walk are uploaded here). Returns the +/// residency-ready row-major `[row*NCOLS+col]` buffer. `row_index` is the identity +/// (every input row is a real MEMW_R row). +#[allow(clippy::too_many_arguments)] +pub fn gpu_fill_memw_register( + reg_addr: &[u32], + ts: &[u64], + value: &[u64], + is_read: &[u8], + old_value: &[u64], + old_ts: &[u64], + num_rows: usize, +) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let buf = fill_memw_register_on(&stream, reg_addr, ts, value, is_read, old_value, old_ts, num_rows)?; + stream.synchronize()?; + Ok(buf) +} + +/// Host-returning wrapper over [`gpu_fill_memw_register`] for byte-parity tests. +#[allow(clippy::too_many_arguments)] +pub fn gpu_fill_memw_register_host( + reg_addr: &[u32], + ts: &[u64], + value: &[u64], + is_read: &[u8], + old_value: &[u64], + old_ts: &[u64], + num_rows: usize, +) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let buf = fill_memw_register_on(&stream, reg_addr, ts, value, is_read, old_value, old_ts, num_rows)?; + let host = stream.clone_dtoh(&buf)?; + stream.synchronize()?; + Ok(host) +} + +#[allow(clippy::too_many_arguments)] +fn fill_memw_register_on( + stream: &std::sync::Arc, + reg_addr: &[u32], + ts: &[u64], + value: &[u64], + is_read: &[u8], + old_value: &[u64], + old_ts: &[u64], + num_rows: usize, +) -> Result> { + let n = reg_addr.len(); + debug_assert_eq!(ts.len(), n); + debug_assert_eq!(value.len(), n); + debug_assert_eq!(is_read.len(), n); + debug_assert_eq!(old_value.len(), n); + debug_assert_eq!(old_ts.len(), n); + let be = backend()?; + let mut buf = stream.alloc_zeros::(num_rows * MEMW_REGISTER_NCOLS)?; + if n == 0 { + return Ok(buf); + } + let keys_d = stream.clone_htod(reg_addr)?; + let ts_d = stream.clone_htod(ts)?; + let value_d = stream.clone_htod(value)?; + let is_read_d = stream.clone_htod(is_read)?; + let old_value_d = stream.clone_htod(old_value)?; + let old_ts_d = stream.clone_htod(old_ts)?; + let row_index: Vec = (0..n as i64).collect(); + let row_index_d = stream.clone_htod(&row_index)?; + let n_u64 = n as u64; + let ncols_u32 = MEMW_REGISTER_NCOLS as u32; + unsafe { + stream + .launch_builder(&be.memw_register_fill) + .arg(&n_u64) + .arg(&keys_d) + .arg(&ts_d) + .arg(&value_d) + .arg(&is_read_d) + .arg(&row_index_d) + .arg(&old_value_d) + .arg(&old_ts_d) + .arg(&ncols_u32) + .arg(&mut buf) + .launch(LaunchConfig::for_num_elems(n as u32))?; + } + Ok(buf) +} diff --git a/crypto/math-cuda/tests/lde_dev_parity.rs b/crypto/math-cuda/tests/lde_dev_parity.rs new file mode 100644 index 000000000..e358b2c5b --- /dev/null +++ b/crypto/math-cuda/tests/lde_dev_parity.rs @@ -0,0 +1,86 @@ +//! Parity for the on-GPU trace generator's LDE seam: the device-input keep path +//! [`coset_lde_row_major_with_merkle_tree_keep_dev`] must produce a +//! byte-identical Merkle root and row-major LDE as the host-input keep path +//! [`coset_lde_row_major_with_merkle_tree_keep`] for the same matrix. This is +//! the "Step 1a" isolation test: it validates the seam without any trace-fill +//! kernel — upload a host matrix, run both paths, compare. +//! +//! Requires a GPU (skips cleanly if the CUDA backend is unavailable). + +use math::field::element::FieldElement; +use math::field::goldilocks::GoldilocksField; +use math::field::traits::{IsField, IsPrimeField}; +use math_cuda::device::backend; +use rand::{Rng, SeedableRng}; +use rand_chacha::ChaCha8Rng; + +/// Coset weights `[1/N, g/N, ...]` — the layout `crypto/stark/src/prover.rs` uses. +fn coset_weights(n: usize, coset_offset: u64) -> Vec { + let inv_n = FieldElement::::from(n as u64) + .inv() + .expect("n non-zero"); + let mut w = Vec::with_capacity(n); + let mut cur = *inv_n.value(); + for _ in 0..n { + w.push(cur); + cur = GoldilocksField::mul(&cur, &coset_offset); + } + w +} + +fn assert_dev_matches_host(log_n: u64, m: usize, blowup: usize, seed: u64) { + // Skip cleanly when no GPU/CUDA backend is present. + if backend().is_err() { + eprintln!("skipping lde_dev_parity: no CUDA backend"); + return; + } + + let n = 1usize << log_n; + let mut rng = ChaCha8Rng::seed_from_u64(seed); + // Row-major [row*m + col], the layout the trace generator emits. + let row_major: Vec = (0..n * m).map(|_| rng.r#gen::()).collect(); + + let coset_offset = 7u64; + let weights = coset_weights(n, coset_offset); + + // Host path: uploads `row_major` internally. + let (h_handle, h_lde) = + math_cuda::lde::coset_lde_row_major_with_merkle_tree_keep(&row_major, n, m, blowup, &weights) + .expect("host keep"); + + // Device path: pre-upload the SAME matrix, then run the device-input LDE. + let be = backend().unwrap(); + let stream = be.next_stream(); + let dev = stream.clone_htod(&row_major).expect("upload matrix"); + stream.synchronize().expect("sync upload"); + let (d_handle, d_lde) = math_cuda::lde::coset_lde_row_major_with_merkle_tree_keep_dev( + &dev, n, m, blowup, &weights, + ) + .expect("dev keep"); + + assert_eq!( + h_handle.tree.as_ref().unwrap().root, + d_handle.tree.as_ref().unwrap().root, + "Merkle root mismatch (n={n}, m={m}, blowup={blowup})" + ); + let hc: Vec = h_lde.iter().map(GoldilocksField::canonical).collect(); + let dc: Vec = d_lde.iter().map(GoldilocksField::canonical).collect(); + assert_eq!( + hc, dc, + "row-major LDE mismatch (n={n}, m={m}, blowup={blowup})" + ); +} + +#[test] +fn dev_matches_host_cpu_table_width() { + // CPU table is 38 columns — the P1 target. + assert_dev_matches_host(10, 38, 2, 0xC0DE); + assert_dev_matches_host(12, 38, 4, 0xBEEF); +} + +#[test] +fn dev_matches_host_various_widths() { + assert_dev_matches_host(12, 10, 2, 0x11); // memw_register width + assert_dev_matches_host(14, 16, 2, 0x22); // store width + assert_dev_matches_host(13, 49, 2, 0x33); // memw width +} diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 6a81162a5..59e5af3e5 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -508,6 +508,70 @@ where Some((tree, handle, lde_out)) } +/// Device-input analog of [`try_expand_leaf_and_tree_row_major_keep`]: the main +/// trace is already resident on device (row-major `[row*m + col]`, produced by +/// the on-GPU trace generator), so there is NO host upload — the LDE copies it +/// device-to-device. Returns the same (root-only host tree, resident LDE handle +/// with the trace-domain snapshot for the aux build, row-major host LDE) triple. +pub(crate) fn try_expand_leaf_and_tree_row_major_keep_dev( + input_dev: &math_cuda::CudaSlice, + n: usize, + m: usize, + blowup_factor: usize, + weights: &[FieldElement], +) -> Option<( + MerkleTree, + math_cuda::lde::GpuLdeBase, + Vec>, +)> +where + F: IsField + 'static, + E: IsField + 'static, + B: IsMerkleTreeBackend, +{ + // NOTE: no `lde_size < gpu_lde_threshold()` decline here. The threshold is a + // host-side heuristic (small tables are faster to LDE on the CPU, avoiding the + // H2D). For a DEVICE-resident input there is no valid CPU fallback — the host + // trace is a zeroed placeholder — so declining would silently commit zeros + // (bus imbalance / wrong proof). The input is already on-device, so we always + // run the GPU LDE regardless of size. + let _lde_size = n.saturating_mul(blowup_factor); + if TypeId::of::() != TypeId::of::() { + return None; + } + if TypeId::of::() != TypeId::of::() { + return None; + } + if input_dev.len() != n * m || m == 0 || n == 0 { + return None; + } + + let weights_u64 = unsafe { weights_to_u64::(weights) }; + + GPU_LDE_CALLS.fetch_add(m as u64, Ordering::Relaxed); + GPU_LEAF_HASH_CALLS.fetch_add(1, Ordering::Relaxed); + GPU_MERKLE_TREE_CALLS.fetch_add(1, Ordering::Relaxed); + + let (handle, lde_u64) = math_cuda::lde::coset_lde_row_major_with_merkle_tree_keep_dev( + input_dev, + n, + m, + blowup_factor, + &weights_u64, + ) + .ok()?; + + // Transmute Vec → Vec> (zero-copy, E == GoldilocksField). + let lde_out: Vec> = unsafe { + let mut v = std::mem::ManuallyDrop::new(lde_u64); + Vec::from_raw_parts(v.as_mut_ptr() as *mut FieldElement, v.len(), v.capacity()) + }; + + let root = handle.tree.as_ref()?.root; + let tree = MerkleTree::::from_root(root); + Some((tree, handle, lde_out)) +} + /// Row-major ext3 GPU path: single H2D → row-major NTT (m*3 base-field cols) → /// row-major Keccak → Merkle → single D2H → transpose to GpuLdeExt3 handle. /// Same optimization as the base-field path: no extract_columns, no CPU transpose. diff --git a/crypto/stark/src/instruments.rs b/crypto/stark/src/instruments.rs index 96bf6ffae..d0dee46e9 100644 --- a/crypto/stark/src/instruments.rs +++ b/crypto/stark/src/instruments.rs @@ -238,6 +238,14 @@ static AUX_FINGERPRINT_US: AtomicU64 = AtomicU64::new(0); static AUX_INVERT_US: AtomicU64 = AtomicU64::new(0); static AUX_TERM_US: AtomicU64 = AtomicU64::new(0); static AUX_ACCUM_US: AtomicU64 = AtomicU64::new(0); +// GPU trace-generation transfer accounting (the bytes we aim to eliminate). +// `MAIN_H2D_BYTES` sums every host->device upload of a main-trace matrix in +// `commit_main_trace`; `MAIN_DEV_BUILDS` counts tables that instead fed the LDE +// from an already-device-resident buffer (no upload) — 0 until the on-GPU +// trace generator lands, then the target is "all main tables". +static MAIN_H2D_BYTES: AtomicU64 = AtomicU64::new(0); +static MAIN_H2D_CALLS: AtomicU64 = AtomicU64::new(0); +static MAIN_DEV_BUILDS: AtomicU64 = AtomicU64::new(0); thread_local! { static TIMING_DATA: RefCell> = const { RefCell::new(None) }; @@ -281,6 +289,27 @@ pub fn accum_aux_accumulate(d: Duration) { AUX_ACCUM_US.fetch_add(d.as_micros() as u64, Ordering::Relaxed); } +/// Record a host->device upload of a main-trace matrix (`bytes`) at the LDE +/// commit boundary — the transfer the on-GPU trace generator removes. +pub fn accum_main_h2d(bytes: usize) { + MAIN_H2D_BYTES.fetch_add(bytes as u64, Ordering::Relaxed); + MAIN_H2D_CALLS.fetch_add(1, Ordering::Relaxed); +} + +/// Record a main table whose LDE input was already device-resident (no upload). +pub fn accum_main_dev_build() { + MAIN_DEV_BUILDS.fetch_add(1, Ordering::Relaxed); +} + +/// `(uploaded_bytes, upload_calls, device_resident_builds)` since the last reset. +pub fn take_main_transfer() -> (u64, u64, u64) { + ( + MAIN_H2D_BYTES.swap(0, Ordering::Relaxed), + MAIN_H2D_CALLS.swap(0, Ordering::Relaxed), + MAIN_DEV_BUILDS.swap(0, Ordering::Relaxed), + ) +} + pub fn take_r1_sub() -> Round1SubOps { Round1SubOps { main_lde: Duration::from_micros(R1_MAIN_LDE_US.swap(0, Ordering::Relaxed)), @@ -311,6 +340,9 @@ pub fn reset_all() { AUX_INVERT_US.store(0, Ordering::Relaxed); AUX_TERM_US.store(0, Ordering::Relaxed); AUX_ACCUM_US.store(0, Ordering::Relaxed); + MAIN_H2D_BYTES.store(0, Ordering::Relaxed); + MAIN_H2D_CALLS.store(0, Ordering::Relaxed); + MAIN_DEV_BUILDS.store(0, Ordering::Relaxed); TIMING_DATA.with(|cell| { cell.borrow_mut().take(); }); diff --git a/crypto/stark/src/logup_gpu.rs b/crypto/stark/src/logup_gpu.rs index bc9e88302..1cf8cd51e 100644 --- a/crypto/stark/src/logup_gpu.rs +++ b/crypto/stark/src/logup_gpu.rs @@ -431,7 +431,15 @@ where { return None; } - if trace_len < GPU_LOGUP_MIN_ROWS || main_cols.is_empty() || interactions.is_empty() { + // The `trace_len < GPU_LOGUP_MIN_ROWS` decline is a host-side heuristic (small + // tables are cheaper on the CPU aux path). It must NOT apply when the main is + // device-resident (`main_dev.is_some()`): that table's host trace is a zeroed + // placeholder, so a CPU fallback would build the aux from zeros → LogUp bus + // imbalance. Device-resident tables always build the aux on-device here. + if (trace_len < GPU_LOGUP_MIN_ROWS && main_dev.is_none()) + || main_cols.is_empty() + || interactions.is_empty() + { return None; } if std::env::var_os("LAMBDA_VM_NO_GPU_LOGUP").is_some() { @@ -479,13 +487,7 @@ where None => math_cuda::logup::ResidentMain::Host(&main_flat), }; let ra = math_cuda::logup::logup_aux_resident( - main, - trace_len, - &md, - &alpha_flat, - z_arr, - inv_n, - &stream, + main, trace_len, &md, &alpha_flat, z_arr, inv_n, &stream, ) .ok()?; crate::gpu_lde::GPU_LOGUP_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed); diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 5c03292da..87f7098cc 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -777,6 +777,42 @@ pub trait IsStarkProver< // Fused GPU path (cuda only): row-major NTT — single H2D from the // already-row-major trace, no column extraction, no transpose. // Falls back to CPU if GPU path returns None. + // Device-born main trace (on-GPU trace generator): feed the LDE from the + // resident buffer directly — no host upload. Falls through to the host + // path if the handle is absent or the GPU path declines. + #[cfg(feature = "cuda")] + if precomputed.is_none() + && let Some(dev) = trace.main_input_dev() + { + let num_cols = trace.num_main_columns; + let n = trace.num_rows(); + #[cfg(feature = "instruments")] + let t_sub = Instant::now(); + if let Some((tree, handle, main_data)) = + crate::gpu_lde::try_expand_leaf_and_tree_row_major_keep_dev::< + Field, + Field, + BatchedMerkleTreeBackend, + >( + dev, n, num_cols, domain.blowup_factor, &twiddles.coset_weights + ) + { + #[cfg(feature = "instruments")] + crate::instruments::accum_r1_main(t_sub.elapsed(), std::time::Duration::ZERO); + // No host->device upload happened for this table. + #[cfg(feature = "instruments")] + crate::instruments::accum_main_dev_build(); + let root = tree.root; + return Ok(( + TableCommit::plain(tree, root), + (main_data, num_cols), + Some(handle), + )); + } + } + + // Fused GPU host path (cuda only): row-major NTT — single H2D from the + // already-row-major trace, no column extraction, no transpose. #[cfg(feature = "cuda")] if precomputed.is_none() { let (trace_slice, num_cols) = trace.main_data_row_major(); @@ -805,6 +841,12 @@ pub trait IsStarkProver< let root = tree.root; #[cfg(feature = "instruments")] crate::instruments::accum_r1_main(main_lde_dur, std::time::Duration::ZERO); + // This table's main matrix was uploaded host->device for the LDE. + // Counts the transfer the on-GPU trace generator will remove. + #[cfg(feature = "instruments")] + crate::instruments::accum_main_h2d( + trace_slice.len() * std::mem::size_of::(), + ); return Ok(( TableCommit::plain(tree, root), (main_data, num_cols), @@ -2191,6 +2233,16 @@ pub trait IsStarkProver< heap_snaps.push(s); } + // Free the device-born main input buffers now that every main trace is + // committed. They are only the D2D source for `commit_main_trace`; the + // resident LDE handle (`main_gpu_handles`) carries what R2-4 needs. Holding + // them through the aux/DEEP/FRI VRAM peak makes large blocks OOM where the + // CPU-trace stream-and-free path fits. + #[cfg(feature = "cuda")] + for (_, trace, _) in air_trace_pairs.iter_mut() { + trace.clear_main_input_dev(); + } + // ===================================================================== // Round 1, Phase B: Sample shared LogUp challenges // ===================================================================== diff --git a/crypto/stark/src/trace.rs b/crypto/stark/src/trace.rs index 6d40425b7..b4647bad4 100644 --- a/crypto/stark/src/trace.rs +++ b/crypto/stark/src/trace.rs @@ -45,8 +45,41 @@ where /// LDE did not run for this table. #[cfg(feature = "cuda")] pub(crate) main_trace_dev: Option, + /// Device-born main trace produced by the on-GPU trace generator, row-major + /// `[row*num_main_columns + col]`, `num_rows * num_main_columns` u64s. When + /// present, `commit_main_trace` feeds the LDE from this buffer directly + /// (device-to-device, no host upload). None on the CPU trace-gen path. + #[cfg(feature = "cuda")] + pub(crate) main_input_dev: Option, +} + +/// Device-born main trace (row-major `[row*num_main_columns + col]`) produced by +/// the on-GPU trace generator, feeding `commit_main_trace` directly. GPU-only and +/// transient; excluded from logical trace equality and opaque in `Debug`, +/// matching `ResidentMainTrace`. +#[cfg(feature = "cuda")] +#[derive(Clone)] +pub(crate) struct DeviceMainInput { + pub(crate) buf: std::sync::Arc>, +} + +#[cfg(feature = "cuda")] +impl core::fmt::Debug for DeviceMainInput { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("DeviceMainInput").finish_non_exhaustive() + } +} + +#[cfg(feature = "cuda")] +impl PartialEq for DeviceMainInput { + fn eq(&self, _other: &Self) -> bool { + true + } } +#[cfg(feature = "cuda")] +impl Eq for DeviceMainInput {} + /// Device-resident trace-domain main columns (column-major `[col*rows + row]`), /// retained from the R1 main LDE for the aux fingerprint kernel. GPU-only and /// transient; the device buffer is excluded from logical trace equality (only @@ -105,6 +138,8 @@ where resident_aux_ok: true, #[cfg(feature = "cuda")] main_trace_dev: None, + #[cfg(feature = "cuda")] + main_input_dev: None, } } @@ -133,6 +168,8 @@ where resident_aux_ok: true, #[cfg(feature = "cuda")] main_trace_dev: None, + #[cfg(feature = "cuda")] + main_input_dev: None, } } @@ -154,6 +191,8 @@ where resident_aux_ok: true, #[cfg(feature = "cuda")] main_trace_dev: None, + #[cfg(feature = "cuda")] + main_input_dev: None, } } @@ -213,6 +252,32 @@ where self.main_trace_dev = None; } + /// Attach a device-born main trace (row-major `[row*num_main_columns + col]`, + /// `num_rows * num_main_columns` u64s) from the on-GPU trace generator, so + /// `commit_main_trace` feeds the LDE from it directly (no host upload). + #[cfg(feature = "cuda")] + pub fn set_main_input_dev(&mut self, buf: std::sync::Arc>) { + self.main_input_dev = Some(DeviceMainInput { buf }); + } + + /// The device-born main trace buffer, if the on-GPU trace generator produced + /// one for this table. + #[cfg(feature = "cuda")] + pub fn main_input_dev(&self) -> Option<&math_cuda::CudaSlice> { + self.main_input_dev.as_ref().map(|b| b.buf.as_ref()) + } + + /// Drop the device-born main input buffer. It is only read by `commit_main_trace` + /// (the D2D source for the LDE); once the main commit is done nothing else uses + /// it, so freeing it after Round-1 Phase A reclaims the main-trace-sized device + /// buffers (≈ the whole main trace) before the aux LDE / DEEP / FRI VRAM peak — + /// without this the on-GPU trace generator holds them to end-of-proof and large + /// blocks (e.g. ethrex 10-tx) OOM where the CPU-trace (stream-and-free) path fits. + #[cfg(feature = "cuda")] + pub fn clear_main_input_dev(&mut self) { + self.main_input_dev = None; + } + pub fn num_steps(&self) -> usize { debug_assert!(self.main_table.height.is_multiple_of(self.step_size)); self.main_table.height / self.step_size diff --git a/prover/Cargo.toml b/prover/Cargo.toml index 3695689d6..ba3a32977 100644 --- a/prover/Cargo.toml +++ b/prover/Cargo.toml @@ -7,7 +7,7 @@ license.workspace = true [features] default = ["parallel"] parallel = ["stark/parallel", "math/parallel", "crypto/parallel", "dep:rayon"] -cuda = ["stark/cuda"] +cuda = ["stark/cuda", "dep:math-cuda"] test-cuda-faults = ["cuda", "stark/test-cuda-faults"] debug-checks = ["stark/debug-checks"] instruments = ["stark/instruments"] @@ -20,6 +20,7 @@ crypto = { path = "../crypto/crypto" } math = { path = "../crypto/math" } executor = { path = "../executor" } ecsm = { path = "../crypto/ecsm" } +math-cuda = { path = "../crypto/math-cuda", optional = true } serde = { version = "1.0", features = ["derive"] } rayon = { version = "1.8.0", optional = true } sysinfo = { version = "0.31", default-features = false, features = ["system"] } diff --git a/prover/src/lib.rs b/prover/src/lib.rs index 9c315faf1..eb6399914 100644 --- a/prover/src/lib.rs +++ b/prover/src/lib.rs @@ -972,6 +972,16 @@ pub fn prove_with_options_and_inputs( drop(__root); let spans = stark::instruments::take_timeline(); print!("{}", stark::instruments::format_timeline(&spans)); + // Main-trace CPU->GPU transfer accounting: the bytes the on-GPU trace + // generator aims to eliminate, and how many tables already skipped the + // upload by feeding the LDE a device-resident buffer. + let (h2d_bytes, h2d_calls, dev_builds) = stark::instruments::take_main_transfer(); + if h2d_calls > 0 || dev_builds > 0 { + println!( + "[transfer] main-trace H2D: {:.1} MiB over {h2d_calls} table(s); device-resident builds: {dev_builds}", + h2d_bytes as f64 / (1024.0 * 1024.0), + ); + } if let Ok(path) = std::env::var("LAMBDA_VM_TIMELINE_JSON") { let _ = std::fs::write(&path, stark::instruments::timeline_json(&spans)); println!("[timeline] wrote {path}"); diff --git a/prover/src/tables/gpu_trace.rs b/prover/src/tables/gpu_trace.rs new file mode 100644 index 000000000..420210f86 --- /dev/null +++ b/prover/src/tables/gpu_trace.rs @@ -0,0 +1,482 @@ +//! On-GPU trace generation: build trace tables directly in device memory so +//! they feed the already-GPU LDE/commit without a host round-trip. +//! +//! This module is compiled only under the `cuda` feature. It hosts the +//! device-build dispatch (added table-by-table) plus the kill-switch used to +//! A/B the GPU path against the CPU trace generator. +//! +//! Design: `reports/tracegen/GPU-TRACEGEN-DESIGN-V2.md`. +#![cfg(feature = "cuda")] + +use std::sync::{Arc, OnceLock}; + +use stark::trace::TraceTable; + +use std::collections::HashMap; + +use super::cpu::{self, CpuOperation}; +use super::load::{self, LoadOperation}; +use super::lt::{self, LtOperation}; +use super::memw::MemwOperation; +use super::memw_aligned; +use super::memw_register::{self, RegRow}; +use super::shift::{self, ShiftOperation}; +use super::store::{self, StoreOperation}; +use super::types::{GoldilocksExtension, GoldilocksField}; + +/// When set (`LAMBDA_VM_CPU_TRACE=1`), all GPU trace-build dispatchers return +/// `None` so callers fall back to the CPU trace generator. This is the one-flag +/// A/B switch: same binary, `LAMBDA_VM_CPU_TRACE=1` runs the CPU baseline, +/// unset runs the GPU path. Read once and cached. +pub(crate) fn gpu_trace_disabled() -> bool { + static DISABLED: OnceLock = OnceLock::new(); + *DISABLED.get_or_init(|| { + std::env::var("LAMBDA_VM_CPU_TRACE") + .map(|v| v != "0" && !v.is_empty()) + .unwrap_or(false) + }) +} + +// ============================================================================= +// CPU table (the first table built on device — see GPU-TRACEGEN-DESIGN-V2 §P1) +// ============================================================================= + +/// Marshal one chunk of `CpuOperation`s into the packed layout the `trace_cpu` +/// kernel consumes (stride `CPU_OP_STRIDE` u64/op). The kernel does the same +/// bit-slicing as `cpu::generate_cpu_trace`, so this only copies fields — no +/// per-column encoding on the host. +fn pack_cpu_ops(chunk: &[CpuOperation]) -> Vec { + let stride = math_cuda::trace_cpu::CPU_OP_STRIDE; + let mut packed = vec![0u64; chunk.len() * stride]; + for (i, op) in chunk.iter().enumerate() { + let f = &op.decode.fields; + let flags = (f.word_instr as u64) + | ((f.read_register1 as u64) << 1) + | ((f.read_register2 as u64) << 2) + | ((f.write_register as u64) << 3) + | ((f.alu as u64) << 4) + | ((f.add as u64) << 5) + | ((f.sub as u64) << 6) + | ((f.memory as u64) << 7) + | ((f.branch as u64) << 8) + | ((f.ecall as u64) << 9) + | ((op.branch_cond as u64) << 10); + let bytes = (f.rs1 as u64) + | ((f.rs2 as u64) << 8) + | ((f.rd as u64) << 16) + | ((f.half_instruction_length as u64) << 24) + | ((f.alu_flags as u64) << 32) + | ((f.mem_flags as u64) << 40); + let b = i * stride; + packed[b] = op.timestamp; + packed[b + 1] = op.decode.pc; + packed[b + 2] = op.decode.imm; + packed[b + 3] = op.next_pc; + packed[b + 4] = op.rvd; + packed[b + 5] = op.rv1; + packed[b + 6] = op.rv2; + packed[b + 7] = op.arg2; + packed[b + 8] = op.res; + packed[b + 9] = flags; + packed[b + 10] = bytes; + } + packed +} + +/// Build one CPU trace-table chunk on device: pack ops → GPU fill → a +/// `TraceTable` whose main matrix is resident on device (fed to the LDE with no +/// upload). The host main table is a zeroed placeholder sized for the correct +/// `num_rows`; it is never read on the GPU commit path (commit consumes the +/// device buffer, the aux build reads the resident snapshot, queries gather from +/// the device tree). Returns `None` if the GPU build fails, so the caller can +/// fall back to the CPU generator. +fn build_cpu_chunk(chunk: &[CpuOperation]) -> Option> { + let n = chunk.len(); + let num_rows = n.next_power_of_two().max(4); + let last_ts = chunk.last().map(|op| op.timestamp).unwrap_or(0); + let packed = pack_cpu_ops(chunk); + let dev = math_cuda::trace_cpu::gpu_build_cpu_trace(&packed, n, num_rows, last_ts).ok()?; + let mut trace = TraceTable::new_main( + crate::tables::types::zeroed_fe_vec(num_rows * cpu::cols::NUM_COLUMNS), + cpu::cols::NUM_COLUMNS, + 1, + ); + trace.set_main_input_dev(Arc::new(dev)); + Some(trace) +} + +/// Build all CPU trace-table chunks on device, mirroring `chunk_and_generate`'s +/// chunking (`max_rows`, one empty chunk when there are no ops). Returns `None` +/// when the kill-switch is set or any chunk fails to build, so the caller falls +/// back to the CPU generator. +pub(crate) fn gpu_build_cpu_trace_tables( + cpu_ops: &[CpuOperation], + max_rows: usize, +) -> Option>> { + if gpu_trace_disabled() { + return None; + } + let chunks: Vec<&[CpuOperation]> = if cpu_ops.is_empty() { + vec![&[][..]] + } else { + cpu_ops.chunks(max_rows).collect() + }; + let mut tables = Vec::with_capacity(chunks.len()); + for chunk in chunks { + tables.push(build_cpu_chunk(chunk)?); + } + Some(tables) +} + +// ============================================================================= +// MEMW_R (register fast path — the biggest table, ~15M rows on ethrex) +// ============================================================================= + +/// Build one MEMW_R trace-table chunk on device: marshal the walked `RegRow`s into +/// the SoA the `memw_register_fill` kernel consumes, fill the 10 columns row-major +/// on device, and leave the matrix RESIDENT (fed to the LDE with no full-column +/// upload). The host main table is a zeroed placeholder sized to `num_rows`; it is +/// never read on the GPU commit path. The `old_*` come from the (correct, +/// precompile-inclusive) sequential walk — so this is program-agnostic — with only +/// the compact `RegRow` fields uploaded, not the full column matrix. Returns `None` +/// on GPU failure so the caller can fall back to the CPU fill. +fn build_memw_register_chunk( + chunk: &[RegRow], +) -> Option> { + let n = chunk.len(); + let num_rows = n.next_power_of_two().max(4); + + let mut reg_addr = Vec::with_capacity(n); + let mut ts = Vec::with_capacity(n); + let mut value = Vec::with_capacity(n); + let mut is_read = Vec::with_capacity(n); + let mut old_value = Vec::with_capacity(n); + let mut old_ts = Vec::with_capacity(n); + for r in chunk { + let (ra, t, v, ir, ov, ot) = r.fill_soa(); + reg_addr.push(ra); + ts.push(t); + value.push(v); + is_read.push(ir); + old_value.push(ov); + old_ts.push(ot); + } + + let dev = math_cuda::trace_cpu::gpu_fill_memw_register( + ®_addr, &ts, &value, &is_read, &old_value, &old_ts, num_rows, + ) + .ok()?; + + let mut trace = TraceTable::new_main( + crate::tables::types::zeroed_fe_vec(num_rows * memw_register::cols::NUM_COLUMNS), + memw_register::cols::NUM_COLUMNS, + 1, + ); + trace.set_main_input_dev(Arc::new(dev)); + Some(trace) +} + +/// Build all MEMW_R trace-table chunks on device, mirroring `chunk_and_generate`'s +/// chunking (`max_rows`, one empty chunk when there are no rows). Returns `None` +/// when the kill-switch is set or any chunk fails to build, so the caller falls +/// back to the CPU fill. +pub(crate) fn gpu_build_memw_register_tables( + rows: &[RegRow], + max_rows: usize, +) -> Option>> { + if gpu_trace_disabled() { + return None; + } + let chunks: Vec<&[RegRow]> = if rows.is_empty() { + vec![&[][..]] + } else { + rows.chunks(max_rows).collect() + }; + let mut tables = Vec::with_capacity(chunks.len()); + for chunk in chunks { + tables.push(build_memw_register_chunk(chunk)?); + } + Some(tables) +} + +// ============================================================================= +// MEMW_A (aligned memory — the biggest remaining uploader, ~2M rows on ethrex) +// ============================================================================= + +/// Pack one aligned `MemwOperation` into the stride-`MEMW_ALIGNED_STRIDE` layout +/// the `memw_aligned_fill` kernel consumes (see `trace_cpu.cu`). The op is already +/// walked (old_value/old_timestamp filled), so this only copies fields; value/old +/// (`[u32; 8]` each) pack two-per-u64. +pub(crate) fn pack_memw_aligned_op( + op: &MemwOperation, +) -> [u64; math_cuda::trace_cpu::MEMW_ALIGNED_STRIDE] { + let flags = (op.is_register as u64) | ((op.is_read as u64) << 1) | ((op.width as u64) << 8); + let v = &op.value; + let o = &op.old; + [ + flags, + op.base_address, + op.timestamp, + op.old_timestamp[0], + v[0] as u64 | ((v[1] as u64) << 32), + v[2] as u64 | ((v[3] as u64) << 32), + v[4] as u64 | ((v[5] as u64) << 32), + v[6] as u64 | ((v[7] as u64) << 32), + o[0] as u64 | ((o[1] as u64) << 32), + o[2] as u64 | ((o[3] as u64) << 32), + o[4] as u64 | ((o[5] as u64) << 32), + o[6] as u64 | ((o[7] as u64) << 32), + ] +} + +/// Build one MEMW_A trace-table chunk on device: pack the walked ops → GPU fill → +/// a resident matrix fed to the LDE with no full-column upload (only the compact +/// packed ops are H2D'd). Returns `None` on GPU failure so the caller falls back. +fn build_memw_aligned_chunk( + chunk: &[MemwOperation], +) -> Option> { + let n = chunk.len(); + let num_rows = n.next_power_of_two().max(4); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::MEMW_ALIGNED_STRIDE); + for op in chunk { + packed.extend_from_slice(&pack_memw_aligned_op(op)); + } + let dev = math_cuda::trace_cpu::gpu_build_memw_aligned_trace(&packed, n, num_rows).ok()?; + let mut trace = TraceTable::new_main( + crate::tables::types::zeroed_fe_vec(num_rows * memw_aligned::cols::NUM_COLUMNS), + memw_aligned::cols::NUM_COLUMNS, + 1, + ); + trace.set_main_input_dev(Arc::new(dev)); + Some(trace) +} + +/// Build all MEMW_A trace-table chunks on device, mirroring `chunk_and_generate`'s +/// chunking. Returns `None` when the kill-switch is set or any chunk fails to build. +pub(crate) fn gpu_build_memw_aligned_tables( + ops: &[MemwOperation], + max_rows: usize, +) -> Option>> { + if gpu_trace_disabled() { + return None; + } + let chunks: Vec<&[MemwOperation]> = if ops.is_empty() { + vec![&[][..]] + } else { + ops.chunks(max_rows).collect() + }; + let mut tables = Vec::with_capacity(chunks.len()); + for chunk in chunks { + tables.push(build_memw_aligned_chunk(chunk)?); + } + Some(tables) +} + +// ============================================================================= +// LOAD / STORE (per-row-map memory tables — same shape as MEMW_A) +// ============================================================================= + +/// Pack one `LoadOperation` into the `load_fill` stride (see `trace_cpu.cu`). +pub(crate) fn pack_load_op(op: &LoadOperation) -> [u64; math_cuda::trace_cpu::LOAD_STRIDE] { + let flags = (op.signed as u64) | ((op.width as u64) << 8); + let r = &op.res; + [ + flags, + op.base_address, + op.timestamp, + r[0] | (r[1] << 32), + r[2] | (r[3] << 32), + r[4] | (r[5] << 32), + r[6] | (r[7] << 32), + ] +} + +/// Pack one `StoreOperation` into the `store_fill` stride (see `trace_cpu.cu`). +pub(crate) fn pack_store_op(op: &StoreOperation) -> [u64; math_cuda::trace_cpu::STORE_STRIDE] { + let flags = + (op.write2 as u64) | ((op.write4 as u64) << 1) | ((op.write8 as u64) << 2); + [flags, op.base_address, op.timestamp, op.value] +} + +fn build_load_chunk(chunk: &[LoadOperation]) -> Option> { + let n = chunk.len(); + let num_rows = n.next_power_of_two().max(4); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::LOAD_STRIDE); + for op in chunk { + packed.extend_from_slice(&pack_load_op(op)); + } + let dev = math_cuda::trace_cpu::gpu_build_load_trace(&packed, n, num_rows).ok()?; + let mut trace = TraceTable::new_main( + crate::tables::types::zeroed_fe_vec(num_rows * load::cols::NUM_COLUMNS), + load::cols::NUM_COLUMNS, + 1, + ); + trace.set_main_input_dev(Arc::new(dev)); + Some(trace) +} + +pub(crate) fn gpu_build_load_tables( + ops: &[LoadOperation], + max_rows: usize, +) -> Option>> { + if gpu_trace_disabled() { + return None; + } + let chunks: Vec<&[LoadOperation]> = if ops.is_empty() { + vec![&[][..]] + } else { + ops.chunks(max_rows).collect() + }; + let mut tables = Vec::with_capacity(chunks.len()); + for chunk in chunks { + tables.push(build_load_chunk(chunk)?); + } + Some(tables) +} + +fn build_store_chunk( + chunk: &[StoreOperation], +) -> Option> { + let n = chunk.len(); + let num_rows = n.next_power_of_two().max(4); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::STORE_STRIDE); + for op in chunk { + packed.extend_from_slice(&pack_store_op(op)); + } + let dev = math_cuda::trace_cpu::gpu_build_store_trace(&packed, n, num_rows).ok()?; + let mut trace = TraceTable::new_main( + crate::tables::types::zeroed_fe_vec(num_rows * store::cols::NUM_COLUMNS), + store::cols::NUM_COLUMNS, + 1, + ); + trace.set_main_input_dev(Arc::new(dev)); + Some(trace) +} + +pub(crate) fn gpu_build_store_tables( + ops: &[StoreOperation], + max_rows: usize, +) -> Option>> { + if gpu_trace_disabled() { + return None; + } + let chunks: Vec<&[StoreOperation]> = if ops.is_empty() { + vec![&[][..]] + } else { + ops.chunks(max_rows).collect() + }; + let mut tables = Vec::with_capacity(chunks.len()); + for chunk in chunks { + tables.push(build_store_chunk(chunk)?); + } + Some(tables) +} + +// ============================================================================= +// SHIFT (ALU table, no dedup — the kernel recomputes the shift aux on device) +// ============================================================================= + +/// Pack one `ShiftOperation` into the `shift_fill` stride (see `trace_cpu.cu`): +/// value (4×u16 in_halves), full shift_amount, and the flag bits. The kernel +/// recomputes bit_shift/zbs/x/y/limb_shift/out, so only 3 u64/op upload. +pub(crate) fn pack_shift_op(op: &ShiftOperation) -> [u64; math_cuda::trace_cpu::SHIFT_STRIDE] { + let h = &op.in_halves; + let value = (h[0] as u64) | ((h[1] as u64) << 16) | ((h[2] as u64) << 32) | ((h[3] as u64) << 48); + let flags = (op.direction as u64) | ((op.signed as u64) << 1) | ((op.word_instr as u64) << 2); + [value, op.shift_amount, flags] +} + +fn build_shift_chunk( + chunk: &[ShiftOperation], +) -> Option> { + let n = chunk.len(); + let num_rows = n.next_power_of_two().max(4); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::SHIFT_STRIDE); + for op in chunk { + packed.extend_from_slice(&pack_shift_op(op)); + } + let dev = math_cuda::trace_cpu::gpu_build_shift_trace(&packed, n, num_rows).ok()?; + let mut trace = TraceTable::new_main( + crate::tables::types::zeroed_fe_vec(num_rows * shift::cols::NUM_COLUMNS), + shift::cols::NUM_COLUMNS, + 1, + ); + trace.set_main_input_dev(Arc::new(dev)); + Some(trace) +} + +pub(crate) fn gpu_build_shift_tables( + ops: &[ShiftOperation], + max_rows: usize, +) -> Option>> { + if gpu_trace_disabled() { + return None; + } + let chunks: Vec<&[ShiftOperation]> = if ops.is_empty() { + vec![&[][..]] + } else { + ops.chunks(max_rows).collect() + }; + let mut tables = Vec::with_capacity(chunks.len()); + for chunk in chunks { + tables.push(build_shift_chunk(chunk)?); + } + Some(tables) +} + +// ============================================================================= +// LT (ALU dedup table): host per-chunk HashMap dedup → device fill (compute) +// ============================================================================= + +/// Pack one unique `LtOperation` + its multiplicity into the `lt_fill` stride. +pub(crate) fn pack_lt_op(op: &LtOperation, mult: u64) -> [u64; math_cuda::trace_cpu::LT_STRIDE] { + let flags = (op.signed as u64) | ((op.invert as u64) << 1); + [op.lhs, op.rhs, flags, mult] +} + +/// Build one LT trace-table chunk on device. Dedup happens HERE on the host (the +/// same per-chunk HashMap `generate_lt_trace` uses), then the unique ops + summed +/// multiplicities are filled on device. LT rides the permutation-invariant ALU +/// bus, so any row order is valid (validated by multiset/prove, not byte order). +fn build_lt_chunk(chunk: &[LtOperation]) -> Option> { + let mut map: HashMap = HashMap::new(); + for op in chunk { + *map.entry(op.clone()).or_insert(0) += 1; + } + let unique: Vec<(LtOperation, u64)> = map.into_iter().collect(); + let n = unique.len(); + let num_rows = n.next_power_of_two().max(4); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::LT_STRIDE); + for (op, mult) in &unique { + packed.extend_from_slice(&pack_lt_op(op, *mult)); + } + let dev = math_cuda::trace_cpu::gpu_build_lt_trace(&packed, n, num_rows).ok()?; + let mut trace = TraceTable::new_main( + crate::tables::types::zeroed_fe_vec(num_rows * lt::cols::NUM_COLUMNS), + lt::cols::NUM_COLUMNS, + 1, + ); + trace.set_main_input_dev(Arc::new(dev)); + Some(trace) +} + +pub(crate) fn gpu_build_lt_tables( + ops: &[LtOperation], + max_rows: usize, +) -> Option>> { + if gpu_trace_disabled() { + return None; + } + // Chunk the RAW ops exactly like `chunk_and_generate`; each chunk dedups + // independently (matching `generate_lt_trace` per chunk). + let chunks: Vec<&[LtOperation]> = if ops.is_empty() { + vec![&[][..]] + } else { + ops.chunks(max_rows).collect() + }; + let mut tables = Vec::with_capacity(chunks.len()); + for chunk in chunks { + tables.push(build_lt_chunk(chunk)?); + } + Some(tables) +} diff --git a/prover/src/tables/memw.rs b/prover/src/tables/memw.rs index 338b85467..d786c2839 100644 --- a/prover/src/tables/memw.rs +++ b/prover/src/tables/memw.rs @@ -101,7 +101,7 @@ pub mod cols { // ========================================================================= /// A single MEMW operation to be added to the trace. -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq, Eq)] pub struct MemwOperation { /// Whether this is a register access (true) or memory access (false) pub is_register: bool, diff --git a/prover/src/tables/memw_register.rs b/prover/src/tables/memw_register.rs index 590a55100..12f4c0d32 100644 --- a/prover/src/tables/memw_register.rs +++ b/prover/src/tables/memw_register.rs @@ -103,7 +103,7 @@ pub mod cols { /// - `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)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) struct RegRow { /// Register index 0..=255 (`base_address / 2`); u16 keeps the struct at /// 32 bytes — it is the largest persisted array of the walk. @@ -158,6 +158,25 @@ impl RegRow { } } + /// Marshal to the SoA the on-device MEMW_R fill (`memw_register_fill`) + /// consumes: `(reg_addr = 2*address, timestamp, value, is_read, old_value, + /// old_ts)`. The old_timestamp upper limb is shared with TIMESTAMP_1 + /// (`timestamp >> 32`), matching the column encoding. + #[cfg(feature = "cuda")] + pub(crate) fn fill_soa(&self) -> (u32, u64, u64, u8, u64, u64) { + let value = (self.val0 as u64) | ((self.val1 as u64) << 32); + let old_value = (self.old0 as u64) | ((self.old1 as u64) << 32); + let old_ts = (self.old_ts_lo as u64) | ((self.timestamp >> 32) << 32); + ( + 2 * self.address as u32, + self.timestamp, + value, + self.is_read as u8, + old_value, + old_ts, + ) + } + /// 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 diff --git a/prover/src/tables/mod.rs b/prover/src/tables/mod.rs index 78d5f9a6a..73fbc1b1d 100644 --- a/prover/src/tables/mod.rs +++ b/prover/src/tables/mod.rs @@ -34,6 +34,8 @@ pub mod ecdas; pub mod ecsm; pub mod eq; pub mod global_memory; +#[cfg(feature = "cuda")] +pub mod gpu_trace; pub mod halt; pub mod keccak; pub mod keccak_rc; diff --git a/prover/src/tables/trace_builder.rs b/prover/src/tables/trace_builder.rs index 36fd107c3..1c350fba0 100644 --- a/prover/src/tables/trace_builder.rs +++ b/prover/src/tables/trace_builder.rs @@ -3184,6 +3184,29 @@ fn build_traces( // generate→spill order keeps trace memory bounded. let cpu_ops_ref = &cpu_ops; let gen_cpus = || { + // On-GPU CPU-table build (cuda, kill-switch off): the matrix is born + // resident on device and feeds the LDE with no host upload. Falls back + // to the CPU generator otherwise. Disk-spill needs the host matrix to + // spill, so it always uses the CPU path. + #[cfg(feature = "cuda")] + { + let use_gpu = { + #[cfg(feature = "disk-spill")] + { + storage_mode != StorageMode::Disk + } + #[cfg(not(feature = "disk-spill"))] + { + true + } + }; + if use_gpu + && let Some(tables) = + crate::tables::gpu_trace::gpu_build_cpu_trace_tables(cpu_ops_ref, max_rows.cpu) + { + return Ok(tables); + } + } chunk_and_generate( cpu_ops_ref, max_rows.cpu, @@ -3202,6 +3225,31 @@ fn build_traces( ) }; let gen_memw_aligneds = || { + // On-GPU MEMW_A build (cuda, kill-switch off): the biggest remaining + // uploader (~2M rows on ethrex). Fill columns on device from the walked + // ops and leave the matrix resident (only the compact packed ops upload). + // Falls back to the CPU fill otherwise (and for disk-spill). + #[cfg(feature = "cuda")] + { + let use_gpu = { + #[cfg(feature = "disk-spill")] + { + storage_mode != StorageMode::Disk + } + #[cfg(not(feature = "disk-spill"))] + { + true + } + }; + if use_gpu + && let Some(tables) = crate::tables::gpu_trace::gpu_build_memw_aligned_tables( + &memw_aligned_ops, + max_rows.memw_aligned, + ) + { + return Ok(tables); + } + } chunk_and_generate( &memw_aligned_ops, max_rows.memw_aligned, @@ -3211,6 +3259,31 @@ fn build_traces( ) }; let gen_memw_registers = || { + // On-GPU MEMW_R build (cuda, kill-switch off): fill the columns on device + // from the walked RegRows and leave the matrix resident so it feeds the LDE + // with no full-column upload. Falls back to the CPU fill otherwise (and for + // disk-spill, which needs the host matrix to spill). + #[cfg(feature = "cuda")] + { + let use_gpu = { + #[cfg(feature = "disk-spill")] + { + storage_mode != StorageMode::Disk + } + #[cfg(not(feature = "disk-spill"))] + { + true + } + }; + if use_gpu + && let Some(tables) = crate::tables::gpu_trace::gpu_build_memw_register_tables( + &memw_register_rows, + max_rows.memw_register, + ) + { + return Ok(tables); + } + } // Direct-to-column fill from compact RegRows — the register fast path never // materializes a `Vec`. chunk_and_generate( @@ -3222,6 +3295,25 @@ fn build_traces( ) }; let gen_loads = || { + #[cfg(feature = "cuda")] + { + let use_gpu = { + #[cfg(feature = "disk-spill")] + { + storage_mode != StorageMode::Disk + } + #[cfg(not(feature = "disk-spill"))] + { + true + } + }; + if use_gpu + && let Some(tables) = + crate::tables::gpu_trace::gpu_build_load_tables(&load_ops, max_rows.load) + { + return Ok(tables); + } + } chunk_and_generate( &load_ops, max_rows.load, @@ -3231,6 +3323,25 @@ fn build_traces( ) }; let gen_lts = || { + #[cfg(feature = "cuda")] + { + let use_gpu = { + #[cfg(feature = "disk-spill")] + { + storage_mode != StorageMode::Disk + } + #[cfg(not(feature = "disk-spill"))] + { + true + } + }; + if use_gpu + && let Some(tables) = + crate::tables::gpu_trace::gpu_build_lt_tables(<_ops, max_rows.lt) + { + return Ok(tables); + } + } chunk_and_generate( <_ops, max_rows.lt, @@ -3240,6 +3351,25 @@ fn build_traces( ) }; let gen_shifts = || { + #[cfg(feature = "cuda")] + { + let use_gpu = { + #[cfg(feature = "disk-spill")] + { + storage_mode != StorageMode::Disk + } + #[cfg(not(feature = "disk-spill"))] + { + true + } + }; + if use_gpu + && let Some(tables) = + crate::tables::gpu_trace::gpu_build_shift_tables(&shift_ops, max_rows.shift) + { + return Ok(tables); + } + } chunk_and_generate( &shift_ops, max_rows.shift, @@ -3297,6 +3427,25 @@ fn build_traces( ) }; let gen_stores = || { + #[cfg(feature = "cuda")] + { + let use_gpu = { + #[cfg(feature = "disk-spill")] + { + storage_mode != StorageMode::Disk + } + #[cfg(not(feature = "disk-spill"))] + { + true + } + }; + if use_gpu + && let Some(tables) = + crate::tables::gpu_trace::gpu_build_store_tables(&store_ops, max_rows.store) + { + return Ok(tables); + } + } chunk_and_generate::( &store_ops, max_rows.store, @@ -4393,3 +4542,355 @@ impl Traces { ) } } + +#[cfg(test)] +mod gpu_fill_tests { + //! Byte-parity (and multiset, for order-independent LT) between each table's + //! device fill and the CPU `generate_*_trace`, so the on-GPU trace tables are + //! bit-identical to the CPU path. Each test skips cleanly with no CUDA backend. + + use super::*; + + /// MEMW_R (register fast-path) device fill must be byte-identical to the CPU + /// `generate_memw_register_trace_from_rows`. The rows are the walked register + /// rows; here we build synthetic `RegRow`s (read + write, varied values/old/ts) + /// and fill both ways. + #[cfg(feature = "cuda")] + #[test] + fn gpu_memw_register_fill_matches_cpu() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_memw_register_fill_matches_cpu: no CUDA backend"); + return; + } + let mut rows = Vec::new(); + for i in 0..500u64 { + let reg_addr = 2 * (i % 40); // even word-address = 2*reg_index + let ts = 4 * i + 100; + let old_ts = 4 * i + 50; + let val = i.wrapping_mul(2_654_435_761); + let old = i.wrapping_mul(40_503) ^ 0xABCD_1234; + let is_read = i % 3 == 0; + rows.push(memw_register::RegRow::new( + reg_addr, + ts, + val as u32, + (val >> 32) as u32, + old as u32, + (old >> 32) as u32, + old_ts, + is_read, + )); + } + let num_rows = rows.len().next_power_of_two().max(4); + + let cpu_table = memw_register::generate_memw_register_trace_from_rows(&rows); + let (cpu_fe, w) = cpu_table.main_data_row_major(); + assert_eq!(w, math_cuda::trace_cpu::MEMW_REGISTER_NCOLS); + let cpu_u64: Vec = cpu_fe + .iter() + .map(|e| unsafe { *(e.value() as *const _ as *const u64) }) + .collect(); + + let mut reg_addr = Vec::new(); + let mut ts = Vec::new(); + let mut value = Vec::new(); + let mut is_read = Vec::new(); + let mut old_value = Vec::new(); + let mut old_tsv = Vec::new(); + for r in &rows { + let (ra, t, v, ir, ov, ot) = r.fill_soa(); + reg_addr.push(ra); + ts.push(t); + value.push(v); + is_read.push(ir); + old_value.push(ov); + old_tsv.push(ot); + } + let gpu_u64 = math_cuda::trace_cpu::gpu_fill_memw_register_host( + ®_addr, &ts, &value, &is_read, &old_value, &old_tsv, num_rows, + ) + .expect("device MEMW_R fill must run on a box with a CUDA backend"); + + assert_eq!(gpu_u64.len(), num_rows * math_cuda::trace_cpu::MEMW_REGISTER_NCOLS); + assert_eq!(gpu_u64, cpu_u64, "device MEMW_R fill must be byte-identical to the CPU fill"); + } + + /// MEMW_A (aligned memory — the biggest remaining uploader) device fill must be + /// byte-identical to the CPU `generate_memw_aligned_trace`. Exercises memory + /// (widths 2/4/8), register fallback (`is_register`), high address bits, and + /// the 2×u32-per-u64 value/old packing. Skips cleanly with no GPU. + #[cfg(feature = "cuda")] + #[test] + fn gpu_memw_aligned_fill_matches_cpu() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_memw_aligned_fill_matches_cpu: no CUDA backend"); + return; + } + let mut ops = Vec::new(); + for i in 0..500u64 { + let width = [2u8, 4, 8][(i % 3) as usize]; + let is_read = i % 2 == 0; + let is_register = i % 7 == 0; + // Vary the high address word too (whh split has a 32-bit high word). + let base = (0x1_0000_0000u64 * (i % 4)) + 0x8000 + i * 8; + let value = [ + (i as u32).wrapping_mul(2654435761), + (i as u32) ^ 0xABCD_1234, + i as u32, + 0, + 7, + 0, + 0, + (i as u32) & 0xFF, + ]; + let old = [ + (i as u32).wrapping_add(99), + (i as u32) ^ 0x0BAD_F00D, + 0, + 3, + 0, + 0, + 0, + 0, + ]; + let ts = 4 * i + 100; + let old_ts = [4 * i + 50; 8]; + ops.push( + MemwOperation::new(is_register, base, value, ts, width, is_read) + .with_old(old, old_ts), + ); + } + let n = ops.len(); + let num_rows = n.next_power_of_two().max(4); + + let cpu_table = memw_aligned::generate_memw_aligned_trace(&ops); + let (cpu_fe, width_cols) = cpu_table.main_data_row_major(); + assert_eq!(width_cols, math_cuda::trace_cpu::MEMW_ALIGNED_NCOLS); + let cpu_u64: Vec = cpu_fe + .iter() + .map(|e| unsafe { *(e.value() as *const _ as *const u64) }) + .collect(); + + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::MEMW_ALIGNED_STRIDE); + for op in &ops { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_memw_aligned_op(op)); + } + let gpu_u64 = math_cuda::trace_cpu::gpu_build_memw_aligned_trace_host(&packed, n, num_rows) + .expect("device MEMW_A build must run on a box with a CUDA backend"); + + assert_eq!( + gpu_u64.len(), + num_rows * math_cuda::trace_cpu::MEMW_ALIGNED_NCOLS + ); + assert_eq!( + gpu_u64, cpu_u64, + "device MEMW_A table must be byte-identical to the CPU fill" + ); + } + + /// LOAD device fill byte-parity (widths 1/2/4/8, signed/unsigned, sign-bit, + /// high address bits, res-byte packing). Skips cleanly with no GPU. + #[cfg(feature = "cuda")] + #[test] + fn gpu_load_fill_matches_cpu() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_load_fill_matches_cpu: no CUDA backend"); + return; + } + let mut ops = Vec::new(); + for i in 0..400u64 { + let width = [1u8, 2, 4, 8][(i % 4) as usize]; + let signed = i % 2 == 0; + let base = (0x1_0000_0000u64 * (i % 3)) + 0x400 + i * 8; + let ts = 4 * i + 7; + let mut res = [0u64; 8]; + for (j, r) in res.iter_mut().enumerate() { + *r = (i.wrapping_mul(31) + j as u64) & 0xFF; + } + ops.push(load::LoadOperation::new(base, ts, width, signed, res)); + } + let n = ops.len(); + let num_rows = n.next_power_of_two().max(4); + let cpu = load::generate_load_trace(&ops); + let (fe, w) = cpu.main_data_row_major(); + assert_eq!(w, math_cuda::trace_cpu::LOAD_NCOLS); + let cpu_u64: Vec = fe + .iter() + .map(|e| unsafe { *(e.value() as *const _ as *const u64) }) + .collect(); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::LOAD_STRIDE); + for op in &ops { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_load_op(op)); + } + let gpu = math_cuda::trace_cpu::gpu_build_load_trace_host(&packed, n, num_rows) + .expect("device LOAD build must run on a box with a CUDA backend"); + assert_eq!(gpu, cpu_u64, "device LOAD table must be byte-identical to the CPU fill"); + } + + /// STORE device fill byte-parity (widths 1/2/4/8 via `bytes`, full-value + /// DWordBL split, high address bits). Skips cleanly with no GPU. + #[cfg(feature = "cuda")] + #[test] + fn gpu_store_fill_matches_cpu() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_store_fill_matches_cpu: no CUDA backend"); + return; + } + let mut ops = Vec::new(); + for i in 0..400u64 { + let bytes = [1u8, 2, 4, 8][(i % 4) as usize]; + let base = (0x1_0000_0000u64 * (i % 3)) + 0x800 + i * 8; + let ts = 4 * i + 9; + let value = i.wrapping_mul(0x9E37_79B9_7F4A_7C15); + ops.push(store::StoreOperation::new(base, ts, value, bytes)); + } + let n = ops.len(); + let num_rows = n.next_power_of_two().max(4); + let cpu = store::generate_store_trace(&ops); + let (fe, w) = cpu.main_data_row_major(); + assert_eq!(w, math_cuda::trace_cpu::STORE_NCOLS); + let cpu_u64: Vec = fe + .iter() + .map(|e| unsafe { *(e.value() as *const _ as *const u64) }) + .collect(); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::STORE_STRIDE); + for op in &ops { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_store_op(op)); + } + let gpu = math_cuda::trace_cpu::gpu_build_store_trace_host(&packed, n, num_rows) + .expect("device STORE build must run on a box with a CUDA backend"); + assert_eq!(gpu, cpu_u64, "device STORE table must be byte-identical to the CPU fill"); + } + + /// SHIFT device fill byte-parity: the kernel recomputes the full aux + /// (bit_shift/zbs/x/y/limb_shift/out) — covers left/right, signed/unsigned, + /// word/64-bit, shift amounts spanning 0 / >16 / >32 / >64, negative MSB inputs, + /// and the padding rows (ZBS=1). SHIFT is per-row (μ=1), so byte-parity holds. + #[cfg(feature = "cuda")] + #[test] + fn gpu_shift_fill_matches_cpu() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_shift_fill_matches_cpu: no CUDA backend"); + return; + } + let mut ops = Vec::new(); + // Exhaustive over shift byte 0..256 × all flags × MSB-set / edge values + // (so the signed arithmetic-right-shift extension path is fully covered). + let values: [u64; 8] = [ + 0, + 1, + 0x8000_0000_0000_0000, + 0xFFFF_FFFF_FFFF_FFFF, + 0x1234_5678_9ABC_DEF0, + 0xFEDC_BA98_7654_3210, + 0x0000_0000_FFFF_FFFF, + 0xFFFF_FFFF_0000_0000, + ]; + for &value in &values { + for shift_amount in 0u64..256 { + for &direction in &[false, true] { + for &signed in &[false, true] { + for &word_instr in &[false, true] { + ops.push(shift::ShiftOperation::new( + value, + shift_amount, + direction, + signed, + word_instr, + )); + } + } + } + } + } + // Also a few large shift_amounts (high SHIFT_B1/H1/HIGH limbs) — real ops + // carry the full rv2 on the ALU bus. + for &sa in &[0x1_0000u64, 0xFFFF_FFFFu64, 0x1234_5678_9ABCu64, u64::MAX] { + ops.push(shift::ShiftOperation::new(0xDEAD_BEEF_CAFE_1234, sa, false, true, false)); + } + let n = ops.len(); + let num_rows = n.next_power_of_two().max(4); + let cpu = shift::generate_shift_trace(&ops); + let (fe, w) = cpu.main_data_row_major(); + assert_eq!(w, math_cuda::trace_cpu::SHIFT_NCOLS); + let cpu_u64: Vec = fe + .iter() + .map(|e| unsafe { *(e.value() as *const _ as *const u64) }) + .collect(); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::SHIFT_STRIDE); + for op in &ops { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_shift_op(op)); + } + let gpu = math_cuda::trace_cpu::gpu_build_shift_trace_host(&packed, n, num_rows) + .expect("device SHIFT build must run on a box with a CUDA backend"); + assert_eq!(gpu, cpu_u64, "device SHIFT table must be byte-identical to the CPU fill"); + } + + /// LT device fill: dedup rides the permutation-invariant ALU bus, so the row + /// ORDER is non-deterministic (HashMap iteration). Validate as a MULTISET — + /// the set of real rows (μ>0), incl. summed multiplicities, must match the CPU + /// `generate_lt_trace`. Covers signed/unsigned, invert, MSBs, and duplicates + /// (μ>1). Skips cleanly with no GPU. + #[cfg(feature = "cuda")] + #[test] + fn gpu_lt_fill_matches_cpu_multiset() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_lt_fill_matches_cpu_multiset: no CUDA backend"); + return; + } + // Raw ops with deliberate duplicates (some pushed twice → μ=2). + let mut raw = Vec::new(); + for i in 0..800u64 { + let lhs = i.wrapping_mul(0x9E37_79B9_7F4A_7C15) ^ (i << 50); + let rhs = i.wrapping_mul(0x1234_5678_9ABC_DEF1).rotate_left(17); + let signed = i % 2 == 0; + let invert = i % 3 == 0; + let op = lt::LtOperation::new_with_invert(lhs, rhs, signed, invert); + raw.push(op.clone()); + if i % 4 == 0 { + raw.push(op); // duplicate → multiplicity 2 + } + } + + let ncols = math_cuda::trace_cpu::LT_NCOLS; + // Extract the real (μ>0) rows of a row-major u64 buffer as a sorted multiset. + let real_rows = |flat: &[u64]| -> Vec> { + let mut rows: Vec> = flat + .chunks(ncols) + .filter(|row| row[lt::cols::MU] > 0) + .map(|row| row.to_vec()) + .collect(); + rows.sort(); + rows + }; + + let cpu_table = lt::generate_lt_trace(&raw); + let (cpu_fe, w) = cpu_table.main_data_row_major(); + assert_eq!(w, ncols); + let cpu_u64: Vec = cpu_fe + .iter() + .map(|e| unsafe { *(e.value() as *const _ as *const u64) }) + .collect(); + + // Dedup on the host exactly like `gpu_build_lt_tables`, then device-fill. + let mut map: std::collections::HashMap = HashMap::new(); + for op in &raw { + *map.entry(op.clone()).or_insert(0) += 1; + } + let unique: Vec<(lt::LtOperation, u64)> = map.into_iter().collect(); + let n = unique.len(); + let num_rows = n.next_power_of_two().max(4); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::LT_STRIDE); + for (op, mult) in &unique { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_lt_op(op, *mult)); + } + let gpu_u64 = math_cuda::trace_cpu::gpu_build_lt_trace_host(&packed, n, num_rows) + .expect("device LT build must run on a box with a CUDA backend"); + + assert_eq!( + real_rows(&gpu_u64), + real_rows(&cpu_u64), + "device LT rows must match the CPU fill as a multiset" + ); + } +} From fd95330def7a19b4339d4a2f92334514886a707b Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Mon, 13 Jul 2026 14:54:59 -0300 Subject: [PATCH 18/21] fmt --- crypto/math-cuda/src/trace_cpu.rs | 108 +++++++++++++++++++---- crypto/math-cuda/tests/lde_dev_parity.rs | 14 +-- crypto/stark/src/gpu_lde.rs | 6 +- crypto/stark/src/logup_gpu.rs | 8 +- crypto/stark/src/prover.rs | 10 ++- prover/src/tables/gpu_trace.rs | 18 ++-- prover/src/tables/trace_builder.rs | 33 +++++-- 7 files changed, 156 insertions(+), 41 deletions(-) diff --git a/crypto/math-cuda/src/trace_cpu.rs b/crypto/math-cuda/src/trace_cpu.rs index b469efb3e..6879b561e 100644 --- a/crypto/math-cuda/src/trace_cpu.rs +++ b/crypto/math-cuda/src/trace_cpu.rs @@ -170,7 +170,13 @@ pub fn gpu_build_lt_trace(packed_ops: &[u64], n: usize, num_rows: usize) -> Resu let be = backend()?; let stream = be.next_stream(); let out = build_interleaved_on( - &stream, lt_kernel(be), packed_ops, n, num_rows, LT_NCOLS, LT_STRIDE, + &stream, + lt_kernel(be), + packed_ops, + n, + num_rows, + LT_NCOLS, + LT_STRIDE, )?; stream.synchronize()?; Ok(out) @@ -181,7 +187,13 @@ pub fn gpu_build_lt_trace_host(packed_ops: &[u64], n: usize, num_rows: usize) -> let be = backend()?; let stream = be.next_stream(); let out = build_interleaved_on( - &stream, lt_kernel(be), packed_ops, n, num_rows, LT_NCOLS, LT_STRIDE, + &stream, + lt_kernel(be), + packed_ops, + n, + num_rows, + LT_NCOLS, + LT_STRIDE, )?; let host = stream.clone_dtoh(&out)?; stream.synchronize()?; @@ -190,22 +202,42 @@ pub fn gpu_build_lt_trace_host(packed_ops: &[u64], n: usize, num_rows: usize) -> /// Build one SHIFT trace-table chunk on device (residency-ready row-major buffer). /// The kernel recomputes the shift aux from the packed inputs (see `trace_cpu.cu`). -pub fn gpu_build_shift_trace(packed_ops: &[u64], n: usize, num_rows: usize) -> Result> { +pub fn gpu_build_shift_trace( + packed_ops: &[u64], + n: usize, + num_rows: usize, +) -> Result> { let be = backend()?; let stream = be.next_stream(); let out = build_interleaved_on( - &stream, shift_kernel(be), packed_ops, n, num_rows, SHIFT_NCOLS, SHIFT_STRIDE, + &stream, + shift_kernel(be), + packed_ops, + n, + num_rows, + SHIFT_NCOLS, + SHIFT_STRIDE, )?; stream.synchronize()?; Ok(out) } /// Host-returning SHIFT build for byte-parity tests. -pub fn gpu_build_shift_trace_host(packed_ops: &[u64], n: usize, num_rows: usize) -> Result> { +pub fn gpu_build_shift_trace_host( + packed_ops: &[u64], + n: usize, + num_rows: usize, +) -> Result> { let be = backend()?; let stream = be.next_stream(); let out = build_interleaved_on( - &stream, shift_kernel(be), packed_ops, n, num_rows, SHIFT_NCOLS, SHIFT_STRIDE, + &stream, + shift_kernel(be), + packed_ops, + n, + num_rows, + SHIFT_NCOLS, + SHIFT_STRIDE, )?; let host = stream.clone_dtoh(&out)?; stream.synchronize()?; @@ -213,22 +245,42 @@ pub fn gpu_build_shift_trace_host(packed_ops: &[u64], n: usize, num_rows: usize) } /// Build one LOAD trace-table chunk on device (residency-ready row-major buffer). -pub fn gpu_build_load_trace(packed_ops: &[u64], n: usize, num_rows: usize) -> Result> { +pub fn gpu_build_load_trace( + packed_ops: &[u64], + n: usize, + num_rows: usize, +) -> Result> { let be = backend()?; let stream = be.next_stream(); let out = build_interleaved_on( - &stream, load_kernel(be), packed_ops, n, num_rows, LOAD_NCOLS, LOAD_STRIDE, + &stream, + load_kernel(be), + packed_ops, + n, + num_rows, + LOAD_NCOLS, + LOAD_STRIDE, )?; stream.synchronize()?; Ok(out) } /// Host-returning LOAD build for byte-parity tests. -pub fn gpu_build_load_trace_host(packed_ops: &[u64], n: usize, num_rows: usize) -> Result> { +pub fn gpu_build_load_trace_host( + packed_ops: &[u64], + n: usize, + num_rows: usize, +) -> Result> { let be = backend()?; let stream = be.next_stream(); let out = build_interleaved_on( - &stream, load_kernel(be), packed_ops, n, num_rows, LOAD_NCOLS, LOAD_STRIDE, + &stream, + load_kernel(be), + packed_ops, + n, + num_rows, + LOAD_NCOLS, + LOAD_STRIDE, )?; let host = stream.clone_dtoh(&out)?; stream.synchronize()?; @@ -236,22 +288,42 @@ pub fn gpu_build_load_trace_host(packed_ops: &[u64], n: usize, num_rows: usize) } /// Build one STORE trace-table chunk on device (residency-ready row-major buffer). -pub fn gpu_build_store_trace(packed_ops: &[u64], n: usize, num_rows: usize) -> Result> { +pub fn gpu_build_store_trace( + packed_ops: &[u64], + n: usize, + num_rows: usize, +) -> Result> { let be = backend()?; let stream = be.next_stream(); let out = build_interleaved_on( - &stream, store_kernel(be), packed_ops, n, num_rows, STORE_NCOLS, STORE_STRIDE, + &stream, + store_kernel(be), + packed_ops, + n, + num_rows, + STORE_NCOLS, + STORE_STRIDE, )?; stream.synchronize()?; Ok(out) } /// Host-returning STORE build for byte-parity tests. -pub fn gpu_build_store_trace_host(packed_ops: &[u64], n: usize, num_rows: usize) -> Result> { +pub fn gpu_build_store_trace_host( + packed_ops: &[u64], + n: usize, + num_rows: usize, +) -> Result> { let be = backend()?; let stream = be.next_stream(); let out = build_interleaved_on( - &stream, store_kernel(be), packed_ops, n, num_rows, STORE_NCOLS, STORE_STRIDE, + &stream, + store_kernel(be), + packed_ops, + n, + num_rows, + STORE_NCOLS, + STORE_STRIDE, )?; let host = stream.clone_dtoh(&out)?; stream.synchronize()?; @@ -316,7 +388,9 @@ pub fn gpu_fill_memw_register( ) -> Result> { let be = backend()?; let stream = be.next_stream(); - let buf = fill_memw_register_on(&stream, reg_addr, ts, value, is_read, old_value, old_ts, num_rows)?; + let buf = fill_memw_register_on( + &stream, reg_addr, ts, value, is_read, old_value, old_ts, num_rows, + )?; stream.synchronize()?; Ok(buf) } @@ -334,7 +408,9 @@ pub fn gpu_fill_memw_register_host( ) -> Result> { let be = backend()?; let stream = be.next_stream(); - let buf = fill_memw_register_on(&stream, reg_addr, ts, value, is_read, old_value, old_ts, num_rows)?; + let buf = fill_memw_register_on( + &stream, reg_addr, ts, value, is_read, old_value, old_ts, num_rows, + )?; let host = stream.clone_dtoh(&buf)?; stream.synchronize()?; Ok(host) diff --git a/crypto/math-cuda/tests/lde_dev_parity.rs b/crypto/math-cuda/tests/lde_dev_parity.rs index e358b2c5b..064653760 100644 --- a/crypto/math-cuda/tests/lde_dev_parity.rs +++ b/crypto/math-cuda/tests/lde_dev_parity.rs @@ -44,19 +44,19 @@ fn assert_dev_matches_host(log_n: u64, m: usize, blowup: usize, seed: u64) { let weights = coset_weights(n, coset_offset); // Host path: uploads `row_major` internally. - let (h_handle, h_lde) = - math_cuda::lde::coset_lde_row_major_with_merkle_tree_keep(&row_major, n, m, blowup, &weights) - .expect("host keep"); + let (h_handle, h_lde) = math_cuda::lde::coset_lde_row_major_with_merkle_tree_keep( + &row_major, n, m, blowup, &weights, + ) + .expect("host keep"); // Device path: pre-upload the SAME matrix, then run the device-input LDE. let be = backend().unwrap(); let stream = be.next_stream(); let dev = stream.clone_htod(&row_major).expect("upload matrix"); stream.synchronize().expect("sync upload"); - let (d_handle, d_lde) = math_cuda::lde::coset_lde_row_major_with_merkle_tree_keep_dev( - &dev, n, m, blowup, &weights, - ) - .expect("dev keep"); + let (d_handle, d_lde) = + math_cuda::lde::coset_lde_row_major_with_merkle_tree_keep_dev(&dev, n, m, blowup, &weights) + .expect("dev keep"); assert_eq!( h_handle.tree.as_ref().unwrap().root, diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 59e5af3e5..923b2457d 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -564,7 +564,11 @@ where // Transmute Vec → Vec> (zero-copy, E == GoldilocksField). let lde_out: Vec> = unsafe { let mut v = std::mem::ManuallyDrop::new(lde_u64); - Vec::from_raw_parts(v.as_mut_ptr() as *mut FieldElement, v.len(), v.capacity()) + Vec::from_raw_parts( + v.as_mut_ptr() as *mut FieldElement, + v.len(), + v.capacity(), + ) }; let root = handle.tree.as_ref()?.root; diff --git a/crypto/stark/src/logup_gpu.rs b/crypto/stark/src/logup_gpu.rs index 1cf8cd51e..19da57766 100644 --- a/crypto/stark/src/logup_gpu.rs +++ b/crypto/stark/src/logup_gpu.rs @@ -487,7 +487,13 @@ where None => math_cuda::logup::ResidentMain::Host(&main_flat), }; let ra = math_cuda::logup::logup_aux_resident( - main, trace_len, &md, &alpha_flat, z_arr, inv_n, &stream, + main, + trace_len, + &md, + &alpha_flat, + z_arr, + inv_n, + &stream, ) .ok()?; crate::gpu_lde::GPU_LOGUP_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed); diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 87f7098cc..c501e080e 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -794,7 +794,11 @@ pub trait IsStarkProver< Field, BatchedMerkleTreeBackend, >( - dev, n, num_cols, domain.blowup_factor, &twiddles.coset_weights + dev, + n, + num_cols, + domain.blowup_factor, + &twiddles.coset_weights, ) { #[cfg(feature = "instruments")] @@ -844,9 +848,7 @@ pub trait IsStarkProver< // This table's main matrix was uploaded host->device for the LDE. // Counts the transfer the on-GPU trace generator will remove. #[cfg(feature = "instruments")] - crate::instruments::accum_main_h2d( - trace_slice.len() * std::mem::size_of::(), - ); + crate::instruments::accum_main_h2d(trace_slice.len() * std::mem::size_of::()); return Ok(( TableCommit::plain(tree, root), (main_data, num_cols), diff --git a/prover/src/tables/gpu_trace.rs b/prover/src/tables/gpu_trace.rs index 420210f86..80af2358c 100644 --- a/prover/src/tables/gpu_trace.rs +++ b/prover/src/tables/gpu_trace.rs @@ -90,7 +90,9 @@ fn pack_cpu_ops(chunk: &[CpuOperation]) -> Vec { /// device buffer, the aux build reads the resident snapshot, queries gather from /// the device tree). Returns `None` if the GPU build fails, so the caller can /// fall back to the CPU generator. -fn build_cpu_chunk(chunk: &[CpuOperation]) -> Option> { +fn build_cpu_chunk( + chunk: &[CpuOperation], +) -> Option> { let n = chunk.len(); let num_rows = n.next_power_of_two().max(4); let last_ts = chunk.last().map(|op| op.timestamp).unwrap_or(0); @@ -293,12 +295,13 @@ pub(crate) fn pack_load_op(op: &LoadOperation) -> [u64; math_cuda::trace_cpu::LO /// Pack one `StoreOperation` into the `store_fill` stride (see `trace_cpu.cu`). pub(crate) fn pack_store_op(op: &StoreOperation) -> [u64; math_cuda::trace_cpu::STORE_STRIDE] { - let flags = - (op.write2 as u64) | ((op.write4 as u64) << 1) | ((op.write8 as u64) << 2); + let flags = (op.write2 as u64) | ((op.write4 as u64) << 1) | ((op.write8 as u64) << 2); [flags, op.base_address, op.timestamp, op.value] } -fn build_load_chunk(chunk: &[LoadOperation]) -> Option> { +fn build_load_chunk( + chunk: &[LoadOperation], +) -> Option> { let n = chunk.len(); let num_rows = n.next_power_of_two().max(4); let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::LOAD_STRIDE); @@ -381,7 +384,8 @@ pub(crate) fn gpu_build_store_tables( /// recomputes bit_shift/zbs/x/y/limb_shift/out, so only 3 u64/op upload. pub(crate) fn pack_shift_op(op: &ShiftOperation) -> [u64; math_cuda::trace_cpu::SHIFT_STRIDE] { let h = &op.in_halves; - let value = (h[0] as u64) | ((h[1] as u64) << 16) | ((h[2] as u64) << 32) | ((h[3] as u64) << 48); + let value = + (h[0] as u64) | ((h[1] as u64) << 16) | ((h[2] as u64) << 32) | ((h[3] as u64) << 48); let flags = (op.direction as u64) | ((op.signed as u64) << 1) | ((op.word_instr as u64) << 2); [value, op.shift_amount, flags] } @@ -438,7 +442,9 @@ pub(crate) fn pack_lt_op(op: &LtOperation, mult: u64) -> [u64; math_cuda::trace_ /// same per-chunk HashMap `generate_lt_trace` uses), then the unique ops + summed /// multiplicities are filled on device. LT rides the permutation-invariant ALU /// bus, so any row order is valid (validated by multiset/prove, not byte order). -fn build_lt_chunk(chunk: &[LtOperation]) -> Option> { +fn build_lt_chunk( + chunk: &[LtOperation], +) -> Option> { let mut map: HashMap = HashMap::new(); for op in chunk { *map.entry(op.clone()).or_insert(0) += 1; diff --git a/prover/src/tables/trace_builder.rs b/prover/src/tables/trace_builder.rs index ee71f4d2e..0bc302428 100644 --- a/prover/src/tables/trace_builder.rs +++ b/prover/src/tables/trace_builder.rs @@ -4599,8 +4599,14 @@ mod gpu_fill_tests { ) .expect("device MEMW_R fill must run on a box with a CUDA backend"); - assert_eq!(gpu_u64.len(), num_rows * math_cuda::trace_cpu::MEMW_REGISTER_NCOLS); - assert_eq!(gpu_u64, cpu_u64, "device MEMW_R fill must be byte-identical to the CPU fill"); + assert_eq!( + gpu_u64.len(), + num_rows * math_cuda::trace_cpu::MEMW_REGISTER_NCOLS + ); + assert_eq!( + gpu_u64, cpu_u64, + "device MEMW_R fill must be byte-identical to the CPU fill" + ); } /// MEMW_A (aligned memory — the biggest remaining uploader) device fill must be @@ -4712,7 +4718,10 @@ mod gpu_fill_tests { } let gpu = math_cuda::trace_cpu::gpu_build_load_trace_host(&packed, n, num_rows) .expect("device LOAD build must run on a box with a CUDA backend"); - assert_eq!(gpu, cpu_u64, "device LOAD table must be byte-identical to the CPU fill"); + assert_eq!( + gpu, cpu_u64, + "device LOAD table must be byte-identical to the CPU fill" + ); } /// STORE device fill byte-parity (widths 1/2/4/8 via `bytes`, full-value @@ -4747,7 +4756,10 @@ mod gpu_fill_tests { } let gpu = math_cuda::trace_cpu::gpu_build_store_trace_host(&packed, n, num_rows) .expect("device STORE build must run on a box with a CUDA backend"); - assert_eq!(gpu, cpu_u64, "device STORE table must be byte-identical to the CPU fill"); + assert_eq!( + gpu, cpu_u64, + "device STORE table must be byte-identical to the CPU fill" + ); } /// SHIFT device fill byte-parity: the kernel recomputes the full aux @@ -4794,7 +4806,13 @@ mod gpu_fill_tests { // Also a few large shift_amounts (high SHIFT_B1/H1/HIGH limbs) — real ops // carry the full rv2 on the ALU bus. for &sa in &[0x1_0000u64, 0xFFFF_FFFFu64, 0x1234_5678_9ABCu64, u64::MAX] { - ops.push(shift::ShiftOperation::new(0xDEAD_BEEF_CAFE_1234, sa, false, true, false)); + ops.push(shift::ShiftOperation::new( + 0xDEAD_BEEF_CAFE_1234, + sa, + false, + true, + false, + )); } let n = ops.len(); let num_rows = n.next_power_of_two().max(4); @@ -4811,7 +4829,10 @@ mod gpu_fill_tests { } let gpu = math_cuda::trace_cpu::gpu_build_shift_trace_host(&packed, n, num_rows) .expect("device SHIFT build must run on a box with a CUDA backend"); - assert_eq!(gpu, cpu_u64, "device SHIFT table must be byte-identical to the CPU fill"); + assert_eq!( + gpu, cpu_u64, + "device SHIFT table must be byte-identical to the CPU fill" + ); } /// LT device fill: dedup rides the permutation-invariant ALU bus, so the row From 9ee20635fd1f1b967527895687a1cdceb2e8f6ba Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Mon, 13 Jul 2026 15:03:22 -0300 Subject: [PATCH 19/21] fmt --- prover/src/tables/trace_builder.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/prover/src/tables/trace_builder.rs b/prover/src/tables/trace_builder.rs index 0bc302428..0de6cabd1 100644 --- a/prover/src/tables/trace_builder.rs +++ b/prover/src/tables/trace_builder.rs @@ -4531,7 +4531,7 @@ impl Traces { } } -#[cfg(test)] +#[cfg(all(test, feature = "cuda"))] mod gpu_fill_tests { //! Byte-parity (and multiset, for order-independent LT) between each table's //! device fill and the CPU `generate_*_trace`, so the on-GPU trace tables are From 219be6ded8a6cf84637ee983e7b6f8dcb74e7998 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Mon, 13 Jul 2026 15:26:01 -0300 Subject: [PATCH 20/21] add tests --- crypto/math-cuda/src/trace_cpu.rs | 34 +++++++++ prover/src/tables/gpu_trace.rs | 2 +- prover/src/tables/trace_builder.rs | 109 +++++++++++++++++++++++++++-- 3 files changed, 138 insertions(+), 7 deletions(-) diff --git a/crypto/math-cuda/src/trace_cpu.rs b/crypto/math-cuda/src/trace_cpu.rs index 6879b561e..1b5ca6420 100644 --- a/crypto/math-cuda/src/trace_cpu.rs +++ b/crypto/math-cuda/src/trace_cpu.rs @@ -57,6 +57,40 @@ pub fn gpu_build_cpu_trace( Ok(out) } +/// Host-returning wrapper over [`gpu_build_cpu_trace`] for byte-parity tests: +/// builds the row-major CPU table on device and copies it back on the same stream. +pub fn gpu_build_cpu_trace_host( + packed_ops: &[u64], + n: usize, + num_rows: usize, + last_ts: u64, +) -> Result> { + debug_assert_eq!(packed_ops.len(), n * CPU_OP_STRIDE); + let be = backend()?; + let stream = be.next_stream(); + + let ops_dev = if packed_ops.is_empty() { + stream.alloc_zeros::(CPU_OP_STRIDE)? + } else { + stream.clone_htod(packed_ops)? + }; + let mut out = stream.alloc_zeros::(num_rows * CPU_NCOLS)?; + + unsafe { + stream + .launch_builder(&be.trace_cpu_fill) + .arg(&ops_dev) + .arg(&(n as u64)) + .arg(&(num_rows as u64)) + .arg(&last_ts) + .arg(&mut out) + .launch(LaunchConfig::for_num_elems(num_rows as u32))?; + } + let host = stream.clone_dtoh(&out)?; + stream.synchronize()?; + Ok(host) +} + /// MEMW_A table width (`prover::tables::memw_aligned::cols::NUM_COLUMNS`). pub const MEMW_ALIGNED_NCOLS: usize = 29; /// Packed MEMW_A input stride, u64 per op (must match `memw_aligned_fill` in `trace_cpu.cu`). diff --git a/prover/src/tables/gpu_trace.rs b/prover/src/tables/gpu_trace.rs index 80af2358c..e587423ee 100644 --- a/prover/src/tables/gpu_trace.rs +++ b/prover/src/tables/gpu_trace.rs @@ -45,7 +45,7 @@ pub(crate) fn gpu_trace_disabled() -> bool { /// kernel consumes (stride `CPU_OP_STRIDE` u64/op). The kernel does the same /// bit-slicing as `cpu::generate_cpu_trace`, so this only copies fields — no /// per-column encoding on the host. -fn pack_cpu_ops(chunk: &[CpuOperation]) -> Vec { +pub(crate) fn pack_cpu_ops(chunk: &[CpuOperation]) -> Vec { let stride = math_cuda::trace_cpu::CPU_OP_STRIDE; let mut packed = vec![0u64; chunk.len() * stride]; for (i, op) in chunk.iter().enumerate() { diff --git a/prover/src/tables/trace_builder.rs b/prover/src/tables/trace_builder.rs index 0de6cabd1..7b90d6fa0 100644 --- a/prover/src/tables/trace_builder.rs +++ b/prover/src/tables/trace_builder.rs @@ -4539,6 +4539,103 @@ mod gpu_fill_tests { use super::*; + /// CPU-table device fill must be byte-identical to `cpu::generate_cpu_trace`. + /// The CPU kernel is the most intricate of the seven (word-delegate column + /// masking, `PC_DOUBLE_READ`/`PREV_PC_TIMESTAMP_BORROW`, and the `+4` padding + /// cadence with PC=1), so this guards it the same way its six siblings are + /// guarded. The fill is a pure function of the packed op fields, so the synthetic + /// ops need only be diverse (not a valid execution): word-delegate rows, x255 PC + /// reads (`pc_double_read`), `ts_lo < 3` rows (`prev_pc_timestamp_borrow`), x0 + /// registers, and high-word values (DWordWL/DWordHL splits). `n = 300 < 512` + /// forces padding rows, exercising the `+4` cadence off `last_ts`. Skips cleanly + /// with no GPU. + #[cfg(feature = "cuda")] + #[test] + fn gpu_cpu_fill_matches_cpu() { + use crate::tables::types::{DecodeEntry, ShrunkDecode}; + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_cpu_fill_matches_cpu: no CUDA backend"); + return; + } + let mut ops = Vec::new(); + for i in 0..300u64 { + let word = i % 5 == 0; + // rs1 cycles through x255 (PC register), x0, and normal registers so the + // pc_double_read and register-zero-suppression paths are both hit. + let rs1 = match i % 4 { + 0 => 255u8, + 1 => 0, + _ => (i % 31 + 1) as u8, + }; + let fields = ShrunkDecode { + read_register1: i % 3 != 0, + read_register2: i % 2 == 0, + write_register: i % 3 == 0, + word_instr: word, + alu: i % 6 == 0, + add: i % 6 == 1, + sub: i % 6 == 2, + memory: i % 6 == 3, + branch: i % 6 == 4, + ecall: i % 6 == 5, + rs1, + rs2: (i % 32) as u8, + rd: ((i + 7) % 32) as u8, + half_instruction_length: if i % 2 == 0 { 1 } else { 2 }, + alu_flags: (i & 0xFF) as u8, + mem_flags: ((i >> 1) & 0xFF) as u8, + }; + // A few timestamps with ts_lo < 3 (0,1,2) trip prev_pc_timestamp_borrow; + // the rest are large and continue a +4 cadence. + let timestamp = match i { + 0 => 0, + 1 => 1, + 2 => 2, + _ => 4 * i + 100, + }; + // PC/imm/next_pc carry high words (>32 bits) to exercise the DWordWL split + // into PC_1/NEXT_PC_1/IMM_1. + let decode = DecodeEntry { + pc: 0x1000 + i * 4 + ((i % 3) << 33), + imm: i.wrapping_mul(0x9E37_79B9) | ((i % 4) << 40), + fields, + }; + ops.push(CpuOperation { + decode, + timestamp, + next_pc: 0x2000 + i * 4 + ((i % 2) << 34), + rvd: i.wrapping_mul(0x1234_5678_9ABC), + rv1: i.wrapping_mul(0xDEAD_BEEF).rotate_left(13), + rv2: (i ^ 0xFFFF_0000_1111) << 3, + arg2: i.wrapping_add(0x7777_8888_9999), + res: i.wrapping_mul(0xABCD_1234_5678) ^ (i << 48), + branch_cond: i % 3 == 1, + ..Default::default() + }); + } + let n = ops.len(); + let num_rows = n.next_power_of_two().max(4); + let last_ts = ops.last().map(|op| op.timestamp).unwrap_or(0); + + let cpu_table = cpu::generate_cpu_trace(&ops); + let (cpu_fe, w) = cpu_table.main_data_row_major(); + assert_eq!(w, math_cuda::trace_cpu::CPU_NCOLS); + let cpu_u64: Vec = cpu_fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + + let packed = crate::tables::gpu_trace::pack_cpu_ops(&ops); + let gpu_u64 = math_cuda::trace_cpu::gpu_build_cpu_trace_host(&packed, n, num_rows, last_ts) + .expect("device CPU build must run on a box with a CUDA backend"); + + assert_eq!(gpu_u64.len(), num_rows * math_cuda::trace_cpu::CPU_NCOLS); + assert_eq!( + gpu_u64, cpu_u64, + "device CPU table must be byte-identical to the CPU fill" + ); + } + /// MEMW_R (register fast-path) device fill must be byte-identical to the CPU /// `generate_memw_register_trace_from_rows`. The rows are the walked register /// rows; here we build synthetic `RegRow`s (read + write, varied values/old/ts) @@ -4576,7 +4673,7 @@ mod gpu_fill_tests { assert_eq!(w, math_cuda::trace_cpu::MEMW_REGISTER_NCOLS); let cpu_u64: Vec = cpu_fe .iter() - .map(|e| unsafe { *(e.value() as *const _ as *const u64) }) + .map(|e| unsafe { *(e.value() as *const u64) }) .collect(); let mut reg_addr = Vec::new(); @@ -4662,7 +4759,7 @@ mod gpu_fill_tests { assert_eq!(width_cols, math_cuda::trace_cpu::MEMW_ALIGNED_NCOLS); let cpu_u64: Vec = cpu_fe .iter() - .map(|e| unsafe { *(e.value() as *const _ as *const u64) }) + .map(|e| unsafe { *(e.value() as *const u64) }) .collect(); let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::MEMW_ALIGNED_STRIDE); @@ -4710,7 +4807,7 @@ mod gpu_fill_tests { assert_eq!(w, math_cuda::trace_cpu::LOAD_NCOLS); let cpu_u64: Vec = fe .iter() - .map(|e| unsafe { *(e.value() as *const _ as *const u64) }) + .map(|e| unsafe { *(e.value() as *const u64) }) .collect(); let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::LOAD_STRIDE); for op in &ops { @@ -4748,7 +4845,7 @@ mod gpu_fill_tests { assert_eq!(w, math_cuda::trace_cpu::STORE_NCOLS); let cpu_u64: Vec = fe .iter() - .map(|e| unsafe { *(e.value() as *const _ as *const u64) }) + .map(|e| unsafe { *(e.value() as *const u64) }) .collect(); let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::STORE_STRIDE); for op in &ops { @@ -4821,7 +4918,7 @@ mod gpu_fill_tests { assert_eq!(w, math_cuda::trace_cpu::SHIFT_NCOLS); let cpu_u64: Vec = fe .iter() - .map(|e| unsafe { *(e.value() as *const _ as *const u64) }) + .map(|e| unsafe { *(e.value() as *const u64) }) .collect(); let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::SHIFT_STRIDE); for op in &ops { @@ -4878,7 +4975,7 @@ mod gpu_fill_tests { assert_eq!(w, ncols); let cpu_u64: Vec = cpu_fe .iter() - .map(|e| unsafe { *(e.value() as *const _ as *const u64) }) + .map(|e| unsafe { *(e.value() as *const u64) }) .collect(); // Dedup on the host exactly like `gpu_build_lt_tables`, then device-fill. From 9e3c2e3f0f3dcddef441f5b1059e283215fc5716 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Mon, 13 Jul 2026 18:22:37 -0300 Subject: [PATCH 21/21] move tests --- prover/src/tables/gpu_trace.rs | 128 +++----- prover/src/tables/trace_builder.rs | 470 ----------------------------- prover/src/tests/gpu_fill_tests.rs | 465 ++++++++++++++++++++++++++++ prover/src/tests/mod.rs | 2 + 4 files changed, 502 insertions(+), 563 deletions(-) create mode 100644 prover/src/tests/gpu_fill_tests.rs diff --git a/prover/src/tables/gpu_trace.rs b/prover/src/tables/gpu_trace.rs index e587423ee..c93c2b523 100644 --- a/prover/src/tables/gpu_trace.rs +++ b/prover/src/tables/gpu_trace.rs @@ -37,6 +37,33 @@ pub(crate) fn gpu_trace_disabled() -> bool { }) } +/// Shared GPU table-build dispatcher. Mirrors `chunk_and_generate`'s chunking +/// (`ops.chunks(max_rows)`, with one empty chunk when there are no ops so the +/// table still emits its padded minimum), builds each chunk on device via +/// `build_chunk`, and collects the resident tables. Returns `None` when the +/// kill-switch is set (`LAMBDA_VM_CPU_TRACE`) or any chunk fails to build, so the +/// caller falls back to the CPU generator. Every `gpu_build_*_tables` entry point +/// is a thin wrapper over this. +fn gpu_build_tables( + ops: &[T], + max_rows: usize, + build_chunk: impl Fn(&[T]) -> Option>, +) -> Option>> { + if gpu_trace_disabled() { + return None; + } + let chunks: Vec<&[T]> = if ops.is_empty() { + vec![&[][..]] + } else { + ops.chunks(max_rows).collect() + }; + let mut tables = Vec::with_capacity(chunks.len()); + for chunk in chunks { + tables.push(build_chunk(chunk)?); + } + Some(tables) +} + // ============================================================================= // CPU table (the first table built on device — see GPU-TRACEGEN-DESIGN-V2 §P1) // ============================================================================= @@ -115,19 +142,7 @@ pub(crate) fn gpu_build_cpu_trace_tables( cpu_ops: &[CpuOperation], max_rows: usize, ) -> Option>> { - if gpu_trace_disabled() { - return None; - } - let chunks: Vec<&[CpuOperation]> = if cpu_ops.is_empty() { - vec![&[][..]] - } else { - cpu_ops.chunks(max_rows).collect() - }; - let mut tables = Vec::with_capacity(chunks.len()); - for chunk in chunks { - tables.push(build_cpu_chunk(chunk)?); - } - Some(tables) + gpu_build_tables(cpu_ops, max_rows, build_cpu_chunk) } // ============================================================================= @@ -186,19 +201,7 @@ pub(crate) fn gpu_build_memw_register_tables( rows: &[RegRow], max_rows: usize, ) -> Option>> { - if gpu_trace_disabled() { - return None; - } - let chunks: Vec<&[RegRow]> = if rows.is_empty() { - vec![&[][..]] - } else { - rows.chunks(max_rows).collect() - }; - let mut tables = Vec::with_capacity(chunks.len()); - for chunk in chunks { - tables.push(build_memw_register_chunk(chunk)?); - } - Some(tables) + gpu_build_tables(rows, max_rows, build_memw_register_chunk) } // ============================================================================= @@ -259,19 +262,7 @@ pub(crate) fn gpu_build_memw_aligned_tables( ops: &[MemwOperation], max_rows: usize, ) -> Option>> { - if gpu_trace_disabled() { - return None; - } - let chunks: Vec<&[MemwOperation]> = if ops.is_empty() { - vec![&[][..]] - } else { - ops.chunks(max_rows).collect() - }; - let mut tables = Vec::with_capacity(chunks.len()); - for chunk in chunks { - tables.push(build_memw_aligned_chunk(chunk)?); - } - Some(tables) + gpu_build_tables(ops, max_rows, build_memw_aligned_chunk) } // ============================================================================= @@ -322,19 +313,7 @@ pub(crate) fn gpu_build_load_tables( ops: &[LoadOperation], max_rows: usize, ) -> Option>> { - if gpu_trace_disabled() { - return None; - } - let chunks: Vec<&[LoadOperation]> = if ops.is_empty() { - vec![&[][..]] - } else { - ops.chunks(max_rows).collect() - }; - let mut tables = Vec::with_capacity(chunks.len()); - for chunk in chunks { - tables.push(build_load_chunk(chunk)?); - } - Some(tables) + gpu_build_tables(ops, max_rows, build_load_chunk) } fn build_store_chunk( @@ -360,19 +339,7 @@ pub(crate) fn gpu_build_store_tables( ops: &[StoreOperation], max_rows: usize, ) -> Option>> { - if gpu_trace_disabled() { - return None; - } - let chunks: Vec<&[StoreOperation]> = if ops.is_empty() { - vec![&[][..]] - } else { - ops.chunks(max_rows).collect() - }; - let mut tables = Vec::with_capacity(chunks.len()); - for chunk in chunks { - tables.push(build_store_chunk(chunk)?); - } - Some(tables) + gpu_build_tables(ops, max_rows, build_store_chunk) } // ============================================================================= @@ -413,19 +380,7 @@ pub(crate) fn gpu_build_shift_tables( ops: &[ShiftOperation], max_rows: usize, ) -> Option>> { - if gpu_trace_disabled() { - return None; - } - let chunks: Vec<&[ShiftOperation]> = if ops.is_empty() { - vec![&[][..]] - } else { - ops.chunks(max_rows).collect() - }; - let mut tables = Vec::with_capacity(chunks.len()); - for chunk in chunks { - tables.push(build_shift_chunk(chunk)?); - } - Some(tables) + gpu_build_tables(ops, max_rows, build_shift_chunk) } // ============================================================================= @@ -470,19 +425,6 @@ pub(crate) fn gpu_build_lt_tables( ops: &[LtOperation], max_rows: usize, ) -> Option>> { - if gpu_trace_disabled() { - return None; - } - // Chunk the RAW ops exactly like `chunk_and_generate`; each chunk dedups - // independently (matching `generate_lt_trace` per chunk). - let chunks: Vec<&[LtOperation]> = if ops.is_empty() { - vec![&[][..]] - } else { - ops.chunks(max_rows).collect() - }; - let mut tables = Vec::with_capacity(chunks.len()); - for chunk in chunks { - tables.push(build_lt_chunk(chunk)?); - } - Some(tables) + // Each chunk dedups independently (matching `generate_lt_trace` per chunk). + gpu_build_tables(ops, max_rows, build_lt_chunk) } diff --git a/prover/src/tables/trace_builder.rs b/prover/src/tables/trace_builder.rs index 7b90d6fa0..1c864dfc8 100644 --- a/prover/src/tables/trace_builder.rs +++ b/prover/src/tables/trace_builder.rs @@ -4530,473 +4530,3 @@ impl Traces { ) } } - -#[cfg(all(test, feature = "cuda"))] -mod gpu_fill_tests { - //! Byte-parity (and multiset, for order-independent LT) between each table's - //! device fill and the CPU `generate_*_trace`, so the on-GPU trace tables are - //! bit-identical to the CPU path. Each test skips cleanly with no CUDA backend. - - use super::*; - - /// CPU-table device fill must be byte-identical to `cpu::generate_cpu_trace`. - /// The CPU kernel is the most intricate of the seven (word-delegate column - /// masking, `PC_DOUBLE_READ`/`PREV_PC_TIMESTAMP_BORROW`, and the `+4` padding - /// cadence with PC=1), so this guards it the same way its six siblings are - /// guarded. The fill is a pure function of the packed op fields, so the synthetic - /// ops need only be diverse (not a valid execution): word-delegate rows, x255 PC - /// reads (`pc_double_read`), `ts_lo < 3` rows (`prev_pc_timestamp_borrow`), x0 - /// registers, and high-word values (DWordWL/DWordHL splits). `n = 300 < 512` - /// forces padding rows, exercising the `+4` cadence off `last_ts`. Skips cleanly - /// with no GPU. - #[cfg(feature = "cuda")] - #[test] - fn gpu_cpu_fill_matches_cpu() { - use crate::tables::types::{DecodeEntry, ShrunkDecode}; - if math_cuda::device::backend().is_err() { - eprintln!("skipping gpu_cpu_fill_matches_cpu: no CUDA backend"); - return; - } - let mut ops = Vec::new(); - for i in 0..300u64 { - let word = i % 5 == 0; - // rs1 cycles through x255 (PC register), x0, and normal registers so the - // pc_double_read and register-zero-suppression paths are both hit. - let rs1 = match i % 4 { - 0 => 255u8, - 1 => 0, - _ => (i % 31 + 1) as u8, - }; - let fields = ShrunkDecode { - read_register1: i % 3 != 0, - read_register2: i % 2 == 0, - write_register: i % 3 == 0, - word_instr: word, - alu: i % 6 == 0, - add: i % 6 == 1, - sub: i % 6 == 2, - memory: i % 6 == 3, - branch: i % 6 == 4, - ecall: i % 6 == 5, - rs1, - rs2: (i % 32) as u8, - rd: ((i + 7) % 32) as u8, - half_instruction_length: if i % 2 == 0 { 1 } else { 2 }, - alu_flags: (i & 0xFF) as u8, - mem_flags: ((i >> 1) & 0xFF) as u8, - }; - // A few timestamps with ts_lo < 3 (0,1,2) trip prev_pc_timestamp_borrow; - // the rest are large and continue a +4 cadence. - let timestamp = match i { - 0 => 0, - 1 => 1, - 2 => 2, - _ => 4 * i + 100, - }; - // PC/imm/next_pc carry high words (>32 bits) to exercise the DWordWL split - // into PC_1/NEXT_PC_1/IMM_1. - let decode = DecodeEntry { - pc: 0x1000 + i * 4 + ((i % 3) << 33), - imm: i.wrapping_mul(0x9E37_79B9) | ((i % 4) << 40), - fields, - }; - ops.push(CpuOperation { - decode, - timestamp, - next_pc: 0x2000 + i * 4 + ((i % 2) << 34), - rvd: i.wrapping_mul(0x1234_5678_9ABC), - rv1: i.wrapping_mul(0xDEAD_BEEF).rotate_left(13), - rv2: (i ^ 0xFFFF_0000_1111) << 3, - arg2: i.wrapping_add(0x7777_8888_9999), - res: i.wrapping_mul(0xABCD_1234_5678) ^ (i << 48), - branch_cond: i % 3 == 1, - ..Default::default() - }); - } - let n = ops.len(); - let num_rows = n.next_power_of_two().max(4); - let last_ts = ops.last().map(|op| op.timestamp).unwrap_or(0); - - let cpu_table = cpu::generate_cpu_trace(&ops); - let (cpu_fe, w) = cpu_table.main_data_row_major(); - assert_eq!(w, math_cuda::trace_cpu::CPU_NCOLS); - let cpu_u64: Vec = cpu_fe - .iter() - .map(|e| unsafe { *(e.value() as *const u64) }) - .collect(); - - let packed = crate::tables::gpu_trace::pack_cpu_ops(&ops); - let gpu_u64 = math_cuda::trace_cpu::gpu_build_cpu_trace_host(&packed, n, num_rows, last_ts) - .expect("device CPU build must run on a box with a CUDA backend"); - - assert_eq!(gpu_u64.len(), num_rows * math_cuda::trace_cpu::CPU_NCOLS); - assert_eq!( - gpu_u64, cpu_u64, - "device CPU table must be byte-identical to the CPU fill" - ); - } - - /// MEMW_R (register fast-path) device fill must be byte-identical to the CPU - /// `generate_memw_register_trace_from_rows`. The rows are the walked register - /// rows; here we build synthetic `RegRow`s (read + write, varied values/old/ts) - /// and fill both ways. - #[cfg(feature = "cuda")] - #[test] - fn gpu_memw_register_fill_matches_cpu() { - if math_cuda::device::backend().is_err() { - eprintln!("skipping gpu_memw_register_fill_matches_cpu: no CUDA backend"); - return; - } - let mut rows = Vec::new(); - for i in 0..500u64 { - let reg_addr = 2 * (i % 40); // even word-address = 2*reg_index - let ts = 4 * i + 100; - let old_ts = 4 * i + 50; - let val = i.wrapping_mul(2_654_435_761); - let old = i.wrapping_mul(40_503) ^ 0xABCD_1234; - let is_read = i % 3 == 0; - rows.push(memw_register::RegRow::new( - reg_addr, - ts, - val as u32, - (val >> 32) as u32, - old as u32, - (old >> 32) as u32, - old_ts, - is_read, - )); - } - let num_rows = rows.len().next_power_of_two().max(4); - - let cpu_table = memw_register::generate_memw_register_trace_from_rows(&rows); - let (cpu_fe, w) = cpu_table.main_data_row_major(); - assert_eq!(w, math_cuda::trace_cpu::MEMW_REGISTER_NCOLS); - let cpu_u64: Vec = cpu_fe - .iter() - .map(|e| unsafe { *(e.value() as *const u64) }) - .collect(); - - let mut reg_addr = Vec::new(); - let mut ts = Vec::new(); - let mut value = Vec::new(); - let mut is_read = Vec::new(); - let mut old_value = Vec::new(); - let mut old_tsv = Vec::new(); - for r in &rows { - let (ra, t, v, ir, ov, ot) = r.fill_soa(); - reg_addr.push(ra); - ts.push(t); - value.push(v); - is_read.push(ir); - old_value.push(ov); - old_tsv.push(ot); - } - let gpu_u64 = math_cuda::trace_cpu::gpu_fill_memw_register_host( - ®_addr, &ts, &value, &is_read, &old_value, &old_tsv, num_rows, - ) - .expect("device MEMW_R fill must run on a box with a CUDA backend"); - - assert_eq!( - gpu_u64.len(), - num_rows * math_cuda::trace_cpu::MEMW_REGISTER_NCOLS - ); - assert_eq!( - gpu_u64, cpu_u64, - "device MEMW_R fill must be byte-identical to the CPU fill" - ); - } - - /// MEMW_A (aligned memory — the biggest remaining uploader) device fill must be - /// byte-identical to the CPU `generate_memw_aligned_trace`. Exercises memory - /// (widths 2/4/8), register fallback (`is_register`), high address bits, and - /// the 2×u32-per-u64 value/old packing. Skips cleanly with no GPU. - #[cfg(feature = "cuda")] - #[test] - fn gpu_memw_aligned_fill_matches_cpu() { - if math_cuda::device::backend().is_err() { - eprintln!("skipping gpu_memw_aligned_fill_matches_cpu: no CUDA backend"); - return; - } - let mut ops = Vec::new(); - for i in 0..500u64 { - let width = [2u8, 4, 8][(i % 3) as usize]; - let is_read = i % 2 == 0; - let is_register = i % 7 == 0; - // Vary the high address word too (whh split has a 32-bit high word). - let base = (0x1_0000_0000u64 * (i % 4)) + 0x8000 + i * 8; - let value = [ - (i as u32).wrapping_mul(2654435761), - (i as u32) ^ 0xABCD_1234, - i as u32, - 0, - 7, - 0, - 0, - (i as u32) & 0xFF, - ]; - let old = [ - (i as u32).wrapping_add(99), - (i as u32) ^ 0x0BAD_F00D, - 0, - 3, - 0, - 0, - 0, - 0, - ]; - let ts = 4 * i + 100; - let old_ts = [4 * i + 50; 8]; - ops.push( - MemwOperation::new(is_register, base, value, ts, width, is_read) - .with_old(old, old_ts), - ); - } - let n = ops.len(); - let num_rows = n.next_power_of_two().max(4); - - let cpu_table = memw_aligned::generate_memw_aligned_trace(&ops); - let (cpu_fe, width_cols) = cpu_table.main_data_row_major(); - assert_eq!(width_cols, math_cuda::trace_cpu::MEMW_ALIGNED_NCOLS); - let cpu_u64: Vec = cpu_fe - .iter() - .map(|e| unsafe { *(e.value() as *const u64) }) - .collect(); - - let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::MEMW_ALIGNED_STRIDE); - for op in &ops { - packed.extend_from_slice(&crate::tables::gpu_trace::pack_memw_aligned_op(op)); - } - let gpu_u64 = math_cuda::trace_cpu::gpu_build_memw_aligned_trace_host(&packed, n, num_rows) - .expect("device MEMW_A build must run on a box with a CUDA backend"); - - assert_eq!( - gpu_u64.len(), - num_rows * math_cuda::trace_cpu::MEMW_ALIGNED_NCOLS - ); - assert_eq!( - gpu_u64, cpu_u64, - "device MEMW_A table must be byte-identical to the CPU fill" - ); - } - - /// LOAD device fill byte-parity (widths 1/2/4/8, signed/unsigned, sign-bit, - /// high address bits, res-byte packing). Skips cleanly with no GPU. - #[cfg(feature = "cuda")] - #[test] - fn gpu_load_fill_matches_cpu() { - if math_cuda::device::backend().is_err() { - eprintln!("skipping gpu_load_fill_matches_cpu: no CUDA backend"); - return; - } - let mut ops = Vec::new(); - for i in 0..400u64 { - let width = [1u8, 2, 4, 8][(i % 4) as usize]; - let signed = i % 2 == 0; - let base = (0x1_0000_0000u64 * (i % 3)) + 0x400 + i * 8; - let ts = 4 * i + 7; - let mut res = [0u64; 8]; - for (j, r) in res.iter_mut().enumerate() { - *r = (i.wrapping_mul(31) + j as u64) & 0xFF; - } - ops.push(load::LoadOperation::new(base, ts, width, signed, res)); - } - let n = ops.len(); - let num_rows = n.next_power_of_two().max(4); - let cpu = load::generate_load_trace(&ops); - let (fe, w) = cpu.main_data_row_major(); - assert_eq!(w, math_cuda::trace_cpu::LOAD_NCOLS); - let cpu_u64: Vec = fe - .iter() - .map(|e| unsafe { *(e.value() as *const u64) }) - .collect(); - let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::LOAD_STRIDE); - for op in &ops { - packed.extend_from_slice(&crate::tables::gpu_trace::pack_load_op(op)); - } - let gpu = math_cuda::trace_cpu::gpu_build_load_trace_host(&packed, n, num_rows) - .expect("device LOAD build must run on a box with a CUDA backend"); - assert_eq!( - gpu, cpu_u64, - "device LOAD table must be byte-identical to the CPU fill" - ); - } - - /// STORE device fill byte-parity (widths 1/2/4/8 via `bytes`, full-value - /// DWordBL split, high address bits). Skips cleanly with no GPU. - #[cfg(feature = "cuda")] - #[test] - fn gpu_store_fill_matches_cpu() { - if math_cuda::device::backend().is_err() { - eprintln!("skipping gpu_store_fill_matches_cpu: no CUDA backend"); - return; - } - let mut ops = Vec::new(); - for i in 0..400u64 { - let bytes = [1u8, 2, 4, 8][(i % 4) as usize]; - let base = (0x1_0000_0000u64 * (i % 3)) + 0x800 + i * 8; - let ts = 4 * i + 9; - let value = i.wrapping_mul(0x9E37_79B9_7F4A_7C15); - ops.push(store::StoreOperation::new(base, ts, value, bytes)); - } - let n = ops.len(); - let num_rows = n.next_power_of_two().max(4); - let cpu = store::generate_store_trace(&ops); - let (fe, w) = cpu.main_data_row_major(); - assert_eq!(w, math_cuda::trace_cpu::STORE_NCOLS); - let cpu_u64: Vec = fe - .iter() - .map(|e| unsafe { *(e.value() as *const u64) }) - .collect(); - let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::STORE_STRIDE); - for op in &ops { - packed.extend_from_slice(&crate::tables::gpu_trace::pack_store_op(op)); - } - let gpu = math_cuda::trace_cpu::gpu_build_store_trace_host(&packed, n, num_rows) - .expect("device STORE build must run on a box with a CUDA backend"); - assert_eq!( - gpu, cpu_u64, - "device STORE table must be byte-identical to the CPU fill" - ); - } - - /// SHIFT device fill byte-parity: the kernel recomputes the full aux - /// (bit_shift/zbs/x/y/limb_shift/out) — covers left/right, signed/unsigned, - /// word/64-bit, shift amounts spanning 0 / >16 / >32 / >64, negative MSB inputs, - /// and the padding rows (ZBS=1). SHIFT is per-row (μ=1), so byte-parity holds. - #[cfg(feature = "cuda")] - #[test] - fn gpu_shift_fill_matches_cpu() { - if math_cuda::device::backend().is_err() { - eprintln!("skipping gpu_shift_fill_matches_cpu: no CUDA backend"); - return; - } - let mut ops = Vec::new(); - // Exhaustive over shift byte 0..256 × all flags × MSB-set / edge values - // (so the signed arithmetic-right-shift extension path is fully covered). - let values: [u64; 8] = [ - 0, - 1, - 0x8000_0000_0000_0000, - 0xFFFF_FFFF_FFFF_FFFF, - 0x1234_5678_9ABC_DEF0, - 0xFEDC_BA98_7654_3210, - 0x0000_0000_FFFF_FFFF, - 0xFFFF_FFFF_0000_0000, - ]; - for &value in &values { - for shift_amount in 0u64..256 { - for &direction in &[false, true] { - for &signed in &[false, true] { - for &word_instr in &[false, true] { - ops.push(shift::ShiftOperation::new( - value, - shift_amount, - direction, - signed, - word_instr, - )); - } - } - } - } - } - // Also a few large shift_amounts (high SHIFT_B1/H1/HIGH limbs) — real ops - // carry the full rv2 on the ALU bus. - for &sa in &[0x1_0000u64, 0xFFFF_FFFFu64, 0x1234_5678_9ABCu64, u64::MAX] { - ops.push(shift::ShiftOperation::new( - 0xDEAD_BEEF_CAFE_1234, - sa, - false, - true, - false, - )); - } - let n = ops.len(); - let num_rows = n.next_power_of_two().max(4); - let cpu = shift::generate_shift_trace(&ops); - let (fe, w) = cpu.main_data_row_major(); - assert_eq!(w, math_cuda::trace_cpu::SHIFT_NCOLS); - let cpu_u64: Vec = fe - .iter() - .map(|e| unsafe { *(e.value() as *const u64) }) - .collect(); - let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::SHIFT_STRIDE); - for op in &ops { - packed.extend_from_slice(&crate::tables::gpu_trace::pack_shift_op(op)); - } - let gpu = math_cuda::trace_cpu::gpu_build_shift_trace_host(&packed, n, num_rows) - .expect("device SHIFT build must run on a box with a CUDA backend"); - assert_eq!( - gpu, cpu_u64, - "device SHIFT table must be byte-identical to the CPU fill" - ); - } - - /// LT device fill: dedup rides the permutation-invariant ALU bus, so the row - /// ORDER is non-deterministic (HashMap iteration). Validate as a MULTISET — - /// the set of real rows (μ>0), incl. summed multiplicities, must match the CPU - /// `generate_lt_trace`. Covers signed/unsigned, invert, MSBs, and duplicates - /// (μ>1). Skips cleanly with no GPU. - #[cfg(feature = "cuda")] - #[test] - fn gpu_lt_fill_matches_cpu_multiset() { - if math_cuda::device::backend().is_err() { - eprintln!("skipping gpu_lt_fill_matches_cpu_multiset: no CUDA backend"); - return; - } - // Raw ops with deliberate duplicates (some pushed twice → μ=2). - let mut raw = Vec::new(); - for i in 0..800u64 { - let lhs = i.wrapping_mul(0x9E37_79B9_7F4A_7C15) ^ (i << 50); - let rhs = i.wrapping_mul(0x1234_5678_9ABC_DEF1).rotate_left(17); - let signed = i % 2 == 0; - let invert = i % 3 == 0; - let op = lt::LtOperation::new_with_invert(lhs, rhs, signed, invert); - raw.push(op.clone()); - if i % 4 == 0 { - raw.push(op); // duplicate → multiplicity 2 - } - } - - let ncols = math_cuda::trace_cpu::LT_NCOLS; - // Extract the real (μ>0) rows of a row-major u64 buffer as a sorted multiset. - let real_rows = |flat: &[u64]| -> Vec> { - let mut rows: Vec> = flat - .chunks(ncols) - .filter(|row| row[lt::cols::MU] > 0) - .map(|row| row.to_vec()) - .collect(); - rows.sort(); - rows - }; - - let cpu_table = lt::generate_lt_trace(&raw); - let (cpu_fe, w) = cpu_table.main_data_row_major(); - assert_eq!(w, ncols); - let cpu_u64: Vec = cpu_fe - .iter() - .map(|e| unsafe { *(e.value() as *const u64) }) - .collect(); - - // Dedup on the host exactly like `gpu_build_lt_tables`, then device-fill. - let mut map: std::collections::HashMap = HashMap::new(); - for op in &raw { - *map.entry(op.clone()).or_insert(0) += 1; - } - let unique: Vec<(lt::LtOperation, u64)> = map.into_iter().collect(); - let n = unique.len(); - let num_rows = n.next_power_of_two().max(4); - let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::LT_STRIDE); - for (op, mult) in &unique { - packed.extend_from_slice(&crate::tables::gpu_trace::pack_lt_op(op, *mult)); - } - let gpu_u64 = math_cuda::trace_cpu::gpu_build_lt_trace_host(&packed, n, num_rows) - .expect("device LT build must run on a box with a CUDA backend"); - - assert_eq!( - real_rows(&gpu_u64), - real_rows(&cpu_u64), - "device LT rows must match the CPU fill as a multiset" - ); - } -} diff --git a/prover/src/tests/gpu_fill_tests.rs b/prover/src/tests/gpu_fill_tests.rs new file mode 100644 index 000000000..ddc11ff34 --- /dev/null +++ b/prover/src/tests/gpu_fill_tests.rs @@ -0,0 +1,465 @@ +//! GPU device-fill parity tests: byte-parity (and multiset, for the +//! order-independent dedup ALU tables) between each table's on-device fill and its +//! CPU `generate_*_trace`, so the GPU trace tables are bit-identical to the CPU +//! path. Each test skips cleanly with no CUDA backend. Grouped here (rather than +//! inline in the table modules) so production code carries no test code; the whole +//! module is `cuda`-gated at its `mod.rs` registration. + +use std::collections::HashMap; + +use crate::tables::cpu::CpuOperation; +use crate::tables::memw::MemwOperation; +use crate::tables::{cpu, load, lt, memw_aligned, memw_register, shift, store}; + +/// CPU-table device fill must be byte-identical to `cpu::generate_cpu_trace`. +/// The CPU kernel is the most intricate of the seven (word-delegate column +/// masking, `PC_DOUBLE_READ`/`PREV_PC_TIMESTAMP_BORROW`, and the `+4` padding +/// cadence with PC=1), so this guards it the same way its six siblings are +/// guarded. The fill is a pure function of the packed op fields, so the synthetic +/// ops need only be diverse (not a valid execution): word-delegate rows, x255 PC +/// reads (`pc_double_read`), `ts_lo < 3` rows (`prev_pc_timestamp_borrow`), x0 +/// registers, and high-word values (DWordWL/DWordHL splits). `n = 300 < 512` +/// forces padding rows, exercising the `+4` cadence off `last_ts`. Skips cleanly +/// with no GPU. +#[test] +fn gpu_cpu_fill_matches_cpu() { + use crate::tables::types::{DecodeEntry, ShrunkDecode}; + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_cpu_fill_matches_cpu: no CUDA backend"); + return; + } + let mut ops = Vec::new(); + for i in 0..300u64 { + let word = i % 5 == 0; + // rs1 cycles through x255 (PC register), x0, and normal registers so the + // pc_double_read and register-zero-suppression paths are both hit. + let rs1 = match i % 4 { + 0 => 255u8, + 1 => 0, + _ => (i % 31 + 1) as u8, + }; + let fields = ShrunkDecode { + read_register1: i % 3 != 0, + read_register2: i % 2 == 0, + write_register: i % 3 == 0, + word_instr: word, + alu: i % 6 == 0, + add: i % 6 == 1, + sub: i % 6 == 2, + memory: i % 6 == 3, + branch: i % 6 == 4, + ecall: i % 6 == 5, + rs1, + rs2: (i % 32) as u8, + rd: ((i + 7) % 32) as u8, + half_instruction_length: if i % 2 == 0 { 1 } else { 2 }, + alu_flags: (i & 0xFF) as u8, + mem_flags: ((i >> 1) & 0xFF) as u8, + }; + // A few timestamps with ts_lo < 3 (0,1,2) trip prev_pc_timestamp_borrow; + // the rest are large and continue a +4 cadence. + let timestamp = match i { + 0 => 0, + 1 => 1, + 2 => 2, + _ => 4 * i + 100, + }; + // PC/imm/next_pc carry high words (>32 bits) to exercise the DWordWL split + // into PC_1/NEXT_PC_1/IMM_1. + let decode = DecodeEntry { + pc: 0x1000 + i * 4 + ((i % 3) << 33), + imm: i.wrapping_mul(0x9E37_79B9) | ((i % 4) << 40), + fields, + }; + ops.push(CpuOperation { + decode, + timestamp, + next_pc: 0x2000 + i * 4 + ((i % 2) << 34), + rvd: i.wrapping_mul(0x1234_5678_9ABC), + rv1: i.wrapping_mul(0xDEAD_BEEF).rotate_left(13), + rv2: (i ^ 0xFFFF_0000_1111) << 3, + arg2: i.wrapping_add(0x7777_8888_9999), + res: i.wrapping_mul(0xABCD_1234_5678) ^ (i << 48), + branch_cond: i % 3 == 1, + ..Default::default() + }); + } + let n = ops.len(); + let num_rows = n.next_power_of_two().max(4); + let last_ts = ops.last().map(|op| op.timestamp).unwrap_or(0); + + let cpu_table = cpu::generate_cpu_trace(&ops); + let (cpu_fe, w) = cpu_table.main_data_row_major(); + assert_eq!(w, math_cuda::trace_cpu::CPU_NCOLS); + let cpu_u64: Vec = cpu_fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + + let packed = crate::tables::gpu_trace::pack_cpu_ops(&ops); + let gpu_u64 = math_cuda::trace_cpu::gpu_build_cpu_trace_host(&packed, n, num_rows, last_ts) + .expect("device CPU build must run on a box with a CUDA backend"); + + assert_eq!(gpu_u64.len(), num_rows * math_cuda::trace_cpu::CPU_NCOLS); + assert_eq!( + gpu_u64, cpu_u64, + "device CPU table must be byte-identical to the CPU fill" + ); +} + +/// MEMW_R (register fast-path) device fill must be byte-identical to the CPU +/// `generate_memw_register_trace_from_rows`. The rows are the walked register +/// rows; here we build synthetic `RegRow`s (read + write, varied values/old/ts) +/// and fill both ways. +#[test] +fn gpu_memw_register_fill_matches_cpu() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_memw_register_fill_matches_cpu: no CUDA backend"); + return; + } + let mut rows = Vec::new(); + for i in 0..500u64 { + let reg_addr = 2 * (i % 40); // even word-address = 2*reg_index + let ts = 4 * i + 100; + let old_ts = 4 * i + 50; + let val = i.wrapping_mul(2_654_435_761); + let old = i.wrapping_mul(40_503) ^ 0xABCD_1234; + let is_read = i % 3 == 0; + rows.push(memw_register::RegRow::new( + reg_addr, + ts, + val as u32, + (val >> 32) as u32, + old as u32, + (old >> 32) as u32, + old_ts, + is_read, + )); + } + let num_rows = rows.len().next_power_of_two().max(4); + + let cpu_table = memw_register::generate_memw_register_trace_from_rows(&rows); + let (cpu_fe, w) = cpu_table.main_data_row_major(); + assert_eq!(w, math_cuda::trace_cpu::MEMW_REGISTER_NCOLS); + let cpu_u64: Vec = cpu_fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + + let mut reg_addr = Vec::new(); + let mut ts = Vec::new(); + let mut value = Vec::new(); + let mut is_read = Vec::new(); + let mut old_value = Vec::new(); + let mut old_tsv = Vec::new(); + for r in &rows { + let (ra, t, v, ir, ov, ot) = r.fill_soa(); + reg_addr.push(ra); + ts.push(t); + value.push(v); + is_read.push(ir); + old_value.push(ov); + old_tsv.push(ot); + } + let gpu_u64 = math_cuda::trace_cpu::gpu_fill_memw_register_host( + ®_addr, &ts, &value, &is_read, &old_value, &old_tsv, num_rows, + ) + .expect("device MEMW_R fill must run on a box with a CUDA backend"); + + assert_eq!( + gpu_u64.len(), + num_rows * math_cuda::trace_cpu::MEMW_REGISTER_NCOLS + ); + assert_eq!( + gpu_u64, cpu_u64, + "device MEMW_R fill must be byte-identical to the CPU fill" + ); +} + +/// MEMW_A (aligned memory — the biggest remaining uploader) device fill must be +/// byte-identical to the CPU `generate_memw_aligned_trace`. Exercises memory +/// (widths 2/4/8), register fallback (`is_register`), high address bits, and +/// the 2×u32-per-u64 value/old packing. Skips cleanly with no GPU. +#[test] +fn gpu_memw_aligned_fill_matches_cpu() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_memw_aligned_fill_matches_cpu: no CUDA backend"); + return; + } + let mut ops = Vec::new(); + for i in 0..500u64 { + let width = [2u8, 4, 8][(i % 3) as usize]; + let is_read = i % 2 == 0; + let is_register = i % 7 == 0; + // Vary the high address word too (whh split has a 32-bit high word). + let base = (0x1_0000_0000u64 * (i % 4)) + 0x8000 + i * 8; + let value = [ + (i as u32).wrapping_mul(2654435761), + (i as u32) ^ 0xABCD_1234, + i as u32, + 0, + 7, + 0, + 0, + (i as u32) & 0xFF, + ]; + let old = [ + (i as u32).wrapping_add(99), + (i as u32) ^ 0x0BAD_F00D, + 0, + 3, + 0, + 0, + 0, + 0, + ]; + let ts = 4 * i + 100; + let old_ts = [4 * i + 50; 8]; + ops.push( + MemwOperation::new(is_register, base, value, ts, width, is_read).with_old(old, old_ts), + ); + } + let n = ops.len(); + let num_rows = n.next_power_of_two().max(4); + + let cpu_table = memw_aligned::generate_memw_aligned_trace(&ops); + let (cpu_fe, width_cols) = cpu_table.main_data_row_major(); + assert_eq!(width_cols, math_cuda::trace_cpu::MEMW_ALIGNED_NCOLS); + let cpu_u64: Vec = cpu_fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::MEMW_ALIGNED_STRIDE); + for op in &ops { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_memw_aligned_op(op)); + } + let gpu_u64 = math_cuda::trace_cpu::gpu_build_memw_aligned_trace_host(&packed, n, num_rows) + .expect("device MEMW_A build must run on a box with a CUDA backend"); + + assert_eq!( + gpu_u64.len(), + num_rows * math_cuda::trace_cpu::MEMW_ALIGNED_NCOLS + ); + assert_eq!( + gpu_u64, cpu_u64, + "device MEMW_A table must be byte-identical to the CPU fill" + ); +} + +/// LOAD device fill byte-parity (widths 1/2/4/8, signed/unsigned, sign-bit, +/// high address bits, res-byte packing). Skips cleanly with no GPU. +#[test] +fn gpu_load_fill_matches_cpu() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_load_fill_matches_cpu: no CUDA backend"); + return; + } + let mut ops = Vec::new(); + for i in 0..400u64 { + let width = [1u8, 2, 4, 8][(i % 4) as usize]; + let signed = i % 2 == 0; + let base = (0x1_0000_0000u64 * (i % 3)) + 0x400 + i * 8; + let ts = 4 * i + 7; + let mut res = [0u64; 8]; + for (j, r) in res.iter_mut().enumerate() { + *r = (i.wrapping_mul(31) + j as u64) & 0xFF; + } + ops.push(load::LoadOperation::new(base, ts, width, signed, res)); + } + let n = ops.len(); + let num_rows = n.next_power_of_two().max(4); + let cpu = load::generate_load_trace(&ops); + let (fe, w) = cpu.main_data_row_major(); + assert_eq!(w, math_cuda::trace_cpu::LOAD_NCOLS); + let cpu_u64: Vec = fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::LOAD_STRIDE); + for op in &ops { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_load_op(op)); + } + let gpu = math_cuda::trace_cpu::gpu_build_load_trace_host(&packed, n, num_rows) + .expect("device LOAD build must run on a box with a CUDA backend"); + assert_eq!( + gpu, cpu_u64, + "device LOAD table must be byte-identical to the CPU fill" + ); +} + +/// STORE device fill byte-parity (widths 1/2/4/8 via `bytes`, full-value +/// DWordBL split, high address bits). Skips cleanly with no GPU. +#[test] +fn gpu_store_fill_matches_cpu() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_store_fill_matches_cpu: no CUDA backend"); + return; + } + let mut ops = Vec::new(); + for i in 0..400u64 { + let bytes = [1u8, 2, 4, 8][(i % 4) as usize]; + let base = (0x1_0000_0000u64 * (i % 3)) + 0x800 + i * 8; + let ts = 4 * i + 9; + let value = i.wrapping_mul(0x9E37_79B9_7F4A_7C15); + ops.push(store::StoreOperation::new(base, ts, value, bytes)); + } + let n = ops.len(); + let num_rows = n.next_power_of_two().max(4); + let cpu = store::generate_store_trace(&ops); + let (fe, w) = cpu.main_data_row_major(); + assert_eq!(w, math_cuda::trace_cpu::STORE_NCOLS); + let cpu_u64: Vec = fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::STORE_STRIDE); + for op in &ops { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_store_op(op)); + } + let gpu = math_cuda::trace_cpu::gpu_build_store_trace_host(&packed, n, num_rows) + .expect("device STORE build must run on a box with a CUDA backend"); + assert_eq!( + gpu, cpu_u64, + "device STORE table must be byte-identical to the CPU fill" + ); +} + +/// SHIFT device fill byte-parity: the kernel recomputes the full aux +/// (bit_shift/zbs/x/y/limb_shift/out) — covers left/right, signed/unsigned, +/// word/64-bit, shift amounts spanning 0 / >16 / >32 / >64, negative MSB inputs, +/// and the padding rows (ZBS=1). SHIFT is per-row (μ=1), so byte-parity holds. +#[test] +fn gpu_shift_fill_matches_cpu() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_shift_fill_matches_cpu: no CUDA backend"); + return; + } + let mut ops = Vec::new(); + // Exhaustive over shift byte 0..256 × all flags × MSB-set / edge values + // (so the signed arithmetic-right-shift extension path is fully covered). + let values: [u64; 8] = [ + 0, + 1, + 0x8000_0000_0000_0000, + 0xFFFF_FFFF_FFFF_FFFF, + 0x1234_5678_9ABC_DEF0, + 0xFEDC_BA98_7654_3210, + 0x0000_0000_FFFF_FFFF, + 0xFFFF_FFFF_0000_0000, + ]; + for &value in &values { + for shift_amount in 0u64..256 { + for &direction in &[false, true] { + for &signed in &[false, true] { + for &word_instr in &[false, true] { + ops.push(shift::ShiftOperation::new( + value, + shift_amount, + direction, + signed, + word_instr, + )); + } + } + } + } + } + // Also a few large shift_amounts (high SHIFT_B1/H1/HIGH limbs) — real ops + // carry the full rv2 on the ALU bus. + for &sa in &[0x1_0000u64, 0xFFFF_FFFFu64, 0x1234_5678_9ABCu64, u64::MAX] { + ops.push(shift::ShiftOperation::new( + 0xDEAD_BEEF_CAFE_1234, + sa, + false, + true, + false, + )); + } + let n = ops.len(); + let num_rows = n.next_power_of_two().max(4); + let cpu = shift::generate_shift_trace(&ops); + let (fe, w) = cpu.main_data_row_major(); + assert_eq!(w, math_cuda::trace_cpu::SHIFT_NCOLS); + let cpu_u64: Vec = fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::SHIFT_STRIDE); + for op in &ops { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_shift_op(op)); + } + let gpu = math_cuda::trace_cpu::gpu_build_shift_trace_host(&packed, n, num_rows) + .expect("device SHIFT build must run on a box with a CUDA backend"); + assert_eq!( + gpu, cpu_u64, + "device SHIFT table must be byte-identical to the CPU fill" + ); +} + +/// LT device fill: dedup rides the permutation-invariant ALU bus, so the row +/// ORDER is non-deterministic (HashMap iteration). Validate as a MULTISET — +/// the set of real rows (μ>0), incl. summed multiplicities, must match the CPU +/// `generate_lt_trace`. Covers signed/unsigned, invert, MSBs, and duplicates +/// (μ>1). Skips cleanly with no GPU. +#[test] +fn gpu_lt_fill_matches_cpu_multiset() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_lt_fill_matches_cpu_multiset: no CUDA backend"); + return; + } + // Raw ops with deliberate duplicates (some pushed twice → μ=2). + let mut raw = Vec::new(); + for i in 0..800u64 { + let lhs = i.wrapping_mul(0x9E37_79B9_7F4A_7C15) ^ (i << 50); + let rhs = i.wrapping_mul(0x1234_5678_9ABC_DEF1).rotate_left(17); + let signed = i % 2 == 0; + let invert = i % 3 == 0; + let op = lt::LtOperation::new_with_invert(lhs, rhs, signed, invert); + raw.push(op.clone()); + if i % 4 == 0 { + raw.push(op); // duplicate → multiplicity 2 + } + } + + let ncols = math_cuda::trace_cpu::LT_NCOLS; + // Extract the real (μ>0) rows of a row-major u64 buffer as a sorted multiset. + let real_rows = |flat: &[u64]| -> Vec> { + let mut rows: Vec> = flat + .chunks(ncols) + .filter(|row| row[lt::cols::MU] > 0) + .map(|row| row.to_vec()) + .collect(); + rows.sort(); + rows + }; + + let cpu_table = lt::generate_lt_trace(&raw); + let (cpu_fe, w) = cpu_table.main_data_row_major(); + assert_eq!(w, ncols); + let cpu_u64: Vec = cpu_fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + + // Dedup on the host exactly like `gpu_build_lt_tables`, then device-fill. + let mut map: std::collections::HashMap = HashMap::new(); + for op in &raw { + *map.entry(op.clone()).or_insert(0) += 1; + } + let unique: Vec<(lt::LtOperation, u64)> = map.into_iter().collect(); + let n = unique.len(); + let num_rows = n.next_power_of_two().max(4); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::LT_STRIDE); + for (op, mult) in &unique { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_lt_op(op, *mult)); + } + let gpu_u64 = math_cuda::trace_cpu::gpu_build_lt_trace_host(&packed, n, num_rows) + .expect("device LT build must run on a box with a CUDA backend"); + + assert_eq!( + real_rows(&gpu_u64), + real_rows(&cpu_u64), + "device LT rows must match the CPU fill as a multiset" + ); +} diff --git a/prover/src/tests/mod.rs b/prover/src/tests/mod.rs index c3600c359..e43337ec9 100644 --- a/prover/src/tests/mod.rs +++ b/prover/src/tests/mod.rs @@ -44,6 +44,8 @@ pub mod ecdas_tests; pub mod ecsm_tests; #[cfg(test)] pub mod eq_tests; +#[cfg(all(test, feature = "cuda"))] +pub mod gpu_fill_tests; #[cfg(test)] pub mod keccak_rnd_tests; #[cfg(test)]