perf: remove pc_base from pc_idx calc

jonathanpwang · jonathanpwang · commit e066de4b4f35 · 2025-08-21T20:55:52.000-07:00
diff --git a/crates/vm/derive/src/tco.rs b/crates/vm/derive/src/tco.rs
@@ -61,7 +61,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();
diff --git a/crates/vm/src/arch/execution.rs b/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")]
diff --git a/crates/vm/src/arch/interpreter.rs b/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,11 +86,7 @@ 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")]
             {
@@ -151,21 +149,19 @@ 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(
@@ -191,7 +187,6 @@ where
             system_config: inventory.config().clone(),
             pre_compute_buf,
             pre_compute_insns,
-            pc_base,
             pc_start,
             init_memory,
             #[cfg(feature = "tco")]
@@ -209,7 +204,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,14 +223,6 @@ 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>> {
@@ -261,7 +248,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,13 +258,11 @@ 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(
@@ -305,7 +290,6 @@ where
             system_config: inventory.config().clone(),
             pre_compute_buf,
             pre_compute_insns,
-            pc_base,
             pc_start,
             init_memory,
             #[cfg(feature = "tco")]
@@ -448,8 +432,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 +444,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 +462,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 +473,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 +491,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 +629,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 +665,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 +685,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 +725,7 @@ where
                 }
             } else {
                 PreComputeInstruction {
-                    handler: |_, vm_state| {
-                        vm_state.exit_code = Err(ExecutionError::Unreachable(vm_state.pc));
-                    },
+                    handler: unreachable_handler,
                     pre_compute: buf,
                 }
             };
diff --git a/crates/vm/src/arch/interpreter_preflight.rs b/crates/vm/src/arch/interpreter_preflight.rs
@@ -1,4 +1,4 @@
-use std::sync::Arc;
+use std::{iter::repeat_n, sync::Arc};
 
 use openvm_instructions::{instruction::Instruction, program::Program, LocalOpcode, SystemOpcode};
 use openvm_stark_backend::{
@@ -36,6 +36,7 @@ pub struct PreflightInterpretedInstance<F, E> {
 }
 
 #[repr(C)]
+#[derive(Clone)]
 pub struct PcEntry<F> {
     // NOTE[jpw]: revisit storing only smaller `precompute` for better cache locality. Currently
     // VmOpcode is usize so align=8 and there are 7 u32 operands so we store ExecutorId(u32) after
@@ -60,7 +61,10 @@ impl<F: Field, E> PreflightInterpretedInstance<F, E> {
             return Err(StaticProgramError::TooManyExecutors);
         }
         let len = program.instructions_and_debug_infos.len();
-        let mut pc_handler = Vec::with_capacity(len);
+        let pc_base = program.pc_base;
+        let base_idx = get_pc_index(pc_base);
+        let mut pc_handler = Vec::with_capacity(base_idx + len);
+        pc_handler.extend(repeat_n(PcEntry::undefined(), base_idx));
         for insn_and_debug_info in &program.instructions_and_debug_infos {
             if let Some((insn, _)) = insn_and_debug_info {
                 let insn = insn.clone();
@@ -86,9 +90,9 @@ impl<F: Field, E> PreflightInterpretedInstance<F, E> {
         }
         Ok(Self {
             inventory,
-            execution_frequencies: vec![0u32; len],
+            execution_frequencies: vec![0u32; base_idx + len],
+            pc_base,
             pc_handler,
-            pc_base: program.pc_base,
             executor_idx_to_air_idx,
         })
     }
@@ -101,9 +105,11 @@ impl<F: Field, E> PreflightInterpretedInstance<F, E> {
     where
         E: Send + Sync,
     {
+        let base_idx = get_pc_index(self.pc_base);
         self.pc_handler
             .par_iter()
             .enumerate()
+            .skip(base_idx)
             .filter(|(_, entry)| entry.is_some())
             .map(|(i, _)| self.execution_frequencies[i])
             .collect()
@@ -157,15 +163,11 @@ impl<F: PrimeField32, E> PreflightInterpretedInstance<F, E> {
         E: PreflightExecutor<F, RA>,
     {
         let pc = state.pc;
-        let pc_idx = get_pc_index(self.pc_base, pc);
-        let pc_entry =
-            self.pc_handler
-                .get(pc_idx)
-                .ok_or_else(|| ExecutionError::PcOutOfBounds {
-                    pc,
-                    pc_base: self.pc_base,
-                    program_len: self.pc_handler.len(),
-                })?;
+        let pc_idx = get_pc_index(pc);
+        let pc_entry = self
+            .pc_handler
+            .get(pc_idx)
+            .ok_or_else(|| ExecutionError::PcOutOfBounds(pc))?;
         // SAFETY: `execution_frequencies` has the same length as `pc_handler` so `get_pc_entry`
         // already does the bounds check
         unsafe {

Original file line number	Diff line number	Diff line change
`@@ -61,7 +61,7 @@ pub fn tco_impl(item: TokenStream) -> TokenStream {`
`61`	`61`	`// exec_state.pc should have been updated by execute_impl at this point`
`62`	`62`	`let next_handler = interpreter.get_handler(exec_state.vm_state.pc);`
`63`	`63`	`if next_handler.is_none() {`
`64`		`- exec_state.exit_code = Err(interpreter.pc_out_of_bounds_err(exec_state.vm_state.pc));`
	`64`	`+ exec_state.exit_code = Err(ExecutionError::PcOutOfBounds (exec_state.vm_state.pc));`
`65`	`65`	`return;`
`66`	`66`	`}`
`67`	`67`	`let next_handler = next_handler.unwrap_unchecked();`