fix: runtime sharding

core/Cargo.toml

@@ -53,6 +53,7 @@ serde_json = { version = "1.0.113", default-features = false, features = [
 ] }
 k256 = { version = "0.13.3", features = ["expose-field"] }
 elliptic-curve = "0.13.8"
+serial_test = "3.0.0"
 
 [dev-dependencies]
 criterion = "0.5.1"

core/src/cpu/air/mod.rs

@@ -20,6 +20,7 @@ impl<AB> Air<AB> for CpuChip
 where
     AB: CurtaAirBuilder,
 {
+    #[inline(never)]
     fn eval(&self, builder: &mut AB) {
         let main = builder.main();
         let local: &CpuCols<AB::Var> = main.row_slice(0).borrow();

core/src/runtime/mod.rs

@@ -414,6 +414,14 @@ impl Runtime {
         self.syscall_map.get(&code)
     }
 
+    fn max_syscall_cycles(&self) -> u32 {
+        self.syscall_map
+            .values()
+            .map(|syscall| syscall.num_extra_cycles())
+            .max()
+            .unwrap_or(0)
+    }
+
     /// Execute the given instruction over the current state of the runtime.
     fn execute(&mut self, instruction: Instruction) {
         let pc = self.state.pc;
@@ -738,6 +746,8 @@ impl Runtime {
             self.state.memory.insert(*addr, (*value, 0, 0));
         }
 
+        let max_syscall_cycles = self.max_syscall_cycles();
+
         self.state.clk += 1;
         while self.state.pc.wrapping_sub(self.program.pc_base)
             < (self.program.instructions.len() * 4) as u32
@@ -779,8 +789,9 @@ impl Runtime {
             self.state.global_clk += 1;
             self.state.clk += 4;
 
-            // Every `shard_size` cycles, increment shard and reset clk within the shard.
-            if self.state.global_clk % self.shard_size == 0 && !self.unconstrained {
+            // If there's not enough cycles left for another instruction, move to the next shard.
+            // We multiply by 4 because clk is incremented by 4 for each normal instruction.
+            if !self.unconstrained && max_syscall_cycles + self.state.clk >= self.shard_size * 4 {
                 self.state.current_shard += 1;
                 self.state.clk = 0;
             }
@@ -862,7 +873,10 @@ impl Runtime {
 #[cfg(test)]
 pub mod tests {
 
-    use crate::{runtime::Register, utils::tests::FIBONACCI_ELF};
+    use crate::{
+        runtime::Register,
+        utils::tests::{FIBONACCI_ELF, SSZ_WITHDRAWALS_ELF},
+    };
 
     use super::{Instruction, Opcode, Program, Runtime};
 
@@ -879,6 +893,10 @@ pub mod tests {
         Program::from(FIBONACCI_ELF)
     }
 
+    pub fn ssz_withdrawals_program() -> Program {
+        Program::from(SSZ_WITHDRAWALS_ELF)
+    }
+
     pub fn ecall_lwa_program() -> Program {
         let instructions = vec![
             Instruction::new(Opcode::ADD, 5, 0, 101, false, true),

core/src/stark/debug.rs

@@ -1,3 +1,5 @@
+use std::panic::{catch_unwind, AssertUnwindSafe};
+
 use p3_air::{
     Air, AirBuilder, BaseAir, ExtensionBuilder, PairBuilder, PermutationAirBuilder,
     TwoRowMatrixView,
@@ -75,7 +77,14 @@ pub fn debug_constraints<F: PrimeField, EF: ExtensionField<F>, A>(
             builder.is_transition = F::zero();
         }
 
-        air.eval(&mut builder);
+        let result = catch_unwind(AssertUnwindSafe(|| {
+            air.eval(&mut builder);
+        }));
+        if result.is_err() {
+            println!("local: {:?}", main_local);
+            println!("next:  {:?}", main_local);
+            panic!("failed at row {} of chip {}", i, air.name());
+        }
         eval_permutation_constraints(air, &mut builder, cumulative_sum);
     });
 }
@@ -176,7 +185,11 @@ where
     }
 
     fn assert_zero<I: Into<Self::Expr>>(&mut self, x: I) {
-        assert_eq!(x.into(), F::zero(), "constraints must evaluate to zero");
+        let f: F = x.into();
+        if f != F::zero() {
+            let backtrace = std::backtrace::Backtrace::force_capture();
+            panic!("constraint failed: {}", backtrace);
+        }
     }
 }
 

core/src/stark/runtime.rs

@@ -227,12 +227,14 @@ pub mod tests {
     use crate::runtime::tests::fibonacci_program;
     use crate::runtime::tests::simple_memory_program;
     use crate::runtime::tests::simple_program;
+    use crate::runtime::tests::ssz_withdrawals_program;
     use crate::runtime::Instruction;
     use crate::runtime::Opcode;
     use crate::runtime::Program;
     use crate::utils;
     use crate::utils::prove;
     use crate::utils::setup_logger;
+    use serial_test::serial;
 
     #[test]
     fn test_simple_prove() {
@@ -375,6 +377,14 @@ pub mod tests {
         prove(program);
     }
 
+    #[test]
+    #[serial]
+    fn test_ssz_withdrawals_prove() {
+        setup_logger();
+        let program = ssz_withdrawals_program();
+        prove(program);
+    }
+
     #[test]
     fn test_simple_memory_program_prove() {
         let program = simple_memory_program();

core/src/utils/prove.rs

@@ -5,6 +5,7 @@ use crate::{
     stark::{LocalProver, StarkConfig},
 };
 pub use baby_bear_blake3::BabyBearBlake3;
+use size::Size;
 
 pub trait StarkUtils: StarkConfig {
     type UniConfig: p3_uni_stark::StarkGenericConfig<
@@ -61,11 +62,11 @@ pub fn prove_core(runtime: &mut Runtime) -> crate::stark::Proof<BabyBearBlake3>
     let nb_bytes = bincode::serialize(&proof).unwrap().len();
 
     tracing::info!(
-        "cycles={}, e2e={}, khz={:.2}, proofSize={}kb",
+        "cycles={}, e2e={}, khz={:.2}, proofSize={}",
         cycles,
         time,
         (cycles as f64 / time as f64),
-        nb_bytes / 1000
+        Size::from_bytes(nb_bytes),
     );
 
     #[cfg(not(feature = "perf"))]