E2 supports suspension

Author: nyunyunyunyu

crates/sdk/tests/integration_test.rs

@@ -8,7 +8,10 @@ use eyre::Result;
 use openvm_build::GuestOptions;
 use openvm_circuit::{
     self,
-    arch::{instructions::exe::VmExe, ContinuationVmProof, ExecutionError, VirtualMachineError},
+    arch::{
+        instructions::exe::VmExe, ContinuationVmProof, ExecutionError, VirtualMachine,
+        VirtualMachineError, VmExecState,
+    },
     utils::test_system_config,
 };
 use openvm_continuations::verifier::{
@@ -19,15 +22,16 @@ use openvm_native_circuit::{execute_program_with_config, NativeConfig, NativeCpu
 use openvm_native_compiler::{conversion::CompilerOptions, prelude::*};
 use openvm_sdk::{
     codec::{Decode, Encode},
-    config::{AggregationConfig, AppConfig, SdkSystemConfig, SdkVmConfig},
+    config::{AggregationConfig, AppConfig, SdkSystemConfig, SdkVmBuilder, SdkVmConfig},
     prover::verify_app_proof,
-    Sdk, StdIn,
+    DefaultStarkEngine, Sdk, StdIn,
 };
 use openvm_stark_sdk::{
     config::{
         baby_bear_poseidon2::{BabyBearPoseidon2Config, BabyBearPoseidon2Engine},
         setup_tracing, FriParameters,
     },
+    engine::StarkFriEngine,
     openvm_stark_backend::p3_field::FieldAlgebra,
     p3_baby_bear::BabyBear,
 };
@@ -265,6 +269,35 @@ fn test_public_values_and_leaf_verification() -> eyre::Result<()> {
     Ok(())
 }
 
+#[test]
+fn test_metered_execution_suspension() -> eyre::Result<()> {
+    setup_tracing();
+    let app_log_blowup = 1;
+    let app_config = small_test_app_config(app_log_blowup);
+    let exe = app_exe_for_test();
+
+    let engine = DefaultStarkEngine::new(app_config.app_fri_params.fri_params);
+    let (vm, _) = VirtualMachine::new_with_keygen(engine, SdkVmBuilder, app_config.app_vm_config)?;
+    let executor_idx_to_air_idx = vm.executor_idx_to_air_idx();
+    let interpreter = vm
+        .executor()
+        .metered_instance(&exe, &executor_idx_to_air_idx)?;
+    let metered_ctx = vm.build_metered_ctx(&exe).with_segment_suspend(true);
+    let vm_state = interpreter.create_initial_vm_state(vec![]);
+    let mut vm_exec_state = VmExecState::new(vm_state, metered_ctx);
+    vm_exec_state = interpreter.execute_metered_from_state_until_suspension(vm_exec_state)?;
+    assert!(
+        vm_exec_state.exit_code.is_ok(),
+        "Execution exits with an error: {:?}",
+        vm_exec_state.exit_code
+    );
+    assert!(
+        vm_exec_state.exit_code?.is_none(),
+        "Expect more than 1 segment but the execution terminates."
+    );
+    Ok(())
+}
+
 #[cfg(feature = "evm-verify")]
 #[test]
 #[ignore = "slow"]

crates/vm/src/arch/execution_mode/metered/ctx.rs

@@ -23,6 +23,7 @@ pub struct MeteredCtx<const PAGE_BITS: usize = DEFAULT_PAGE_BITS> {
     pub is_trace_height_constant: Vec<bool>,
     pub memory_ctx: MemoryCtx<PAGE_BITS>,
     pub segmentation_ctx: SegmentationCtx,
+    pub segment_suspend: bool,
 }
 
 impl<const PAGE_BITS: usize> MeteredCtx<PAGE_BITS> {
@@ -75,6 +76,7 @@ impl<const PAGE_BITS: usize> MeteredCtx<PAGE_BITS> {
             is_trace_height_constant,
             memory_ctx,
             segmentation_ctx,
+            segment_suspend: false,
         };
         if !config.continuation_enabled {
             // force single segment
@@ -106,6 +108,11 @@ impl<const PAGE_BITS: usize> MeteredCtx<PAGE_BITS> {
         self
     }
 
+    pub fn with_segment_suspend(mut self, segment_suspend: bool) -> Self {
+        self.segment_suspend = segment_suspend;
+        self
+    }
+
     pub fn segments(&self) -> &[Segment] {
         &self.segmentation_ctx.segments
     }
@@ -121,7 +128,7 @@ impl<const PAGE_BITS: usize> MeteredCtx<PAGE_BITS> {
     }
 
     #[inline(always)]
-    pub fn check_and_segment(&mut self, instret: u64, segment_check_insns: u64) {
+    pub fn check_and_segment(&mut self, instret: u64, segment_check_insns: u64) -> bool {
         let threshold = self
             .segmentation_ctx
             .instret_last_segment_check
@@ -131,7 +138,7 @@ impl<const PAGE_BITS: usize> MeteredCtx<PAGE_BITS> {
             "overflow in segment check threshold calculation"
         );
         if instret < threshold {
-            return;
+            return false;
         }
 
         self.memory_ctx
@@ -145,6 +152,7 @@ impl<const PAGE_BITS: usize> MeteredCtx<PAGE_BITS> {
         if did_segment {
             self.reset_segment();
         }
+        did_segment
     }
 
     #[allow(dead_code)]
@@ -200,12 +208,11 @@ impl<const PAGE_BITS: usize> ExecutionCtxTrait for MeteredCtx<PAGE_BITS> {
         segment_check_insns: u64,
         exec_state: &mut VmExecState<F, GuestMemory, Self>,
     ) -> bool {
-        // E2 always runs until termination. Here we use the function as a hook called every
-        // instruction.
+        // If `segment_suspend` is set, suspend every segment. Otherwise, execute until termination.
         exec_state
             .ctx
-            .check_and_segment(instret, segment_check_insns);
-        false
+            .check_and_segment(instret, segment_check_insns)
+            && exec_state.ctx.segment_suspend
     }
 
     #[inline(always)]

crates/vm/src/arch/interpreter.rs

@@ -201,6 +201,15 @@ where
         })
     }
 
+    pub fn create_initial_vm_state(&self, inputs: impl Into<Streams<F>>) -> VmState<F> {
+        VmState::initial(
+            &self.system_config,
+            &self.init_memory,
+            self.pc_start,
+            inputs,
+        )
+    }
+
     /// # Safety
     /// - This function assumes that the `pc` is within program bounds - this should be the case if
     ///   the pc is checked to be in bounds before jumping to it.
@@ -381,12 +390,7 @@ where
         inputs: impl Into<Streams<F>>,
         ctx: MeteredCtx,
     ) -> Result<(Vec<Segment>, VmState<F, GuestMemory>), ExecutionError> {
-        let vm_state = VmState::initial(
-            &self.system_config,
-            &self.init_memory,
-            self.pc_start,
-            inputs,
-        );
+        let vm_state = self.create_initial_vm_state(inputs);
         self.execute_metered_from_state(vm_state, ctx)
     }
 
@@ -405,22 +409,38 @@ where
     ) -> Result<(Vec<Segment>, VmState<F, GuestMemory>), ExecutionError> {
         let mut exec_state = VmExecState::new(from_state, ctx);
 
-        let instret = exec_state.instret();
-        let pc = exec_state.pc();
-        let segmentation_check_insns = exec_state.ctx.segmentation_ctx.segment_check_insns;
+        loop {
+            exec_state = self.execute_metered_from_state_until_suspension(exec_state)?;
+            // The execution has terminated.
+            if exec_state.exit_code.is_ok() && exec_state.exit_code.as_ref().unwrap().is_some() {
+                break;
+            }
+            if exec_state.exit_code.is_err() {
+                return Err(exec_state.exit_code.unwrap_err());
+            }
+        }
+        check_termination(exec_state.exit_code)?;
+        let VmExecState { vm_state, ctx, .. } = exec_state;
+        Ok((ctx.into_segments(), vm_state))
+    }
+    pub fn execute_metered_from_state_until_suspension(
+        &self,
+        mut from_state: VmExecState<F, GuestMemory, MeteredCtx>,
+    ) -> Result<VmExecState<F, GuestMemory, MeteredCtx>, ExecutionError> {
+        let instret = from_state.instret();
+        let pc = from_state.pc();
+        let segmentation_check_insns = from_state.ctx.segmentation_ctx.segment_check_insns;
         // Start execution
         run!(
             "execute_metered",
             self,
             instret,
             pc,
             segmentation_check_insns,
-            exec_state,
+            from_state,
             MeteredCtx
         );
-        check_termination(exec_state.exit_code)?;
-        let VmExecState { vm_state, ctx, .. } = exec_state;
-        Ok((ctx.into_segments(), vm_state))
+        Ok(from_state)
     }
 }
 
@@ -437,12 +457,7 @@ where
         inputs: impl Into<Streams<F>>,
         ctx: MeteredCostCtx,
     ) -> Result<(u64, VmState<F, GuestMemory>), ExecutionError> {
-        let vm_state = VmState::initial(
-            &self.system_config,
-            &self.init_memory,
-            self.pc_start,
-            inputs,
-        );
+        let vm_state = self.create_initial_vm_state(inputs);
         self.execute_metered_cost_from_state(vm_state, ctx)
     }