Refactor

georgwiese · georgwiese · commit 7851df59f880 · 2025-11-27T15:39:21.000-05:00
diff --git a/openvm/src/customize_exe.rs b/openvm/src/customize_exe.rs
@@ -228,8 +228,7 @@ pub fn customize<'a, P: PgoAdapter<Adapter = BabyBearOpenVmApcAdapter<'a>>>(
         ),
         inputs,
         &blocks,
-    )
-    .unwrap();
+    );
 
     let start = std::time::Instant::now();
     let apcs = pgo.filter_blocks_and_create_apcs_with_pgo(
diff --git a/openvm/src/execution_stats.rs b/openvm/src/execution_stats.rs
@@ -1,21 +1,8 @@
-#![cfg_attr(feature = "tco", allow(internal_features))]
-#![cfg_attr(feature = "tco", allow(incomplete_features))]
-#![cfg_attr(feature = "tco", feature(explicit_tail_calls))]
-#![cfg_attr(feature = "tco", feature(core_intrinsics))]
-
-use eyre::Result;
 use itertools::Itertools;
-use openvm_circuit::arch::execution_mode::Segment;
-use openvm_circuit::arch::{PreflightExecutionOutput, VirtualMachine, VmCircuitConfig, VmInstance};
+use openvm_circuit::arch::VmCircuitConfig;
 use openvm_instructions::exe::VmExe;
-use openvm_sdk::prover::vm::new_local_prover;
-use openvm_sdk::{
-    config::{AppConfig, DEFAULT_APP_LOG_BLOWUP},
-    StdIn,
-};
+use openvm_sdk::StdIn;
 use openvm_stark_backend::p3_matrix::dense::DenseMatrix;
-use openvm_stark_sdk::config::baby_bear_poseidon2::BabyBearPermutationEngine;
-use openvm_stark_sdk::config::FriParameters;
 use openvm_stark_sdk::openvm_stark_backend::p3_field::PrimeField32;
 use openvm_stark_sdk::p3_baby_bear::BabyBear;
 use powdr_autoprecompiles::blocks::BasicBlock;
@@ -24,136 +11,36 @@ use std::collections::hash_map::Entry;
 use std::collections::BTreeMap;
 use std::{collections::HashMap, sync::Arc};
 
-#[cfg(not(feature = "cuda"))]
-use crate::PowdrSdkCpu;
-use crate::{Instr, SpecializedConfig, SpecializedConfigCpuBuilder};
-use tracing::info_span;
+use crate::trace_generation::do_with_trace;
+use crate::{Instr, SpecializedConfig};
 
 use std::collections::HashSet;
 use std::hash::Hash;
 
-// ChatGPT generated code
-fn intersect_partitions<Id>(partitions: &[Vec<Vec<Id>>]) -> Vec<Vec<Id>>
-where
-    Id: Eq + Hash + Copy,
-{
-    if partitions.is_empty() {
-        return Vec::new();
-    }
-
-    // 1) For each partition, build a map: Id -> class_index
-    let mut maps: Vec<HashMap<Id, usize>> = Vec::with_capacity(partitions.len());
-    for part in partitions {
-        let mut m = HashMap::new();
-        for (class_idx, class) in part.iter().enumerate() {
-            for &id in class {
-                m.insert(id, class_idx);
-            }
-        }
-        maps.push(m);
-    }
-
-    // 2) Collect the universe of all Ids
-    let mut universe: HashSet<Id> = HashSet::new();
-    for part in partitions {
-        for class in part {
-            for &id in class {
-                universe.insert(id);
-            }
-        }
-    }
-
-    // 3) For each Id, build its "signature" of class indices across all partitions
-    //    and group by that signature.
-    let mut grouped: HashMap<Vec<usize>, Vec<Id>> = HashMap::new();
-
-    for &id in &universe {
-        let mut signature = Vec::with_capacity(maps.len());
-        let mut is_singleton = false;
-        for m in &maps {
-            let Some(class_idx) = m.get(&id) else {
-                // The element did not appear in one of the partition, so it is its
-                // own equivalence class. We can also omit it in the output partition.
-                is_singleton = true;
-                break;
-            };
-            signature.push(*class_idx);
-        }
-        if !is_singleton {
-            grouped.entry(signature).or_default().push(id);
-        }
-    }
-
-    // 4) Resulting equivalence classes are the grouped values
-    grouped.into_values().collect()
+#[derive(Default)]
+struct Trace {
+    /// Mapping (segment_idx, timestamp) -> Vec<u32>
+    rows_by_time: BTreeMap<(usize, u32), Vec<u32>>,
+    trace_values_by_pc: HashMap<u32, Vec<Vec<u32>>>,
+    air_id_by_pc: HashMap<u32, usize>,
+    column_names_by_air_id: HashMap<usize, Vec<String>>,
 }
 
 pub fn execution_stats(
     exe: Arc<VmExe<BabyBear>>,
     vm_config: SpecializedConfig,
     inputs: StdIn,
     blocks: &[BasicBlock<Instr<BabyBear>>],
-) -> Result<ExecutionStats, Box<dyn std::error::Error>> {
-    // 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());
-
-    // Create the SDK
-    #[cfg(feature = "cuda")]
-    let sdk = PowdrSdkGpu::new(app_config).unwrap();
-    #[cfg(not(feature = "cuda"))]
-    let sdk = PowdrSdkCpu::new(app_config).unwrap();
-    // Build owned vm instance, so we can mutate it later
-    let vm_builder = sdk.app_vm_builder().clone();
-    let vm_pk = sdk.app_pk().app_vm_pk.clone();
-    let exe = sdk.convert_to_exe(exe.clone())?;
-    let mut vm_instance: VmInstance<_, _> = new_local_prover(vm_builder, &vm_pk, exe.clone())?;
-
-    vm_instance.reset_state(inputs.clone());
-    let metered_ctx = vm_instance.vm.build_metered_ctx(&exe);
-    let metered_interpreter = vm_instance.vm.metered_interpreter(vm_instance.exe())?;
-    let (segments, _) = metered_interpreter.execute_metered(inputs.clone(), metered_ctx)?;
-    let mut state = vm_instance.state_mut().take();
-
-    // Get reusable inputs for `debug_proving_ctx`, the mock prover API from OVM.
-    let vm: &mut VirtualMachine<BabyBearPermutationEngine<_>, SpecializedConfigCpuBuilder> =
-        &mut vm_instance.vm;
-
-    // Mapping (segment_idx, timestamp) -> Vec<u32>
-    let mut rows_by_time = BTreeMap::new();
-
-    let mut trace_values_by_pc = HashMap::new();
-    let mut column_names_by_air_id = HashMap::new();
-    let mut air_id_by_pc = HashMap::new();
-
-    for (seg_idx, segment) in segments.into_iter().enumerate() {
-        let _segment_span = info_span!("prove_segment", segment = seg_idx).entered();
-        // We need a separate span so the metric label includes "segment" from _segment_span
-        let _prove_span = info_span!("total_proof").entered();
-        let Segment {
-            instret_start,
-            num_insns,
-            trace_heights,
-        } = segment;
-        assert_eq!(state.as_ref().unwrap().instret(), instret_start);
-        let from_state = Option::take(&mut state).unwrap();
-        vm.transport_init_memory_to_device(&from_state.memory);
-        let PreflightExecutionOutput {
-            system_records,
-            record_arenas,
-            to_state,
-        } = vm.execute_preflight(
-            &mut vm_instance.interpreter,
-            from_state,
-            Some(num_insns),
-            &trace_heights,
-        )?;
-        state = Some(to_state);
+) -> ExecutionStats {
+    let trace = collect_trace(exe, vm_config, inputs);
+    generate_execution_stats(blocks, trace)
+}
 
-        // Generate proving context for each segment
-        let ctx = vm.generate_proving_ctx(system_records, record_arenas)?;
+fn collect_trace(exe: Arc<VmExe<BabyBear>>, vm_config: SpecializedConfig, inputs: StdIn) -> Trace {
+    let mut trace = Trace::default();
+    let mut seg_idx = 0;
 
+    do_with_trace(exe, vm_config, inputs, |vm, _pk, ctx| {
         let global_airs = vm
             .config()
             .create_airs()
@@ -188,26 +75,28 @@ pub fn execution_stats(
                 let row = row.iter().map(|v| v.as_canonical_u32()).collect::<Vec<_>>();
                 let pc_value = row[pc_index];
                 let ts_value = row[ts_index];
-                rows_by_time.insert((seg_idx, ts_value), row.clone());
+                trace.rows_by_time.insert((seg_idx, ts_value), row.clone());
 
                 if pc_value == 0 {
                     // Padding row!
                     continue;
                 }
 
-                if let Entry::Vacant(e) = trace_values_by_pc.entry(pc_value) {
+                if let Entry::Vacant(e) = trace.trace_values_by_pc.entry(pc_value) {
                     // First time we see this PC, initialize the column -> values map
                     e.insert(vec![Vec::new(); row.len()]);
-                    column_names_by_air_id.insert(*air_id, column_names.clone());
-                    air_id_by_pc.insert(pc_value, *air_id);
+                    trace
+                        .column_names_by_air_id
+                        .insert(*air_id, column_names.clone());
+                    trace.air_id_by_pc.insert(pc_value, *air_id);
                 }
-                let values_by_col = trace_values_by_pc.get_mut(&pc_value).unwrap();
+                let values_by_col = trace.trace_values_by_pc.get_mut(&pc_value).unwrap();
                 assert_eq!(
-                    air_id_by_pc[&pc_value],
+                    trace.air_id_by_pc[&pc_value],
                     *air_id,
                     "Mismatched air IDs for PC {}: {} vs {}",
                     pc_value,
-                    global_airs[&air_id_by_pc[&pc_value]].name(),
+                    global_airs[&trace.air_id_by_pc[&pc_value]].name(),
                     air.name()
                 );
                 assert_eq!(values_by_col.len(), row.len());
@@ -217,8 +106,15 @@ pub fn execution_stats(
                 }
             }
         }
-    }
+        seg_idx += 1;
+    });
+    trace
+}
 
+fn generate_execution_stats(
+    blocks: &[BasicBlock<Instr<BabyBear>>],
+    trace: Trace,
+) -> ExecutionStats {
     // Block ID -> instruction count mapping
     let instruction_counts = blocks
         .iter()
@@ -228,7 +124,7 @@ pub fn execution_stats(
     // Block ID -> Vec<Vec<Row>>
     let mut block_rows = BTreeMap::new();
     let mut i = 0;
-    let rows_by_time = rows_by_time.values().collect::<Vec<_>>();
+    let rows_by_time = trace.rows_by_time.values().collect::<Vec<_>>();
     while i < rows_by_time.len() {
         let row = &rows_by_time[i];
         let pc_value = row[0] as u64;
@@ -288,7 +184,8 @@ pub fn execution_stats(
         .collect::<BTreeMap<_, _>>();
 
     // Map all column values to their range (1st and 99th percentile) for each pc
-    let column_ranges_by_pc: HashMap<u32, Vec<(u32, u32)>> = trace_values_by_pc
+    let column_ranges_by_pc: HashMap<u32, Vec<(u32, u32)>> = trace
+        .trace_values_by_pc
         .into_iter()
         .map(|(pc, values_by_col)| {
             let column_ranges = values_by_col
@@ -306,8 +203,8 @@ pub fn execution_stats(
         .collect();
 
     let export = ExecutionStats {
-        air_id_by_pc: air_id_by_pc.into_iter().collect(),
-        column_names_by_air_id: column_names_by_air_id.into_iter().collect(),
+        air_id_by_pc: trace.air_id_by_pc.into_iter().collect(),
+        column_names_by_air_id: trace.column_names_by_air_id.into_iter().collect(),
         column_ranges_by_pc: column_ranges_by_pc.into_iter().collect(),
         equivalence_classes_by_block: intersected_equivalence_classes,
     };
@@ -316,5 +213,61 @@ pub fn execution_stats(
     let json = serde_json::to_string_pretty(&export).unwrap();
     std::fs::write("pgo_range_constraints.json", json).unwrap();
 
-    Ok(export)
+    export
+}
+
+// ChatGPT generated code
+fn intersect_partitions<Id>(partitions: &[Vec<Vec<Id>>]) -> Vec<Vec<Id>>
+where
+    Id: Eq + Hash + Copy,
+{
+    if partitions.is_empty() {
+        return Vec::new();
+    }
+
+    // 1) For each partition, build a map: Id -> class_index
+    let mut maps: Vec<HashMap<Id, usize>> = Vec::with_capacity(partitions.len());
+    for part in partitions {
+        let mut m = HashMap::new();
+        for (class_idx, class) in part.iter().enumerate() {
+            for &id in class {
+                m.insert(id, class_idx);
+            }
+        }
+        maps.push(m);
+    }
+
+    // 2) Collect the universe of all Ids
+    let mut universe: HashSet<Id> = HashSet::new();
+    for part in partitions {
+        for class in part {
+            for &id in class {
+                universe.insert(id);
+            }
+        }
+    }
+
+    // 3) For each Id, build its "signature" of class indices across all partitions
+    //    and group by that signature.
+    let mut grouped: HashMap<Vec<usize>, Vec<Id>> = HashMap::new();
+
+    for &id in &universe {
+        let mut signature = Vec::with_capacity(maps.len());
+        let mut is_singleton = false;
+        for m in &maps {
+            let Some(class_idx) = m.get(&id) else {
+                // The element did not appear in one of the partition, so it is its
+                // own equivalence class. We can also omit it in the output partition.
+                is_singleton = true;
+                break;
+            };
+            signature.push(*class_idx);
+        }
+        if !is_singleton {
+            grouped.entry(signature).or_default().push(id);
+        }
+    }
+
+    // 4) Resulting equivalence classes are the grouped values
+    grouped.into_values().collect()
 }
diff --git a/openvm/src/lib.rs b/openvm/src/lib.rs
@@ -880,9 +880,14 @@ pub fn prove(
     #[cfg(not(feature = "cuda"))]
     let sdk = PowdrSdkCpu::new(app_config).unwrap();
     if mock {
-        do_with_trace(program, inputs, |vm, pk, ctx| {
-            debug_proving_ctx(vm, pk, &ctx);
-        });
+        do_with_trace(
+            program.exe.clone(),
+            program.vm_config.clone(),
+            inputs,
+            |vm, pk, ctx| {
+                debug_proving_ctx(vm, pk, &ctx);
+            },
+        );
     } else {
         let mut app_prover = sdk.app_prover(exe.clone())?;
 
diff --git a/openvm/src/trace_generation.rs b/openvm/src/trace_generation.rs
@@ -1,16 +1,19 @@
+use std::sync::Arc;
+
 use openvm_circuit::arch::{
     execution_mode::Segment, PreflightExecutionOutput, VirtualMachine, VmCircuitConfig, VmInstance,
 };
+use openvm_instructions::exe::VmExe;
 use openvm_sdk::{
     config::{AppConfig, DEFAULT_APP_LOG_BLOWUP},
     prover::vm::new_local_prover,
     StdIn,
 };
 use openvm_stark_backend::{keygen::types::MultiStarkProvingKey, prover::types::ProvingContext};
-use openvm_stark_sdk::{config::FriParameters, engine::StarkEngine};
+use openvm_stark_sdk::{config::FriParameters, engine::StarkEngine, p3_baby_bear::BabyBear};
 use tracing::info_span;
 
-use crate::{BabyBearSC, CompiledProgram, SpecializedConfigCpuBuilder};
+use crate::{BabyBearSC, SpecializedConfig, SpecializedConfigCpuBuilder};
 
 #[cfg(not(feature = "cuda"))]
 use crate::PowdrSdkCpu as PowdrSdk;
@@ -25,17 +28,15 @@ use openvm_stark_sdk::config::baby_bear_poseidon2::BabyBearPoseidon2Engine;
 /// Given a program and input, generates the trace segment by segment and calls the provided
 /// callback with the VM, proving key, and proving context (containing the trace) for each segment.
 pub fn do_with_trace(
-    program: &CompiledProgram,
+    exe: Arc<VmExe<BabyBear>>,
+    vm_config: SpecializedConfig,
     inputs: StdIn,
     mut callback: impl FnMut(
         &VirtualMachine<BabyBearPoseidon2Engine, SpecializedConfigCpuBuilder>,
         &MultiStarkProvingKey<BabyBearSC>,
         ProvingContext<<BabyBearPoseidon2Engine as StarkEngine>::PB>,
     ),
 ) {
-    let exe = &program.exe;
-    let vm_config = program.vm_config.clone();
-
     // Set app configuration
     let app_fri_params =
         FriParameters::standard_with_100_bits_conjectured_security(DEFAULT_APP_LOG_BLOWUP);
