Skip to content

[PGO] collect cycle data #3442

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 4 commits into
base: main
Choose a base branch
from
Draft

[PGO] collect cycle data #3442

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
720b982
first commit
qwang98 Nov 14, 2025
2b79014
MemoryAccess dict
georgwiese Nov 14, 2025
d047389
Store JSON, add notebook
georgwiese Nov 14, 2025
6780bf4
Always store JSON, update notebook
georgwiese Nov 14, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 29 additions & 29 deletions Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,40 +44,40 @@ powdr-openvm-hints-transpiler = { path = "./openvm/extensions/hints-transpiler",
powdr-openvm-hints-circuit = { path = "./openvm/extensions/hints-circuit", version = "0.1.4" }

# openvm
openvm = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr" }
openvm-build = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr" }
openvm-rv32im-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr" }
openvm-rv32im-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr" }
openvm-rv32im-guest = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr", default-features = false }
openvm-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr" }
openvm-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr" }
openvm-circuit-derive = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr" }
openvm-circuit-primitives = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr" }
openvm-circuit-primitives-derive = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr" }
openvm-instructions = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr" }
openvm-instructions-derive = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr" }
openvm-sdk = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr", default-features = false, features = [
openvm = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data" }
openvm-build = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data" }
openvm-rv32im-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data" }
openvm-rv32im-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data" }
openvm-rv32im-guest = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data", default-features = false }
openvm-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data" }
openvm-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data" }
openvm-circuit-derive = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data" }
openvm-circuit-primitives = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data" }
openvm-circuit-primitives-derive = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data" }
openvm-instructions = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data" }
openvm-instructions-derive = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data" }
openvm-sdk = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data", default-features = false, features = [
"parallel",
"jemalloc",
"nightly-features",
"evm-prove",
] }
openvm-ecc-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr" }
openvm-ecc-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr" }
openvm-keccak256-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr" }
openvm-keccak256-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr" }
openvm-sha256-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr" }
openvm-sha256-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr" }
openvm-algebra-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr" }
openvm-algebra-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr" }
openvm-bigint-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr" }
openvm-bigint-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr" }
openvm-pairing-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr" }
openvm-pairing-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr" }
openvm-native-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr", default-features = false }
openvm-native-recursion = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr", default-features = false }
openvm-platform = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr" }
openvm-custom-insn = { git = "https://github.com/powdr-labs/openvm.git", rev = "v1.4.1-powdr" }
openvm-ecc-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data" }
openvm-ecc-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data" }
openvm-keccak256-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data" }
openvm-keccak256-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data" }
openvm-sha256-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data" }
openvm-sha256-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data" }
openvm-algebra-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data" }
openvm-algebra-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data" }
openvm-bigint-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data" }
openvm-bigint-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data" }
openvm-pairing-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data" }
openvm-pairing-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data" }
openvm-native-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data", default-features = false }
openvm-native-recursion = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data", default-features = false }
openvm-platform = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data" }
openvm-custom-insn = { git = "https://github.com/powdr-labs/openvm.git", rev = "cycle-data" }

# stark-backend
openvm-stark-sdk = { git = "https://github.com/powdr-labs/stark-backend.git", rev = "v1.2.1-powdr", default-features = false, features = [
Expand Down
161 changes: 158 additions & 3 deletions autoprecompiles/src/execution_profile.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ use crate::blocks::Program;
use std::collections::HashMap;
use std::sync::atomic::AtomicU32;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::sync::{Arc, Mutex};
use tracing::dispatcher::Dispatch;
use tracing::field::Field as TracingField;
use tracing::{Event, Level, Subscriber};
Expand All @@ -14,8 +14,8 @@ use tracing_subscriber::{
Layer,
};

// Produces execution count by pc
// Used in Pgo::Cell and Pgo::Instruction to help rank basic blocks to create APCs for
/// Produces execution count by pc
/// Used in Pgo::Cell and Pgo::Instruction to help rank basic blocks to create APCs for
pub fn execution_profile<A: Adapter>(
program: &A::Program,
execute_fn: impl FnOnce(),
Expand Down Expand Up @@ -132,3 +132,158 @@ where
}
}
}

/// Produces a rich execution trace with per-cycle memory access information.
pub fn execution_data<A: Adapter>(
_program: &A::Program,
execute_fn: impl FnOnce(),
) -> Vec<Cycle> {
let collector = PgoExecutionCollector::default();

let subscriber = Registry::default().with(collector.clone());
let dispatch = Dispatch::new(subscriber);
tracing::dispatcher::with_default(&dispatch, execute_fn);

collector.into_cycles()
}

// holds basic type fields of execution objects captured in trace by subscriber
#[derive(Default)]
struct CycleData {
pc: Option<u64>,
address_space: Option<u64>,
address: Option<u64>,
value: Option<u64>,
}

impl tracing::field::Visit for CycleData {
// when we receive a u64 field, they are parsed into fields of the pgo data
fn record_u64(&mut self, field: &tracing::field::Field, value: u64) {
if field.name() == "pc" {
self.pc = Some(value);
} else if field.name() == "address_space" {
self.address_space = Some(value);
} else if field.name() == "address" {
self.address = Some(value);
} else if field.name() == "value" {
self.value = Some(value);
}
}

// required for implementation, but in practice we will only receive u64 fields
// the fields we receive are determined by the instruction trace print out of our openvm fork during execution
fn record_debug(&mut self, _: &TracingField, _: &dyn std::fmt::Debug) {}
}

