perf: pass state variables by value in tco exec handlers (#2121)

- pass `pc`, `instret` and
`instret_end`/`max_execution_cost`/`segment_check_insns` by value in
execution handlers to get them to be passed in registers
- add `likely`, `unlikely` hints for suspension/termination in `tco`

[benchmark
comparison](https://github.com/axiom-crypto/openvm-reth-benchmark/actions/runs/17513217695#summary-49747838308)

Towards INT-4921

Author: shuklaayush

crates/sdk/src/lib.rs

@@ -411,7 +411,7 @@ where
         let (cost, final_state) = interpreter
             .execute_metered_cost(inputs, ctx)
             .map_err(VirtualMachineError::from)?;
-        let instret = final_state.instret;
+        let instret = final_state.instret();
 
         let public_values = extract_public_values(
             self.executor.config.as_ref().num_public_values,

crates/vm/derive/src/tco.rs

@@ -30,16 +30,19 @@ pub fn tco_impl(item: TokenStream) -> TokenStream {
     // Build the function call with all the generics
     let generic_args = build_generic_args(generics);
     let execute_call = if generic_args.is_empty() {
-        quote! { #fn_name(pre_compute, exec_state) }
+        quote! { #fn_name(pre_compute, &mut instret, &mut pc, arg, exec_state) }
     } else {
-        quote! { #fn_name::<#(#generic_args),*>(pre_compute, exec_state) }
+        quote! { #fn_name::<#(#generic_args),*>(pre_compute, &mut instret, &mut pc, arg, exec_state) }
     };
 
     // Generate the TCO handler function
     let handler_fn = quote! {
         #[inline(never)]
         unsafe fn #handler_name #handler_generics (
             interpreter: &::openvm_circuit::arch::interpreter::InterpretedInstance<#f_type, #ctx_type>,
+            mut instret: u64,
+            mut pc: u32,
+            arg: u64,
             exec_state: &mut ::openvm_circuit::arch::VmExecState<
                 #f_type,
                 ::openvm_circuit::system::memory::online::GuestMemory,
@@ -50,28 +53,30 @@ pub fn tco_impl(item: TokenStream) -> TokenStream {
         {
             use ::openvm_circuit::arch::ExecutionError;
 
-            let pre_compute = interpreter.get_pre_compute(exec_state.vm_state.pc);
+            let pre_compute = interpreter.get_pre_compute(pc);
             #execute_call;
 
-            if exec_state.exit_code.is_err() {
-                // stop execution
+            if ::core::intrinsics::unlikely(exec_state.exit_code.is_err()) {
+                exec_state.set_instret_and_pc(instret, pc);
                 return;
             }
-            if #ctx_type::should_suspend(exec_state) {
+            if ::core::intrinsics::unlikely(#ctx_type::should_suspend(instret, pc, arg, exec_state)) {
+                exec_state.set_instret_and_pc(instret, pc);
                 return;
             }
-            // exec_state.pc should have been updated by execute_impl at this point
-            let next_handler = interpreter.get_handler(exec_state.vm_state.pc);
-            if next_handler.is_none() {
-                exec_state.exit_code = Err(ExecutionError::PcOutOfBounds (exec_state.vm_state.pc));
+
+            let next_handler = interpreter.get_handler(pc);
+            if ::core::intrinsics::unlikely(next_handler.is_none()) {
+                exec_state.set_instret_and_pc(instret, pc);
+                exec_state.exit_code = Err(ExecutionError::PcOutOfBounds(pc));
                 return;
             }
             let next_handler = next_handler.unwrap_unchecked();
 
             // NOTE: `become` is a keyword that requires Rust Nightly.
             // It is part of the explicit tail calls RFC: <https://github.com/rust-lang/rust/issues/112788>
             // which is still incomplete.
-            become next_handler(interpreter, exec_state)
+            become next_handler(interpreter, instret, pc, arg, exec_state)
         }
     };
 

crates/vm/src/arch/execution.rs

@@ -85,22 +85,40 @@ pub enum StaticProgramError {
     ExecutorNotFound { opcode: VmOpcode },
 }
 
-/// Function pointer for interpreter execution with function signature `(pre_compute, exec_state)`.
-/// The `pre_compute: &[u8]` is a pre-computed buffer of data corresponding to a single instruction.
-/// The contents of `pre_compute` are determined from the program code as specified by the
-/// [Executor] and [MeteredExecutor] traits.
-pub type ExecuteFunc<F, CTX> =
-    unsafe fn(pre_compute: &[u8], exec_state: &mut VmExecState<F, GuestMemory, CTX>);
+/// Function pointer for interpreter execution with function signature `(pre_compute, instret, pc,
+/// arg, exec_state)`. The `pre_compute: &[u8]` is a pre-computed buffer of data
+/// corresponding to a single instruction. The contents of `pre_compute` are determined from the
+/// program code as specified by the [Executor] and [MeteredExecutor] traits.
+/// `arg` is a runtime constant that we want to keep in register:
+/// - For pure execution it is `instret_end`
+/// - For metered cost execution it is the `max_execution_cost`
+/// - For metered execution it is `segment_check_insns`
+pub type ExecuteFunc<F, CTX> = unsafe fn(
+    pre_compute: &[u8],
+    instret: &mut u64,
+    pc: &mut u32,
+    arg: u64,
+    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
+);
 
 /// Handler for tail call elimination. The `CTX` is assumed to contain pointers to the pre-computed
 /// buffer and the function handler table.
 ///
 /// - `pre_compute_buf` is the starting pointer of the pre-computed buffer.
 /// - `handlers` is the starting pointer of the table of function pointers of `Handler` type. The
 ///   pointer is typeless to avoid self-referential types.
+/// - `pc`, `instret`, `instret_end` are passed as separate arguments for efficiency
+///
+/// `arg` is a runtime constant that we want to keep in register:
+/// - For pure execution it is `instret_end`
+/// - For metered cost execution it is the `max_execution_cost`
+/// - For metered execution it is `segment_check_insns`
 #[cfg(feature = "tco")]
 pub type Handler<F, CTX> = unsafe fn(
     interpreter: &InterpretedInstance<F, CTX>,
+    instret: u64,
+    pc: u32,
+    arg: u64,
     exec_state: &mut VmExecState<F, GuestMemory, CTX>,
 );
 

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

@@ -125,11 +125,11 @@ impl<const PAGE_BITS: usize> MeteredCtx<PAGE_BITS> {
     }
 
     #[inline(always)]
-    pub fn check_and_segment(&mut self, instret: u64) {
+    pub fn check_and_segment(&mut self, instret: u64, segment_check_insns: u64) {
         let threshold = self
             .segmentation_ctx
             .instret_last_segment_check
-            .wrapping_add(self.segmentation_ctx.segment_check_insns);
+            .wrapping_add(segment_check_insns);
         debug_assert!(
             threshold >= self.segmentation_ctx.instret_last_segment_check,
             "overflow in segment check threshold calculation"
@@ -195,23 +195,30 @@ impl<const PAGE_BITS: usize> ExecutionCtxTrait for MeteredCtx<PAGE_BITS> {
     }
 
     #[inline(always)]
-    fn should_suspend<F>(vm_state: &mut VmExecState<F, GuestMemory, Self>) -> bool {
+    fn should_suspend<F>(
+        instret: u64,
+        _pc: u32,
+        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.
-        vm_state.ctx.check_and_segment(vm_state.instret);
+        exec_state
+            .ctx
+            .check_and_segment(instret, segment_check_insns);
         false
     }
 
     #[inline(always)]
-    fn on_terminate<F>(vm_state: &mut VmExecState<F, GuestMemory, Self>) {
-        vm_state
+    fn on_terminate<F>(instret: u64, _pc: u32, exec_state: &mut VmExecState<F, GuestMemory, Self>) {
+        exec_state
             .ctx
             .memory_ctx
-            .lazy_update_boundary_heights(&mut vm_state.ctx.trace_heights);
-        vm_state
+            .lazy_update_boundary_heights(&mut exec_state.ctx.trace_heights);
+        exec_state
             .ctx
             .segmentation_ctx
-            .segment(vm_state.instret, &vm_state.ctx.trace_heights);
+            .segment(instret, &exec_state.ctx.trace_heights);
     }
 }
 

crates/vm/src/arch/execution_mode/metered_cost.rs

@@ -123,8 +123,14 @@ impl ExecutionCtxTrait for MeteredCostCtx {
         );
     }
 
-    fn should_suspend<F>(vm_state: &mut VmExecState<F, GuestMemory, Self>) -> bool {
-        vm_state.ctx.cost > vm_state.ctx.max_execution_cost
+    #[inline(always)]
+    fn should_suspend<F>(
+        _instret: u64,
+        _pc: u32,
+        max_execution_cost: u64,
+        exec_state: &mut VmExecState<F, GuestMemory, Self>,
+    ) -> bool {
+        exec_state.ctx.cost > max_execution_cost
     }
 }
 

crates/vm/src/arch/execution_mode/mod.rs

@@ -12,8 +12,20 @@ pub use pure::ExecutionCtx;
 
 pub trait ExecutionCtxTrait: Sized {
     fn on_memory_operation(&mut self, address_space: u32, ptr: u32, size: u32);
-    fn should_suspend<F>(vm_state: &mut VmExecState<F, GuestMemory, Self>) -> bool;
-    fn on_terminate<F>(_vm_state: &mut VmExecState<F, GuestMemory, Self>) {}
+
+    fn should_suspend<F>(
+        instret: u64,
+        pc: u32,
+        _arg: u64,
+        exec_state: &mut VmExecState<F, GuestMemory, Self>,
+    ) -> bool;
+
+    fn on_terminate<F>(
+        _instret: u64,
+        _pc: u32,
+        _exec_state: &mut VmExecState<F, GuestMemory, Self>,
+    ) {
+    }
 }
 
 pub trait MeteredExecutionCtxTrait: ExecutionCtxTrait {

crates/vm/src/arch/execution_mode/pure.rs

@@ -4,7 +4,7 @@ use crate::{
 };
 
 pub struct ExecutionCtx {
-    instret_end: u64,
+    pub instret_end: u64,
 }
 
 impl ExecutionCtx {
@@ -22,8 +22,14 @@ impl ExecutionCtx {
 impl ExecutionCtxTrait for ExecutionCtx {
     #[inline(always)]
     fn on_memory_operation(&mut self, _address_space: u32, _ptr: u32, _size: u32) {}
+
     #[inline(always)]
-    fn should_suspend<F>(vm_state: &mut VmExecState<F, GuestMemory, Self>) -> bool {
-        vm_state.instret >= vm_state.ctx.instret_end
+    fn should_suspend<F>(
+        instret: u64,
+        _pc: u32,
+        instret_end: u64,
+        _exec_state: &mut VmExecState<F, GuestMemory, Self>,
+    ) -> bool {
+        instret >= instret_end
     }
 }