feat: add cycle tracker report

.vscode/settings.json

@@ -24,6 +24,8 @@
         // // Examples.
         // "examples/chess/program/Cargo.toml",
         // "examples/chess/script/Cargo.toml",
+        // "examples/cycle-tracking/program/Cargo.toml",
+        // "examples/cycle-tracking/script/Cargo.toml",
         // "examples/fibonacci/program/Cargo.toml",
         // "examples/fibonacci/script/Cargo.toml",
         // "examples/io/program/Cargo.toml",

book/writing-programs/cycle-tracking.md

@@ -7,7 +7,7 @@ When writing a program, it is useful to know how many RISC-V cycles a portion of
 To track the number of cycles spent in a portion of the program, you can either put `println!("cycle-tracker-start: block name")` + `println!("cycle-tracker-end: block name")` statements (block name must be same between start and end) around the portion of your program you want to profile or use the `#[sp1_derive::cycle_tracker]` macro on a function. An example is shown below:
 
 ```rust,noplayground
-{{#include ../../examples/cycle-tracking/program/src/main.rs}}
+{{#include ../../examples/cycle-tracking/program/bin/normal.rs}}
 ```
 
 Note that to use the macro, you must add the `sp1-derive` crate to your dependencies for your program.
@@ -42,6 +42,19 @@ stdout: result: 2940
 
 Note that we elegantly handle nested cycle tracking, as you can see above.
 
+### Get Tracked Cycle Counts
+To include tracked cycle counts in the `ExecutionReport` when using `ProverClient::execute`, use the following annotations:
+
+```rust,noplayground
+fn main() {
+  println!("cycle-tracker-report-start: block name");
+  // ...
+  println!("cycle-tracker-report-end: block name");
+}
+```
+
+This will log the cycle count for `block name` and include it in the `ExecutionReport` in the `cycle_tracker` map.
+
 ## Tracking Cycles with Tracing
 
 The `cycle-tracker` annotation is a convenient way to track cycles for specific sections of code. However, sometimes it can also be useful to track what functions are taking the most cycles across the entire program, without having to annotate every function individually.

core/src/runtime/report.rs

