perf: remove pc_base from pc_idx calc

Author: jonathanpwang

crates/vm/derive/src/tco.rs

@@ -48,6 +48,8 @@ pub fn tco_impl(item: TokenStream) -> TokenStream {
         )
         #where_clause
         {
+            use ::openvm_circuit::arch::ExecutionError;
+
             let pre_compute = interpreter.get_pre_compute(exec_state.vm_state.pc);
             #execute_call;
 
@@ -61,7 +63,7 @@ pub fn tco_impl(item: TokenStream) -> TokenStream {
             // 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(interpreter.pc_out_of_bounds_err(exec_state.vm_state.pc));
+                exec_state.exit_code = Err(ExecutionError::PcOutOfBounds (exec_state.vm_state.pc));
                 return;
             }
             let next_handler = next_handler.unwrap_unchecked();

crates/vm/src/arch/execution.rs

@@ -28,12 +28,8 @@ use crate::{
 pub enum ExecutionError {
     #[error("execution failed at pc {pc}, err: {msg}")]
     Fail { pc: u32, msg: &'static str },
-    #[error("pc {pc} out of bounds for program of length {program_len}, with pc_base {pc_base}")]
-    PcOutOfBounds {
-        pc: u32,
-        pc_base: u32,
-        program_len: usize,
-    },
+    #[error("pc {0} out of bounds")]
+    PcOutOfBounds(u32),
     #[error("unreachable instruction at pc {0}")]
     Unreachable(u32),
     #[error("at pc {pc}, opcode {opcode} was not enabled")]

crates/vm/src/arch/interpreter.rs

@@ -1,6 +1,7 @@
 use std::{
     alloc::{alloc, dealloc, handle_alloc_error, Layout},
     borrow::{Borrow, BorrowMut},
+    iter::repeat_n,
     ptr::NonNull,
 };
 
@@ -44,15 +45,16 @@ pub struct InterpretedInstance<'a, F, Ctx> {
     #[allow(dead_code)]
     pre_compute_buf: AlignedBuf,
     /// Instruction table of function pointers and pointers to the pre-computed buffer. Indexed by
-    /// `pc_index = (pc - pc_base) / DEFAULT_PC_STEP`.
+    /// `pc_index = pc / DEFAULT_PC_STEP`.
+    /// SAFETY: The first `pc_base / DEFAULT_PC_STEP` entries will be unreachable. We do this to
+    /// avoid needing to subtract `pc_base` during runtime.
     pre_compute_insns: Vec<PreComputeInstruction<'a, F, Ctx>>,
     #[cfg(feature = "tco")]
     pre_compute_max_size: usize,
     /// Handler function pointers for tail call optimization.
     #[cfg(feature = "tco")]
     handlers: Vec<Handler<F, Ctx>>,
 
-    pc_base: u32,
     pc_start: u32,
 
     init_memory: SparseMemoryImage,
@@ -84,22 +86,14 @@ macro_rules! run {
             #[cfg(not(feature = "tco"))]
             unsafe {
                 tracing::debug!("execute_trampoline");
-                execute_trampoline(
-                    $interpreter.pc_base,
-                    &mut $exec_state,
-                    &$interpreter.pre_compute_insns,
-                );
+                execute_trampoline(&mut $exec_state, &$interpreter.pre_compute_insns);
             }
             #[cfg(feature = "tco")]
             {
                 tracing::debug!("execute_tco");
-                let handler = $interpreter.get_handler($exec_state.pc).ok_or(
-                    ExecutionError::PcOutOfBounds {
-                        pc: $exec_state.pc,
-                        pc_base: $interpreter.pc_base,
-                        program_len: $interpreter.handlers.len(),
-                    },
-                )?;
+                let handler = $interpreter
+                    .get_handler($exec_state.pc)
+                    .ok_or(ExecutionError::PcOutOfBounds($exec_state.pc))?;
                 // SAFETY:
                 // - handler is generated by Executor, MeteredExecutor traits
                 // - it is the responsibility of each Executor to ensure handler is safe given a
@@ -151,27 +145,25 @@ where
     {
         let program = &exe.program;
         let pre_compute_max_size = get_pre_compute_max_size(program, inventory);
-        let mut pre_compute_buf = alloc_pre_compute_buf(program.len(), pre_compute_max_size);
+        let mut pre_compute_buf = alloc_pre_compute_buf(program, pre_compute_max_size);
         let mut split_pre_compute_buf =
             split_pre_compute_buf(program, &mut pre_compute_buf, pre_compute_max_size);
         let pre_compute_insns = get_pre_compute_instructions::<F, Ctx, E>(
             program,
             inventory,
             &mut split_pre_compute_buf,
         )?;
-        let pc_base = program.pc_base;
         let pc_start = exe.pc_start;
         let init_memory = exe.init_memory.clone();
         #[cfg(feature = "tco")]
-        let handlers = program
-            .instructions_and_debug_infos
-            .iter()
+        let handlers = repeat_n(&None, get_pc_index(program.pc_base))
+            .chain(program.instructions_and_debug_infos.iter())
             .zip_eq(split_pre_compute_buf.iter_mut())
             .enumerate()
             .map(
                 |(pc_idx, (inst_opt, pre_compute))| -> Result<Handler<F, Ctx>, StaticProgramError> {
                     if let Some((inst, _)) = inst_opt {
-                        let pc = pc_base + pc_idx as u32 * DEFAULT_PC_STEP;
+                        let pc = pc_idx as u32 * DEFAULT_PC_STEP;
                         if get_system_opcode_handler::<F, Ctx>(inst, pre_compute).is_some() {
                             Ok(terminate_execute_e12_tco_handler)
                         } else {
@@ -191,7 +183,6 @@ where
             system_config: inventory.config().clone(),
             pre_compute_buf,
             pre_compute_insns,
-            pc_base,
             pc_start,
             init_memory,
             #[cfg(feature = "tco")]
@@ -209,7 +200,7 @@ where
     #[cfg(feature = "tco")]
     #[inline(always)]
     pub fn get_pre_compute(&self, pc: u32) -> &[u8] {
-        let pc_idx = get_pc_index(self.pc_base, pc);
+        let pc_idx = get_pc_index(pc);
         // SAFETY:
         // - we assume that pc is in bounds
         // - pre_compute_buf is allocated for pre_compute_max_size * program_len bytes, with each
@@ -228,18 +219,10 @@ where
         }
     }
 
-    pub fn pc_out_of_bounds_err(&self, pc: u32) -> ExecutionError {
-        ExecutionError::PcOutOfBounds {
-            pc,
-            pc_base: self.pc_base,
-            program_len: self.pre_compute_insns.len(),
-        }
-    }
-
     #[cfg(feature = "tco")]
     #[inline(always)]
     pub fn get_handler(&self, pc: u32) -> Option<Handler<F, Ctx>> {
-        let pc_idx = get_pc_index(self.pc_base, pc);
+        let pc_idx = get_pc_index(pc);
         self.handlers.get(pc_idx).copied()
     }
 }
@@ -261,7 +244,7 @@ where
     {
         let program = &exe.program;
         let pre_compute_max_size = get_metered_pre_compute_max_size(program, inventory);
-        let mut pre_compute_buf = alloc_pre_compute_buf(program.len(), pre_compute_max_size);
+        let mut pre_compute_buf = alloc_pre_compute_buf(program, pre_compute_max_size);
         let mut split_pre_compute_buf =
             split_pre_compute_buf(program, &mut pre_compute_buf, pre_compute_max_size);
         let pre_compute_insns = get_metered_pre_compute_instructions::<F, Ctx, E>(
@@ -271,19 +254,17 @@ where
             &mut split_pre_compute_buf,
         )?;
 
-        let pc_base = program.pc_base;
         let pc_start = exe.pc_start;
         let init_memory = exe.init_memory.clone();
         #[cfg(feature = "tco")]
-        let handlers = program
-            .instructions_and_debug_infos
-            .iter()
+        let handlers = repeat_n(&None, get_pc_index(program.pc_base))
+            .chain(program.instructions_and_debug_infos.iter())
             .zip_eq(split_pre_compute_buf.iter_mut())
             .enumerate()
             .map(
                 |(pc_idx, (inst_opt, pre_compute))| -> Result<Handler<F, Ctx>, StaticProgramError> {
                     if let Some((inst, _)) = inst_opt {
-                        let pc = pc_base + pc_idx as u32 * DEFAULT_PC_STEP;
+                        let pc = pc_idx as u32 * DEFAULT_PC_STEP;
                         if get_system_opcode_handler::<F, Ctx>(inst, pre_compute).is_some() {
                             Ok(terminate_execute_e12_tco_handler)
                         } else {
@@ -305,7 +286,6 @@ where
             system_config: inventory.config().clone(),
             pre_compute_buf,
             pre_compute_insns,
-            pc_base,
             pc_start,
             init_memory,
             #[cfg(feature = "tco")]
@@ -448,8 +428,10 @@ where
     }
 }
 
-fn alloc_pre_compute_buf(program_len: usize, pre_compute_max_size: usize) -> AlignedBuf {
-    let buf_len = program_len * pre_compute_max_size;
+fn alloc_pre_compute_buf<F>(program: &Program<F>, pre_compute_max_size: usize) -> AlignedBuf {
+    let base_idx = get_pc_index(program.pc_base);
+    let padded_program_len = base_idx + program.instructions_and_debug_infos.len();
+    let buf_len = padded_program_len * pre_compute_max_size;
     AlignedBuf::uninit(buf_len, pre_compute_max_size)
 }
 
@@ -458,8 +440,9 @@ fn split_pre_compute_buf<'a, F>(
     pre_compute_buf: &'a mut AlignedBuf,
     pre_compute_max_size: usize,
 ) -> Vec<&'a mut [u8]> {
-    let program_len = program.instructions_and_debug_infos.len();
-    let buf_len = program_len * pre_compute_max_size;
+    let base_idx = get_pc_index(program.pc_base);
+    let padded_program_len = base_idx + program.instructions_and_debug_infos.len();
+    let buf_len = padded_program_len * pre_compute_max_size;
     // SAFETY:
     // - pre_compute_buf.ptr was allocated with exactly buf_len bytes
     // - lifetime 'a ensures the returned slices don't outlive the AlignedBuf
@@ -475,7 +458,6 @@ fn split_pre_compute_buf<'a, F>(
 /// The `fn_ptrs` pointer to pre-computed buffers that outlive this function.
 #[inline(always)]
 unsafe fn execute_trampoline<F: PrimeField32, Ctx: ExecutionCtxTrait>(
-    pc_base: u32,
     vm_state: &mut VmExecState<F, GuestMemory, Ctx>,
     fn_ptrs: &[PreComputeInstruction<F, Ctx>],
 ) {
@@ -487,16 +469,12 @@ unsafe fn execute_trampoline<F: PrimeField32, Ctx: ExecutionCtxTrait>(
         if Ctx::should_suspend(vm_state) {
             break;
         }
-        let pc_index = get_pc_index(pc_base, vm_state.pc);
+        let pc_index = get_pc_index(vm_state.pc);
         if let Some(inst) = fn_ptrs.get(pc_index) {
             // SAFETY: pre_compute assumed to live long enough
             unsafe { (inst.handler)(inst.pre_compute, vm_state) };
         } else {
-            vm_state.exit_code = Err(ExecutionError::PcOutOfBounds {
-                pc: vm_state.pc,
-                pc_base,
-                program_len: fn_ptrs.len(),
-            });
+            vm_state.exit_code = Err(ExecutionError::PcOutOfBounds(vm_state.pc));
         }
     }
     if vm_state
@@ -509,8 +487,8 @@ unsafe fn execute_trampoline<F: PrimeField32, Ctx: ExecutionCtxTrait>(
 }
 
 #[inline(always)]
-pub fn get_pc_index(pc_base: u32, pc: u32) -> usize {
-    ((pc - pc_base) / DEFAULT_PC_STEP) as usize
+pub fn get_pc_index(pc: u32) -> usize {
+    (pc / DEFAULT_PC_STEP) as usize
 }
 
 /// Bytes allocated according to the given Layout
@@ -647,15 +625,19 @@ where
     Ctx: ExecutionCtxTrait,
     E: Executor<F>,
 {
-    program
-        .instructions_and_debug_infos
-        .iter()
+    let unreachable_handler: ExecuteFunc<F, Ctx> = |_, vm_state| {
+        vm_state.exit_code = Err(ExecutionError::Unreachable(vm_state.pc));
+    };
+
+    repeat_n(&None, get_pc_index(program.pc_base))
+        .chain(program.instructions_and_debug_infos.iter())
         .zip_eq(pre_compute.iter_mut())
         .enumerate()
         .map(|(i, (inst_opt, buf))| {
-            // SAFETY: we cast to raw pointer and then borrow to remove the lifetime. This is safe
-            // only in the current context because `buf` comes from `pre_compute_buf` which will
-            // outlive the returned `PreComputeInstruction`s.
+            // SAFETY: we cast to raw pointer and then borrow to remove the lifetime. This
+            // is safe only in the current context because `buf` comes
+            // from `pre_compute_buf` which will outlive the returned
+            // `PreComputeInstruction`s.
             let buf: &mut [u8] = unsafe { &mut *(*buf as *mut [u8]) };
             let pre_inst = if let Some((inst, _)) = inst_opt {
                 tracing::trace!("get_pre_compute_instruction {inst:?}");
@@ -679,9 +661,7 @@ where
             } else {
                 // Dead instruction at this pc
                 PreComputeInstruction {
-                    handler: |_, vm_state| {
-                        vm_state.exit_code = Err(ExecutionError::Unreachable(vm_state.pc));
-                    },
+                    handler: unreachable_handler,
                     pre_compute: buf,
                 }
             };
@@ -701,15 +681,18 @@ where
     Ctx: MeteredExecutionCtxTrait,
     E: MeteredExecutor<F>,
 {
-    program
-        .instructions_and_debug_infos
-        .iter()
+    let unreachable_handler: ExecuteFunc<F, Ctx> = |_, vm_state| {
+        vm_state.exit_code = Err(ExecutionError::Unreachable(vm_state.pc));
+    };
+    repeat_n(&None, get_pc_index(program.pc_base))
+        .chain(program.instructions_and_debug_infos.iter())
         .zip_eq(pre_compute.iter_mut())
         .enumerate()
         .map(|(i, (inst_opt, buf))| {
-            // SAFETY: we cast to raw pointer and then borrow to remove the lifetime. This is safe
-            // only in the current context because `buf` comes from `pre_compute_buf` which will
-            // outlive the returned `PreComputeInstruction`s.
+            // SAFETY: we cast to raw pointer and then borrow to remove the lifetime. This
+            // is safe only in the current context because `buf` comes
+            // from `pre_compute_buf` which will outlive the returned
+            // `PreComputeInstruction`s.
             let buf: &mut [u8] = unsafe { &mut *(*buf as *mut [u8]) };
             let pre_inst = if let Some((inst, _)) = inst_opt {
                 tracing::trace!("get_metered_pre_compute_instruction {inst:?}");
@@ -738,9 +721,7 @@ where
                 }
             } else {
                 PreComputeInstruction {
-                    handler: |_, vm_state| {
-                        vm_state.exit_code = Err(ExecutionError::Unreachable(vm_state.pc));
-                    },
+                    handler: unreachable_handler,
                     pre_compute: buf,
                 }
             };