feat(openvm_circuit): store span info in cycle tracker for convenient debug info

Author: roynalnaruto

crates/vm/src/arch/segment.rs

@@ -281,16 +281,18 @@ impl<F: PrimeField32, VC: VmConfig<F>> ExecutionSegment<F, VC> {
                         Some(SysPhantom::CtStart) =>
                         {
                             #[cfg(feature = "bench-metrics")]
-                            metrics
-                                .cycle_tracker
-                                .start(dsl_instr.cloned().unwrap_or("Default".to_string()))
+                            metrics.cycle_tracker.start(
+                                dsl_instr.cloned().unwrap_or("Default".to_string()),
+                                metrics.cycle_count,
+                            )
                         }
                         Some(SysPhantom::CtEnd) =>
                         {
                             #[cfg(feature = "bench-metrics")]
-                            metrics
-                                .cycle_tracker
-                                .end(dsl_instr.cloned().unwrap_or("Default".to_string()))
+                            metrics.cycle_tracker.end(
+                                dsl_instr.cloned().unwrap_or("Default".to_string()),
+                                metrics.cycle_count,
+                            )
                         }
                         _ => {}
                     }

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

@@ -1,7 +1,18 @@
+/// 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.
+    tag: String,
+    /// The cycle count at which the span starts.
+    start: usize,
+}
+
 #[derive(Clone, Debug, Default)]
 pub struct CycleTracker {
     /// Stack of span names, with most recent at the end
-    stack: Vec<String>,
+    stack: Vec<SpanInfo>,
+    /// Depth of the stack.
+    depth: usize,
 }
 
 impl CycleTracker {
@@ -11,23 +22,33 @@ 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) {
+    pub fn start(&mut self, mut name: String, cycles_count: usize) {
         // hack to remove "CT-" prefix
         if name.starts_with("CT-") {
             name = name.split_off(3);
         }
-        self.stack.push(name);
+        self.stack.push(SpanInfo {
+            tag: name.clone(),
+            start: cycles_count,
+        });
+        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) {
+    pub fn end(&mut self, mut name: String, cycles_count: usize) {
         // hack to remove "CT-" prefix
         if name.starts_with("CT-") {
             name = name.split_off(3);
         }
-        let stack_top = self.stack.pop();
-        assert_eq!(stack_top.unwrap(), name, "Stack top does not match name");
+        let SpanInfo { tag, start } = self.stack.pop().unwrap();
+        assert_eq!(tag, name, "Stack top does not match name");
+        self.depth -= 1;
+        let padding = "│ ".repeat(self.depth);
+        let span_cycles = cycles_count - start;
+        tracing::info!("{}└╴{} cycles", padding, span_cycles);
     }
 
     /// Ends the current cycle tracker span.
@@ -37,7 +58,11 @@ impl CycleTracker {
 
     /// Get full name of span with all parent names separated by ";" in flamegraph format
     pub fn get_full_name(&self) -> String {
-        self.stack.join(";")
+        self.stack
+            .iter()
+            .map(|span_info| span_info.tag.clone())
+            .collect::<Vec<String>>()
+            .join(";")
     }
 }
 

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

@@ -492,11 +492,11 @@ impl<C: Config + Debug> Halo2ConstraintCompiler<C> {
                     }
                     DslIr::CycleTrackerStart(_name) => {
                         #[cfg(feature = "bench-metrics")]
-                        cell_tracker.start(_name);
+                        cell_tracker.start(_name, 0);
                     }
                     DslIr::CycleTrackerEnd(_name) => {
                         #[cfg(feature = "bench-metrics")]
-                        cell_tracker.end(_name);
+                        cell_tracker.end(_name, 0);
                     }
                     DslIr::CircuitPublish(val, index) => {
                         public_values[index] = vars[&val.0];