@@ -10,6 +10,7 @@ use super::*;
 pub struct ExecutionReport {
     pub opcode_counts: HashMap<Opcode, u64>,
     pub syscall_counts: HashMap<SyscallCode, u64>,
+    pub cycle_tracker: HashMap<String, u64>,
 }
 
 impl ExecutionReport {

core/src/syscall/write.rs

@@ -1,5 +1,5 @@
 use crate::{
-    runtime::{Register, Syscall, SyscallContext},
+    runtime::{Register, Runtime, Syscall, SyscallContext},
     utils::num_to_comma_separated,
 };
 
@@ -12,6 +12,24 @@ impl SyscallWrite {
 }
 
 impl Syscall for SyscallWrite {
+    /// Handle writes to file descriptors during execution.
+    ///
+    /// If stdout (fd = 1):
+    /// - If the stream is a cycle tracker, either log the cycle tracker or accumulate it in the report.
+    /// - Else, print the stream to stdout.
+    ///
+    /// If stderr (fd = 2):
+    /// - Print the stream to stderr.
+    ///
+    /// If fd = 3:
+    /// - Update the public value stream.
+    ///
+    /// If fd = 4:
+    /// - Update the input stream.
+    ///
+    /// If the fd matches a hook in the hook registry, invoke the hook.
+    ///
+    /// Else, log a warning.
     fn execute(&self, ctx: &mut SyscallContext, arg1: u32, arg2: u32) -> Option<u32> {
         let a2 = Register::X12;
         let rt = &mut ctx.rt;
@@ -25,39 +43,16 @@ impl Syscall for SyscallWrite {
         let slice = bytes.as_slice();
         if fd == 1 {
             let s = core::str::from_utf8(slice).unwrap();
-            if s.contains("cycle-tracker-start:") {
-                let fn_name = s
-                    .split("cycle-tracker-start:")
-                    .last()
-                    .unwrap()
-                    .trim_end()
-                    .trim_start();
-                let depth = rt.cycle_tracker.len() as u32;
-                rt.cycle_tracker
-                    .insert(fn_name.to_string(), (rt.state.global_clk, depth));
-                let padding = (0..depth).map(|_| "│ ").collect::<String>();
-                log::info!("{}┌╴{}", padding, fn_name);
-            } else if s.contains("cycle-tracker-end:") {
-                let fn_name = s
-                    .split("cycle-tracker-end:")
-                    .last()
-                    .unwrap()
-                    .trim_end()
-                    .trim_start();
-                let (start, depth) = rt.cycle_tracker.remove(fn_name).unwrap_or((0, 0));
-                // Leftpad by 2 spaces for each depth.
-                let padding = (0..depth).map(|_| "│ ").collect::<String>();
-                log::info!(
-                    "{}└╴{} cycles",
-                    padding,
-                    num_to_comma_separated(rt.state.global_clk - start as u64)
-                );
-            } else {
-                let flush_s = update_io_buf(ctx, fd, s);
-                if !flush_s.is_empty() {
-                    flush_s
-                        .into_iter()
-                        .for_each(|line| println!("stdout: {}", line));
+            match parse_cycle_tracker_command(s) {
+                Some(command) => handle_cycle_tracker_command(rt, command),
+                None => {
+                    // If the string does not match any known command, print it to stdout.
+                    let flush_s = update_io_buf(ctx, fd, s);
+                    if !flush_s.is_empty() {
+                        flush_s
+                            .into_iter()
+                            .for_each(|line| println!("stdout: {}", line));
+                    }
                 }
             }
         } else if fd == 2 {
@@ -84,7 +79,79 @@ impl Syscall for SyscallWrite {
     }
 }
 
-pub fn update_io_buf(ctx: &mut SyscallContext, fd: u32, s: &str) -> Vec<String> {
+/// An enum representing the different cycle tracker commands.
+#[derive(Clone)]
+enum CycleTrackerCommand {
+    Start(String),
+    End(String),
+    ReportStart(String),
+    ReportEnd(String),
+}
+
+/// Parse a cycle tracker command from a string. If the string does not match any known command, returns None.
+fn parse_cycle_tracker_command(s: &str) -> Option<CycleTrackerCommand> {
+    let (command, fn_name) = s.split_once(':')?;
+    let trimmed_name = fn_name.trim().to_string();
+
+    match command {
+        "cycle-tracker-start" => Some(CycleTrackerCommand::Start(trimmed_name)),
+        "cycle-tracker-end" => Some(CycleTrackerCommand::End(trimmed_name)),
+        "cycle-tracker-report-start" => Some(CycleTrackerCommand::ReportStart(trimmed_name)),
+        "cycle-tracker-report-end" => Some(CycleTrackerCommand::ReportEnd(trimmed_name)),
+        _ => None,
+    }
+}
+
+/// Handle a cycle tracker command.
+fn handle_cycle_tracker_command(rt: &mut Runtime, command: CycleTrackerCommand) {
+    match command {
+        CycleTrackerCommand::Start(name) | CycleTrackerCommand::ReportStart(name) => {
+            start_cycle_tracker(rt, &name)
+        }
+        CycleTrackerCommand::End(name) => {
+            end_cycle_tracker(rt, &name);
+        }
+        CycleTrackerCommand::ReportEnd(name) => {
+            // Attempt to end the cycle tracker and accumulate the total cycles in the fn_name's
+            // entry in the ExecutionReport.
+            if let Some(total_cycles) = end_cycle_tracker(rt, &name) {
+                rt.report
+                    .cycle_tracker
+                    .entry(name.to_string())
+                    .and_modify(|cycles| *cycles += total_cycles)
+                    .or_insert(total_cycles);
+            }
+        }
+    }
+}
+
+/// Start tracking cycles for the given name at the specific depth and print out the log.
+fn start_cycle_tracker(rt: &mut Runtime, name: &str) {
+    let depth = rt.cycle_tracker.len() as u32;
+    rt.cycle_tracker
+        .insert(name.to_string(), (rt.state.global_clk, depth));
+    let padding = "│ ".repeat(depth as usize);
+    log::info!("{}┌╴{}", padding, name);
+}
+
+/// End tracking cycles for the given name, print out the log, and return the total number of cycles in the span.
+/// If the name is not found in the cycle tracker cache, returns None.
+fn end_cycle_tracker(rt: &mut Runtime, name: &str) -> Option<u64> {
+    if let Some((start, depth)) = rt.cycle_tracker.remove(name) {
+        let padding = "│ ".repeat(depth as usize);
+        let total_cycles = rt.state.global_clk - start as u64;
+        log::info!(
+            "{}└╴{} cycles",
+            padding,
+            num_to_comma_separated(total_cycles)
+        );
+        return Some(total_cycles);
+    }
+    None
+}
+
+/// Update the io buffer for the given file descriptor with the given string.
+fn update_io_buf(ctx: &mut SyscallContext, fd: u32, s: &str) -> Vec<String> {
     let rt = &mut ctx.rt;
     let entry = rt.io_buf.entry(fd).or_default();
     entry.push_str(s);

examples/cycle-tracking/program/bin/normal.rs

@@ -0,0 +1,32 @@
+#![no_main]
+sp1_zkvm::entrypoint!(main);
+
+#[sp1_derive::cycle_tracker]
+pub fn expensive_function(x: usize) -> usize {
+    let mut y = 1;
+    for _ in 0..100 {
+        y *= x;
+        y %= 7919;
+    }
+    y
+}
+
+pub fn main() {
+    let mut nums = vec![1, 1];
+
+    // Setup a large vector with Fibonacci-esque numbers.
+    println!("cycle-tracker-start: setup");
+    for _ in 0..100 {
+        let mut c = nums[nums.len() - 1] + nums[nums.len() - 2];
+        c %= 7919;
+        nums.push(c);
+    }
+    println!("cycle-tracker-end: setup");
+
+    println!("cycle-tracker-start: main-body");
+    for i in 0..2 {
+        let result = expensive_function(nums[nums.len() - i - 1]);
+        println!("result: {}", result);
+    }
+    println!("cycle-tracker-end: main-body");
+}

examples/cycle-tracking/program/bin/report.rs

@@ -0,0 +1,25 @@
+#![no_main]
+sp1_zkvm::entrypoint!(main);
+
+#[sp1_derive::cycle_tracker]
+pub fn expensive_function(x: usize) -> usize {
+    let mut y = 1;
+    for _ in 0..100 {
+        y *= x;
+        y %= 7919;
+    }
+    y
+}
+
+pub fn main() {
+    let mut nums = vec![1, 1];
+
+    // Setup a large vector with Fibonacci-esque numbers.
+    println!("cycle-tracker-report-start: setup");
+    for _ in 0..100 {
+        let mut c = nums[nums.len() - 1] + nums[nums.len() - 2];
+        c %= 7919;
+        nums.push(c);
+    }
+    println!("cycle-tracker-report-end: setup");
+}