report (dsl,opcode) usage in cycle tracker

Author: kunxian-xia

crates/vm/src/metrics/cycle_tracker/mod.rs

@@ -1,23 +1,44 @@
+use std::collections::BTreeMap;
+
 /// Stats for a nested span in the execution segment that is tracked by the [`CycleTracker`].
 #[derive(Clone, Debug, Default)]
 pub struct SpanInfo {
     /// The name of the span.
     pub tag: String,
     /// The cycle count at which the span starts.
     pub start: usize,
+    /// Maps (dsl_ir, opcode) to number of times opcode was executed
+    pub counts: BTreeMap<(Option<String>, String), usize>,
 }
 
-#[derive(Clone, Debug, Default)]
+#[derive(Clone, Debug)]
 pub struct CycleTracker {
     /// Stack of span names, with most recent at the end
     stack: Vec<SpanInfo>,
     /// Depth of the stack.
     depth: usize,
+    max_depth: usize,
+}
+
+impl Default for CycleTracker {
+    fn default() -> Self {
+        Self {
+            stack: Vec::new(),
+            depth: 0,
+            max_depth: std::env::var("CYCLE_TRACKER_MAX_DEPTH")
+                .ok()
+                .and_then(|s| s.parse().ok())
+                .unwrap_or(2),
+        }
+    }
 }
 
 impl CycleTracker {
-    pub fn new() -> Self {
-        Self::default()
+    pub fn new(max_depth: usize) -> Self {
+        Self {
+            max_depth,
+            ..Default::default()
+        }
     }
 
     pub fn top(&self) -> Option<&String> {
@@ -30,33 +51,48 @@ impl CycleTracker {
     /// Starts a new cycle tracker span for the given name.
     /// If a span already exists for the given name, it ends the existing span and pushes a new one
     /// to the vec.
-    pub fn start(&mut self, mut name: String, cycles_count: usize) {
+    pub fn start(&mut self, mut name: String, cycles_count: usize, num_insns_by_dsl: &BTreeMap<(Option<String>, String), usize>) {
         // hack to remove "CT-" prefix
         if name.starts_with("CT-") {
             name = name.split_off(3);
         }
+        self.depth += 1;
+        if self.depth > self.max_depth {
+            return;
+        }
         self.stack.push(SpanInfo {
             tag: name.clone(),
             start: cycles_count,
+            counts: num_insns_by_dsl.clone(),
         });
         let padding = "│ ".repeat(self.depth);
         tracing::info!("{}┌╴{}", padding, name);
-        self.depth += 1;
     }
 
     /// Ends the cycle tracker span for the given name.
     /// If no span exists for the given name, it panics.
-    pub fn end(&mut self, mut name: String, cycles_count: usize) {
+    pub fn end(&mut self, mut name: String, cycles_count: usize, num_insns_by_dsl: &BTreeMap<(Option<String>, String), usize>) {
         // hack to remove "CT-" prefix
         if name.starts_with("CT-") {
             name = name.split_off(3);
         }
-        let SpanInfo { tag, start } = self.stack.pop().unwrap();
-        assert_eq!(tag, name, "Stack top does not match name");
         self.depth -= 1;
+        if self.depth >= self.max_depth {
+            return;
+        }
+        let SpanInfo { tag, start, counts: num_insns_start } = self.stack.pop().unwrap();
+        assert_eq!(tag, name, "Stack top does not match name");
         let padding = "│ ".repeat(self.depth);
         let span_cycles = cycles_count - start;
-        tracing::info!("{}└╴{} cycles", padding, span_cycles);
+        for (dsl_opcode, num_insns) in num_insns_by_dsl {
+            let start_count = num_insns_start.get(dsl_opcode).cloned().unwrap_or(0);
+            let span_count = num_insns - start_count;
+            if span_count > 0 {
+                tracing::info!("{}│   ({:?},{}): {} instructions", padding, dsl_opcode.0, dsl_opcode.1, span_count);
+
+            }
+        }
+        tracing::info!("{}└╴({}) {} cycles", padding, name, span_cycles);
     }
 
     /// Ends the current cycle tracker span.

crates/vm/src/metrics/mod.rs

@@ -40,6 +40,7 @@ pub struct VmMetrics {
     /// Metric collection tools. Only collected when "perf-metrics" feature is enabled.
     pub cycle_tracker: CycleTracker,
 
+    pub(crate) num_insns_executed: usize,
     pub(crate) current_trace_cells: Vec<usize>,
 
     /// Backtrace for guest debug panic display
@@ -68,6 +69,7 @@ pub fn update_instruction_metrics<F, RA, Executor>(
     {
         let pc = state.pc();
         state.metrics.update_backtrace(pc);
+        state.metrics.num_insns_executed += 1;
     }
 
     #[cfg(feature = "perf-metrics")]
@@ -224,13 +226,13 @@ impl VmMetrics {
                 .map(|(_, func)| (*func).clone())
                 .unwrap();
             if pc == self.current_fn.start {
-                self.cycle_tracker.start(self.current_fn.name.clone(), 0);
+                // self.cycle_tracker.start(self.current_fn.name.clone(), 0);
             } else {
                 while let Some(name) = self.cycle_tracker.top() {
                     if name == &self.current_fn.name {
                         break;
                     }
-                    self.cycle_tracker.force_end();
+                    // self.cycle_tracker.force_end();
                 }
             }
         };

crates/vm/src/system/phantom/mod.rs

@@ -157,14 +157,14 @@ where
                 SysPhantom::CtStart => {
                     let metrics = state.metrics;
                     if let Some(info) = metrics.debug_infos.get(pc) {
-                        metrics.cycle_tracker.start(info.dsl_instruction.clone());
+                        metrics.cycle_tracker.start(info.dsl_instruction.clone(), metrics.num_insns_executed, &metrics.counts);
                     }
                 }
                 #[cfg(feature = "perf-metrics")]
                 SysPhantom::CtEnd => {
                     let metrics = state.metrics;
                     if let Some(info) = metrics.debug_infos.get(pc) {
-                        metrics.cycle_tracker.end(info.dsl_instruction.clone());
+                        metrics.cycle_tracker.end(info.dsl_instruction.clone(), metrics.num_insns_executed, &metrics.counts);
                     }
                 }
                 _ => {}

extensions/native/compiler/src/constraints/halo2/compiler.rs

@@ -136,7 +136,10 @@ impl<C: Config + Debug> Halo2ConstraintCompiler<C> {
         C: Config<N = Bn254Fr, F = BabyBear, EF = BabyBearExt4>,
     {
         #[cfg(feature = "metrics")]
-        let mut cell_tracker = CycleTracker::new();
+        let mut cell_tracker = CycleTracker::new(std::env::var("CYCLE_TRACKER_MAX_DEPTH")
+            .ok()
+            .and_then(|s| s.parse().ok())
+            .unwrap_or(2));
         let range = Arc::new(halo2_state.builder.range_chip());
         let f_chip = Arc::new(BabyBearChip::new(range.clone()));
         let ext_chip = BabyBearExt4Chip::new(Arc::clone(&f_chip));