fix: update the ctx.instret_left in E1 from the register (#2276)

Author: GunaDD

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

@@ -7,6 +7,7 @@ use openvm_stark_backend::p3_field::PrimeField32;
 use crate::{
     arch::{
         aot::common::{sync_gpr_to_xmm, sync_xmm_to_gpr, REG_AS2_PTR},
+        execution_mode::ExecutionCtx,
         interpreter::{AlignedBuf, PreComputeInstruction},
         ExecutionCtxTrait, StaticProgramError, Streams, SystemConfig, VmExecState, VmState,
     },
@@ -300,6 +301,12 @@ unsafe extern "C" fn set_pc_shim<F, Ctx: ExecutionCtxTrait>(state_ptr: *mut c_vo
     state.vm_state.set_pc(pc);
 }
 
+// only needed for pure execution
+unsafe extern "C" fn set_instret_left_shim<F>(state_ptr: *mut c_void, instret_left: u64) {
+    let state = &mut *(state_ptr as *mut VmExecState<F, GuestMemory, ExecutionCtx>);
+    state.ctx.instret_left = instret_left;
+}
+
 pub(crate) extern "C" fn extern_handler<F, Ctx: ExecutionCtxTrait, const E1: bool>(
     state_ptr: *mut c_void,
     pre_compute_insns_ptr: *const c_void,

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

@@ -6,7 +6,10 @@ use openvm_stark_backend::p3_field::PrimeField32;
 use super::{common::*, AotInstance, AsmRunFn};
 use crate::{
     arch::{
-        aot::{asm_to_lib, extern_handler, get_vm_address_space_addr, get_vm_pc_ptr, set_pc_shim},
+        aot::{
+            asm_to_lib, extern_handler, get_vm_address_space_addr, get_vm_pc_ptr,
+            set_instret_left_shim, set_pc_shim,
+        },
         execution_mode::ExecutionCtx,
         interpreter::{
             alloc_pre_compute_buf, get_pre_compute_instructions, get_pre_compute_max_size,
@@ -112,6 +115,7 @@ where
             format!("{:p}", extern_handler::<F, ExecutionCtx, true> as *const ());
         let set_pc_ptr = format!("{:p}", set_pc_shim::<F, ExecutionCtx> as *const ());
         let pre_compute_insns_ptr = format!("{:p}", pre_compute_insns_ptr as *const ());
+        let instret_left_ptr = format!("{:p}", set_instret_left_shim::<F> as *const ());
 
         for (pc, instruction, _) in exe.program.enumerate_by_pc() {
             /* Preprocessing step, to check if we should suspend or not */
@@ -129,6 +133,11 @@ where
                 asm_str += &Self::push_address_space_start();
                 asm_str += &Self::push_internal_registers();
 
+                asm_str += &format!("   mov {REG_FIRST_ARG}, {REG_EXEC_STATE_PTR}\n");
+                asm_str += &format!("   mov {REG_SECOND_ARG}, {REG_INSTRET_END}\n");
+                asm_str += &format!("   mov {REG_D}, {instret_left_ptr}\n");
+                asm_str += &format!("   call {REG_D}\n");
+
                 asm_str += &format!("   mov {REG_FIRST_ARG}, {REG_EXEC_STATE_PTR}\n");
                 asm_str += &format!("   mov {REG_SECOND_ARG}, {pre_compute_insns_ptr}\n");
                 asm_str += &format!("   mov {REG_THIRD_ARG}, {pc}\n");
@@ -199,6 +208,10 @@ where
         // asm_run_end part
         for (pc, _instruction, _) in exe.program.enumerate_by_pc() {
             asm_str += &format!("asm_run_end_{pc}:\n");
+            asm_str += &format!("   mov {REG_FIRST_ARG}, {REG_EXEC_STATE_PTR}\n");
+            asm_str += &format!("   mov {REG_SECOND_ARG}, {REG_INSTRET_END}\n");
+            asm_str += &format!("   mov {REG_D}, {instret_left_ptr}\n");
+            asm_str += &format!("   call {REG_D}\n");
             asm_str += &Self::xmm_to_rv32_regs();
             asm_str += &format!("    mov {REG_FIRST_ARG}, rbx\n");
             asm_str += &format!("    mov {REG_SECOND_ARG}, {pc}\n");