line level and instruction level profiling

Author: morganthomas

crates/lean_vm/src/diagnostics/profiler.rs

@@ -21,7 +21,7 @@ pub(crate) fn profiling_report(
     let mut call_stack: Vec<String> = Vec::new();
     let mut prev_function_name = String::new();
 
-    for (&line_num, &cycle_count) in instructions.lines.iter().zip(&instructions.cycles) {
+    for (&line_num, &cycle_count) in instructions.lines.iter().zip(&instructions.lines_cycles) {
         let (_, current_function_name) = find_function_for_line(line_num, function_locations);
 
         if prev_function_name != current_function_name {

crates/lean_vm/src/execution/context.rs

@@ -1,10 +1,10 @@
-use crate::core::SourceLineNumber;
+use crate::core::{SourceLineNumber};
 use std::collections::BTreeMap;
 
 #[derive(Debug, Clone, Default)]
 pub struct ExecutionHistory {
     pub lines: Vec<SourceLineNumber>,
-    pub cycles: Vec<usize>, // for each line, how many cycles it took
+    pub lines_cycles: Vec<usize>, // for each line, how many cycles it took
 }
 
 impl ExecutionHistory {
@@ -14,11 +14,11 @@ impl ExecutionHistory {
 
     pub fn add_line(&mut self, location: SourceLineNumber, cycles: usize) {
         self.lines.push(location);
-        self.cycles.push(cycles);
+        self.lines_cycles.push(cycles);
     }
 
     pub fn total_cycles(&self) -> usize {
-        self.cycles.iter().sum()
+        self.lines_cycles.iter().sum()
     }
 
     pub const fn len(&self) -> usize {

crates/lean_vm/src/execution/runner.rs

@@ -1,6 +1,6 @@
 //! VM execution runner
 
-use crate::HintExecutionContext;
+use crate::{CodeAddress, SourceLineNumber, HintExecutionContext};
 use crate::core::{
     DIMENSION, F, ONE_VEC_PTR, POSEIDON_16_NULL_HASH_PTR, POSEIDON_24_NULL_HASH_PTR,
     PUBLIC_INPUT_START, VECTOR_LEN, ZERO_VEC_PTR,
@@ -72,36 +72,82 @@ pub fn execute_bytecode(
 ) -> ExecutionResult {
     let mut std_out = String::new();
     let mut instruction_history = ExecutionHistory::new();
-    execute_bytecode_helper(
-        bytecode,
-        public_input,
-        private_input,
-        no_vec_runtime_memory,
-        &mut std_out,
-        &mut instruction_history,
-        profiler,
-        function_locations,
-    )
-    .unwrap_or_else(|err| {
-        let lines_history = &instruction_history.lines;
-        let latest_instructions =
-            &lines_history[lines_history.len().saturating_sub(STACK_TRACE_INSTRUCTIONS)..];
-        println!(
-            "\n{}",
-            crate::diagnostics::pretty_stack_trace(
-                source_code,
-                latest_instructions,
-                function_locations
-            )
-        );
-        if !std_out.is_empty() {
-            println!("╔══════════════════════════════════════════════════════════════╗");
-            println!("║                         STD-OUT                              ║");
-            println!("╚══════════════════════════════════════════════════════════════╝\n");
-            print!("{std_out}");
+    let result =
+        execute_bytecode_helper(
+            bytecode,
+            public_input,
+            private_input,
+            no_vec_runtime_memory,
+            &mut std_out,
+            &mut instruction_history,
+            profiler,
+            function_locations,
+        )
+        .unwrap_or_else(|err| {
+            let lines_history = &instruction_history.lines;
+            let latest_instructions =
+                &lines_history[lines_history.len().saturating_sub(STACK_TRACE_INSTRUCTIONS)..];
+            println!(
+                "\n{}",
+                crate::diagnostics::pretty_stack_trace(
+                    source_code,
+                    latest_instructions,
+                    function_locations
+                )
+            );
+            if !std_out.is_empty() {
+                println!("╔══════════════════════════════════════════════════════════════╗");
+                println!("║                         STD-OUT                              ║");
+                println!("╚══════════════════════════════════════════════════════════════╝\n");
+                print!("{std_out}");
+            }
+            panic!("Error during bytecode execution: {err}");
+        });
+    if profiler {
+        print_line_cycle_counts(instruction_history);
+        print_instruction_cycle_counts(bytecode, result.pcs.clone());
+    }
+    result
+}
+
+fn print_line_cycle_counts(history: ExecutionHistory) {
+    println!("Line by line cycle counts");
+    println!("=========================\n");
+
+    let mut gross_cycle_counts: BTreeMap<SourceLineNumber, usize> = BTreeMap::new();
+    for (line, cycle_count) in history.lines.iter().zip(history.lines_cycles.iter()) {
+        let prev_count = gross_cycle_counts.get(line).unwrap_or(&0);
+        gross_cycle_counts.insert(*line, *prev_count + cycle_count);
+    }
+    for (line, cycle_count) in gross_cycle_counts.iter() {
+        println!("line {line}: {cycle_count} cycles");
+    }
+    println!("");
+}
+
+fn print_instruction_cycle_counts(bytecode: &Bytecode, pcs: Vec<CodeAddress>) {
+    println!("Instruction level cycle counts");
+    println!("==============================");
+
+    let mut gross_cycle_counts: BTreeMap<CodeAddress, usize> = BTreeMap::new();
+    for pc in pcs.iter() {
+        let prev_count = gross_cycle_counts.get(pc).unwrap_or(&0);
+        gross_cycle_counts.insert(*pc, *prev_count + 1);
+    }
+    for (pc, cycle_count) in gross_cycle_counts.iter() {
+        let instruction = &bytecode.instructions[*pc];
+        let hints = bytecode.hints.get(pc);
+        match hints {
+            Some(hints) => {
+                for hint in hints {
+                    println!("hint: {hint}");
+                }
+            },
+            None => {},
         }
-        panic!("Error during bytecode execution: {err}");
-    })
+        println!("pc {pc}: {cycle_count} cycles: {instruction}");
+    }
+    println!("");
 }
 
 /// Helper function that performs the actual bytecode execution

crates/lean_vm/src/isa/hint.rs

@@ -187,7 +187,7 @@ impl Hint {
             Self::LocationReport { location } => {
                 ctx.instruction_history.lines.push(*location);
                 ctx.instruction_history
-                    .cycles
+                    .lines_cycles
                     .push(*ctx.cpu_cycles_before_new_line);
                 *ctx.cpu_cycles_before_new_line = 0;
             }