Call execution_stats from customize

Author: georgwiese

cli-openvm/src/main.rs

@@ -146,8 +146,14 @@ fn run_command(command: Commands) {
                 stdin_from(input),
             );
             let pgo_config = pgo_config(pgo, max_columns, execution_profile);
-            let program =
-                powdr_openvm::compile_guest(&guest, guest_opts, powdr_config, pgo_config).unwrap();
+            let program = powdr_openvm::compile_guest(
+                &guest,
+                guest_opts,
+                powdr_config,
+                pgo_config,
+                stdin_from(input),
+            )
+            .unwrap();
             write_program_to_file(program, &format!("{guest}_compiled.cbor")).unwrap();
         }
 
@@ -172,9 +178,14 @@ fn run_command(command: Commands) {
             );
             let pgo_config = pgo_config(pgo, max_columns, execution_profile);
             let compile_and_exec = || {
-                let program =
-                    powdr_openvm::compile_guest(&guest, guest_opts, powdr_config, pgo_config)
-                        .unwrap();
+                let program = powdr_openvm::compile_guest(
+                    &guest,
+                    guest_opts,
+                    powdr_config,
+                    pgo_config,
+                    stdin_from(input),
+                )
+                .unwrap();
                 powdr_openvm::execute(program, stdin_from(input)).unwrap();
             };
             if let Some(metrics_path) = metrics {
@@ -203,17 +214,31 @@ fn run_command(command: Commands) {
             if let Some(apc_candidates_dir) = &apc_candidates_dir {
                 powdr_config = powdr_config.with_apc_candidates_dir(apc_candidates_dir);
             }
+            let inputs = stdin_from(input);
             let execution_profile = powdr_openvm::execution_profile_from_guest(
                 &guest,
                 guest_opts.clone(),
-                stdin_from(input),
+                inputs.clone(),
             );
             let pgo_config = pgo_config(pgo, max_columns, execution_profile);
             let compile_and_prove = || {
-                let program =
-                    powdr_openvm::compile_guest(&guest, guest_opts, powdr_config, pgo_config)
-                        .unwrap();
-                powdr_openvm::prove(&program, mock, recursion, stdin_from(input), None, apc_candidates_dir).unwrap()
+                let program = powdr_openvm::compile_guest(
+                    &guest,
+                    guest_opts,
+                    powdr_config,
+                    pgo_config,
+                    inputs.clone(),
+                )
+                .unwrap();
+                powdr_openvm::prove(
+                    &program,
+                    mock,
+                    recursion,
+                    inputs.clone(),
+                    None,
+                    apc_candidates_dir,
+                )
+                .unwrap()
             };
             if let Some(metrics_path) = metrics {
                 run_with_metric_collection_to_file(

openvm/src/customize_exe.rs

@@ -13,12 +13,13 @@ use crate::memory_bus_interaction::OpenVmMemoryBusInteraction;
 use crate::opcode::branch_opcodes_bigint_set;
 use crate::powdr_extension::chip::PowdrAir;
 use crate::utils::UnsupportedOpenVmReferenceError;
-use crate::OriginalCompiledProgram;
+use crate::{execution_stats, OriginalCompiledProgram};
 use crate::{CompiledProgram, SpecializedConfig};
 use itertools::Itertools;
 use openvm_instructions::instruction::Instruction as OpenVmInstruction;
 use openvm_instructions::program::{Program as OpenVmProgram, DEFAULT_PC_STEP};
 use openvm_instructions::VmOpcode;
+use openvm_sdk::StdIn;
 use openvm_stark_backend::{
     interaction::SymbolicInteraction,
     p3_field::{FieldAlgebra, PrimeField32},
@@ -147,6 +148,7 @@ pub fn customize<'a, P: PgoAdapter<Adapter = BabyBearOpenVmApcAdapter<'a>>>(
     }: OriginalCompiledProgram,
     config: PowdrConfig,
     pgo: P,
+    inputs: StdIn,
 ) -> CompiledProgram {
     let labels = elf.text_labels();
     let debug_info = elf.debug_info();
@@ -217,6 +219,18 @@ pub fn customize<'a, P: PgoAdapter<Adapter = BabyBearOpenVmApcAdapter<'a>>>(
         })
         .collect();
 
+    execution_stats(
+        exe.clone(),
+        SpecializedConfig::new(
+            original_config.clone(),
+            vec![],
+            config.degree_bound.identities,
+        ),
+        inputs,
+        &blocks,
+    )
+    .unwrap();
+
     let start = std::time::Instant::now();
     let apcs = pgo.filter_blocks_and_create_apcs_with_pgo(blocks, &config, vm_config, labels);
     metrics::gauge!("total_apc_gen_time_ms").set(start.elapsed().as_millis() as f64);

openvm/src/execution_stats.rs

@@ -5,9 +5,9 @@
 
 use eyre::Result;
 use itertools::Itertools;
-use openvm_circuit::arch::execution_mode::metered::segment_ctx::SegmentationLimits;
 use openvm_circuit::arch::execution_mode::Segment;
 use openvm_circuit::arch::{PreflightExecutionOutput, VirtualMachine, VmCircuitConfig, VmInstance};
+use openvm_instructions::exe::VmExe;
 use openvm_sdk::prover::vm::new_local_prover;
 use openvm_sdk::{
     config::{AppConfig, DEFAULT_APP_LOG_BLOWUP},
@@ -18,14 +18,15 @@ 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;
 use powdr_autoprecompiles::JsonExport;
 use std::collections::hash_map::Entry;
 use std::collections::BTreeMap;
-use std::{collections::HashMap, path::PathBuf, sync::Arc};
+use std::{collections::HashMap, sync::Arc};
 
 #[cfg(not(feature = "cuda"))]
 use crate::PowdrSdkCpu;
-use crate::{CompiledProgram, SpecializedConfigCpuBuilder};
+use crate::{Instr, SpecializedConfig, SpecializedConfigCpuBuilder};
 use tracing::info_span;
 
 use std::collections::HashSet;
@@ -80,27 +81,11 @@ where
 }
 
 pub fn execution_stats(
-    program: &CompiledProgram,
+    exe: Arc<VmExe<BabyBear>>,
+    vm_config: SpecializedConfig,
     inputs: StdIn,
-    segment_height: Option<usize>, // uses the default height if None
-    apc_candidates_dir: Option<PathBuf>,
+    blocks: &[BasicBlock<Instr<BabyBear>>],
 ) -> Result<(), Box<dyn std::error::Error>> {
-    let exe = &program.exe;
-    let mut vm_config = program.vm_config.clone();
-
-    // DefaultSegmentationStrategy { max_segment_len: 4194204, max_cells_per_chip_in_segment: 503304480 }
-    if let Some(segment_height) = segment_height {
-        vm_config
-            .sdk
-            .config_mut()
-            .sdk
-            .system
-            .config
-            .segmentation_limits =
-            SegmentationLimits::default().with_max_trace_height(segment_height as u32);
-        tracing::debug!("Setting max segment len to {}", segment_height);
-    }
-
     // Set app configuration
     let app_fri_params =
         FriParameters::standard_with_100_bits_conjectured_security(DEFAULT_APP_LOG_BLOWUP);
@@ -226,23 +211,10 @@ pub fn execution_stats(
         }
     }
 
-    let apc_candidates_dir = apc_candidates_dir.unwrap();
-    let apc_candiates: powdr_autoprecompiles::pgo::JsonExport = {
-        let json_str =
-            std::fs::read_to_string(apc_candidates_dir.join("apc_candidates.json")).unwrap();
-        serde_json::from_str(&json_str).unwrap()
-    };
-    let apcs = apc_candiates.apcs;
-
     // Block ID -> instruction count mapping
-    let instruction_counts = apcs
+    let instruction_counts = blocks
         .iter()
-        .map(|apc| {
-            (
-                apc.original_block.start_pc,
-                apc.original_block.statements.len(),
-            )
-        })
+        .map(|block| (block.start_pc, block.statements.len()))
         .collect::<HashMap<_, _>>();
 
     // Block ID -> Vec<Vec<Row>>

openvm/src/lib.rs

@@ -534,6 +534,7 @@ pub fn compile_guest(
     guest_opts: GuestOptions,
     config: PowdrConfig,
     pgo_config: PgoConfig,
+    inputs: StdIn,
 ) -> Result<CompiledProgram, Box<dyn std::error::Error>> {
     let original_program = compile_openvm(guest, guest_opts.clone())?;
 
@@ -542,7 +543,7 @@ pub fn compile_guest(
         tally_opcode_frequency(&pgo_config, &original_program.exe);
     }
 
-    compile_exe(original_program, config, pgo_config)
+    compile_exe(original_program, config, pgo_config, inputs)
 }
 
 fn instruction_index_to_pc(program: &Program<BabyBear>, idx: usize) -> u64 {
@@ -590,6 +591,7 @@ pub fn compile_exe(
     original_program: OriginalCompiledProgram,
     config: PowdrConfig,
     pgo_config: PgoConfig,
+    inputs: StdIn,
 ) -> Result<CompiledProgram, Box<dyn std::error::Error>> {
     let compiled = match pgo_config {
         PgoConfig::Cell(pgo_data, max_total_columns) => {
@@ -611,14 +613,16 @@ pub fn compile_exe(
                     pgo_data,
                     max_total_apc_columns,
                 ),
+                inputs,
             )
         }
         PgoConfig::Instruction(pgo_data) => customize(
             original_program,
             config,
             InstructionPgo::with_pgo_data(pgo_data),
+            inputs,
         ),
-        PgoConfig::None => customize(original_program, config, NonePgo::default()),
+        PgoConfig::None => customize(original_program, config, NonePgo::default(), inputs),
     };
     // Export the compiled program to a PIL file for debugging purposes.
     export_pil(
@@ -848,10 +852,8 @@ pub fn prove(
     recursion: bool,
     inputs: StdIn,
     segment_height: Option<usize>, // uses the default height if None
-    apc_candidates_dir: Option<PathBuf>,
+    _apc_candidates_dir: Option<PathBuf>,
 ) -> Result<(), Box<dyn std::error::Error>> {
-    execution_stats(program, inputs.clone(), segment_height, apc_candidates_dir)?;
-
     let exe = &program.exe;
     let mut vm_config = program.vm_config.clone();
 
@@ -1001,7 +1003,14 @@ mod tests {
         pgo_config: PgoConfig,
         segment_height: Option<usize>,
     ) -> Result<(), Box<dyn std::error::Error>> {
-        let program = compile_guest(guest, GuestOptions::default(), config, pgo_config).unwrap();
+        let program = compile_guest(
+            guest,
+            GuestOptions::default(),
+            config,
+            pgo_config,
+            stdin.clone(),
+        )
+        .unwrap();
         prove(&program, mock, recursion, stdin, segment_height, None)
     }
 
@@ -1120,8 +1129,14 @@ mod tests {
         let pgo_data = execution_profile_from_guest(GUEST, GuestOptions::default(), stdin.clone());
 
         let config = default_powdr_openvm_config(GUEST_APC, GUEST_SKIP_NO_APC_EXECUTED);
-        let program =
-            compile_guest(GUEST, GuestOptions::default(), config, PgoConfig::None).unwrap();
+        let program = compile_guest(
+            GUEST,
+            GuestOptions::default(),
+            config,
+            PgoConfig::None,
+            stdin.clone(),
+        )
+        .unwrap();
 
         // Assert that all APCs aren't executed
         program
@@ -1170,9 +1185,14 @@ mod tests {
     fn matmul_compile() {
         let guest = "guest-matmul";
         let config = default_powdr_openvm_config(1, 0);
-        assert!(
-            compile_guest(guest, GuestOptions::default(), config, PgoConfig::default()).is_ok()
-        );
+        assert!(compile_guest(
+            guest,
+            GuestOptions::default(),
+            config,
+            PgoConfig::default(),
+            StdIn::default()
+        )
+        .is_ok());
     }
 
     #[test]
@@ -1654,6 +1674,8 @@ mod tests {
             GuestOptions::default(),
             config,
             guest.pgo_config,
+            // TODO
+            StdIn::default(),
         )
         .unwrap();