feat: rv32im tco without become keyword

`become` keyword was causing some corruption and not passing references
properly between calls.

Author: jonathanpwang

crates/vm/derive/src/tco.rs

@@ -49,19 +49,25 @@ pub fn tco_impl(item: TokenStream) -> TokenStream {
         ) -> Result<(), ::openvm_circuit::arch::ExecutionError>
         #where_clause
         {
-            let pre_compute = interpreter.get_pre_compute(exec_state.pc);
+            let pre_compute = interpreter.get_pre_compute(exec_state.vm_state.pc);
             #execute_call;
 
             if std::hint::unlikely(exec_state.exit_code.is_err()) {
                 return Err(::openvm_circuit::arch::ExecutionError::ExecStateError);
             }
             if std::hint::unlikely(exec_state.exit_code.as_ref().unwrap().is_some()) {
+                #ctx_type::on_terminate(exec_state);
                 // terminate
                 return Ok(());
             }
+            if #ctx_type::should_suspend(exec_state) {
+                return Ok(());
+            }
             // exec_state.pc should have been updated by execute_impl at this point
-            let next_handler = interpreter.get_handler(exec_state.pc)?;
-            become next_handler(interpreter, exec_state)
+            let next_handler = interpreter.get_handler(exec_state.vm_state.pc)?;
+
+            // The `become` keyword has a bug that is not re-passing the `interpreter`, `exec_state` references properly. But llvm seems to almost always guarantee tail call elimination when the function signature is the same as the current function.
+            next_handler(interpreter, exec_state)
         }
     };
 

crates/vm/src/arch/interpreter.rs

@@ -139,7 +139,7 @@ where
                             // unwrap because get_pre_compute_instructions would have errored
                             // already on DisabledOperation
                             let executor = inventory.get_executor(inst.opcode).unwrap();
