diff --git a/.vscode/settings.json b/.vscode/settings.json index ad9ab59b7b..7e4fd0ca18 100644 --- a/.vscode/settings.json +++ b/.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", diff --git a/book/writing-programs/cycle-tracking.md b/book/writing-programs/cycle-tracking.md index 171a36d3ff..176249b45e 100644 --- a/book/writing-programs/cycle-tracking.md +++ b/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. diff --git a/core/src/runtime/report.rs b/core/src/runtime/report.rs index 62e85f108b..36ca39616f 100644 --- a/core/src/runtime/report.rs +++ b/core/src/runtime/report.rs @@ -10,6 +10,7 @@ use super::*; pub struct ExecutionReport { pub opcode_counts: HashMap, pub syscall_counts: HashMap, + pub cycle_tracker: HashMap, } impl ExecutionReport { diff --git a/core/src/syscall/write.rs b/core/src/syscall/write.rs index 8e9d389cbd..9c42252b93 100644 --- a/core/src/syscall/write.rs +++ b/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 { 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::(); - 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::(); - 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 { +/// 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 { + 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 { + 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 { let rt = &mut ctx.rt; let entry = rt.io_buf.entry(fd).or_default(); entry.push_str(s); diff --git a/examples/cycle-tracking/program/bin/normal.rs b/examples/cycle-tracking/program/bin/normal.rs new file mode 100644 index 0000000000..cfffe0efe7 --- /dev/null +++ b/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"); +} diff --git a/examples/cycle-tracking/program/bin/report.rs b/examples/cycle-tracking/program/bin/report.rs new file mode 100644 index 0000000000..e60ce2e2af --- /dev/null +++ b/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"); +} diff --git a/examples/cycle-tracking/program/elf/normal b/examples/cycle-tracking/program/elf/normal new file mode 100755 index 0000000000..812a0f97a3 Binary files /dev/null and b/examples/cycle-tracking/program/elf/normal differ diff --git a/examples/cycle-tracking/program/elf/report b/examples/cycle-tracking/program/elf/report new file mode 100755 index 0000000000..afcc118fd0 Binary files /dev/null and b/examples/cycle-tracking/program/elf/report differ