Collect empirical constraints

Author: georgwiese

autoprecompiles/src/empirical_constraints.rs

@@ -0,0 +1,165 @@
+use std::collections::btree_map::Entry;
+use std::collections::{BTreeMap, BTreeSet, HashMap};
+use std::fmt::Debug;
+use std::hash::Hash;
+
+use itertools::Itertools;
+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,);
+            }
+        }
+    }
+}
+
+/// Intersects multiple partitions of the same universe into a single partition.
+/// In other words, two elements are in the same equivalence class in the resulting partition
+/// if and only if they are in the same equivalence class in all input partitions.
+/// Singleton equivalence classes are omitted from the result.
+pub fn intersect_partitions<Id>(partitions: &[BTreeSet<BTreeSet<Id>>]) -> BTreeSet<BTreeSet<Id>>
+where
+    Id: Eq + Hash + Copy + Ord,
+{
+    // For each partition, build a map: Id -> class_index
+    let class_ids: Vec<HashMap<Id, usize>> = partitions
+        .iter()
+        .map(|partition| {
+            partition
+                .iter()
+                .enumerate()
+                .flat_map(|(class_idx, class)| class.iter().map(move |&id| (id, class_idx)))
+                .collect()
+        })
+        .collect();
+
+    // Iterate over all elements in the universe
+    partitions
+        .iter()
+        .flat_map(|partition| partition.iter())
+        .flat_map(|class| class.iter().copied())
+        .unique()
+        .filter_map(|id| {
+            // Build the signature of the element: the list of class indices it belongs to
+            // (one index per partition)
+            class_ids
+                .iter()
+                .map(|m| m.get(&id).cloned())
+                // If an element did not appear in any one of the partitions, it is
+                // a singleton and we skip it.
+                .collect::<Option<Vec<usize>>>()
+                .map(|signature| (signature, id))
+        })
+        // Group elements by their signatures
+        .into_group_map()
+        .into_values()
+        // Remove singletons and convert to Set
+        .filter_map(|ids| (ids.len() > 1).then_some(ids.into_iter().collect()))
+        .collect()
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::BTreeSet;
+
+    fn partition(sets: Vec<Vec<u32>>) -> BTreeSet<BTreeSet<u32>> {
+        sets.into_iter().map(|s| s.into_iter().collect()).collect()
+    }
+
+    #[test]
+    fn test_intersect_partitions() {
+        let partition1 = partition(vec![
+            // Two classes: 1-4 and 5-9
+            vec![1, 2, 3, 4],
+            vec![5, 6, 7, 8, 9],
+        ]);
+        let partition2 = partition(vec![
+            // Four classes: 1, 2-3, 4-5, 6-8, 9 (implicit)
+            vec![1],
+            vec![2, 3],
+            vec![4, 5],
+            vec![6, 7, 8],
+        ]);
+
+        let result = super::intersect_partitions(&[partition1, partition2]);
+
+        let expected = partition(vec![vec![2, 3], vec![6, 7, 8]]);
+
+        assert_eq!(result, expected);
+    }
+}

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();
+    }
+}