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;
@@ -53,6 +54,8 @@ pub struct PowdrConfig {
     pub degree_bound: DegreeBound,
     /// The path to the APC candidates dir, if any.
     pub apc_candidates_dir_path: Option<PathBuf>,
+    /// Whether to use optimistic precompiles.
+    pub optimistic_precompiles: bool,
 }
 
 impl PowdrConfig {
@@ -62,13 +65,19 @@ impl PowdrConfig {
             skip_autoprecompiles,
             degree_bound,
             apc_candidates_dir_path: None,
+            optimistic_precompiles: false,
         }
     }
 
     pub fn with_apc_candidates_dir<P: AsRef<Path>>(mut self, path: P) -> Self {
         self.apc_candidates_dir_path = Some(path.as_ref().to_path_buf());
         self
     }
+
+    pub fn with_optimistic_precompiles(mut self, optimistic: bool) -> Self {
+        self.optimistic_precompiles = optimistic;
+        self
+    }
 }
 
 #[derive(Debug, Clone, PartialEq, Hash, Eq, Serialize, Deserialize)]

cli-openvm/src/main.rs

@@ -3,8 +3,11 @@ use metrics_tracing_context::{MetricsLayer, TracingContextLayer};
 use metrics_util::{debugging::DebuggingRecorder, layers::Layer};
 use openvm_sdk::StdIn;
 use openvm_stark_sdk::bench::serialize_metric_snapshot;
+use powdr_autoprecompiles::empirical_constraints::EmpiricalConstraintsJson;
 use powdr_autoprecompiles::pgo::{pgo_config, PgoType};
+use powdr_autoprecompiles::PowdrConfig;
 use powdr_openvm::{compile_openvm, default_powdr_openvm_config, CompiledProgram, GuestOptions};
+use powdr_openvm::{detect_empirical_constraints, OriginalCompiledProgram};
 
 #[cfg(feature = "metrics")]
 use openvm_stark_sdk::metrics_tracing::TimingMetricsLayer;
@@ -46,6 +49,12 @@ enum Commands {
         /// When `--pgo-mode cell`, the directory to persist all APC candidates + a metrics summary
         #[arg(long)]
         apc_candidates_dir: Option<PathBuf>,
+
+        /// If active, generates "optimistic" precompiles. Optimistic precompiles are smaller in size
+        /// but may fail at runtime if the assumptions they make are violated.
+        #[arg(long)]
+        #[arg(default_value_t = false)]
+        optimistic_precompiles: bool,
     },
 
     Execute {
@@ -73,6 +82,12 @@ enum Commands {
         /// When `--pgo-mode cell`, the directory to persist all APC candidates + a metrics summary
         #[arg(long)]
         apc_candidates_dir: Option<PathBuf>,
+
+        /// If active, generates "optimistic" precompiles. Optimistic precompiles are smaller in size
+        /// but may fail at runtime if the assumptions they make are violated.
+        #[arg(long)]
+        #[arg(default_value_t = false)]
+        optimistic_precompiles: bool,
     },
 
     Prove {
@@ -108,6 +123,12 @@ enum Commands {
         /// When `--pgo-mode cell`, the directory to persist all APC candidates + a metrics summary
         #[arg(long)]
         apc_candidates_dir: Option<PathBuf>,
+
+        /// If active, generates "optimistic" precompiles. Optimistic precompiles are smaller in size
+        /// but may fail at runtime if the assumptions they make are violated.
+        #[arg(long)]
+        #[arg(default_value_t = false)]
+        optimistic_precompiles: bool,
     },
 }
 
@@ -135,15 +156,20 @@ fn run_command(command: Commands) {
             max_columns,
             input,
             apc_candidates_dir,
+            optimistic_precompiles,
         } => {
             let mut powdr_config = default_powdr_openvm_config(autoprecompiles as u64, skip as u64);
             if let Some(apc_candidates_dir) = apc_candidates_dir {
                 powdr_config = powdr_config.with_apc_candidates_dir(apc_candidates_dir);
             }
+            if optimistic_precompiles {
+                powdr_config = powdr_config.with_optimistic_precompiles(true);
+            }
             let guest_program = compile_openvm(&guest, guest_opts.clone()).unwrap();
             let execution_profile =
                 powdr_openvm::execution_profile_from_guest(&guest_program, stdin_from(input));
 
+            maybe_compute_empirical_constraints(&guest_program, &powdr_config, stdin_from(input));
             let pgo_config = pgo_config(pgo, max_columns, execution_profile);
             let program =
                 powdr_openvm::compile_exe(guest_program, powdr_config, pgo_config).unwrap();
@@ -159,12 +185,17 @@ fn run_command(command: Commands) {
             input,
             metrics,
             apc_candidates_dir,
+            optimistic_precompiles,
         } => {
             let mut powdr_config = default_powdr_openvm_config(autoprecompiles as u64, skip as u64);
             if let Some(apc_candidates_dir) = apc_candidates_dir {
                 powdr_config = powdr_config.with_apc_candidates_dir(apc_candidates_dir);
             }
+            if optimistic_precompiles {
+                powdr_config = powdr_config.with_optimistic_precompiles(true);
+            }
             let guest_program = compile_openvm(&guest, guest_opts.clone()).unwrap();
+            maybe_compute_empirical_constraints(&guest_program, &powdr_config, stdin_from(input));
             let execution_profile =
                 powdr_openvm::execution_profile_from_guest(&guest_program, stdin_from(input));
             let pgo_config = pgo_config(pgo, max_columns, execution_profile);
@@ -194,12 +225,17 @@ fn run_command(command: Commands) {
             input,
             metrics,
             apc_candidates_dir,
+            optimistic_precompiles,
         } => {
             let mut powdr_config = default_powdr_openvm_config(autoprecompiles as u64, skip as u64);
             if let Some(apc_candidates_dir) = apc_candidates_dir {
                 powdr_config = powdr_config.with_apc_candidates_dir(apc_candidates_dir);
             }
+            if optimistic_precompiles {
+                powdr_config = powdr_config.with_optimistic_precompiles(true);
+            }
             let guest_program = compile_openvm(&guest, guest_opts).unwrap();
+            maybe_compute_empirical_constraints(&guest_program, &powdr_config, stdin_from(input));
 
             let execution_profile =
                 powdr_openvm::execution_profile_from_guest(&guest_program, stdin_from(input));
@@ -261,3 +297,33 @@ pub fn run_with_metric_collection_to_file<R>(file: std::fs::File, f: impl FnOnce
         .unwrap();
     res
 }
+
+fn maybe_compute_empirical_constraints(
+    guest_program: &OriginalCompiledProgram,
+    powdr_config: &PowdrConfig,
+    stdin: StdIn,
+) {
+    if !powdr_config.optimistic_precompiles {
+        return;
+    }
+
+    tracing::warn!(
+        "Optimistic precompiles are not implemented yet. Computing empirical constraints..."
+    );
+
+    let (empirical_constraints, debug_info) =
+        detect_empirical_constraints(&guest_program, powdr_config.degree_bound, stdin);
+
+    if let Some(path) = &powdr_config.apc_candidates_dir_path {
+        tracing::info!(
+            "Saving empirical constraints debug info to {}/empirical_constraints.json",
+            path.display()
+        );
+        let export = EmpiricalConstraintsJson {
+            empirical_constraints: empirical_constraints.clone(),
+            debug_info,
+        };
+        let json = serde_json::to_string_pretty(&export).unwrap();
+        std::fs::write(path.join("empirical_constraints.json"), json).unwrap();
+    }
+}