-                            executor.handler(pc, inst, *pre_compute)
+                            executor.handler(pc, inst, pre_compute)
                         }
                     } else {
                         Ok(unreachable_tco_handler)
@@ -177,6 +177,9 @@ where
         //   instruction getting pre_compute_max_size bytes
         // - self.pre_compute_buf.ptr is non-null
         // - initialization of the contents of the slice is the responsibility of each Executor
+        debug_assert!(
+            (pc_idx + 1) * self.pre_compute_max_size <= self.pre_compute_buf.layout.size()
+        );
         unsafe {
             let ptr = self
                 .pre_compute_buf
@@ -190,7 +193,7 @@ where
     #[inline(always)]
     pub fn get_handler(&self, pc: u32) -> Result<Handler<F, Ctx>, ExecutionError> {
         let pc_idx = get_pc_index(self.pc_base, pc);
-        if std::hint::unlikely(pc_idx > self.handlers.len()) {
+        if std::hint::unlikely(pc_idx >= self.handlers.len()) {
             return Err(ExecutionError::PcOutOfBounds {
                 pc,
                 pc_base: self.pc_base,
@@ -287,13 +290,38 @@ where
     ) -> Result<VmState<F, GuestMemory>, ExecutionError> {
         let ctx = ExecutionCtx::new(num_insns);
         let mut exec_state = VmExecState::new(from_state, ctx);
-        // Start execution
-        execute_with_metrics!(
-            "execute_e1",
-            self.pc_base,
-            &mut exec_state,
-            &self.pre_compute_insns
-        );
+
+        #[cfg(feature = "metrics")]
+        let start = std::time::Instant::now();
+        #[cfg(feature = "metrics")]
+        let start_instret = exec_state.instret;
+
+        #[cfg(not(feature = "tco"))]
+        unsafe {
+            execute_trampoline(self.pc_base, &mut exec_state, &self.pre_compute_insns);
+        }
+        #[cfg(feature = "tco")]
+        unsafe {
+            let handler = self.get_handler(exec_state.pc)?;
+            let res = handler(self, &mut exec_state);
+            if let Err(err) = res {
+                match err {
+                    ExecutionError::ExecStateError => {}
+                    _ => {
+                        return Err(err);
+                    }
+                }
+            }
+        }
+
+        #[cfg(feature = "metrics")]
+        {
+            let elapsed = start.elapsed();
+            let insns = exec_state.instret - start_instret;
+            tracing::info!("instructions_executed={insns}");
+            metrics::counter!("execute_e1_insns").absolute(insns);
+            metrics::gauge!("execute_e1_insn_mi/s").set(insns as f64 / elapsed.as_micros() as f64);
+        }
         if num_insns.is_some() {
             check_exit_code(exec_state.exit_code)?;
         } else {
@@ -410,15 +438,10 @@ fn split_pre_compute_buf<'a, F>(
 ) -> Vec<&'a mut [u8]> {
     let program_len = program.instructions_and_debug_infos.len();
     let buf_len = program_len * pre_compute_max_size;
-    let mut pre_compute_buf_ptr =
-        unsafe { std::slice::from_raw_parts_mut(pre_compute_buf.ptr, buf_len) };
-    let mut split_pre_compute_buf = Vec::with_capacity(program_len);
-    for _ in 0..program_len {
-        let (first, last) = pre_compute_buf_ptr.split_at_mut(pre_compute_max_size);
-        pre_compute_buf_ptr = last;
-        split_pre_compute_buf.push(first);
-    }
-    split_pre_compute_buf
+    let pre_compute_buf = unsafe { std::slice::from_raw_parts_mut(pre_compute_buf.ptr, buf_len) };
+    pre_compute_buf
+        .chunks_exact_mut(pre_compute_max_size)
+        .collect()
 }
 
 /// Executes using function pointers with the trampoline (loop) approach.

crates/vm/src/utils/stark_utils.rs

@@ -114,6 +114,8 @@ where
     let exe = exe.into();
     let input = input.into();
     let metered_ctx = vm.build_metered_ctx();
+    // TEMP: for testing
+    vm.interpreter(&exe)?.execute(input.clone(), None)?;
     let (segments, _) = vm
         .metered_interpreter(&exe)?
         .execute_metered(input.clone(), metered_ctx)?;

extensions/rv32im/circuit/Cargo.toml

@@ -33,9 +33,10 @@ openvm-circuit = { workspace = true, features = ["test-utils"] }
 test-case.workspace = true
 
 [features]
-default = ["parallel", "jemalloc"]
+default = ["parallel", "jemalloc", "tco"]
 parallel = ["openvm-circuit/parallel"]
 test-utils = ["openvm-circuit/test-utils", "dep:openvm-stark-sdk"]
+tco = ["openvm-circuit/tco"]
 # performance features:
 mimalloc = ["openvm-circuit/mimalloc"]
 jemalloc = ["openvm-circuit/jemalloc"]

extensions/rv32im/circuit/src/auipc/execution.rs

@@ -3,13 +3,7 @@ use std::{
     mem::size_of,
 };
 
-use openvm_circuit::{
-    arch::{
-        E2PreCompute, ExecuteFunc, ExecutionCtxTrait, Executor, MeteredExecutionCtxTrait,
-        MeteredExecutor, StaticProgramError, VmExecState,
-    },
-    system::memory::online::GuestMemory,
-};
+use openvm_circuit::{arch::*, system::memory::online::GuestMemory};
 use openvm_circuit_primitives_derive::AlignedBytesBorrow;
 use openvm_instructions::{
     instruction::Instruction, program::DEFAULT_PC_STEP, riscv::RV32_REGISTER_AS,
@@ -66,6 +60,21 @@ where
         self.pre_compute_impl(pc, inst, data)?;
         Ok(execute_e1_impl)
     }
+
+    #[cfg(feature = "tco")]
+    fn handler<Ctx>(
+        &self,
+        pc: u32,
+        inst: &Instruction<F>,
+        data: &mut [u8],
+    ) -> Result<Handler<F, Ctx>, StaticProgramError>
+    where
+        Ctx: ExecutionCtxTrait,
+    {
+        let data: &mut AuiPcPreCompute = data.borrow_mut();
+        self.pre_compute_impl(pc, inst, data)?;
+        Ok(execute_e1_tco_handler)
+    }
 }
 
 impl<F, A> MeteredExecutor<F> for Rv32AuipcExecutor<A>
@@ -105,6 +114,7 @@ unsafe fn execute_e12_impl<F: PrimeField32, CTX: ExecutionCtxTrait>(
     vm_state.instret += 1;
 }
 
+#[create_tco_handler]
 unsafe fn execute_e1_impl<F: PrimeField32, CTX: ExecutionCtxTrait>(
     pre_compute: &[u8],
     vm_state: &mut VmExecState<F, GuestMemory, CTX>,

extensions/rv32im/circuit/src/base_alu/execution.rs

@@ -3,13 +3,7 @@ use std::{
     mem::size_of,
 };
 
-use openvm_circuit::{
-    arch::{
-        E2PreCompute, ExecuteFunc, ExecutionCtxTrait, Executor, MeteredExecutionCtxTrait,
-        MeteredExecutor, StaticProgramError, VmExecState,
-    },
-    system::memory::online::GuestMemory,
-};
+use openvm_circuit::{arch::*, system::memory::online::GuestMemory};
 use openvm_circuit_primitives_derive::AlignedBytesBorrow;
 use openvm_instructions::{
     instruction::Instruction,
@@ -102,6 +96,38 @@ where
         };
         Ok(fn_ptr)
     }
+
+    #[cfg(feature = "tco")]
+    fn handler<Ctx>(
+        &self,
+        pc: u32,
+        inst: &Instruction<F>,
+        data: &mut [u8],
+    ) -> Result<Handler<F, Ctx>, StaticProgramError>
+    where
+        Ctx: ExecutionCtxTrait,
+    {
+        let data: &mut BaseAluPreCompute = data.borrow_mut();
+        let is_imm = self.pre_compute_impl(pc, inst, data)?;
+        let opcode = inst.opcode;
+
+        let fn_ptr = match (
+            is_imm,
+            BaseAluOpcode::from_usize(opcode.local_opcode_idx(self.offset)),
+        ) {
+            (true, BaseAluOpcode::ADD) => execute_e1_tco_handler::<_, _, true, AddOp>,
+            (false, BaseAluOpcode::ADD) => execute_e1_tco_handler::<_, _, false, AddOp>,
+            (true, BaseAluOpcode::SUB) => execute_e1_tco_handler::<_, _, true, SubOp>,
+            (false, BaseAluOpcode::SUB) => execute_e1_tco_handler::<_, _, false, SubOp>,
+            (true, BaseAluOpcode::XOR) => execute_e1_tco_handler::<_, _, true, XorOp>,
+            (false, BaseAluOpcode::XOR) => execute_e1_tco_handler::<_, _, false, XorOp>,
+            (true, BaseAluOpcode::OR) => execute_e1_tco_handler::<_, _, true, OrOp>,
+            (false, BaseAluOpcode::OR) => execute_e1_tco_handler::<_, _, false, OrOp>,
+            (true, BaseAluOpcode::AND) => execute_e1_tco_handler::<_, _, true, AndOp>,
+            (false, BaseAluOpcode::AND) => execute_e1_tco_handler::<_, _, false, AndOp>,
+        };
+        Ok(fn_ptr)
+    }
 }
 
 impl<F, A, const LIMB_BITS: usize> MeteredExecutor<F>
@@ -174,6 +200,7 @@ unsafe fn execute_e12_impl<
     vm_state.instret += 1;
 }
 
+#[create_tco_handler]
 #[inline(always)]
 unsafe fn execute_e1_impl<
     F: PrimeField32,

extensions/rv32im/circuit/src/branch_eq/execution.rs

@@ -3,13 +3,7 @@ use std::{
     mem::size_of,
 };
 
-use openvm_circuit::{
-    arch::{
-        E2PreCompute, ExecuteFunc, ExecutionCtxTrait, Executor, MeteredExecutionCtxTrait,
-        MeteredExecutor, StaticProgramError, VmExecState,
-    },
-    system::memory::online::GuestMemory,
-};
+use openvm_circuit::{arch::*, system::memory::online::GuestMemory};
 use openvm_circuit_primitives_derive::AlignedBytesBorrow;
 use openvm_instructions::{
     instruction::Instruction, program::DEFAULT_PC_STEP, riscv::RV32_REGISTER_AS, LocalOpcode,
@@ -84,6 +78,26 @@ where
         };
         Ok(fn_ptr)
     }
+
+    #[cfg(feature = "tco")]
+    fn handler<Ctx>(
+        &self,
+        pc: u32,
+        inst: &Instruction<F>,
+        data: &mut [u8],
+    ) -> Result<Handler<F, Ctx>, StaticProgramError>
+    where
+        Ctx: ExecutionCtxTrait,
+    {
+        let data: &mut BranchEqualPreCompute = data.borrow_mut();
+        let is_bne = self.pre_compute_impl(pc, inst, data)?;
+        let fn_ptr = if is_bne {
+            execute_e1_tco_handler::<_, _, true>
+        } else {
+            execute_e1_tco_handler::<_, _, false>
+        };
+        Ok(fn_ptr)
+    }
 }
 
 impl<F, A, const NUM_LIMBS: usize> MeteredExecutor<F> for BranchEqualExecutor<A, NUM_LIMBS>
@@ -131,6 +145,7 @@ unsafe fn execute_e12_impl<F: PrimeField32, CTX: ExecutionCtxTrait, const IS_NE:
     vm_state.instret += 1;
 }
 
