Collect empirical constraints

Author: georgwiese

autoprecompiles/src/empirical_constraints.rs

@@ -0,0 +1,144 @@
+use std::collections::btree_map::Entry;
+use std::collections::{BTreeMap, BTreeSet, HashMap};
+use std::fmt::Debug;
+use std::hash::Hash;
+
+use serde::{Deserialize, Serialize};
+
+/// "Constraints" that were inferred from execution statistics. They hold empirically
+/// (most of the time), but are not guaranteed to hold in all cases.
+#[derive(Serialize, Deserialize, Clone, Default, Debug)]
+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, BTreeSet<BTreeSet<(usize, usize)>>>,
+}
+
+/// Debug information mapping AIR ids to program counters and column names.
+#[derive(Serialize, Deserialize, Default)]
+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,
+}
+
+impl EmpiricalConstraints {
+    pub fn combine_with(&mut self, other: EmpiricalConstraints) {
+        // Combine column ranges by PC
+        for (pc, ranges) in other.column_ranges_by_pc {
+            self.column_ranges_by_pc
+                .entry(pc)
+                .and_modify(|existing_ranges| {
+                    for (i, (min, max)) in ranges.iter().enumerate() {
+                        if let Some((existing_min, existing_max)) = existing_ranges.get_mut(i) {
+                            *existing_min = (*existing_min).min(*min);
+                            *existing_max = (*existing_max).max(*max);
+                        }
+                    }
+                })
+                .or_insert(ranges);
+        }
+
+        // Combine equivalence classes by block
+        for (block_pc, classes) in other.equivalence_classes_by_block {
+            self.equivalence_classes_by_block
+                .entry(block_pc)
+                .and_modify(|existing_classes| {
+                    let combined =
+                        intersect_partitions(&[existing_classes.clone(), classes.clone()]);
+                    *existing_classes = combined;
+                })
+                .or_insert(classes);
+        }
+    }
+}
+
+impl DebugInfo {
+    pub fn combine_with(&mut self, other: DebugInfo) {
+        merge_maps(&mut self.air_id_by_pc, other.air_id_by_pc);
+        merge_maps(
+            &mut self.column_names_by_air_id,
+            other.column_names_by_air_id,
+        );
+    }
+}
+
+/// Merges two maps, asserting that existing keys map to equal values.
+fn merge_maps<K: Ord, V: Eq + Debug>(map1: &mut BTreeMap<K, V>, map2: BTreeMap<K, V>) {
+    for (key, value) in map2 {
+        match map1.entry(key) {
+            Entry::Vacant(v) => {
+                v.insert(value);
+            }
+            Entry::Occupied(existing) => {
+                assert_eq!(*existing.get(), value,);
+            }
+        }
+    }
+}
+
+// ChatGPT generated code
+pub fn intersect_partitions<Id>(partitions: &[BTreeSet<BTreeSet<Id>>]) -> BTreeSet<BTreeSet<Id>>
+where
+    Id: Eq + Hash + Copy + Ord,
+{
+    if partitions.is_empty() {
+        return BTreeSet::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: BTreeSet<Id> = BTreeSet::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>, BTreeSet<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().insert(id);
+        }
+    }
+
+    // 4) Resulting equivalence classes are the grouped values
+    grouped.into_values().collect()
+}

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, vec![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();
+    }
+}