openvm_circuit/metrics/
mod.rs

Help
1use std::{collections::BTreeMap, mem};
2
3use cycle_tracker::CycleTracker;
4use metrics::counter;
5use openvm_instructions::{
6    exe::{FnBound, FnBounds},
7    VmOpcode,
8};
9use openvm_stark_backend::p3_field::PrimeField32;
10
11use crate::arch::{ExecutionSegment, InstructionExecutor, VmConfig};
12
13pub mod cycle_tracker;
14
15#[derive(Clone, Debug, Default)]
16pub struct VmMetrics {
17    pub cycle_count: usize,
18    pub chip_heights: Vec<(String, usize)>,
19    /// Maps (dsl_ir, opcode) to number of times opcode was executed
20    pub counts: BTreeMap<(Option<String>, String), usize>,
21    /// Maps (dsl_ir, opcode, air_name) to number of trace cells generated by opcode
22    pub trace_cells: BTreeMap<(Option<String>, String, String), usize>,
23    /// Metric collection tools. Only collected when `config.profiling` is true.
24    pub cycle_tracker: CycleTracker,
25    #[allow(dead_code)]
26    pub(crate) fn_bounds: FnBounds,
27    /// Cycle span by function if function start/end addresses are available
28    #[allow(dead_code)]
29    pub(crate) current_fn: FnBound,
30    pub(crate) current_trace_cells: Vec<usize>,
31}
32
33impl<F, VC> ExecutionSegment<F, VC>
34where
35    F: PrimeField32,
36    VC: VmConfig<F>,
37{
38    /// Update metrics that increment per instruction
39    #[allow(unused_variables)]
40    pub fn update_instruction_metrics(
41        &mut self,
42        pc: u32,
43        opcode: VmOpcode,
44        dsl_instr: Option<String>,
45    ) {
46        self.metrics.cycle_count += 1;
47
48        if self.system_config().profiling {
49            let executor = self.chip_complex.inventory.get_executor(opcode).unwrap();
50            let opcode_name = executor.get_opcode_name(opcode.as_usize());
51            self.metrics.update_trace_cells(
52                &self.air_names,
53                self.current_trace_cells(),
54                opcode_name,
55                dsl_instr,
56            );
57
58            #[cfg(feature = "function-span")]
59            self.metrics.update_current_fn(pc);
60        }
61    }
62}
63
64impl VmMetrics {
65    fn update_trace_cells(
66        &mut self,
67        air_names: &[String],
68        now_trace_cells: Vec<usize>,
69        opcode_name: String,
70        dsl_instr: Option<String>,
71    ) {
72        let key = (dsl_instr, opcode_name);
73        self.cycle_tracker.increment_opcode(&key);
74        *self.counts.entry(key.clone()).or_insert(0) += 1;
75
76        for (air_name, now_value, prev_value) in
77            itertools::izip!(air_names, &now_trace_cells, &self.current_trace_cells)
78        {
79            if prev_value != now_value {
80                let key = (key.0.clone(), key.1.clone(), air_name.to_owned());
81                self.cycle_tracker
82                    .increment_cells_used(&key, now_value - prev_value);
83                *self.trace_cells.entry(key).or_insert(0) += now_value - prev_value;
84            }
85        }
86        self.current_trace_cells = now_trace_cells;
87    }
88
89    /// Take the cycle tracker and fn bounds information for use in
90    /// next segment. Leave the rest of the metrics for recording purposes.
91    pub fn partial_take(&mut self) -> Self {
92        Self {
93            cycle_tracker: mem::take(&mut self.cycle_tracker),
94            fn_bounds: mem::take(&mut self.fn_bounds),
95            current_fn: mem::take(&mut self.current_fn),
96            ..Default::default()
97        }
98    }
99
100    /// Clear statistics that are local to a segment
101    // Important: chip and cycle count metrics should start over for SegmentationStrategy,
102    // but we need to carry over the cycle tracker so spans can cross segments
103    pub fn clear(&mut self) {
104        *self = self.partial_take();
105    }
106
107    #[cfg(feature = "function-span")]
108    fn update_current_fn(&mut self, pc: u32) {
109        if !self.fn_bounds.is_empty() && (pc < self.current_fn.start || pc > self.current_fn.end) {
110            self.current_fn = self
111                .fn_bounds
112                .range(..=pc)
113                .next_back()
114                .map(|(_, func)| (*func).clone())
115                .unwrap();
116            if pc == self.current_fn.start {
117                self.cycle_tracker.start(self.current_fn.name.clone());
118            } else {
119                self.cycle_tracker.force_end();
120            }
121        };
122    }
123    pub fn emit(&self) {
124        for (name, value) in self.chip_heights.iter() {
125            let labels = [("chip_name", name.clone())];
126            counter!("rows_used", &labels).absolute(*value as u64);
127        }
128
129        for ((dsl_ir, opcode), value) in self.counts.iter() {
130            let labels = [
131                ("dsl_ir", dsl_ir.clone().unwrap_or_else(String::new)),
132                ("opcode", opcode.clone()),
133            ];
134            counter!("frequency", &labels).absolute(*value as u64);
135        }
136
137        for ((dsl_ir, opcode, air_name), value) in self.trace_cells.iter() {
138            let labels = [
139                ("dsl_ir", dsl_ir.clone().unwrap_or_else(String::new)),
140                ("opcode", opcode.clone()),
141                ("air_name", air_name.clone()),
142            ];
143            counter!("cells_used", &labels).absolute(*value as u64);
144        }
145    }
146}
openvm_circuit/metrics/mod.rs

openvm_circuit/metrics/
mod.rs
