[Fix] metered execution suspend before guest program completed (#246)

* better guest exec

* lint

* update according to reviews

Author: noel2004

crates/integration/src/lib.rs

@@ -357,8 +357,14 @@ pub fn tester_execute<T: ProverTester>(
             .map(|p| p.as_stark_proof().expect("must be stark proof")),
     )?;
 
+    let app_vm_config = app_config.app_vm_config.clone();
     let sdk = Sdk::new(app_config)?;
-    let ret = scroll_zkvm_prover::utils::vm::execute_guest(&sdk, app_exe, &stdin)?;
+    let ret = scroll_zkvm_prover::utils::vm::execute_guest(
+        &sdk,
+        app_vm_config.as_ref(),
+        app_exe,
+        &stdin,
+    )?;
     Ok(ret)
 }
 

crates/prover/src/prover/mod.rs

@@ -10,11 +10,17 @@ use openvm_native_circuit::NativeGpuBuilder as NativeBuilder;
 
 use openvm_circuit::arch::instructions::exe::VmExe;
 use openvm_sdk::{DefaultStarkEngine, config::SdkVmBuilder};
-use openvm_sdk::{F, Sdk, StdIn, prover::StarkProver};
+use openvm_sdk::{
+    F, Sdk, StdIn,
+    config::{AppConfig, SdkVmConfig},
+    prover::StarkProver,
+};
 use scroll_zkvm_types::{proof::OpenVmEvmProof, types_agg::ProgramCommitment, utils::serialize_vk};
 use scroll_zkvm_verifier::verifier::{AGG_STARK_PROVING_KEY, UniversalVerifier};
 use tracing::instrument;
 
+type SdkAppConfig = AppConfig<SdkVmConfig>;
+
 // Re-export from openvm_sdk.
 pub use openvm_sdk::{self};
 