impl CycleData {
fn memory_access(&self) -> Option<(u64, u64, u64)> {
Some((
self.address_space?,
self.address?,
self.value?,
))
}
}

/// Captures a complete description of a VM cycle.
#[derive(Debug, Clone)]
pub struct Cycle {
pub pc: u64,
pub reads: Vec<(u64, u64, u64)>,
pub writes: Vec<(u64, u64, u64)>,
}

impl Cycle {
fn new(pc: u64) -> Self {
Self {
pc,
reads: Vec::new(),
writes: Vec::new(),
}
}
}

#[derive(Default)]
struct ExecutionCollectorState {
cycles: Vec<Cycle>,
current_cycle: Option<Cycle>,
}

#[derive(Clone, Default)]
struct PgoExecutionCollector {
state: Arc<Mutex<ExecutionCollectorState>>,
}

impl PgoExecutionCollector {
fn start_cycle(&self, pc: u64) {
let mut state = self
.state
.lock()
.expect("execution collector state lock poisoned");
if let Some(cycle) = state.current_cycle.take() {
state.cycles.push(cycle);
}
state.current_cycle = Some(Cycle::new(pc));
}

fn record_read(&self, access: (u64, u64, u64)) {
self.record_access(access, false);
}

fn record_write(&self, access: (u64, u64, u64)) {
self.record_access(access, true);
}

fn record_access(&self, access: (u64, u64, u64), is_write: bool) {
let mut state = self
.state
.lock()
.expect("execution collector state lock poisoned");
if let Some(cycle) = state.current_cycle.as_mut() {
if is_write {
cycle.writes.push(access);
} else {
cycle.reads.push(access);
}
}
}

fn into_cycles(self) -> Vec<Cycle> {
let mut state = self
.state
.lock()
.expect("execution collector state lock poisoned");
if let Some(cycle) = state.current_cycle.take() {
state.cycles.push(cycle);
}
std::mem::take(&mut state.cycles)
}
}

impl<S> Layer<S> for PgoExecutionCollector
where
S: Subscriber + for<'a> LookupSpan<'a>,
{
fn on_event(&self, event: &Event<'_>, _ctx: Context<'_, S>) {
let mut visitor = CycleData::default();
event.record(&mut visitor);

if let Some(pc) = visitor.pc {
self.start_cycle(pc);
}

match event.metadata().target() {
"read" => {
if let Some(access) = visitor.memory_access() {
self.record_read(access);
}
}
"write" => {
if let Some(access) = visitor.memory_access() {
self.record_write(access);
}
}
_ => {}
}
}
}
30 changes: 30 additions & 0 deletions openvm/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ use openvm_stark_sdk::openvm_stark_backend::p3_field::PrimeField32;
use openvm_stark_sdk::p3_baby_bear::BabyBear;
use openvm_transpiler::transpiler::Transpiler;
use powdr_autoprecompiles::evaluation::AirStats;
use powdr_autoprecompiles::execution_profile::{Cycle, execution_data};
use powdr_autoprecompiles::pgo::{CellPgo, InstructionPgo, NonePgo};
use powdr_autoprecompiles::{execution_profile::execution_profile, PowdrConfig};
use powdr_extension::PowdrExtension;
Expand Down Expand Up @@ -979,6 +980,27 @@ pub fn execution_profile_from_guest(
})
}

pub fn execution_data_from_guest(
guest: &str,
guest_opts: GuestOptions,
inputs: StdIn,
) -> Vec<Cycle> {
let OriginalCompiledProgram { exe, vm_config, .. } = compile_openvm(guest, guest_opts).unwrap();
let program = Prog::from(&exe.program);

// Set app configuration
let app_fri_params =
FriParameters::standard_with_100_bits_conjectured_security(DEFAULT_APP_LOG_BLOWUP);
let app_config = AppConfig::new(app_fri_params, vm_config.clone());

// prepare for execute
let sdk = PowdrExecutionProfileSdkCpu::new(app_config).unwrap();

execution_data::<BabyBearOpenVmApcAdapter>(&program, || {
sdk.execute(exe.clone(), inputs.clone()).unwrap();
})
}
Comment on lines 986 to 1013
Copy link
Collaborator Author

@qwang98 qwang98 Nov 14, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

TODO: Top part mostly same as execution_profile_from_guest. Make it a helper.


#[cfg(test)]
mod tests {
use super::*;
Expand Down Expand Up @@ -2132,4 +2154,12 @@ mod tests {
// "Total columns exceeded the limit: {total_columns} > {MAX_TOTAL_COLUMNS}"
// );
}

#[test]
fn guest_execution_data() {
let mut stdin = StdIn::default();
stdin.write(&GUEST_ITER);
let execution_data = execution_data_from_guest(GUEST, GuestOptions::default(), stdin.clone());
println!("Execution data: {execution_data:#?}");
}
Comment on lines 2169 to 2182
Copy link
Collaborator Author

@qwang98 qwang98 Nov 14, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Happy to make this a part of our CLI if we run these frequently.

}
Loading