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,12 @@ 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
+            * (P::ONES - lv.is_kernel_mode - 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),
@@ -140,6 +146,17 @@ pub(crate) fn eval_ext_circuit<F: RichField + Extendable<D>, const D: usize>(
         yield_constr.constraint_transition(builder, kernel_constr);
     }
 
+    // Check the helper value in the general columns.
+    {
+        let not_kernel_mode = builder.sub_extension(one, lv.is_kernel_mode);
+        let diff = builder.sub_extension(
+            not_kernel_mode,
+            lv.general.push().push_prover_input_not_kernel,
+        );
+        let constr = builder.mul_extension(lv.op.push_prover_input, diff);
+        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

@@ -265,7 +265,10 @@ pub(crate) fn ctl_data_byte_packing_push<F: Field>() -> Vec<Column<F>> {
 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),
+            bit_col,
+        )],
         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;