+#[create_tco_handler]
 unsafe fn execute_e1_impl<F: PrimeField32, CTX: ExecutionCtxTrait, const IS_NE: bool>(
     pre_compute: &[u8],
     vm_state: &mut VmExecState<F, GuestMemory, CTX>,

extensions/rv32im/circuit/src/branch_lt/execution.rs

@@ -3,13 +3,7 @@ use std::{
     mem::size_of,
 };
 
-use openvm_circuit::{
-    arch::{
-        E2PreCompute, ExecuteFunc, ExecutionCtxTrait, Executor, MeteredExecutionCtxTrait,
-        MeteredExecutor, StaticProgramError, VmExecState,
-    },
-    system::memory::online::GuestMemory,
-};
+use openvm_circuit::{arch::*, system::memory::online::GuestMemory};
 use openvm_circuit_primitives_derive::AlignedBytesBorrow;
 use openvm_instructions::{
     instruction::Instruction, program::DEFAULT_PC_STEP, riscv::RV32_REGISTER_AS, LocalOpcode,
@@ -86,6 +80,27 @@ where
         };
         Ok(fn_ptr)
     }
+
+    #[cfg(feature = "tco")]
+    fn handler<Ctx>(
+        &self,
+        pc: u32,
+        inst: &Instruction<F>,
+        data: &mut [u8],
+    ) -> Result<Handler<F, Ctx>, StaticProgramError>
+    where
+        Ctx: ExecutionCtxTrait,
+    {
+        let data: &mut BranchLePreCompute = data.borrow_mut();
+        let local_opcode = self.pre_compute_impl(pc, inst, data)?;
+        let fn_ptr = match local_opcode {
+            BranchLessThanOpcode::BLT => execute_e1_tco_handler::<_, _, BltOp>,
+            BranchLessThanOpcode::BLTU => execute_e1_tco_handler::<_, _, BltuOp>,
+            BranchLessThanOpcode::BGE => execute_e1_tco_handler::<_, _, BgeOp>,
+            BranchLessThanOpcode::BGEU => execute_e1_tco_handler::<_, _, BgeuOp>,
+        };
+        Ok(fn_ptr)
+    }
 }
 
 impl<F, A, const NUM_LIMBS: usize, const LIMB_BITS: usize> MeteredExecutor<F>
@@ -136,6 +151,7 @@ unsafe fn execute_e12_impl<F: PrimeField32, CTX: ExecutionCtxTrait, OP: BranchLe
     vm_state.instret += 1;
 }
 
+#[create_tco_handler]
 unsafe fn execute_e1_impl<F: PrimeField32, CTX: ExecutionCtxTrait, OP: BranchLessThanOp>(
     pre_compute: &[u8],
     vm_state: &mut VmExecState<F, GuestMemory, CTX>,