feat: skip range-checking PUSH operations in KERNEL mode

evm_arithmetization/src/byte_packing/byte_packing_stark.rs

@@ -175,8 +175,7 @@ impl<F: RichField + Extendable<D>, const D: usize> BytePackingStark<F, D> {
         ops: Vec<BytePackingOp>,
         min_rows: usize,
     ) -> Vec<[F; NUM_COLUMNS]> {
-        let base_len: usize = ops.iter().map(|op| usize::from(!op.bytes.is_empty())).sum();
-        let num_rows = core::cmp::max(base_len.max(BYTE_RANGE_MAX), min_rows).next_power_of_two();
+        let num_rows = core::cmp::max(ops.len().max(BYTE_RANGE_MAX), min_rows).next_power_of_two();
         let mut rows = Vec::with_capacity(num_rows);
 
         for op in ops {

evm_arithmetization/src/cpu/columns/general.rs

@@ -14,6 +14,7 @@ pub(crate) union CpuGeneralColumnsView<T: Copy> {
     jumps: CpuJumpsView<T>,
     shift: CpuShiftView<T>,
     stack: CpuStackView<T>,
+    push: CpuPushView<T>,
 }
 
 impl<T: Copy> CpuGeneralColumnsView<T> {
@@ -78,6 +79,18 @@ impl<T: Copy> CpuGeneralColumnsView<T> {
     pub(crate) fn stack_mut(&mut self) -> &mut CpuStackView<T> {
         unsafe { &mut self.stack }
     }
+
+    /// View of the columns required for the push operation.
+    /// SAFETY: Each view is a valid interpretation of the underlying array.
+    pub(crate) fn push(&self) -> &CpuPushView<T> {
+        unsafe { &self.push }
+    }
+
+    /// Mutable view of the columns required for the push operation.
+    /// SAFETY: Each view is a valid interpretation of the underlying array.
+    pub(crate) fn push_mut(&mut self) -> &mut CpuPushView<T> {
+        unsafe { &mut self.push }
+    }
 }
 
 impl<T: Copy + PartialEq> PartialEq<Self> for CpuGeneralColumnsView<T> {
@@ -180,6 +193,18 @@ pub(crate) struct CpuStackView<T: Copy> {
     pub(crate) stack_len_bounds_aux: T,
 }
 
+/// View of the first `CpuGeneralColumn` storing the product of the negated
+/// `is_kernel_mode` flag with the `push_prover_input` combined op flag, to
+/// filter out `PUSH` instructions from being range-checked when happening in
+/// the KERNEL context.
+#[repr(C)]
+#[derive(Copy, Clone)]
+pub(crate) struct CpuPushView<T: Copy> {
+    /// Product of `push_prover_input` with the negated `is_kernel_mode` flag.
+    pub(crate) push_prover_input_not_kernel: T,
+    /// Reserve the unused columns.
+    _padding_columns: [T; NUM_SHARED_COLUMNS - 1],
+}
 /// The number of columns shared by all views of [`CpuGeneralColumnsView`].
 /// This is defined in terms of the largest view in order to determine the
 /// number of padding columns to add to each field without creating a cycle
@@ -195,3 +220,4 @@ const_assert!(size_of::<CpuLogicView<u8>>() == NUM_SHARED_COLUMNS);
 const_assert!(size_of::<CpuJumpsView<u8>>() == NUM_SHARED_COLUMNS);
 const_assert!(size_of::<CpuShiftView<u8>>() == NUM_SHARED_COLUMNS);
 const_assert!(size_of::<CpuStackView<u8>>() == NUM_SHARED_COLUMNS);
+const_assert!(size_of::<CpuPushView<u8>>() == NUM_SHARED_COLUMNS);

evm_arithmetization/src/cpu/control_flow.rs

@@ -78,6 +78,16 @@ pub(crate) fn eval_packed_generic<P: PackedField>(
     );
     yield_constr.constraint_transition(is_prover_input * (lv.is_kernel_mode - nv.is_kernel_mode));
 
+    // Check the helper value in the general columns.
+    yield_constr.constraint(
+        lv.op.push_prover_input
+            * ((lv.is_kernel_mode + lv.general.push().push_prover_input_not_kernel) - P::ONES),
+    );
+
+    // The general column should be not be set when the operation is a
+    // `PROVER_INPUT`.
+    yield_constr.constraint(is_prover_input * lv.general.push().push_prover_input_not_kernel);
+
     // If a non-CPU cycle row is followed by a CPU cycle row, then:
     //  - the `program_counter` of the CPU cycle row is `main` (the entry point of
     //    our kernel),
@@ -105,6 +115,12 @@ pub(crate) fn eval_ext_circuit<F: RichField + Extendable<D>, const D: usize>(
 
     let next_halt_state = builder.sub_extension(one, is_cpu_cycle_next);
 
+    let is_prover_input = builder.mul_sub_extension(
+        lv.op.push_prover_input,
+        lv.opcode_bits[5],
+        lv.op.push_prover_input,
+    );
+
     // Once we start executing instructions, then we continue until the end of the
     // table or we reach dummy padding rows. This, along with the constraints on
     // the first row, enforces that operation flags and the halt flag are
@@ -129,17 +145,33 @@ pub(crate) fn eval_ext_circuit<F: RichField + Extendable<D>, const D: usize>(
         yield_constr.constraint_transition(builder, kernel_constr);
 
         // Same constraints as before, for PROVER_INPUT.
-        let is_prover_input = builder.mul_sub_extension(
-            lv.op.push_prover_input,
-            lv.opcode_bits[5],
-            lv.op.push_prover_input,
-        );
         let pc_constr = builder.mul_add_extension(is_prover_input, pc_diff, is_prover_input);
         yield_constr.constraint_transition(builder, pc_constr);
         let kernel_constr = builder.mul_extension(is_prover_input, kernel_diff);
         yield_constr.constraint_transition(builder, kernel_constr);
     }
 
+    // Check the helper value in the general columns.
+    {
+        let temp = builder.add_extension(
+            lv.is_kernel_mode,
+            lv.general.push().push_prover_input_not_kernel,
+        );
+        let constr =
+            builder.mul_sub_extension(lv.op.push_prover_input, temp, lv.op.push_prover_input);
+        yield_constr.constraint(builder, constr);
+    }
+
+    // The general column should be not be set when the operation is a
+    // `PROVER_INPUT`.
+    {
+        let constr = builder.mul_extension(
+            is_prover_input,
+            lv.general.push().push_prover_input_not_kernel,
+        );
+        yield_constr.constraint(builder, constr);
+    }
+
     // If a non-CPU cycle row is followed by a CPU cycle row, then:
     //  - the `program_counter` of the CPU cycle row is `main` (the entry point of
     //    our kernel),

evm_arithmetization/src/cpu/cpu_stark.rs

@@ -263,9 +263,11 @@ pub(crate) fn ctl_data_byte_packing_push<F: Field>() -> Vec<Column<F>> {
 
 /// CTL filter for the `PUSH` operation.
 pub(crate) fn ctl_filter_byte_packing_push<F: Field>() -> Filter<F> {
-    let bit_col = Column::single(COL_MAP.opcode_bits[5]);
     Filter::new(
-        vec![(Column::single(COL_MAP.op.push_prover_input), bit_col)],
+        vec![(
+            Column::single(COL_MAP.general.push().push_prover_input_not_kernel),
+            Column::single(COL_MAP.op.push_prover_input),
+        )],
         vec![],
     )
 }

evm_arithmetization/src/witness/operation.rs

@@ -3,6 +3,7 @@ use itertools::Itertools;
 use keccak_hash::keccak;
 use plonky2::field::types::Field;
 
+use super::state::KERNEL_CONTEXT;
 use super::transition::Transition;
 use super::util::{
     byte_packing_log, byte_unpacking_log, mem_read_with_log, mem_write_log,
@@ -385,7 +386,15 @@ pub(crate) fn generate_push<F: Field, T: Transition<F>>(
     let val = U256::from_big_endian(&bytes);
     push_with_write(state, &mut row, val)?;
 
-    byte_packing_log(state, base_address, bytes);
+    // This is necessary to filter out PUSH instructions from the BytePackingStark
+    // CTl when happening in the KERNEL context.
+    //
+    // Note that we don't need to
+    row.general.push_mut().push_prover_input_not_kernel = F::ONE - row.is_kernel_mode;
+
+    if code_context != KERNEL_CONTEXT {
+        byte_packing_log(state, base_address, bytes);
+    }
 
     state.push_cpu(row);
 

evm_arithmetization/src/witness/state.rs

@@ -2,7 +2,7 @@ use ethereum_types::U256;
 
 use crate::cpu::kernel::aggregator::KERNEL;
 
-const KERNEL_CONTEXT: usize = 0;
+pub(crate) const KERNEL_CONTEXT: usize = 0;
 
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
 pub struct RegistersState {

evm_arithmetization/src/witness/traces.rs

@@ -65,11 +65,7 @@ impl<T: Copy> Traces<T> {
                     Operation::RangeCheckOperation { .. } => 1,
                 })
                 .sum(),
-            byte_packing_len: self
-                .byte_packing_ops
-                .iter()
-                .map(|op| usize::from(!op.bytes.is_empty()))
-                .sum(),
+            byte_packing_len: self.byte_packing_ops.len(),
             cpu_len: self.cpu.len(),
             keccak_len: self.keccak_inputs.len() * keccak::keccak_stark::NUM_ROUNDS,
             keccak_sponge_len: self

evm_arithmetization/src/witness/util.rs

@@ -344,6 +344,11 @@ pub(crate) fn byte_packing_log<F: Field, T: Transition<F>>(
     base_address: MemoryAddress,
     bytes: Vec<u8>,
 ) {
+    if bytes.is_empty() {
+        // No-op
+        return;
+    }
+
     let clock = state.get_clock();
 
     let mut address = base_address;

zero_bin/tools/prove_stdio.sh

@@ -52,18 +52,28 @@ if [[ $TEST_ONLY == "test_only" ]]; then
     export MEMORY_CIRCUIT_SIZE="17..18"
 else
     if [[ $INPUT_FILE == *"witness_b19240705"* ]]; then
-      # These sizes are configured specifically for block 19240705. Don't use this in other scenarios
+        # These sizes are configured specifically for block 19240705. Don't use this in other scenarios.
         echo "Using specific circuit sizes for witness_b19240705.json"
         export ARITHMETIC_CIRCUIT_SIZE="16..19"
-        export BYTE_PACKING_CIRCUIT_SIZE="16..19"
+        export BYTE_PACKING_CIRCUIT_SIZE="11..15"
         export CPU_CIRCUIT_SIZE="18..21"
         export KECCAK_CIRCUIT_SIZE="15..18"
         export KECCAK_SPONGE_CIRCUIT_SIZE="10..13"
         export LOGIC_CIRCUIT_SIZE="13..17"
         export MEMORY_CIRCUIT_SIZE="20..23"
+    elif [[ $INPUT_FILE == *"witness_b2_b7"* ]]; then
+        # These sizes are configured specifically for custom small blocks. Don't use this in other scenarios.
+        echo "Using specific circuit sizes for witness_b2_b7.json"
+        export ARITHMETIC_CIRCUIT_SIZE="16..17"
+        export BYTE_PACKING_CIRCUIT_SIZE="9..11"
+        export CPU_CIRCUIT_SIZE="16..17"
+        export KECCAK_CIRCUIT_SIZE="14..15"
+        export KECCAK_SPONGE_CIRCUIT_SIZE="9..10"
+        export LOGIC_CIRCUIT_SIZE="12..13"
+        export MEMORY_CIRCUIT_SIZE="18..19"
     else
         export ARITHMETIC_CIRCUIT_SIZE="16..23"
-        export BYTE_PACKING_CIRCUIT_SIZE="9..21"
+        export BYTE_PACKING_CIRCUIT_SIZE="8..21"
         export CPU_CIRCUIT_SIZE="12..25"
         export KECCAK_CIRCUIT_SIZE="14..20"
         export KECCAK_SPONGE_CIRCUIT_SIZE="9..15"