Empirical constraints suggestions

autoprecompiles/src/empirical_constraints.rs

@@ -1,38 +1,35 @@
 use std::collections::btree_map::Entry;
-use std::collections::{BTreeMap, BTreeSet, HashMap};
+use std::collections::{BTreeMap, HashMap};
 use std::fmt::Debug;
 use std::hash::Hash;
 
 use itertools::Itertools;
 use serde::{Deserialize, Serialize};
 
+pub use crate::equivalence_classes::{EquivalenceClass, EquivalenceClasses};
+
 /// "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)]
+#[derive(Serialize, 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)>>>,
+    pub equivalence_classes_by_block: BTreeMap<u64, EquivalenceClasses<BlockCell>>,
+    pub debug_info: DebugInfo,
 }
 
 /// Debug information mapping AIR ids to program counters and column names.
-#[derive(Serialize, Deserialize, Default)]
+#[derive(Serialize, Deserialize, Default, Debug)]
 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
@@ -52,15 +49,14 @@ impl EmpiricalConstraints {
 
         // Combine equivalence classes by block
         for (block_pc, classes) in other.equivalence_classes_by_block {
-            self.equivalence_classes_by_block
+            let existing = 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);
+                .or_default();
+
+            *existing = intersect_partitions(vec![std::mem::take(existing), classes]);
         }
+        self.debug_info.combine_with(other.debug_info);
     }
 }
 
@@ -88,16 +84,32 @@ fn merge_maps<K: Ord, V: Eq + Debug>(map1: &mut BTreeMap<K, V>, map2: BTreeMap<K
     }
 }
 
+#[derive(PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize, Debug, Copy, Clone)]
+pub struct BlockCell {
+    /// The row index which is also the instruction index within the basic block
+    row_idx: usize,
+    /// The column index within the instruction air
+    column_idx: usize,
+}
+
+impl BlockCell {
+    pub fn new(row_idx: usize, column_idx: usize) -> Self {
+        Self {
+            row_idx,
+            column_idx,
+        }
+    }
+}
+
 /// 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,
-{
+pub fn intersect_partitions<T: Eq + Hash + Copy + Ord>(
+    partitions: Vec<EquivalenceClasses<T>>,
+) -> EquivalenceClasses<T> {
     // For each partition, build a map: Id -> class_index
-    let class_ids: Vec<HashMap<Id, usize>> = partitions
+    let class_ids: Vec<HashMap<T, usize>> = partitions
         .iter()
         .map(|partition| {
             partition
@@ -112,14 +124,14 @@ where
     partitions
         .iter()
         .flat_map(|partition| partition.iter())
-        .flat_map(|class| class.iter().copied())
+        .flat_map(|class| class.iter())
         .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())
+                .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>>>()
@@ -128,16 +140,16 @@ where
         // 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()))
+        // Convert to set
+        .map(|ids| ids.into_iter().copied().collect())
         .collect()
 }
 
 #[cfg(test)]
 mod tests {
-    use std::collections::BTreeSet;
+    use crate::empirical_constraints::EquivalenceClasses;
 
-    fn partition(sets: Vec<Vec<u32>>) -> BTreeSet<BTreeSet<u32>> {
+    fn partition(sets: Vec<Vec<u32>>) -> EquivalenceClasses<u32> {
         sets.into_iter().map(|s| s.into_iter().collect()).collect()
     }
 
@@ -156,7 +168,7 @@ mod tests {
             vec![6, 7, 8],
         ]);
 
-        let result = super::intersect_partitions(&[partition1, partition2]);
+        let result = super::intersect_partitions(vec![partition1, partition2]);
 
         let expected = partition(vec![vec![2, 3], vec![6, 7, 8]]);
 

autoprecompiles/src/equivalence_classes.rs

@@ -0,0 +1,36 @@
+use std::collections::BTreeSet;
+
+use serde::Serialize;
+
+/// An equivalence class 
+pub type EquivalenceClass<T> = BTreeSet<T>;
+
+/// A collection of equivalence classes where all classes are guaranteed to have at least two elements
+#[derive(Serialize, Debug, PartialEq, Eq)]
+pub struct EquivalenceClasses<T> {
+    inner: BTreeSet<EquivalenceClass<T>>,
+}
+
+// TODO: derive
+impl<T> Default for EquivalenceClasses<T> {
+    fn default() -> Self {
+        Self {
+            inner: Default::default(),
+        }
+    }
+}
+
+impl<T: Ord> FromIterator<EquivalenceClass<T>> for EquivalenceClasses<T> {
+    fn from_iter<I: IntoIterator<Item = EquivalenceClass<T>>>(iter: I) -> Self {
+        // When collecting, we ignore classes with 0 or 1 elements as they are useless
+        Self {
+            inner: iter.into_iter().filter(|class| class.len() > 1).collect(),
+        }
+    }
+}
+
+impl<T> EquivalenceClasses<T> {
+    pub fn iter(&self) -> impl Iterator<Item = &EquivalenceClass<T>> {
+        self.inner.iter()
+    }
+}

autoprecompiles/src/lib.rs

@@ -41,6 +41,7 @@ mod stats_logger;
 pub mod symbolic_machine_generator;
 pub use pgo::{PgoConfig, PgoType};
 pub use powdr_constraint_solver::inliner::DegreeBound;
+pub mod equivalence_classes;
 pub mod trace_handler;
 
 #[derive(Clone)]

cli-openvm/src/main.rs

@@ -3,7 +3,6 @@ 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};
@@ -311,19 +310,15 @@ fn maybe_compute_empirical_constraints(
         "Optimistic precompiles are not implemented yet. Computing empirical constraints..."
     );
 
-    let (empirical_constraints, debug_info) =
+    let empirical_constraints =
         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();
+        let json = serde_json::to_string_pretty(&empirical_constraints).unwrap();
         std::fs::write(path.join("empirical_constraints.json"), json).unwrap();
     }
 }