diff --git a/Cargo.lock b/Cargo.lock index f4c67bad58..0b9cb306e1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -841,17 +841,37 @@ dependencies = [ [[package]] name = "asm_bridge" -version = "1.4.1-rc.1" +version = "1.4.1-rc.2" dependencies = [ + "asm_bridge_utils", "cc", "openvm-circuit", - "openvm-instructions", "openvm-stark-sdk", ] [[package]] name = "asm_bridge_metered" -version = "1.4.1-rc.1" +version = "1.4.1-rc.2" +dependencies = [ + "asm_bridge_utils", + "cc", + "openvm-circuit", + "openvm-stark-sdk", +] + +[[package]] +name = "asm_bridge_metered_cost" +version = "1.4.1-rc.2" +dependencies = [ + "asm_bridge_utils", + "cc", + "openvm-circuit", + "openvm-stark-sdk", +] + +[[package]] +name = "asm_bridge_utils" +version = "1.4.1-rc.2" dependencies = [ "cc", "openvm-circuit", diff --git a/Cargo.toml b/Cargo.toml index 3d24663806..b55c16cc46 100644 --- a/Cargo.toml +++ b/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", diff --git a/crates/asm/asm_bridge/Cargo.toml b/crates/asm/asm_bridge/Cargo.toml new file mode 100644 index 0000000000..59706b7660 --- /dev/null +++ b/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"] diff --git a/crates/vm/src/arch/asm_bridge/src/asm_x86_run.s b/crates/asm/asm_bridge/src/asm_x86_run.s similarity index 100% rename from crates/vm/src/arch/asm_bridge/src/asm_x86_run.s rename to crates/asm/asm_bridge/src/asm_x86_run.s diff --git a/crates/asm/asm_bridge/src/lib.rs b/crates/asm/asm_bridge/src/lib.rs new file mode 100644 index 0000000000..284f71eece --- /dev/null +++ b/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::(vm_exec_state_ptr, final_pc, final_instret); +} + +/// # Safety +/// - vm_exec_state_ptr must point to a valid VmExecState. +/// - 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::( + 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::(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) }; + + 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` + } +} diff --git a/crates/vm/src/arch/asm_bridge_metered/Cargo.toml b/crates/asm/asm_bridge_metered/Cargo.toml similarity index 72% rename from crates/vm/src/arch/asm_bridge_metered/Cargo.toml rename to crates/asm/asm_bridge_metered/Cargo.toml index 7a634e19c6..1fd28f701f 100644 --- a/crates/vm/src/arch/asm_bridge_metered/Cargo.toml +++ b/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"] diff --git a/crates/vm/src/arch/asm_bridge_metered/src/asm_x86_run.s b/crates/asm/asm_bridge_metered/src/asm_x86_run.s similarity index 100% rename from crates/vm/src/arch/asm_bridge_metered/src/asm_x86_run.s rename to crates/asm/asm_bridge_metered/src/asm_x86_run.s diff --git a/crates/vm/src/arch/asm_bridge_metered/src/lib.rs b/crates/asm/asm_bridge_metered/src/lib.rs similarity index 58% rename from crates/vm/src/arch/asm_bridge_metered/src/lib.rs rename to crates/asm/asm_bridge_metered/src/lib.rs index 19b16d5f31..2e3086a5ee 100644 --- a/crates/vm/src/arch/asm_bridge_metered/src/lib.rs +++ b/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) }; - vm_exec_state_ref - .vm_state - .set_instret_and_pc(final_instret, final_pc); + asm_bridge_utils::set_instret_and_pc_generic::(vm_exec_state_ptr, final_pc, final_instret); } +/// # Safety +/// - vm_exec_state_ptr must point to VmExecState. +/// - 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 = Box::new(cur_instret); // placeholder to call the handler function - let mut pc: Box = Box::new(cur_pc); - - let vm_exec_state_ref = - unsafe { &mut *(vm_exec_state_ptr as *mut VmExecState) }; - - // 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::( + 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::(vm_ptr, pre_ptr, cur_pc, cur_instret, arg) } } diff --git a/crates/asm/asm_bridge_metered_cost/Cargo.toml b/crates/asm/asm_bridge_metered_cost/Cargo.toml new file mode 100644 index 0000000000..a6675cf9d7 --- /dev/null +++ b/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"] diff --git a/crates/asm/asm_bridge_metered_cost/src/asm_x86_run.s b/crates/asm/asm_bridge_metered_cost/src/asm_x86_run.s new file mode 100644 index 0000000000..c123b20b79 --- /dev/null +++ b/crates/asm/asm_bridge_metered_cost/src/asm_x86_run.s @@ -0,0 +1,100 @@ +.intel_syntax noprefix +.code64 +.section .text +.global asm_run_internal + +/* +rbx = vm_exec_state_ptr +rbp = pre_compute_insns_ptr +r13 = cur_pc +r14 = cur_instret +r12 is currently unused + +Assembly code explanation + +asm_run_internal: + push rbp ; push callee saved register + push rbx + push r12 + push r13 + push r14 ; since we push 5 times, this already ensures the stack is 16 bytes aligned + mov rbx, rdi ; rbx = rdi = vm_exec_state_ptr + mov rbp, rsi ; rbp = rsi = pre_compute_insns_ptr + mov r13, rdx ; r13 = rdx = from_state_pc + mov r14, rcx ; r14 = rcx = from_state_instret + +asm_execute: + mov rdi, r14 ; rdi = cur_instret + mov rsi, r13 ; rsi = cur_pc + mov rdx, rbx ; rdx = vm_exec_state_ptr + call should_suspend ; should_suspend(cur_instret, cur_pc, vm_exec_state_ptr) + cmp rax, 1 ; if return value of should_suspend is 1 + je asm_run_end ; jump to asm_run_end + mov rdi, rbx ; rdi = vm_exec_state_ptr + mov rsi, rbp ; rsi = pre_compute_insns_ptr + mov rdx, r13 ; rdx = cur_pc + mov rcx, r14 ; rcx = cur_instret + call metered_cost_extern_handler ; metered_cost_extern_handler(vm_exec_state_ptr, pre_compute_insns_ptr, cur_pc, cur_instret) + add r14, 1 ; cur_instret += 1 + cmp rax, 1 ; if return value of metered_cost_extern_handler is 1 + je asm_run_end ; jump to asm_run_end + mov r13, rax ; cur_pc = return value of metered_cost_extern_handler + jmp asm_execute ; jump to asm_execute + +asm_run_end: + mov rdi, rbx ; rdi = vm_exec_state_ptr + mov rsi, rbp ; rsi = pre_compute_insns_ptr + mov rdx, r13 ; rdx = cur_pc + mov rcx, r14 ; rcx = cur_instret + call metered_cost_set_instret_and_pc ; metered_cost_set_instret_and_pc(vm_exec_state_ptr, pre_compute_insns_ptr, cur_pc, cur_instret) + xor rax, rax ; set return value to 0 + pop r14 ; pop callee saved registers + pop r13 + pop r12 + pop rbx + pop rbp + ret +*/ + +asm_run_internal: + push rbp + push rbx + push r12 + push r13 + push r14 + mov rbx, rdi + mov rbp, rsi + mov r13, rdx + mov r14, rcx + +asm_execute: + mov rdi, r14 + mov rsi, r13 + mov rdx, rbx + call should_suspend + cmp rax, 1 + je asm_run_end + mov rdi, rbx + mov rsi, rbp + mov rdx, r13 + mov rcx, r14 + call metered_cost_extern_handler + add r14, 1 + cmp rax, 1 + je asm_run_end + mov r13, rax + jmp asm_execute + +asm_run_end: + mov rdi, rbx + mov rsi, rbp + mov rdx, r13 + mov rcx, r14 + call metered_cost_set_instret_and_pc + xor rax, rax + pop r14 + pop r13 + pop r12 + pop rbx + pop rbp + ret diff --git a/crates/asm/asm_bridge_metered_cost/src/lib.rs b/crates/asm/asm_bridge_metered_cost/src/lib.rs new file mode 100644 index 0000000000..094e366e65 --- /dev/null +++ b/crates/asm/asm_bridge_metered_cost/src/lib.rs @@ -0,0 +1,98 @@ +use std::ffi::c_void; + +use openvm_circuit::{ + arch::{execution_mode::MeteredCostCtx, VmExecState}, + system::memory::online::GuestMemory, +}; +use openvm_stark_sdk::p3_baby_bear::BabyBear; + +/* +rbx = vm_exec_state +rbp = pre_compute_insns +r13 = from_state_pc +r14 = from_state_instret +*/ + +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 = MeteredCostCtx; + +// 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 +// works for metered cost execution +#[no_mangle] +pub extern "C" fn metered_cost_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::(vm_exec_state_ptr, final_pc, final_instret); +} + +/// # Safety +/// - vm_exec_state_ptr must point to VmExecState. +/// - pre_compute_insns_ptr must be a valid, contiguous array of PreComputeInstruction<'static, F, +/// MeteredCostCtx>. +/// - cur_pc must be a valid PC for the current program. +#[no_mangle] +pub unsafe extern "C" fn metered_cost_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::( + vm_exec_state_ptr, + pre_compute_insns_ptr, + cur_pc, + ); + let arg = (*ctx_ptr).max_execution_cost; + asm_bridge_utils::extern_finish_generic::(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 { + let exec_state_ref = unsafe { &mut *(exec_state_ptr as *mut VmExecState) }; + + let max_execution_cost = exec_state_ref.ctx.max_execution_cost; + + if exec_state_ref.ctx.cost > max_execution_cost { + 1 + } else { + 0 + } +} diff --git a/crates/vm/src/arch/asm_bridge/Cargo.toml b/crates/asm/asm_bridge_utils/Cargo.toml similarity index 80% rename from crates/vm/src/arch/asm_bridge/Cargo.toml rename to crates/asm/asm_bridge_utils/Cargo.toml index ed90c88a39..80da382c36 100644 --- a/crates/vm/src/arch/asm_bridge/Cargo.toml +++ b/crates/asm/asm_bridge_utils/Cargo.toml @@ -1,5 +1,5 @@ [package] -name = "asm_bridge" +name = "asm_bridge_utils" version.workspace = true edition.workspace = true authors.workspace = true @@ -10,10 +10,9 @@ repository.workspace = true ignored = ["cc"] [lib] -crate-type = ["cdylib"] [dependencies] cc = "1.0" openvm-circuit = { workspace = true, features = ["test-utils"] } -openvm-stark-sdk.workspace = true +openvm-stark-sdk = { workspace = true } openvm-instructions.workspace = true \ No newline at end of file diff --git a/crates/asm/asm_bridge_utils/src/lib.rs b/crates/asm/asm_bridge_utils/src/lib.rs new file mode 100644 index 0000000000..8fea4ac4a7 --- /dev/null +++ b/crates/asm/asm_bridge_utils/src/lib.rs @@ -0,0 +1,65 @@ +use std::ffi::c_void; + +use openvm_circuit::{ + arch::{interpreter::PreComputeInstruction, ExecutionCtxTrait, VmExecState}, + system::memory::online::GuestMemory, +}; +use openvm_instructions::program::DEFAULT_PC_STEP; +use openvm_stark_sdk::p3_baby_bear::BabyBear as F; + +pub fn set_instret_and_pc_generic( + vm_exec_state_ptr: *mut c_void, + final_pc: u32, + final_instret: u64, +) { + let vm_exec_state_ref = + unsafe { &mut *(vm_exec_state_ptr as *mut VmExecState) }; + vm_exec_state_ref + .vm_state + .set_instret_and_pc(final_instret, final_pc); +} + +/// # Safety +/// - vm_exec_state_ptr must point to a valid VmExecState. +/// - pre_compute_insns_ptr must be a valid, contiguous array of PreComputeInstruction<'static, F, +/// Ctx>. +pub unsafe fn extern_prep_generic( + vm_exec_state_ptr: *mut c_void, + pre_compute_insns_ptr: *const c_void, + cur_pc: u32, +) -> ( + *mut VmExecState, + *const PreComputeInstruction<'static, F, Ctx>, + *const Ctx, +) { + let vm_ptr = vm_exec_state_ptr as *mut VmExecState; + let base = pre_compute_insns_ptr as *const PreComputeInstruction<'static, F, Ctx>; + let pc_idx = (cur_pc / DEFAULT_PC_STEP) as usize; + let pre_ptr = base.add(pc_idx); + let ctx_ptr = &(*vm_ptr).ctx as *const Ctx; + (vm_ptr, pre_ptr, ctx_ptr) +} + +/// # Safety +/// - vm_exec_state_ptr and pre_compute_insn_ptr must come from extern_prep_generic for the same Ctx +/// and program. +pub unsafe fn extern_finish_generic( + vm_exec_state_ptr: *mut VmExecState, + pre_compute_insn_ptr: *const PreComputeInstruction<'static, F, Ctx>, + cur_pc: u32, + cur_instret: u64, + arg: u64, +) -> u32 { + let vm = &mut *vm_exec_state_ptr; + let pre = &*pre_compute_insn_ptr; + + let mut instret: Box = Box::new(cur_instret); + let mut pc: Box = Box::new(cur_pc); + + (pre.handler)(pre.pre_compute, &mut instret, &mut pc, arg, vm); + + match vm.exit_code { + Ok(None) => *pc, + _ => 1, + } +} diff --git a/crates/toolchain/tests/tests/transpiler_tests.rs b/crates/toolchain/tests/tests/transpiler_tests.rs index 45f7751e87..9ca219e592 100644 --- a/crates/toolchain/tests/tests/transpiler_tests.rs +++ b/crates/toolchain/tests/tests/transpiler_tests.rs @@ -8,6 +8,8 @@ use num_bigint::BigUint; use openvm_algebra_circuit::*; use openvm_algebra_transpiler::{Fp2TranspilerExtension, ModularTranspilerExtension}; use openvm_bigint_circuit::*; +#[cfg(feature = "aot")] +use openvm_circuit::arch::{aot::AotInstance, execution_mode::ExecutionCtx}; use openvm_circuit::{ arch::{InitFileGenerator, SystemConfig, VmExecutor}, derive::VmConfig, @@ -20,6 +22,8 @@ use openvm_platform::memory::MEM_SIZE; use openvm_rv32im_circuit::{ Rv32I, Rv32IExecutor, Rv32ImBuilder, Rv32ImConfig, Rv32Io, Rv32IoExecutor, Rv32M, Rv32MExecutor, }; +#[cfg(feature = "aot")] +use openvm_rv32im_transpiler::Rv32JalrOpcode; use openvm_rv32im_transpiler::{ Rv32ITranspilerExtension, Rv32IoTranspilerExtension, Rv32MTranspilerExtension, }; diff --git a/crates/vm/src/arch/aot.rs b/crates/vm/src/arch/aot.rs index 5431d11b93..98bc0db487 100644 --- a/crates/vm/src/arch/aot.rs +++ b/crates/vm/src/arch/aot.rs @@ -7,7 +7,7 @@ use openvm_stark_backend::p3_field::PrimeField32; use crate::{ arch::{ - execution_mode::{ExecutionCtx, MeteredCtx, Segment}, + execution_mode::{ExecutionCtx, MeteredCostCtx, MeteredCtx, Segment}, interpreter::{ alloc_pre_compute_buf, get_metered_pre_compute_instructions, get_metered_pre_compute_max_size, get_pre_compute_instructions, @@ -82,11 +82,11 @@ where .parent() .unwrap(); - let src_asm_bridge_dir = std::path::Path::new(manifest_dir).join("src/arch/asm_bridge"); + let src_asm_bridge_dir = std::path::Path::new(manifest_dir).join("../asm/asm_bridge"); let src_asm_bridge_dir_str = src_asm_bridge_dir.to_str().unwrap(); // ar rcs libasm_runtime.a asm_run.o - // cargo rustc -- -L /home/ubuntu/openvm/crates/vm/src/arch/asm_bridge -l static=asm_runtime + // cargo rustc -- -L /home/ubuntu/openvm/crates/vm/../asm/asm_bridge -l static=asm_runtime // run the below command from the `src_asm_bridge_dir` directory // as src/asm_run.s -o asm_run.o @@ -181,7 +181,6 @@ where lib, }) } - pub fn create_initial_vm_state(&self, inputs: impl Into>) -> VmState { VmState::initial( &self.system_config, @@ -293,15 +292,47 @@ where E: MeteredExecutor, { let default_name = String::from("asm_x86_run"); - Self::new_metered_with_asm_name(inventory, exe, executor_idx_to_air_idx, &default_name) + let src_asm_bridge_dir_str = String::from("../asm/asm_bridge_metered"); + let asm_so_str = String::from("libasm_bridge_metered.so"); + Self::new_metered_generic_with_asm_name( + inventory, + exe, + executor_idx_to_air_idx, + &default_name, + &src_asm_bridge_dir_str, + &asm_so_str, + ) + } + + pub fn new_metered_cost( + inventory: &'a ExecutorInventory, + exe: &VmExe, + executor_idx_to_air_idx: &[usize], + ) -> Result + where + E: MeteredExecutor, + { + let default_name = String::from("asm_x86_run"); + let src_asm_bridge_dir_str = String::from("../asm/asm_bridge_metered_cost"); + let asm_so_str = String::from("libasm_bridge_metered_cost.so"); + Self::new_metered_generic_with_asm_name( + inventory, + exe, + executor_idx_to_air_idx, + &default_name, + &src_asm_bridge_dir_str, + &asm_so_str, + ) } - /// Creates a new interpreter instance for metered execution. - pub fn new_metered_with_asm_name( + /// Creates a new interpreter instance for metered cost execution. + pub fn new_metered_generic_with_asm_name( inventory: &'a ExecutorInventory, exe: &VmExe, executor_idx_to_air_idx: &[usize], asm_name: &String, + bridge_str: &String, + asm_so_str: &String, ) -> Result where E: MeteredExecutor, @@ -316,12 +347,11 @@ where .parent() .unwrap(); - let src_asm_bridge_dir = - std::path::Path::new(manifest_dir).join("src/arch/asm_bridge_metered"); + let src_asm_bridge_dir = std::path::Path::new(manifest_dir).join(bridge_str); let src_asm_bridge_dir_str = src_asm_bridge_dir.to_str().unwrap(); // ar rcs libasm_runtime.a asm_run.o - // cargo rustc -- -L /home/ubuntu/openvm/crates/vm/src/arch/asm_bridge -l static=asm_runtime + // cargo rustc -- -L /home/ubuntu/openvm/crates/vm/ -l static=asm_runtime // run the below command from the `src_asm_bridge_dir` directory // as src/asm_run.s -o asm_run.o @@ -386,7 +416,7 @@ where .join("target") .join(asm_name) .join("release") - .join("libasm_bridge_metered.so"); + .join(asm_so_str); let lib = unsafe { Library::new(&lib_path).expect("Failed to load library") }; let program = &exe.program; @@ -421,7 +451,7 @@ impl AotInstance<'_, F, MeteredCtx> where F: PrimeField32, { - /// Metered exeecution for the given `inputs`. Execution begins from the initial + /// Metered execution for the given `inputs`. Execution begins from the initial /// state specified by the `VmExe`. This function executes the program until termination. /// /// Returns the segmentation boundary data and the final VM state when execution stops. @@ -483,3 +513,67 @@ where // TODO: implement execute_metered_until_suspend for AOT if needed } + +impl AotInstance<'_, F, MeteredCostCtx> +where + F: PrimeField32, +{ + /// Metered cost execution for the given `inputs`. Execution begins from the initial + /// state specified by the `VmExe`. This function executes the program until termination. + /// + /// Returns the cost and the final VM state when execution stops. + /// + /// Assumes the program doesn't jump to out of bounds pc + pub fn execute_metered_cost( + &mut self, + inputs: impl Into>, + ctx: MeteredCostCtx, + ) -> Result<(u64, VmState), ExecutionError> { + let vm_state = self.create_initial_vm_state(inputs); + self.execute_metered_cost_from_state(vm_state, ctx) + } + + /// Metered cost execution for the given `VmState`. This function executes the program until + /// termination + /// + /// Returns the cost and the final VM state when execution stops. + /// + /// Assume program doesn't jump to out of bounds pc + pub fn execute_metered_cost_from_state( + &self, + from_state: VmState, + ctx: MeteredCostCtx, + ) -> Result<(u64, VmState), ExecutionError> { + let from_state_instret = from_state.instret(); + let from_state_pc = from_state.pc(); + + let mut vm_exec_state: Box> = + Box::new(VmExecState::new(from_state, ctx)); + + unsafe { + let asm_run: libloading::Symbol = self + .lib + .get(b"asm_run") + .expect("Failed to get asm_run symbol"); + + let vm_exec_state_ptr = + &mut *vm_exec_state as *mut VmExecState; + let pre_compute_insns_ptr = self.pre_compute_insns_box.as_ptr(); + + asm_run( + vm_exec_state_ptr as *mut c_void, + pre_compute_insns_ptr as *const c_void, + from_state_pc, + from_state_instret, + ); + } + + // handle execution error + match vm_exec_state.exit_code { + Ok(_) => Ok((vm_exec_state.ctx.cost, vm_exec_state.vm_state)), + Err(e) => Err(e), + } + } + + // TODO: implement execute_metered_cost_until_suspend for AOT if needed +} diff --git a/crates/vm/src/arch/asm_bridge/src/asm_test_name.s b/crates/vm/src/arch/asm_bridge/src/asm_test_name.s deleted file mode 100644 index 410e40157f..0000000000 --- a/crates/vm/src/arch/asm_bridge/src/asm_test_name.s +++ /dev/null @@ -1,98 +0,0 @@ -.intel_syntax noprefix -.code64 -.section .text -.global asm_run_internal - -/* -rbx = vm_exec_state_ptr -rbp = pre_compute_insns_ptr -r13 = cur_pc -r14 = cur_instret -r12 is currently unused - -asm_run_internal: - push rbp ; push callee saved registers - push rbx - push r12 - push r13 - push r14 ; since we push 5 times, this already ensures the stack is 16 bytes aligned - mov rbx, rdi ; rbx = rdi = vm_exec_state_ptr - mov rbp, rsi ; rbp = rsi = pre_compute_insns_ptr - mov r13, rdx ; r13 = rdx = from_state_pc - mov r14, rcx ; r14 = rcx = from_state_instret - -asm_execute: - mov rdi, r14 ; rdi = cur_instret - mov rsi, r13 ; rsi = cur_pc - mov rdx, rbx ; rdx = vm_exec_state_ptr - call should_suspend ; should_suspend(cur_instret, cur_pc, vm_exec_state_ptr) - cmp rax, 1 ; if return value of should_suspend is 1 - je asm_run_end ; jump to asm_run_end - mov rdi, rbx ; rdi = vm_exec_state_ptr - mov rsi, rbp ; rsi = pre_compute_insns_ptr - mov rdx, r13 ; rdx = cur_pc - mov rcx, r14 ; rcx = cur_instret - call extern_handler ; extern_handler(vm_exec_state_ptr, pre_compute_insns_ptr, cur_pc, cur_instret) - add r14, 1 ; cur_instret += 1 - cmp rax, 1 ; if return value of extern_handler is 1 - je asm_run_end ; jump to asm_run_end - mov r13, rax ; cur_pc = return value of extern_handler - jmp asm_execute ; jump to asm_execute - -asm_run_end: - mov rdi, rbx ; rdi = vm_exec_state_ptr - mov rsi, rbp ; rsi = pre_compute_insns_ptr - mov rdx, r13 ; rdx = cur_pc - mov rcx, r14 ; rcx = cur_instret - call set_instret_and_pc ; set_instret_and_pc(vm_exec_state_ptr, pre_compute_insns_ptr, cur_pc, cur_instret) - xor rax, rax ; set return value to 0 - pop r14 ; pop callee saved registers - pop r13 - pop r12 - pop rbx - pop rbp - ret -*/ - -asm_run_internal: - push rbp - push rbx - push r12 - push r13 - push r14 - mov rbx, rdi - mov rbp, rsi - mov r13, rdx - mov r14, rcx - -asm_execute: - mov rdi, r14 - mov rsi, r13 - mov rdx, rbx - call should_suspend - cmp rax, 1 - je asm_run_end - mov rdi, rbx - mov rsi, rbp - mov rdx, r13 - mov rcx, r14 - call extern_handler - add r14, 1 - cmp rax, 1 - je asm_run_end - mov r13, rax - jmp asm_execute - -asm_run_end: - mov rdi, rbx - mov rsi, rbp - mov rdx, r13 - mov rcx, r14 - call set_instret_and_pc - xor rax, rax - pop r14 - pop r13 - pop r12 - pop rbx - pop rbp - ret diff --git a/crates/vm/src/arch/asm_bridge/src/lib.rs b/crates/vm/src/arch/asm_bridge/src/lib.rs deleted file mode 100644 index 0cd7f723bc..0000000000 --- a/crates/vm/src/arch/asm_bridge/src/lib.rs +++ /dev/null @@ -1,131 +0,0 @@ -use std::ffi::c_void; - -use openvm_circuit::{ - arch::{execution_mode::ExecutionCtx, interpreter::PreComputeInstruction, VmExecState}, - system::memory::online::GuestMemory, -}; -use openvm_instructions::program::DEFAULT_PC_STEP; -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 assembly execution, store the instret and pc from the x86 registers -/// to the vm state's pc and instret for the pure execution mode -#[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 - let vm_exec_state_ref = - unsafe { &mut *(vm_exec_state_ptr as *mut VmExecState) }; - vm_exec_state_ref - .vm_state - .set_instret_and_pc(final_instret, final_pc); -} - -/// extern handler for the pure execution mode -/// calls the correct function handler based on `cur_pc` -/// -/// returns 1 if we should terminate and *pc otherwise -/// this is safe because *pc is always a multiple of 4 -#[no_mangle] -pub 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 { - // this is boxed for safety so that when `execute_e12_impl` runs when called by the handler - // it would be able to dereference instret and pc correctly - let mut instret: Box = Box::new(cur_instret); - let mut pc: Box = Box::new(cur_pc); - - // reference to vm_exec_state - let vm_exec_state_ref = - unsafe { &mut *(vm_exec_state_ptr as *mut VmExecState) }; - - // 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 pure execution it is `instret_end` - let arg = ctx.instret_end; - - unsafe { - (pre_compute_insns.handler)( - pre_compute_insns.pre_compute, - &mut instret, - &mut pc, - arg, - vm_exec_state_ref, - ); - }; - - 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 - } - } -} - -#[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) }; - - 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` - } -} diff --git a/crates/vm/src/arch/vm.rs b/crates/vm/src/arch/vm.rs index 77b5156d19..66b11c49c6 100644 --- a/crates/vm/src/arch/vm.rs +++ b/crates/vm/src/arch/vm.rs @@ -308,6 +308,16 @@ where ) -> Result, StaticProgramError> { InterpretedInstance::new_metered(&self.inventory, exe, executor_idx_to_air_idx) } + + // Creates an AOT instance for metered execution of the given `exe`. + #[cfg(feature = "aot")] + pub fn metered_cost_aot_instance( + &self, + exe: &VmExe, + executor_idx_to_air_idx: &[usize], + ) -> Result, StaticProgramError> { + AotInstance::new_metered_cost(&self.inventory, exe, executor_idx_to_air_idx) + } } #[derive(Error, Debug)] @@ -532,6 +542,21 @@ where .metered_cost_instance(exe, &executor_idx_to_air_idx) } + // Metered Cost AOT execution + #[cfg(feature = "aot")] + pub fn get_metered_cost_aot_instance( + &self, + exe: &VmExe>, + ) -> Result, MeteredCostCtx>, StaticProgramError> + where + Val: PrimeField32, + >>::Executor: MeteredExecutor>, + { + let executor_idx_to_air_idx = self.executor_idx_to_air_idx(); + self.executor() + .metered_cost_aot_instance(exe, &executor_idx_to_air_idx) + } + pub fn preflight_interpreter( &self, exe: &VmExe>, @@ -1517,4 +1542,4 @@ mod vm_metrics { } } } -} \ No newline at end of file +} diff --git a/crates/vm/src/utils/stark_utils.rs b/crates/vm/src/utils/stark_utils.rs index 09459b59af..b3112103bb 100644 --- a/crates/vm/src/utils/stark_utils.rs +++ b/crates/vm/src/utils/stark_utils.rs @@ -104,6 +104,7 @@ where Com: AsRef<[Val; CHUNK]> + From<[Val; CHUNK]>, { let metered_ctx = vm.build_metered_ctx(&exe); + let metered_cost_ctx = vm.build_metered_cost_ctx(); /* Assertions for Pure Execution AOT */ @@ -150,6 +151,21 @@ where assert_eq!(segments[i].num_insns, aot_segments[i].num_insns); assert_eq!(segments[i].trace_heights, aot_segments[i].trace_heights); } + + /* + Assertions for Metered Cost AOT + */ + let (aot_cost, aot_state_metered_cost) = vm + .get_metered_cost_aot_instance(&exe)? + .execute_metered_cost(input.clone(), metered_cost_ctx.clone())?; + + let (cost, interp_state_metered_cost) = vm + .metered_cost_interpreter(&exe)? + .execute_metered_cost(input.clone(), metered_cost_ctx.clone())?; + + assert_vm_state_eq(&interp_state_metered_cost, &aot_state_metered_cost); + + assert_eq!(cost, aot_cost); Ok(()) } diff --git a/extensions/native/circuit/tests/integration_test.rs b/extensions/native/circuit/tests/integration_test.rs index d11deefbc1..00c04af90c 100644 --- a/extensions/native/circuit/tests/integration_test.rs +++ b/extensions/native/circuit/tests/integration_test.rs @@ -5,6 +5,10 @@ use std::{ }; use itertools::Itertools; +#[cfg(feature = "aot")] +use openvm_circuit::arch::execution_mode::metered_cost::MeteredCostCtx; +#[cfg(feature = "aot")] +use openvm_circuit::arch::VmState; #[cfg(feature = "cuda")] use openvm_circuit::system::cuda::extensions::SystemGpuBuilder as SystemBuilder; #[cfg(not(feature = "cuda"))] @@ -984,6 +988,61 @@ fn test_single_segment_executor_no_segmentation() { .unwrap(); } +#[cfg(feature = "aot")] +fn compare_vm_states( + vm_state1: &VmState, + vm_state2: &VmState, + memory_dimensions: openvm_circuit::system::memory::dimensions::MemoryDimensions, +) { + assert_eq!(vm_state1.instret(), vm_state2.instret()); + assert_eq!(vm_state1.pc(), vm_state2.pc()); + use openvm_circuit::{ + arch::hasher::poseidon2::vm_poseidon2_hasher, system::memory::merkle::MerkleTree, + }; + + let hasher = vm_poseidon2_hasher::(); + + let tree1 = MerkleTree::from_memory(&vm_state1.memory.memory, &memory_dimensions, &hasher); + let tree2 = MerkleTree::from_memory(&vm_state2.memory.memory, &memory_dimensions, &hasher); + + assert_eq!(tree1.root(), tree2.root(), "Memory states differ"); +} + +// Helper to run AOT metered-cost and compare against interpreter baseline. +#[cfg(feature = "aot")] +fn run_aot_metered_cost_and_compare( + vm: &VirtualMachine, + exe: VmExe, + executor_idx_to_air_idx: Vec, + ctx: MeteredCostCtx, + instructions_len: usize, + baseline_cost: u64, + baseline_state: VmState, + config: NativeConfig, + resume_state: Option>, +) -> (u64, VmState) { + let mut aot_instance = vm + .executor() + .metered_cost_aot_instance(&exe, &executor_idx_to_air_idx) + .unwrap(); + let (mut aot_cost, mut aot_vm_state) = aot_instance + .execute_metered_cost(vec![], ctx.clone()) + .expect("Failed to execute"); + if resume_state.is_some() { + (aot_cost, aot_vm_state) = aot_instance + .execute_metered_cost_from_state(resume_state.clone().unwrap(), ctx.clone()) + .expect("Failed to execute"); + } + assert_eq!(aot_vm_state.instret(), instructions_len as u64); + assert_eq!(baseline_cost, aot_cost); + compare_vm_states( + &aot_vm_state, + &baseline_state, + config.clone().system.memory_config.memory_dimensions(), + ); + (aot_cost, aot_vm_state) +} + #[test] fn test_vm_execute_metered_cost_native_chips() { type F = BabyBear; @@ -993,7 +1052,7 @@ fn test_vm_execute_metered_cost_native_chips() { let engine = TestEngine::new(FriParameters::new_for_testing(3)); let (vm, _) = - VirtualMachine::new_with_keygen(engine, NativeBuilder::default(), config).unwrap(); + VirtualMachine::new_with_keygen(engine, NativeBuilder::default(), config.clone()).unwrap(); let instructions = vec![ // Field Arithmetic operations (FieldArithmeticChip) @@ -1014,11 +1073,26 @@ fn test_vm_execute_metered_cost_native_chips() { .unwrap(); let ctx = vm.build_metered_cost_ctx(); let (cost, vm_state) = instance - .execute_metered_cost(vec![], ctx) + .execute_metered_cost(vec![], ctx.clone()) .expect("Failed to execute"); assert_eq!(vm_state.instret(), instructions.len() as u64); assert!(cost > 0); + + #[cfg(feature = "aot")] + { + let _ = run_aot_metered_cost_and_compare( + &vm, + exe.clone(), + executor_idx_to_air_idx.clone(), + ctx, + instructions.len(), + cost, + vm_state, + config.clone(), + None::>, + ); + } } #[test] @@ -1051,7 +1125,7 @@ fn test_vm_execute_metered_cost_halt() { .unwrap(); let ctx = vm.build_metered_cost_ctx(); let (cost1, vm_state1) = instance1 - .execute_metered_cost(vec![], ctx) + .execute_metered_cost(vec![], ctx.clone()) .expect("Failed to execute"); assert_eq!(vm_state1.instret(), instructions.len() as u64); @@ -1063,9 +1137,137 @@ fn test_vm_execute_metered_cost_halt() { .unwrap(); let ctx2 = vm.build_metered_cost_ctx().with_max_execution_cost(0); let (cost2, vm_state2) = instance2 - .execute_metered_cost(vec![], ctx2) + .execute_metered_cost(vec![], ctx2.clone()) .expect("Failed to execute"); assert_eq!(vm_state2.instret(), 1); assert!(cost2 < cost1); + + let executor_idx_to_air_idx3 = vm.executor_idx_to_air_idx(); + let instance3 = vm + .executor() + .metered_cost_instance(&exe, &executor_idx_to_air_idx3) + .unwrap(); + let ctx3 = vm.build_metered_cost_ctx().with_max_execution_cost(100); + let (cost3, vm_state3) = instance3 + .execute_metered_cost(vec![], ctx3.clone()) + .expect("Failed to execute"); + + assert_eq!(vm_state3.instret(), 3); + assert!(cost2 < cost3); + assert!(cost3 < cost1); + + #[cfg(feature = "aot")] + { + let (aot_cost1, _aot_vm_state1) = run_aot_metered_cost_and_compare( + &vm, + exe.clone(), + executor_idx_to_air_idx.clone(), + ctx, + instructions.len(), + cost1, + vm_state1, + config.clone(), + None::>, + ); + let (aot_cost2, _aot_vm_state2) = run_aot_metered_cost_and_compare( + &vm, + exe.clone(), + executor_idx_to_air_idx2.clone(), + ctx2, + 1usize, + cost2, + vm_state2, + config.clone(), + None::>, + ); + let (aot_cost3, _aot_vm_state3) = run_aot_metered_cost_and_compare( + &vm, + exe.clone(), + executor_idx_to_air_idx3.clone(), + ctx3, + 3usize, + cost3, + vm_state3, + config.clone(), + None::>, + ); + + assert!(aot_cost2 < aot_cost1); + assert!(aot_cost2 < aot_cost3); + assert!(aot_cost3 < aot_cost1); + } +} + +#[test] +fn test_vm_execute_metered_cost_resume_parity() { + type F = BabyBear; + + setup_tracing(); + let config = test_native_config(); + + let engine = TestEngine::new(FriParameters::new_for_testing(3)); + let (vm, _) = + VirtualMachine::new_with_keygen(engine, NativeBuilder::default(), config.clone()).unwrap(); + + // Simple multi-instruction program to ensure we can suspend and then resume to completion. + let instructions = vec![ + Instruction::large_from_isize(ADD.global_opcode(), 0, 0, 1, 4, 0, 0, 0), + Instruction::large_from_isize(SUB.global_opcode(), 1, 10, 2, 4, 0, 0, 0), + Instruction::large_from_isize(MUL.global_opcode(), 2, 3, 4, 4, 0, 0, 0), + Instruction::from_isize(TERMINATE.global_opcode(), 0, 0, 0, 0, 0), + ]; + + let exe = VmExe::new(Program::::from_instructions(&instructions)); + + // Create interpreter instance and suspend after very small budget + let executor_idx_to_air_idx = vm.executor_idx_to_air_idx(); + let interp_instance = vm + .executor() + .metered_cost_instance(&exe, &executor_idx_to_air_idx) + .unwrap(); + let ctx_suspend = vm.build_metered_cost_ctx().with_max_execution_cost(0); + let (cost_suspend_interp, vm_state_suspend_interp) = interp_instance + .execute_metered_cost(vec![], ctx_suspend.clone()) + .expect("Failed to execute"); + assert_eq!(vm_state_suspend_interp.instret(), 1); + + // Resume with unlimited budget + let ctx_unlimited = vm.build_metered_cost_ctx(); + let (cost_interp_resume, vm_state_final_interp) = interp_instance + .execute_metered_cost_from_state(vm_state_suspend_interp.clone(), ctx_unlimited.clone()) + .expect("Failed to resume interp"); + + assert!(cost_suspend_interp < cost_interp_resume); + assert_eq!(vm_state_final_interp.instret(), instructions.len() as u64); + + // Do the same with AOT, compare suspended states, then resume both to completion + #[cfg(feature = "aot")] + { + let (cost_suspend_aot, vm_state_suspend_aot) = run_aot_metered_cost_and_compare( + &vm, + exe.clone(), + executor_idx_to_air_idx.clone(), + ctx_suspend, + 1usize, + cost_suspend_interp, + vm_state_suspend_interp, + config.clone(), + None::>, + ); + + let (cost_resume_aot, vm_state_final_aot) = run_aot_metered_cost_and_compare( + &vm, + exe.clone(), + executor_idx_to_air_idx.clone(), + ctx_unlimited, + instructions.len(), + cost_interp_resume, + vm_state_final_interp, + config.clone(), + Some(vm_state_suspend_aot.clone()), + ); + assert_eq!(cost_interp_resume, cost_resume_aot); + assert_eq!(cost_suspend_aot, cost_suspend_interp); + } }