[PGO] collect cycle data

Cargo.toml

@@ -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 = [

autoprecompiles/src/execution_profile.rs

@@ -1,9 +1,10 @@
 use crate::adapter::Adapter;
 use crate::blocks::Program;
+use serde::Serialize;
 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};
@@ -14,8 +15,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(),
@@ -132,3 +133,162 @@ 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<MemoryAccess> {
+        Some(MemoryAccess {
+            address_space: self.address_space?,
+            address: self.address?,
+            value: self.value?,
+        })
+    }
+}
+
+#[derive(Debug, Clone, Copy, Serialize)]
+pub struct MemoryAccess {
+    pub address_space: u64,
+    pub address: u64,
+    pub value: u64,
+}
+
+/// Captures a complete description of a VM cycle.
+#[derive(Debug, Clone, Serialize)]
+pub struct Cycle {
+    pub pc: u64,
+    pub reads: Vec<MemoryAccess>,
+    pub writes: Vec<MemoryAccess>,
+}
+
+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: MemoryAccess) {
+        self.record_access(access, false);
+    }
+
+    fn record_write(&self, access: MemoryAccess) {
+        self.record_access(access, true);
+    }
+
+    fn record_access(&self, access: MemoryAccess, 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);
+                }
+            }
+            _ => {}
+        }
+    }
+}