openvm_circuit/arch/
interpreter_preflight.rs

1#[cfg(feature = "metrics")]
2use std::collections::BTreeMap;
3use std::{iter::repeat_n, sync::Arc};
4
5#[cfg(not(feature = "parallel"))]
6use itertools::Itertools;
7use openvm_instructions::{instruction::Instruction, program::Program, LocalOpcode, SystemOpcode};
8use openvm_stark_backend::{
9    p3_field::{Field, PrimeField32},
10    p3_maybe_rayon::prelude::*,
11};
12
13use crate::{
14    arch::{
15        execution_mode::PreflightCtx, interpreter::get_pc_index, Arena, ExecutionError, ExecutorId,
16        ExecutorInventory, PreflightExecutor, StaticProgramError, VmExecState,
17    },
18    system::memory::online::TracingMemory,
19};
20
21/// VM preflight executor (E3 executor) for use with trace generation.
22/// Note: This executor doesn't hold any VM state and can be used for multiple execution.
23pub struct PreflightInterpretedInstance<F, E> {
24    // NOTE[jpw]: we use an Arc so that VmInstance can hold both VirtualMachine and
25    // PreflightInterpretedInstance. All we really need is to borrow `executors: &'a [E]`.
26    inventory: Arc<ExecutorInventory<E>>,
27
28    /// This is a map from (pc - pc_base) / pc_step -> [PcEntry].
29    /// We will set `executor_idx` to `u32::MAX` in the [PcEntry] if the program has no instruction
30    /// at that pc.
31    // PERF[jpw/ayush]: We could map directly to the raw pointer(u64) for executor, but storing the
32    // u32 may be better for cache efficiency.
33    pc_handler: Vec<PcEntry<F>>,
34    // pc_handler, execution_frequencies will all have the same length, which equals
35    // `Program::len()`
36    execution_frequencies: Vec<u32>,
37    pc_base: u32,
38
39    pub(super) executor_idx_to_air_idx: Vec<usize>,
40}
41
42#[repr(C)]
43#[derive(Clone)]
44pub struct PcEntry<F> {
45    // NOTE[jpw]: revisit storing only smaller `precompute` for better cache locality. Currently
46    // VmOpcode is usize so align=8 and there are 7 u32 operands so we store ExecutorId(u32) after
47    // to avoid padding. This means PcEntry has align=8 and size=40 bytes, which is too big
48    pub insn: Instruction<F>,
49    pub executor_idx: ExecutorId,
50}
51
52impl<F: Field, E> PreflightInterpretedInstance<F, E> {
53    /// Creates a new interpreter instance for preflight execution.
54    /// Rewrites the program into an internal table specialized for enum dispatch.
55    ///
56    /// ## Assumption
57    /// There are less than `u32::MAX` total AIRs.
58    pub fn new(
59        program: &Program<F>,
60        inventory: Arc<ExecutorInventory<E>>,
61        executor_idx_to_air_idx: Vec<usize>,
62    ) -> Result<Self, StaticProgramError> {
63        if inventory.executors().len() > u32::MAX as usize {
64            // This would mean we cannot use u32::MAX as an "undefined" executor index
65            return Err(StaticProgramError::TooManyExecutors);
66        }
67        let len = program.instructions_and_debug_infos.len();
68        let pc_base = program.pc_base;
69        let base_idx = get_pc_index(pc_base);
70        let mut pc_handler = Vec::with_capacity(base_idx + len);
71        pc_handler.extend(repeat_n(PcEntry::undefined(), base_idx));
72        for insn_and_debug_info in &program.instructions_and_debug_infos {
73            if let Some((insn, _)) = insn_and_debug_info {
74                let insn = insn.clone();
75                let executor_idx = if insn.opcode == SystemOpcode::TERMINATE.global_opcode() {
76                    // The execution loop will always branch to terminate before using this executor
77                    0
78                } else {
79                    *inventory.instruction_lookup.get(&insn.opcode).ok_or(
80                        StaticProgramError::ExecutorNotFound {
81                            opcode: insn.opcode,
82                        },
83                    )?
84                };
85                assert!(
86                    (executor_idx as usize) < inventory.executors.len(),
87                    "ExecutorInventory ensures executor_idx is in bounds"
88                );
89                let pc_entry = PcEntry { insn, executor_idx };
90                pc_handler.push(pc_entry);
91            } else {
92                pc_handler.push(PcEntry::undefined());
93            }
94        }
95        Ok(Self {
96            inventory,
97            execution_frequencies: vec![0u32; base_idx + len],
98            pc_base,
99            pc_handler,
100            executor_idx_to_air_idx,
101        })
102    }
103
104    pub fn executors(&self) -> &[E] {
105        &self.inventory.executors
106    }
107
108    pub fn filtered_execution_frequencies(&self) -> Vec<u32> {
109        let base_idx = get_pc_index(self.pc_base);
110        self.pc_handler
111            .par_iter()
112            .zip_eq(&self.execution_frequencies)
113            .skip(base_idx)
114            .filter_map(|(entry, freq)| entry.is_some().then_some(*freq))
115            .collect()
116    }
117
118    pub fn reset_execution_frequencies(&mut self) {
119        self.execution_frequencies.fill(0);
120    }
121}
122
123impl<F: PrimeField32, E> PreflightInterpretedInstance<F, E> {
124    #[cfg(feature = "metrics")]
125    pub fn opcode_counts_by_air<RA>(&self) -> BTreeMap<(usize, String), u64>
126    where
127        RA: Arena,
128        E: PreflightExecutor<F, RA>,
129    {
130        let mut counts = BTreeMap::new();
131        for (entry, &freq) in self.pc_handler.iter().zip(&self.execution_frequencies) {
132            if freq == 0
133                || !entry.is_some()
134                || entry.insn.opcode == SystemOpcode::TERMINATE.global_opcode()
135            {
136                continue;
137            }
138            let executor_idx = entry.executor_idx as usize;
139            let air_idx = unsafe {
140                // SAFETY: `entry.executor_idx` was produced by `ExecutorInventory`, and
141                // `executor_idx_to_air_idx` has one entry per executor.
142                *self.executor_idx_to_air_idx.get_unchecked(executor_idx)
143            };
144            let executor = unsafe {
145                // SAFETY: same invariant as in `execute_instruction`.
146                self.inventory.executors.get_unchecked(executor_idx)
147            };
148            let opcode = executor.get_opcode_name(entry.insn.opcode.as_usize());
149            *counts.entry((air_idx, opcode)).or_insert(0) += freq as u64;
150        }
151        counts
152    }
153
154    /// Stopping is triggered by should_stop() or if VM is terminated
155    pub fn execute_from_state<RA>(
156        &mut self,
157        state: &mut VmExecState<F, TracingMemory, PreflightCtx<RA>>,
158    ) -> Result<(), ExecutionError>
159    where
160        RA: Arena,
161        E: PreflightExecutor<F, RA>,
162    {
163        loop {
164            if let Ok(Some(_)) = state.exit_code {
165                // should terminate
166                break;
167            }
168            if state.ctx.instret_left == 0 {
169                // should suspend
170                break;
171            }
172
173            // Fetch, decode and execute single instruction
174            self.execute_instruction(state)?;
175            state.ctx.instret_left -= 1;
176        }
177
178        Ok(())
179    }
180
181    /// Executes a single instruction and updates VM state
182    #[inline(always)]
183    fn execute_instruction<RA>(
184        &mut self,
185        state: &mut VmExecState<F, TracingMemory, PreflightCtx<RA>>,
186    ) -> Result<(), ExecutionError>
187    where
188        RA: Arena,
189        E: PreflightExecutor<F, RA>,
190    {
191        let pc = state.pc();
192        let pc_idx = get_pc_index(pc);
193        let pc_entry = self
194            .pc_handler
195            .get(pc_idx)
196            .ok_or_else(|| ExecutionError::PcOutOfBounds(pc))?;
197        // SAFETY: `execution_frequencies` has the same length as `pc_handler` so `get_pc_entry`
198        // already does the bounds check
199        unsafe {
200            *self.execution_frequencies.get_unchecked_mut(pc_idx) += 1;
201        };
202        tracing::trace!("pc: {pc:#x} | {:?}", pc_entry.insn);
203
204        if !pc_entry.is_some() {
205            return Err(ExecutionError::Unreachable(pc));
206        }
207
208        let opcode = pc_entry.insn.opcode;
209        let c = pc_entry.insn.c;
210        // Handle termination instruction
211        if opcode == SystemOpcode::TERMINATE.global_opcode() {
212            state.exit_code = Ok(Some(c.as_canonical_u32()));
213            return Ok(());
214        }
215
216        // SAFETY: non-system `executor_idx` values come from `ExecutorInventory`, which ensures
217        // that `executor_idx` is within bounds.
218        let executor = unsafe {
219            self.inventory
220                .executors
221                .get_unchecked(pc_entry.executor_idx as usize)
222        };
223
224        // Execute the instruction using the control implementation
225        tracing::trace!(
226            "opcode: {} | timestamp: {}",
227            executor.get_opcode_name(pc_entry.insn.opcode.as_usize()),
228            state.memory.timestamp()
229        );
230        let arena = unsafe {
231            // SAFETY: executor_idx is guarantee to be within bounds by ProgramHandler constructor
232            let air_idx = *self
233                .executor_idx_to_air_idx
234                .get_unchecked(pc_entry.executor_idx as usize);
235            // SAFETY: air_idx is a valid AIR index in the vkey, and always construct arenas with
236            // length equal to num_airs
237            state.ctx.arenas.get_unchecked_mut(air_idx)
238        };
239        let vm_state_mut = state.vm_state.into_mut(arena);
240        executor.execute(vm_state_mut, &pc_entry.insn)?;
241
242        #[cfg(feature = "metrics")]
243        {
244            crate::metrics::update_instruction_metrics(state, executor, pc, pc_entry);
245        }
246
247        Ok(())
248    }
249}
250
251impl<F> PcEntry<F> {
252    pub fn is_some(&self) -> bool {
253        self.executor_idx != u32::MAX
254    }
255}
256
257impl<F: Default> PcEntry<F> {
258    fn undefined() -> Self {
259        Self {
260            insn: Instruction::default(),
261            executor_idx: u32::MAX,
262        }
263    }
264}
265
266/// Macro for executing and emitting metrics for instructions/s and number of instructions executed.
267/// Does not include any tracing span.
268#[macro_export]
269macro_rules! execute_spanned {
270    ($name:literal, $executor:expr, $state:expr) => {{
271        #[cfg(feature = "metrics")]
272        let start = std::time::Instant::now();
273        #[cfg(feature = "metrics")]
274        let start_instret_left = $state.ctx.instret_left;
275
276        let result = $executor.execute_from_state($state);
277
278        #[cfg(feature = "metrics")]
279        {
280            let elapsed = start.elapsed();
281            let insns = start_instret_left - $state.ctx.instret_left;
282            tracing::info!("instructions_executed={insns}");
283            metrics::counter!(concat!($name, "_insns")).absolute(insns);
284            metrics::gauge!(concat!($name, "_insn_mi/s"))
285                .set(insns as f64 / elapsed.as_micros() as f64);
286        }
287        result
288    }};
289}
openvm_circuit/arch/interpreter_preflight.rs

openvm_circuit/arch/
interpreter_preflight.rs
