Collect empirical constraints

Author: georgwiese

autoprecompiles/src/empirical_constraints.rs

@@ -0,0 +1,29 @@
+use std::collections::BTreeMap;
+
+use serde::{Deserialize, Serialize};
+
+/// "Constraints" that were inferred from execution statistics.
+#[derive(Serialize, Deserialize, Clone, Default)]
+pub struct EmpiricalConstraints {
+    /// For each program counter, the range constraints for each column.
+    /// The range might not hold in 100% of cases.
+    pub column_ranges_by_pc: BTreeMap<u32, Vec<(u32, u32)>>,
+    /// For each basic block (identified by its starting PC), the equivalence classes of columns.
+    /// Each equivalence class is a list of (instruction index in block, column index).
+    pub equivalence_classes_by_block: BTreeMap<u64, Vec<Vec<(usize, usize)>>>,
+}
+
+/// Debug information mapping AIR ids to program counters and column names.
+#[derive(Serialize, Deserialize)]
+pub struct DebugInfo {
+    /// Mapping from program counter to AIR id.
+    pub air_id_by_pc: BTreeMap<u32, usize>,
+    /// Mapping from AIR id to column names.
+    pub column_names_by_air_id: BTreeMap<usize, Vec<String>>,
+}
+
+#[derive(Serialize, Deserialize)]
+pub struct EmpiricalConstraintsJson {
+    pub empirical_constraints: EmpiricalConstraints,
+    pub debug_info: DebugInfo,
+}

autoprecompiles/src/lib.rs

@@ -26,6 +26,7 @@ pub mod adapter;
 pub mod blocks;
 pub mod bus_map;
 pub mod constraint_optimizer;
+pub mod empirical_constraints;
 pub mod evaluation;
 pub mod execution_profile;
 pub mod expression;

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 {
@@ -210,9 +221,14 @@ fn run_command(command: Commands) {
             );
             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();
+                let program = powdr_openvm::compile_guest(
+                    &guest,
+                    guest_opts,
+                    powdr_config,
+                    pgo_config,
+                    stdin_from(input),
+                )
+                .unwrap();
                 powdr_openvm::prove(&program, mock, recursion, stdin_from(input), None).unwrap()
             };
             if let Some(metrics_path) = metrics {

openvm/Cargo.toml

@@ -61,6 +61,7 @@ derive_more = { version = "2.0.1", default-features = false, features = [
   "from",
 ] }
 itertools = "0.14.0"
+serde_json = "1.0.140"
 
 tracing = "0.1.40"
 tracing-subscriber = { version = "0.3.17", features = ["std", "env-filter"] }

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::{empirical_constraints::detect_empirical_constraints, 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,17 @@ pub fn customize<'a, P: PgoAdapter<Adapter = BabyBearOpenVmApcAdapter<'a>>>(
         })
         .collect();
 
+    let _empirical_constraints = detect_empirical_constraints(
+        exe.clone(),
+        SpecializedConfig::new(
+            original_config.clone(),
+            vec![],
+            config.degree_bound.identities,
+        ),
+        inputs,
+        &blocks,
+    );
+
     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);