Pass execution stats directly

Author: georgwiese

autoprecompiles/src/adapter.rs

@@ -6,6 +6,7 @@ use std::{fmt::Display, sync::Arc};
 use powdr_number::FieldElement;
 use serde::{Deserialize, Serialize};
 
+use crate::ExecutionStats;
 use crate::{
     blocks::{BasicBlock, Instruction, Program},
     constraint_optimizer::IsBusStateful,
@@ -45,12 +46,13 @@ pub trait PgoAdapter {
         config: &PowdrConfig,
         vm_config: AdapterVmConfig<Self::Adapter>,
         labels: BTreeMap<u64, Vec<String>>,
+        execution_stats: ExecutionStats,
     ) -> Vec<AdapterApcWithStats<Self::Adapter>> {
         let filtered_blocks = blocks
             .into_iter()
             .filter(|block| !Self::Adapter::should_skip_block(block))
             .collect();
-        self.create_apcs_with_pgo(filtered_blocks, config, vm_config, labels)
+        self.create_apcs_with_pgo(filtered_blocks, config, vm_config, labels, execution_stats)
     }
 
     fn create_apcs_with_pgo(
@@ -59,6 +61,7 @@ pub trait PgoAdapter {
         config: &PowdrConfig,
         vm_config: AdapterVmConfig<Self::Adapter>,
         labels: BTreeMap<u64, Vec<String>>,
+        execution_stats: ExecutionStats,
     ) -> Vec<AdapterApcWithStats<Self::Adapter>>;
 
     fn pc_execution_count(&self, _pc: u64) -> Option<u32> {

autoprecompiles/src/lib.rs

@@ -410,7 +410,7 @@ impl ColumnAllocator {
 }
 
 #[derive(Serialize, Deserialize)]
-pub struct JsonExport {
+pub struct ExecutionStats {
     pub air_id_by_pc: BTreeMap<u32, usize>,
     pub column_names_by_air_id: BTreeMap<usize, Vec<String>>,
     pub column_ranges_by_pc: BTreeMap<u32, Vec<(u32, u32)>>,
@@ -422,6 +422,7 @@ pub fn build<A: Adapter>(
     vm_config: AdapterVmConfig<A>,
     degree_bound: DegreeBound,
     apc_candidates_dir_path: Option<&Path>,
+    execution_stats: &ExecutionStats,
 ) -> Result<AdapterApc<A>, crate::constraint_optimizer::Error> {
     let start = std::time::Instant::now();
 
@@ -430,14 +431,8 @@ pub fn build<A: Adapter>(
         vm_config.instruction_handler,
         &vm_config.bus_map,
     );
-
-    let pgo_range_constraints_json = std::fs::read_to_string("pgo_range_constraints.json")
-        .expect("Failed to read pgo range constraints json file");
-
-    let pgo_range_constraints: JsonExport = serde_json::from_str(&pgo_range_constraints_json)
-        .expect("JSON format does not match JsonExport struct");
-    let range_constraints = pgo_range_constraints.column_ranges_by_pc;
-    let equivalence_classes_by_block = pgo_range_constraints.equivalence_classes_by_block;
+    let range_constraints = &execution_stats.column_ranges_by_pc;
+    let equivalence_classes_by_block = &execution_stats.equivalence_classes_by_block;
 
     // Mapping (instruction index, column index) -> AlgebraicReference
     let reverse_subs = column_allocator

autoprecompiles/src/pgo/cell/mod.rs

@@ -13,7 +13,7 @@ use crate::{
     blocks::BasicBlock,
     evaluation::EvaluationResult,
     pgo::cell::selection::parallel_fractional_knapsack,
-    PowdrConfig,
+    ExecutionStats, PowdrConfig,
 };
 
 mod selection;
@@ -78,8 +78,8 @@ impl<A, C> CellPgo<A, C> {
 }
 
 #[derive(Serialize, Deserialize)]
-pub struct JsonExport {
-    pub apcs: Vec<ApcCandidateJsonExport>,
+struct JsonExport {
+    apcs: Vec<ApcCandidateJsonExport>,
     labels: BTreeMap<u64, Vec<String>>,
 }
 
@@ -92,6 +92,7 @@ impl<A: Adapter + Send + Sync, C: Candidate<A> + Send + Sync> PgoAdapter for Cel
         config: &PowdrConfig,
         vm_config: AdapterVmConfig<Self::Adapter>,
         labels: BTreeMap<u64, Vec<String>>,
+        execution_stats: ExecutionStats,
     ) -> Vec<AdapterApcWithStats<Self::Adapter>> {
         tracing::info!(
             "Generating autoprecompiles with cell PGO for {} blocks",
@@ -131,6 +132,7 @@ impl<A: Adapter + Send + Sync, C: Candidate<A> + Send + Sync> PgoAdapter for Cel
                     vm_config.clone(),
                     config.degree_bound,
                     config.apc_candidates_dir_path.as_deref(),
+                    &execution_stats,
                 )
                 .ok()?;
                 let candidate = C::create(

autoprecompiles/src/pgo/instruction.rs

@@ -4,7 +4,7 @@ use crate::{
     adapter::{Adapter, AdapterApcWithStats, AdapterVmConfig, PgoAdapter},
     blocks::BasicBlock,
     pgo::create_apcs_for_all_blocks,
-    PowdrConfig,
+    ExecutionStats, PowdrConfig,
 };
 
 pub struct InstructionPgo<A> {
@@ -30,6 +30,7 @@ impl<A: Adapter> PgoAdapter for InstructionPgo<A> {
         config: &PowdrConfig,
         vm_config: AdapterVmConfig<Self::Adapter>,
         _labels: BTreeMap<u64, Vec<String>>,
+        execution_stats: ExecutionStats,
     ) -> Vec<AdapterApcWithStats<Self::Adapter>> {
         tracing::info!(
             "Generating autoprecompiles with instruction PGO for {} blocks",
@@ -70,7 +71,7 @@ impl<A: Adapter> PgoAdapter for InstructionPgo<A> {
                 );
         }
 
-        create_apcs_for_all_blocks::<Self::Adapter>(blocks, config, vm_config)
+        create_apcs_for_all_blocks::<Self::Adapter>(blocks, config, vm_config, execution_stats)
     }
 
     fn pc_execution_count(&self, pc: u64) -> Option<u32> {

autoprecompiles/src/pgo/mod.rs

@@ -6,15 +6,15 @@ use strum::{Display, EnumString};
 use crate::{
     adapter::{Adapter, AdapterApcWithStats, AdapterVmConfig, ApcWithStats},
     blocks::BasicBlock,
-    PowdrConfig,
+    ExecutionStats, PowdrConfig,
 };
 
 mod cell;
 mod instruction;
 mod none;
 
 pub use {
-    cell::{ApcCandidateJsonExport, Candidate, CellPgo, JsonExport, KnapsackItem},
+    cell::{ApcCandidateJsonExport, Candidate, CellPgo, KnapsackItem},
     instruction::InstructionPgo,
     none::NonePgo,
 };
@@ -77,6 +77,7 @@ fn create_apcs_for_all_blocks<A: Adapter>(
     blocks: Vec<BasicBlock<A::Instruction>>,
     config: &PowdrConfig,
     vm_config: AdapterVmConfig<A>,
+    execution_stats: ExecutionStats,
 ) -> Vec<AdapterApcWithStats<A>> {
     let n_acc = config.autoprecompiles as usize;
     tracing::info!("Generating {n_acc} autoprecompiles in parallel");
@@ -97,6 +98,7 @@ fn create_apcs_for_all_blocks<A: Adapter>(
                 vm_config.clone(),
                 config.degree_bound,
                 config.apc_candidates_dir_path.as_deref(),
+                &execution_stats,
             )
             .unwrap()
         })

autoprecompiles/src/pgo/none.rs

@@ -4,7 +4,7 @@ use crate::{
     adapter::{Adapter, AdapterApcWithStats, AdapterVmConfig, PgoAdapter},
     blocks::BasicBlock,
     pgo::create_apcs_for_all_blocks,
-    PowdrConfig,
+    ExecutionStats, PowdrConfig,
 };
 
 pub struct NonePgo<A> {
@@ -29,6 +29,7 @@ impl<A: Adapter> PgoAdapter for NonePgo<A> {
         config: &PowdrConfig,
         vm_config: AdapterVmConfig<Self::Adapter>,
         _labels: BTreeMap<u64, Vec<String>>,
+        execution_stats: ExecutionStats,
     ) -> Vec<AdapterApcWithStats<Self::Adapter>> {
         // cost = number_of_original_instructions
         blocks.sort_by(|a, b| b.statements.len().cmp(&a.statements.len()));
@@ -42,6 +43,6 @@ impl<A: Adapter> PgoAdapter for NonePgo<A> {
             );
         }
 
-        create_apcs_for_all_blocks::<Self::Adapter>(blocks, config, vm_config)
+        create_apcs_for_all_blocks::<Self::Adapter>(blocks, config, vm_config, execution_stats)
     }
 }

openvm/src/customize_exe.rs

@@ -219,7 +219,7 @@ pub fn customize<'a, P: PgoAdapter<Adapter = BabyBearOpenVmApcAdapter<'a>>>(
         })
         .collect();
 
-    execution_stats(
+    let execution_stats = execution_stats(
         exe.clone(),
         SpecializedConfig::new(
             original_config.clone(),
@@ -232,7 +232,13 @@ pub fn customize<'a, P: PgoAdapter<Adapter = BabyBearOpenVmApcAdapter<'a>>>(
     .unwrap();
 
     let start = std::time::Instant::now();
-    let apcs = pgo.filter_blocks_and_create_apcs_with_pgo(blocks, &config, vm_config, labels);
+    let apcs = pgo.filter_blocks_and_create_apcs_with_pgo(
+        blocks,
+        &config,
+        vm_config,
+        labels,
+        execution_stats,
+    );
     metrics::gauge!("total_apc_gen_time_ms").set(start.elapsed().as_millis() as f64);
 
     let pc_base = exe.program.pc_base;

openvm/src/execution_stats.rs

@@ -19,7 +19,7 @@ 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;
-use powdr_autoprecompiles::JsonExport;
+use powdr_autoprecompiles::ExecutionStats;
 use std::collections::hash_map::Entry;
 use std::collections::BTreeMap;
 use std::{collections::HashMap, sync::Arc};
@@ -85,7 +85,7 @@ pub fn execution_stats(
     vm_config: SpecializedConfig,
     inputs: StdIn,
     blocks: &[BasicBlock<Instr<BabyBear>>],
-) -> Result<(), Box<dyn std::error::Error>> {
+) -> 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);
@@ -297,7 +297,7 @@ pub fn execution_stats(
         })
         .collect();
 
-    let export = JsonExport {
+    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(),
         column_ranges_by_pc: column_ranges_by_pc.into_iter().collect(),
@@ -308,5 +308,5 @@ pub fn execution_stats(
     let json = serde_json::to_string_pretty(&export).unwrap();
     std::fs::write("pgo_range_constraints.json", json).unwrap();
 
-    Ok(())
+    Ok(export)
 }