openvm_circuit/metrics/
mod.rs

1use std::{collections::BTreeMap, mem};
2
3use backtrace::Backtrace;
4use cycle_tracker::CycleTracker;
5#[cfg(feature = "perf-metrics")]
6use itertools::Itertools;
7use metrics::counter;
8use openvm_instructions::{
9    exe::{FnBound, FnBounds},
10    program::ProgramDebugInfo,
11};
12use openvm_stark_backend::prover::{DeviceMultiStarkProvingKey, ProverBackend};
13
14use crate::{
15    arch::{
16        execution_mode::PreflightCtx, interpreter_preflight::PcEntry, Arena, PreflightExecutor,
17        VmExecState,
18    },
19    system::memory::online::TracingMemory,
20};
21
22pub mod cycle_tracker;
23
24#[derive(Clone, Debug, Default)]
25pub struct VmMetrics {
26    // Static info
27    pub air_names: Vec<String>,
28    pub debug_infos: ProgramDebugInfo,
29    #[cfg(feature = "perf-metrics")]
30    pub(crate) num_sys_airs: usize,
31    #[cfg(feature = "perf-metrics")]
32    pub(crate) main_widths: Vec<usize>,
33    #[cfg(feature = "perf-metrics")]
34    pub(crate) total_widths: Vec<usize>,
35
36    // Dynamic stats
37    /// Maps (dsl_ir, opcode) to number of times opcode was executed
38    pub counts: BTreeMap<(Option<String>, String), usize>,
39    /// Maps (dsl_ir, opcode, air_name) to number of trace cells generated by opcode
40    pub trace_cells: BTreeMap<(Option<String>, String, String), usize>,
41    /// Metric collection tools. Only collected when "perf-metrics" feature is enabled.
42    pub cycle_tracker: CycleTracker,
43
44    #[cfg(feature = "perf-metrics")]
45    pub(crate) current_trace_cells: Vec<usize>,
46
47    /// Backtrace for guest debug panic display
48    pub prev_backtrace: Option<Backtrace>,
49    #[allow(dead_code)]
50    pub(crate) fn_bounds: FnBounds,
51    /// Cycle span by function if function start/end addresses are available
52    #[allow(dead_code)]
53    pub(crate) current_fn: FnBound,
54}
55
56/// We assume this will be called after execute_instruction, so less error-handling is needed.
57#[allow(unused_variables)]
58#[inline(always)]
59pub fn update_instruction_metrics<F, RA, Executor>(
60    state: &mut VmExecState<F, TracingMemory, PreflightCtx<RA>>,
61    executor: &Executor,
62    prev_pc: u32, // the pc of the instruction executed, state.pc is next pc
63    pc_entry: &PcEntry<F>,
64) where
65    F: Clone + Send + Sync,
66    RA: Arena,
67    Executor: PreflightExecutor<F, RA>,
68{
69    #[cfg(all(feature = "metrics", any(debug_assertions, feature = "perf-metrics")))]
70    {
71        let pc = state.pc();
72        state.metrics.update_backtrace(pc);
73    }
74
75    #[cfg(feature = "perf-metrics")]
76    {
77        use std::iter::zip;
78
79        let pc = state.pc();
80        let opcode = pc_entry.insn.opcode;
81        let opcode_name = executor.get_opcode_name(opcode.as_usize());
82
83        let debug_info = state.metrics.debug_infos.get(prev_pc);
84        let dsl_instr = debug_info.as_ref().map(|info| info.dsl_instruction.clone());
85
86        let now_trace_heights: Vec<usize> = state
87            .ctx
88            .arenas
89            .iter()
90            .map(|arena| arena.current_trace_height())
91            .collect();
92        let now_trace_cells = zip(&state.metrics.main_widths, &now_trace_heights)
93            .map(|(main_width, h)| main_width * h)
94            .collect_vec();
95        state
96            .metrics
97            .update_trace_cells(now_trace_cells, opcode_name, dsl_instr);
98
99        state.metrics.update_current_fn(pc);
100    }
101}
102
103// We clear the current trace cell counts so there aren't negative diffs at the start of the next
104// segment.
105#[cfg(feature = "perf-metrics")]
106pub fn end_segment_metrics<F, RA>(state: &mut VmExecState<F, TracingMemory, PreflightCtx<RA>>)
107where
108    F: Clone + Send + Sync,
109    RA: Arena,
110{
111    state.metrics.current_trace_cells.fill(0);
112}
113
114#[cfg(feature = "metrics")]
115pub fn emit_opcode_counts(metrics: &VmMetrics, counts: BTreeMap<(usize, String), u64>) {
116    for ((air_idx, opcode), count) in counts {
117        let Some(air_name) = metrics.air_names.get(air_idx) else {
118            continue;
119        };
120        let labels = [
121            ("air_name", air_name.clone()),
122            ("air_id", air_idx.to_string()),
123            ("opcode", opcode),
124        ];
125        counter!("opcode_count", &labels).absolute(count);
126    }
127}
128
129impl VmMetrics {
130    #[cfg(feature = "metrics")]
131    pub fn set_pk_air_names<PB: ProverBackend>(&mut self, pk: &DeviceMultiStarkProvingKey<PB>) {
132        self.air_names = pk.per_air.iter().map(|pk| pk.air_name.clone()).collect();
133    }
134
135    #[cfg(feature = "perf-metrics")]
136    pub fn set_pk_trace_info<PB: ProverBackend>(&mut self, pk: &DeviceMultiStarkProvingKey<PB>) {
137        let (main_widths, total_widths): (Vec<_>, Vec<_>) = pk
138            .per_air
139            .iter()
140            .map(|pk| {
141                let width = &pk.vk.params.width;
142                (width.main_width(), width.total_width())
143            })
144            .unzip();
145        self.main_widths = main_widths;
146        self.total_widths = total_widths;
147        self.current_trace_cells = vec![0; self.air_names.len()];
148    }
149
150    #[cfg(feature = "perf-metrics")]
151    pub fn update_trace_cells(
152        &mut self,
153        now_trace_cells: Vec<usize>,
154        opcode_name: String,
155        dsl_instr: Option<String>,
156    ) {
157        let key = (dsl_instr, opcode_name.clone());
158        self.cycle_tracker.increment_opcode(&key);
159        *self.counts.entry(key.clone()).or_insert(0) += 1;
160
161        for (air_name, now_value, prev_value) in
162            itertools::izip!(&self.air_names, &now_trace_cells, &self.current_trace_cells)
163        {
164            if prev_value != now_value {
165                let cells_used = now_value - prev_value;
166                let key = (key.0.clone(), key.1.clone(), air_name.to_owned());
167                self.cycle_tracker.increment_cells_used(&key, cells_used);
168                *self.trace_cells.entry(key).or_insert(0) += cells_used;
169            }
170        }
171        self.current_trace_cells = now_trace_cells;
172    }
173
174    /// Take the cycle tracker and fn bounds information for use in
175    /// next segment. Leave the rest of the metrics for recording purposes.
176    pub fn partial_take(&mut self) -> Self {
177        Self {
178            cycle_tracker: mem::take(&mut self.cycle_tracker),
179            fn_bounds: mem::take(&mut self.fn_bounds),
180            current_fn: mem::take(&mut self.current_fn),
181            ..Default::default()
182        }
183    }
184
185    /// Clear statistics that are local to a segment
186    // Important: chip and cycle count metrics should start over for SegmentationStrategy,
187    // but we need to carry over the cycle tracker so spans can cross segments
188    pub fn clear(&mut self) {
189        *self = self.partial_take();
190    }
191
192    #[cfg(all(feature = "metrics", any(debug_assertions, feature = "perf-metrics")))]
193    pub fn update_backtrace(&mut self, pc: u32) {
194        if let Some(info) = self.debug_infos.get(pc) {
195            if let Some(trace) = &info.trace {
196                self.prev_backtrace = Some(trace.clone());
197            }
198        }
199    }
200
201    #[cfg(feature = "perf-metrics")]
202    pub(super) fn update_current_fn(&mut self, pc: u32) {
203        if self.fn_bounds.is_empty() {
204            return;
205        }
206        if pc < self.current_fn.start || pc > self.current_fn.end {
207            self.current_fn = self
208                .fn_bounds
209                .range(..=pc)
210                .next_back()
211                .map(|(_, func)| (*func).clone())
212                .unwrap();
213            if pc == self.current_fn.start {
214                self.cycle_tracker.start(self.current_fn.name.clone());
215            } else {
216                while let Some(name) = self.cycle_tracker.top() {
217                    if name == &self.current_fn.name {
218                        break;
219                    }
220                    self.cycle_tracker.force_end();
221                }
222            }
223        };
224    }
225
226    pub fn emit(&self) {
227        for ((dsl_ir, opcode), value) in self.counts.iter() {
228            let labels = [
229                ("dsl_ir", dsl_ir.clone().unwrap_or_else(String::new)),
230                ("opcode", opcode.clone()),
231            ];
232            counter!("frequency", &labels).absolute(*value as u64);
233        }
234
235        for ((dsl_ir, opcode, air_name), value) in self.trace_cells.iter() {
236            let labels = [
237                ("dsl_ir", dsl_ir.clone().unwrap_or_else(String::new)),
238                ("opcode", opcode.clone()),
239                ("air_name", air_name.clone()),
240            ];
241            counter!("cells_used", &labels).absolute(*value as u64);
242        }
243    }
244}
openvm_circuit/metrics/mod.rs

openvm_circuit/metrics/
mod.rs
