openvm_circuit/metrics/
mod.rs

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

openvm_circuit/metrics/
mod.rs