@@ -31,6 +37,8 @@ pub struct Prover {
     pub app_exe: Arc<VmExe<F>>,
     /// Prover configuration.
     pub config: ProverConfig,
+    /// SDKConfig
+    app_config: SdkAppConfig,
     /// Lazily initialized SDK
     sdk: OnceLock<Sdk>,
     /// Lazily initialized stark prover
@@ -54,11 +62,18 @@ impl Prover {
     /// Setup the [`Prover`] given paths to the application's exe and proving key.
     #[instrument("Prover::setup")]
     pub fn setup(config: ProverConfig, name: Option<&str>) -> Result<Self, Error> {
+        let mut app_config = read_app_config(&config.path_app_config)?;
+        let segment_len = config.segment_len.unwrap_or(DEFAULT_SEGMENT_SIZE);
+        let segmentation_limits = &mut app_config.app_vm_config.system.config.segmentation_limits;
+        segmentation_limits.max_trace_height = segment_len as u32;
+        segmentation_limits.max_cells = 1_200_000_000_usize; // For 24G vram
+
         let app_exe = read_app_exe(&config.path_app_exe)?;
         Ok(Self {
             app_exe: Arc::new(app_exe),
             config,
             prover_name: name.unwrap_or("universal").to_string(),
+            app_config,
             sdk: OnceLock::new(),
             prover: OnceLock::new(),
         })
@@ -74,14 +89,7 @@ impl Prover {
     fn get_sdk(&self) -> Result<&Sdk, Error> {
         self.sdk.get_or_try_init(|| {
             tracing::info!("Lazy initializing SDK...");
-            let mut app_config = read_app_config(&self.config.path_app_config)?;
-            let segment_len = self.config.segment_len.unwrap_or(DEFAULT_SEGMENT_SIZE);
-            let segmentation_limits =
-                &mut app_config.app_vm_config.system.config.segmentation_limits;
-            segmentation_limits.max_trace_height = segment_len as u32;
-            segmentation_limits.max_cells = 1_200_000_000_usize; // For 24G vram
-
-            let sdk = Sdk::new(app_config).expect("sdk init failed");
+            let sdk = Sdk::new(self.app_config.clone()).expect("sdk init failed");
 
             // 45s for first time
             let sdk = sdk.with_agg_pk(AGG_STARK_PROVING_KEY.clone());
@@ -157,7 +165,13 @@ impl Prover {
     ) -> Result<crate::utils::vm::ExecutionResult, Error> {
         let sdk = self.get_sdk()?;
         let t = std::time::Instant::now();
-        let exec_result = crate::utils::vm::execute_guest(sdk, self.app_exe.clone(), stdin)?;
+        let exec_result = crate::utils::vm::execute_guest(
+            sdk,
+            self.app_config.app_vm_config.as_ref(),
+            self.app_exe.clone(),
+            stdin,
+        )
+        .map_err(|e| Error::GenProof(e.to_string()))?;
         let execution_time_mills = t.elapsed().as_millis() as u64;
         let execution_time_s = execution_time_mills as f32 / 1000.0f32;
         let exec_speed = (exec_result.total_cycle as f32 / 1_000_000.0f32) / execution_time_s; // MHz

crates/prover/src/utils/vm.rs

@@ -1,27 +1,78 @@
-use openvm_sdk::{Sdk, StdIn, types::ExecutableFormat};
-
-use crate::Error;
+use openvm_circuit::{
+    arch::{SystemConfig, VirtualMachineError},
+    system::memory::merkle::public_values::extract_public_values,
+};
+use openvm_sdk::{Sdk, SdkError, StdIn, types::ExecutableFormat};
 
 pub struct ExecutionResult {
     pub total_cycle: u64,
     #[allow(dead_code)]
     pub public_values: Vec<u8>,
 }
 
+// the 100* current cost / ratio for chunk circuit, use to estimated the max cycles
+// which an metered execution can handle
+const COST_CYCLE_RATIO: u64 = 87u64;
+
+// Execute the guest program using the metered executor first to measure actual cycles.
+// If the execution exceeds the maximum cost allowed by the metered executor,
+// we re-execute the program using the normal executor (execute_e1), which has no limitations
+// on the size of the execution process.
 pub fn execute_guest(
     sdk: &Sdk,
+    sys_config: &SystemConfig,
     exe: impl Into<ExecutableFormat>,
-    stdin: &StdIn,
-) -> Result<ExecutionResult, Error> {
-    let (public_values, (_cost, total_cycle)) = sdk
-        .execute_metered_cost(exe, stdin.clone())
-        .map_err(|e| Error::GenProof(e.to_string()))?;
+    inputs: &StdIn,
+) -> Result<ExecutionResult, SdkError> {
+    let app_prover = sdk.app_prover(exe)?;
+
+    let vm = app_prover.vm();
+    let exe = app_prover.exe();
+
+    let ctx = vm.build_metered_cost_ctx();
+    let preset_max_cost = ctx.max_execution_cost * 2;
+    let estimated_max_cycles = preset_max_cost / COST_CYCLE_RATIO;
+    tracing::info!("Double preset max cost to ({preset_max_cost}) for metering execution");
+    let ctx = ctx.with_max_execution_cost(preset_max_cost);
+    let interpreter = vm
+        .metered_cost_interpreter(&exe)
+        .map_err(VirtualMachineError::from)?;
+
+    let (ctx, final_state) = interpreter
+        .execute_metered_cost(inputs.clone(), ctx)
+        .map_err(VirtualMachineError::from)?;
+    let mut total_cycle = ctx.instret;
+
+    let mut public_values =
+        extract_public_values(sys_config.num_public_values, &final_state.memory.memory);
 
-    tracing::debug!(name: "public_values after guest execution", ?public_values);
     if public_values.iter().all(|x| *x == 0) {
-        return Err(Error::GenProof("public_values are all 0s".to_string()));
+        if ctx.cost < ctx.max_execution_cost {
+            return Err(SdkError::Other(eyre::eyre!(
+                "public_values are all 0s for unexpected reason"
+            )));
+        }
+
+        tracing::warn!(
+            "Large execution exceed limit of metered execution, cycle is expected to >{estimated_max_cycles}"
+        );
+        let exe = sdk.convert_to_exe(exe)?;
+        let instance = vm.interpreter(&exe).map_err(VirtualMachineError::from)?;
+        let final_memory = instance
+            .execute(inputs.clone(), None)
+            .map_err(VirtualMachineError::from)?
+            .memory;
+        public_values = extract_public_values(sys_config.num_public_values, &final_memory.memory);
+        total_cycle = estimated_max_cycles;
+
+        if public_values.iter().all(|x| *x == 0) {
+            return Err(SdkError::Other(eyre::eyre!(
+                "public_values are all 0s upon execute_e1"
+            )));
+        }
     }
 
+    tracing::debug!(name: "public_values after guest execution: {:?}", ?public_values);
     Ok(ExecutionResult {
         total_cycle,
         public_values,