feat: aot metered cost

Cargo.toml

@@ -29,8 +29,9 @@ members = [
     "crates/toolchain/tests",
     "crates/continuations",
     "crates/vm",
-    "crates/vm/src/arch/asm_bridge",
-    "crates/vm/src/arch/asm_bridge_metered",
+    "crates/asm/asm_bridge",
+    "crates/asm/asm_bridge_metered",
+    "crates/asm/asm_bridge_metered_cost",
     "extensions/rv32im/circuit",
     "extensions/rv32im/transpiler",
     "extensions/rv32im/guest",

crates/asm/asm_bridge/Cargo.toml

@@ -0,0 +1,19 @@
+[package]
+name = "asm_bridge"
+version.workspace = true
+edition.workspace = true
+authors.workspace = true
+homepage.workspace = true
+repository.workspace = true
+
+[lib]
+crate-type = ["cdylib"]
+
+[dependencies]
+cc = "1.0"
+openvm-circuit = { workspace = true, features = ["test-utils"] }
+openvm-stark-sdk = { workspace = true }
+asm_bridge_utils = { path = "../asm_bridge_utils" }
+
+[package.metadata.cargo-shear]
+ignored = ["cc"]

crates/asm/asm_bridge/src/lib.rs

@@ -0,0 +1,95 @@
+use std::ffi::c_void;
+
+use openvm_circuit::{
+    arch::{execution_mode::ExecutionCtx, VmExecState},
+    system::memory::online::GuestMemory,
+};
+use openvm_stark_sdk::p3_baby_bear::BabyBear;
+
+extern "C" {
+    fn asm_run_internal(
+        vm_exec_state_ptr: *mut c_void,       // rdi = vm_exec_state
+        pre_compute_insns_ptr: *const c_void, // rsi = pre_compute_insns
+        from_state_pc: u32,                   // rdx = from_state.pc
+        from_state_instret: u64,              // rcx = from_state.instret
+    );
+}
+
+/// Runs the VM execution from assembly
+///
+/// # Safety
+///
+///
+/// This function is unsafe because:
+/// - `vm_exec_state_ptr` must be valid
+/// - `pre_compute_insns` must point to valid pre-compute instructions
+#[no_mangle]
+pub unsafe extern "C" fn asm_run(
+    vm_exec_state_ptr: *mut c_void,
+    pre_compute_insns_ptr: *const c_void, // rsi = pre_compute_insns
+    from_state_pc: u32,
+    from_state_instret: u64,
+) {
+    asm_run_internal(
+        vm_exec_state_ptr,
+        pre_compute_insns_ptr,
+        from_state_pc,
+        from_state_instret,
+    );
+}
+
+type F = BabyBear;
+type Ctx = ExecutionCtx;
+
+// at the end of the execution, you want to store the instret and pc from the x86 registers
+// to update the vm state's pc and instret
+#[no_mangle]
+pub extern "C" fn set_instret_and_pc(
+    vm_exec_state_ptr: *mut c_void,        // rdi = vm_exec_state
+    _pre_compute_insns_ptr: *const c_void, // rsi = pre_compute_insns
+    final_pc: u32,                         // rdx = final_pc
+    final_instret: u64,                    // rcx = final_instret
+) {
+    // reference to vm_exec_state
+    asm_bridge_utils::set_instret_and_pc_generic::<Ctx>(vm_exec_state_ptr, final_pc, final_instret);
+}
+
+/// # Safety
+/// - vm_exec_state_ptr must point to a valid VmExecState<F, GuestMemory, Ctx>.
+/// - pre_compute_insns_ptr must be a valid, contiguous array of PreComputeInstruction<'static, F,
+///   Ctx>.
+/// - cur_pc must be a valid PC for the current program.
+#[no_mangle]
+pub unsafe extern "C" fn extern_handler(
+    vm_exec_state_ptr: *mut c_void,
+    pre_compute_insns_ptr: *const c_void,
+    cur_pc: u32,
+    cur_instret: u64,
+) -> u32 {
+    unsafe {
+        let (vm_ptr, pre_ptr, ctx_ptr) = asm_bridge_utils::extern_prep_generic::<Ctx>(
+            vm_exec_state_ptr,
+            pre_compute_insns_ptr,
+            cur_pc,
+        );
+        // `arg` is a runtime constant that we want to keep in register
+        // - For pure execution it is `instret_end`
+        let arg = (*ctx_ptr).instret_end;
+        asm_bridge_utils::extern_finish_generic::<Ctx>(vm_ptr, pre_ptr, cur_pc, cur_instret, arg)
+    }
+}
+
+#[no_mangle]
+pub extern "C" fn should_suspend(instret: u64, _pc: u32, exec_state_ptr: *mut c_void) -> u32 {
+    // reference to vm_exec_state
+    let vm_exec_state_ref =
+        unsafe { &mut *(exec_state_ptr as *mut VmExecState<F, GuestMemory, Ctx>) };
+
+    let instret_end = vm_exec_state_ref.ctx.instret_end;
+
+    if instret >= instret_end {
+        1 // should suspend is `true`
+    } else {
+        0 // should suspend is `false`
+    }
+}

crates/vm/src/arch/asm_bridge_metered/Cargo.toml → crates/asm/asm_bridge_metered/Cargo.toml

@@ -1,19 +1,19 @@
 [package]
 name = "asm_bridge_metered"
-version.workspace = true 
+version.workspace = true
 edition.workspace = true
 authors.workspace = true
 homepage.workspace = true
 repository.workspace = true
 
-[package.metadata.cargo-shear]
-ignored = ["cc"]
-
 [lib]
 crate-type = ["cdylib"]
 
 [dependencies]
 cc = "1.0"
 openvm-circuit = { workspace = true, features = ["test-utils"] }
-openvm-stark-sdk.workspace = true
-openvm-instructions.workspace = true
+openvm-stark-sdk = { workspace = true }
+asm_bridge_utils = { path = "../asm_bridge_utils" }
+
+[package.metadata.cargo-shear]
+ignored = ["cc"]

crates/vm/src/arch/asm_bridge_metered/src/lib.rs → crates/asm/asm_bridge_metered/src/lib.rs

@@ -1,10 +1,9 @@
 use std::ffi::c_void;
 
 use openvm_circuit::{
-    arch::{execution_mode::MeteredCtx, interpreter::PreComputeInstruction, VmExecState},
+    arch::{execution_mode::MeteredCtx, VmExecState},
     system::memory::online::GuestMemory,
 };
-use openvm_instructions::program::DEFAULT_PC_STEP;
 use openvm_stark_sdk::p3_baby_bear::BabyBear;
 
 /*
@@ -27,7 +26,6 @@ extern "C" {
 ///
 /// # Safety
 ///
-///
 /// This function is unsafe because:
 /// - `vm_exec_state_ptr` must be valid
 /// - `pre_compute_insns` must point to valid pre-compute instructions
@@ -60,58 +58,31 @@ pub extern "C" fn metered_set_instret_and_pc(
     final_instret: u64,                    // rcx = final_instret
 ) {
     // reference to vm_exec_state
-    let vm_exec_state_ref =
-        unsafe { &mut *(vm_exec_state_ptr as *mut VmExecState<F, GuestMemory, Ctx>) };
-    vm_exec_state_ref
-        .vm_state
-        .set_instret_and_pc(final_instret, final_pc);
+    asm_bridge_utils::set_instret_and_pc_generic::<Ctx>(vm_exec_state_ptr, final_pc, final_instret);
 }
 
+/// # Safety
+/// - vm_exec_state_ptr must point to VmExecState<F, GuestMemory, MeteredCtx>.
+/// - pre_compute_insns_ptr must be a valid, contiguous array of PreComputeInstruction<'static, F,
+///   MeteredCtx>.
+/// - cur_pc must be a valid PC for the current program.
 #[no_mangle]
-pub extern "C" fn metered_extern_handler(
+pub unsafe extern "C" fn metered_extern_handler(
     vm_exec_state_ptr: *mut c_void,
     pre_compute_insns_ptr: *const c_void,
     cur_pc: u32,
     cur_instret: u64,
 ) -> u32 {
-    let mut instret: Box<u64> = Box::new(cur_instret); // placeholder to call the handler function
-    let mut pc: Box<u32> = Box::new(cur_pc);
-
-    let vm_exec_state_ref =
-        unsafe { &mut *(vm_exec_state_ptr as *mut VmExecState<F, GuestMemory, Ctx>) };
-
-    // pointer to the first element of `pre_compute_insns`
-    let pre_compute_insns_base_ptr =
-        pre_compute_insns_ptr as *const PreComputeInstruction<'static, F, Ctx>;
-    let pc_idx = (cur_pc / DEFAULT_PC_STEP) as usize;
-
-    let pre_compute_insns = unsafe { &*pre_compute_insns_base_ptr.add(pc_idx) };
-
-    let ctx = &vm_exec_state_ref.ctx;
-    // `arg` is a runtime constant that we want to keep in register
-    // - For metered execution it is `segment_check_insns`
-    let arg = ctx.segmentation_ctx.segment_check_insns;
-
     unsafe {
-        (pre_compute_insns.handler)(
-            pre_compute_insns.pre_compute,
-            &mut instret,
-            &mut pc,
-            arg,
-            vm_exec_state_ref,
+        let (vm_ptr, pre_ptr, ctx_ptr) = asm_bridge_utils::extern_prep_generic::<Ctx>(
+            vm_exec_state_ptr,
+            pre_compute_insns_ptr,
+            cur_pc,
         );
-    };
-
-    match vm_exec_state_ref.exit_code {
-        Ok(None) => {
-            // execution continues
-            *pc
-        }
-        _ => {
-            // special indicator that we must terminate
-            // this won't collide with actual pc value because pc values are always multiple of 4
-            1
-        }
+        // `arg` is a runtime constant that we want to keep in register
+        // - For metered execution it is `segment_check_insns`
+        let arg = (*ctx_ptr).segmentation_ctx.segment_check_insns;
+        asm_bridge_utils::extern_finish_generic::<Ctx>(vm_ptr, pre_ptr, cur_pc, cur_instret, arg)
     }
 }
 

crates/asm/asm_bridge_metered_cost/Cargo.toml

@@ -0,0 +1,19 @@
+[package]
+name = "asm_bridge_metered_cost"
+version.workspace = true
+edition.workspace = true
+authors.workspace = true
+homepage.workspace = true
+repository.workspace = true
+
+[lib]
+crate-type = ["cdylib"]
+
+[dependencies]
+cc = "1.0"
+openvm-circuit = { workspace = true, features = ["test-utils"] }
+openvm-stark-sdk = { workspace = true }
+asm_bridge_utils = { path = "../asm_bridge_utils" }
+
+[package.metadata.cargo-shear]
+ignored = ["cc"]