Build optimistic precompiles

georgwiese · georgwiese · commit 2b586293d2ca · 2025-11-27T21:02:39.000-05:00
diff --git a/autoprecompiles/src/adapter.rs b/autoprecompiles/src/adapter.rs
@@ -6,6 +6,7 @@ use std::{fmt::Display, sync::Arc};
 use powdr_number::FieldElement;
 use serde::{Deserialize, Serialize};
 
+use crate::EmpiricalConstraints;
 use crate::{
     blocks::{BasicBlock, Instruction, Program},
     constraint_optimizer::IsBusStateful,
@@ -45,12 +46,19 @@ pub trait PgoAdapter {
         config: &PowdrConfig,
         vm_config: AdapterVmConfig<Self::Adapter>,
         labels: BTreeMap<u64, Vec<String>>,
+        empirical_constraints: EmpiricalConstraints,
     ) -> Vec<AdapterApcWithStats<Self::Adapter>> {
         let filtered_blocks = blocks
             .into_iter()
             .filter(|block| !Self::Adapter::should_skip_block(block))
             .collect();
-        self.create_apcs_with_pgo(filtered_blocks, config, vm_config, labels)
+        self.create_apcs_with_pgo(
+            filtered_blocks,
+            config,
+            vm_config,
+            labels,
+            empirical_constraints,
+        )
     }
 
     fn create_apcs_with_pgo(
@@ -59,6 +67,7 @@ pub trait PgoAdapter {
         config: &PowdrConfig,
         vm_config: AdapterVmConfig<Self::Adapter>,
         labels: BTreeMap<u64, Vec<String>>,
+        empirical_constraints: EmpiricalConstraints,
     ) -> Vec<AdapterApcWithStats<Self::Adapter>>;
 
     fn pc_execution_count(&self, _pc: u64) -> Option<u32> {
diff --git a/autoprecompiles/src/empirical_constraints.rs b/autoprecompiles/src/empirical_constraints.rs
@@ -2,6 +2,13 @@ use std::collections::BTreeMap;
 
 use serde::{Deserialize, Serialize};
 
+use crate::{
+    adapter::Adapter,
+    blocks::BasicBlock,
+    expression::{AlgebraicExpression, AlgebraicReference},
+    SymbolicConstraint,
+};
+
 /// "Constraints" that were inferred from execution statistics.
 #[derive(Serialize, Deserialize, Clone, Default)]
 pub struct EmpiricalConstraints {
@@ -27,3 +34,84 @@ pub struct EmpiricalConstraintsJson {
     pub empirical_constraints: EmpiricalConstraints,
     pub debug_info: DebugInfo,
 }
+
+pub fn add_empirical_constraints<A: Adapter>(
+    empirical_constraints: &EmpiricalConstraints,
+    subs: &[Vec<u64>],
+    block: &BasicBlock<A::Instruction>,
+    columns: impl Iterator<Item = AlgebraicReference>,
+) -> (
+    Vec<SymbolicConstraint<A::PowdrField>>,
+    Vec<SymbolicConstraint<A::PowdrField>>,
+) {
+    let range_constraints = &empirical_constraints.column_ranges_by_pc;
+    let equivalence_classes_by_block = &empirical_constraints.equivalence_classes_by_block;
+
+    let mut range_analyzer_constraints = Vec::new();
+    let mut equivalence_analyzer_constraints = Vec::new();
+
+    // Mapping (instruction index, column index) -> AlgebraicReference
+    let reverse_subs = subs
+        .iter()
+        .enumerate()
+        .flat_map(|(instr_index, subs)| {
+            subs.iter()
+                .enumerate()
+                .map(move |(col_index, &poly_id)| (poly_id, (instr_index, col_index)))
+        })
+        .collect::<BTreeMap<_, _>>();
+    let algebraic_references = columns
+        .map(|r| (reverse_subs.get(&r.id).unwrap().clone(), r.clone()))
+        .collect::<BTreeMap<_, _>>();
+
+    for i in 0..block.statements.len() {
+        let pc = (block.start_pc + (i * 4) as u64) as u32;
+        let Some(range_constraints) = range_constraints.get(&pc) else {
+            continue;
+        };
+        for (col_index, range) in range_constraints.iter().enumerate() {
+            if range.0 == range.1 {
+                let value = A::PowdrField::from(range.0 as u64);
+                let Some(reference) = algebraic_references.get(&(i, col_index)).cloned() else {
+                    panic!(
+                        "Missing reference for (i: {}, col_index: {}, block_id: {})",
+                        i, col_index, block.start_pc
+                    );
+                };
+                let constraint =
+                    AlgebraicExpression::Reference(reference) - AlgebraicExpression::Number(value);
+
+                range_analyzer_constraints.push(SymbolicConstraint { expr: constraint });
+            }
+        }
+    }
+
+    if let Some(equivalence_classes) = equivalence_classes_by_block.get(&block.start_pc) {
+        for equivalence_class in equivalence_classes {
+            let first = equivalence_class.first().unwrap();
+            let Some(first_ref) = algebraic_references.get(first).cloned() else {
+                // TODO: This fails in some blocks. For now, just return no extra constraints.
+                println!(
+                    "Missing reference for (i: {}, col_index: {}, block_id: {})",
+                    first.0, first.1, block.start_pc
+                );
+                return (range_analyzer_constraints, vec![]);
+            };
+            for other in equivalence_class.iter().skip(1) {
+                let Some(other_ref) = algebraic_references.get(other).cloned() else {
+                    // TODO: This fails in some blocks. For now, just return no extra constraints.
+                    println!(
+                        "Missing reference for (i: {}, col_index: {}, block_id: {})",
+                        other.0, other.1, block.start_pc
+                    );
+                    return (range_analyzer_constraints, vec![]);
+                };
+                let constraint = AlgebraicExpression::Reference(first_ref.clone())
+                    - AlgebraicExpression::Reference(other_ref.clone());
+                equivalence_analyzer_constraints.push(SymbolicConstraint { expr: constraint });
+            }
+        }
+    }
+
+    (range_analyzer_constraints, equivalence_analyzer_constraints)
+}
diff --git a/autoprecompiles/src/lib.rs b/autoprecompiles/src/lib.rs
@@ -1,6 +1,7 @@
 use crate::adapter::{Adapter, AdapterApc, AdapterVmConfig};
 use crate::blocks::BasicBlock;
 use crate::bus_map::{BusMap, BusType};
+use crate::empirical_constraints::{add_empirical_constraints, EmpiricalConstraints};
 use crate::evaluation::AirStats;
 use crate::expression_conversion::algebraic_to_grouped_expression;
 use crate::symbolic_machine_generator::convert_machine_field_type;
@@ -394,9 +395,10 @@ impl<T, I> Apc<T, I> {
 }
 
 /// Allocates global poly_ids and keeps track of substitutions
+#[derive(Debug)]
 pub struct ColumnAllocator {
     /// For each original air, for each original column index, the associated poly_id in the APC air
-    subs: Vec<Vec<u64>>,
+    pub subs: Vec<Vec<u64>>,
     /// The next poly_id to issue
     next_poly_id: u64,
 }
@@ -421,6 +423,7 @@ pub fn build<A: Adapter>(
     vm_config: AdapterVmConfig<A>,
     degree_bound: DegreeBound,
     apc_candidates_dir_path: Option<&Path>,
+    empirical_constraints: &EmpiricalConstraints,
 ) -> Result<AdapterApc<A>, crate::constraint_optimizer::Error> {
     let start = std::time::Instant::now();
 
@@ -430,6 +433,14 @@ pub fn build<A: Adapter>(
         &vm_config.bus_map,
     );
 
+    let (range_analyzer_constraints, _equivalence_analyzer_constraints) =
+        add_empirical_constraints::<A>(
+            empirical_constraints,
+            &column_allocator.subs,
+            &block,
+            machine.main_columns(),
+        );
+
     let labels = [("apc_start_pc", block.start_pc.to_string())];
     metrics::counter!("before_opt_cols", &labels)
         .absolute(machine.unique_references().count() as u64);
@@ -438,13 +449,33 @@ pub fn build<A: Adapter>(
     metrics::counter!("before_opt_interactions", &labels)
         .absolute(machine.unique_references().count() as u64);
 
+    let mut baseline = machine;
+
+    let (machine, column_allocator) = optimizer::optimize::<A>(
+        baseline.clone(),
+        vm_config.bus_interaction_handler.clone(),
+        degree_bound,
+        &vm_config.bus_map,
+        column_allocator,
+    )
+    .unwrap();
+    let dumb_precompile = machine.render(&vm_config.bus_map);
+
+    baseline.constraints.extend(range_analyzer_constraints);
+    // TODO: Appears to be buggy
+    // baseline
+    //     .constraints
+    //     .extend(equivalence_analyzer_constraints);
+
     let (machine, column_allocator) = optimizer::optimize::<A>(
-        machine,
+        baseline,
         vm_config.bus_interaction_handler,
         degree_bound,
         &vm_config.bus_map,
         column_allocator,
-    )?;
+    )
+    .unwrap();
+    let ai_precompile = machine.render(&vm_config.bus_map);
 
     // add guards to constraints that are not satisfied by zeroes
     let (machine, column_allocator) = add_guards(machine, column_allocator);
@@ -469,6 +500,16 @@ pub fn build<A: Adapter>(
             std::fs::File::create(&ser_path).expect("Failed to create file for APC candidate");
         let writer = BufWriter::new(file);
         serde_cbor::to_writer(writer, &apc).expect("Failed to write APC candidate to file");
+
+        let dumb_path = path
+            .join(format!("apc_candidate_{}_dumb.txt", apc.start_pc()))
+            .with_extension("txt");
+        std::fs::write(dumb_path, dumb_precompile).unwrap();
+
+        let ai_path = path
+            .join(format!("apc_candidate_{}_ai.txt", apc.start_pc()))
+            .with_extension("txt");
+        std::fs::write(ai_path, ai_precompile).unwrap();
     }
 
     metrics::gauge!("apc_gen_time_ms", &labels).set(start.elapsed().as_millis() as f64);
@@ -556,11 +597,16 @@ fn add_guards<T: FieldElement>(
 
     machine.constraints.extend(is_valid_mults);
 
-    assert_eq!(
-        pre_degree,
-        machine.degree(),
-        "Degree should not change after adding guards"
-    );
+    // TODO: Why do we need this?
+    if pre_degree != 0 {
+        assert_eq!(
+            pre_degree,
+            machine.degree(),
+            "Degree should not change after adding guards, but changed from {} to {}",
+            pre_degree,
+            machine.degree(),
+        );
+    }
 
     // This needs to be added after the assertion above because it's a quadratic constraint
     // so it may increase the degree of the machine.
diff --git a/autoprecompiles/src/memory_optimizer.rs b/autoprecompiles/src/memory_optimizer.rs
@@ -187,6 +187,16 @@ fn redundant_memory_interactions_indices<
                 // In that case, we can replace both bus interactions with equality constraints
                 // between the data that would have been sent and received.
                 if let Some((previous_send, existing_values)) = memory_contents.remove(&addr) {
+                    if existing_values.len() != mem_int.data().len() {
+                        log::error!(
+                            "Memory interaction data length mismatch: existing values = {}, new values = {}. Resetting memory.",
+                            existing_values.iter().map(ToString::to_string).join(", "),
+                            mem_int.data().iter().map(ToString::to_string).join(", ")
+                        );
+                        memory_contents.clear();
+                        continue;
+                    }
+
                     for (existing, new) in existing_values.iter().zip_eq(mem_int.data().iter()) {
                         new_constraints.push(AlgebraicConstraint::assert_zero(
                             existing.clone() - new.clone(),
diff --git a/autoprecompiles/src/pgo/cell/mod.rs b/autoprecompiles/src/pgo/cell/mod.rs
@@ -13,7 +13,7 @@ use crate::{
     blocks::BasicBlock,
     evaluation::EvaluationResult,
     pgo::cell::selection::parallel_fractional_knapsack,
-    PowdrConfig,
+    EmpiricalConstraints, PowdrConfig,
 };
 
 mod selection;
@@ -92,6 +92,7 @@ impl<A: Adapter + Send + Sync, C: Candidate<A> + Send + Sync> PgoAdapter for Cel
         config: &PowdrConfig,
         vm_config: AdapterVmConfig<Self::Adapter>,
         labels: BTreeMap<u64, Vec<String>>,
+        empirical_constraints: EmpiricalConstraints,
     ) -> Vec<AdapterApcWithStats<Self::Adapter>> {
         tracing::info!(
             "Generating autoprecompiles with cell PGO for {} blocks",
@@ -131,6 +132,7 @@ impl<A: Adapter + Send + Sync, C: Candidate<A> + Send + Sync> PgoAdapter for Cel
                     vm_config.clone(),
                     config.degree_bound,
                     config.apc_candidates_dir_path.as_deref(),
+                    &empirical_constraints,
                 )
                 .ok()?;
                 let candidate = C::create(
diff --git a/autoprecompiles/src/pgo/instruction.rs b/autoprecompiles/src/pgo/instruction.rs
@@ -4,7 +4,7 @@ use crate::{
     adapter::{Adapter, AdapterApcWithStats, AdapterVmConfig, PgoAdapter},
     blocks::BasicBlock,
     pgo::create_apcs_for_all_blocks,
-    PowdrConfig,
+    EmpiricalConstraints, PowdrConfig,
 };
 
 pub struct InstructionPgo<A> {
@@ -30,6 +30,7 @@ impl<A: Adapter> PgoAdapter for InstructionPgo<A> {
         config: &PowdrConfig,
         vm_config: AdapterVmConfig<Self::Adapter>,
         _labels: BTreeMap<u64, Vec<String>>,
+        empirical_constraints: EmpiricalConstraints,
     ) -> Vec<AdapterApcWithStats<Self::Adapter>> {
         tracing::info!(
             "Generating autoprecompiles with instruction PGO for {} blocks",
@@ -70,7 +71,7 @@ impl<A: Adapter> PgoAdapter for InstructionPgo<A> {
                 );
         }
 
-        create_apcs_for_all_blocks::<Self::Adapter>(blocks, config, vm_config)
+        create_apcs_for_all_blocks::<Self::Adapter>(blocks, config, vm_config, empirical_constraints)
     }
 
     fn pc_execution_count(&self, pc: u64) -> Option<u32> {
diff --git a/autoprecompiles/src/pgo/mod.rs b/autoprecompiles/src/pgo/mod.rs
@@ -6,7 +6,7 @@ use strum::{Display, EnumString};
 use crate::{
     adapter::{Adapter, AdapterApcWithStats, AdapterVmConfig, ApcWithStats},
     blocks::BasicBlock,
-    PowdrConfig,
+    EmpiricalConstraints, PowdrConfig,
 };
 
 mod cell;
@@ -77,6 +77,7 @@ fn create_apcs_for_all_blocks<A: Adapter>(
     blocks: Vec<BasicBlock<A::Instruction>>,
     config: &PowdrConfig,
     vm_config: AdapterVmConfig<A>,
+    empirical_constraints: EmpiricalConstraints,
 ) -> Vec<AdapterApcWithStats<A>> {
     let n_acc = config.autoprecompiles as usize;
     tracing::info!("Generating {n_acc} autoprecompiles in parallel");
@@ -97,6 +98,7 @@ fn create_apcs_for_all_blocks<A: Adapter>(
                 vm_config.clone(),
                 config.degree_bound,
                 config.apc_candidates_dir_path.as_deref(),
+                &empirical_constraints,
             )
             .unwrap()
         })
diff --git a/autoprecompiles/src/pgo/none.rs b/autoprecompiles/src/pgo/none.rs
@@ -4,7 +4,7 @@ use crate::{
     adapter::{Adapter, AdapterApcWithStats, AdapterVmConfig, PgoAdapter},
     blocks::BasicBlock,
     pgo::create_apcs_for_all_blocks,
-    PowdrConfig,
+    EmpiricalConstraints, PowdrConfig,
 };
 
 pub struct NonePgo<A> {
@@ -29,6 +29,7 @@ impl<A: Adapter> PgoAdapter for NonePgo<A> {
         config: &PowdrConfig,
         vm_config: AdapterVmConfig<Self::Adapter>,
         _labels: BTreeMap<u64, Vec<String>>,
+        empirical_constraints: EmpiricalConstraints,
     ) -> Vec<AdapterApcWithStats<Self::Adapter>> {
         // cost = number_of_original_instructions
         blocks.sort_by(|a, b| b.statements.len().cmp(&a.statements.len()));
@@ -42,6 +43,6 @@ impl<A: Adapter> PgoAdapter for NonePgo<A> {
             );
         }
 
-        create_apcs_for_all_blocks::<Self::Adapter>(blocks, config, vm_config)
+        create_apcs_for_all_blocks::<Self::Adapter>(blocks, config, vm_config, empirical_constraints)
     }
 }
diff --git a/cli-openvm/src/main.rs b/cli-openvm/src/main.rs
@@ -211,13 +211,14 @@ fn run_command(command: Commands) {
             apc_candidates_dir,
         } => {
             let mut powdr_config = default_powdr_openvm_config(autoprecompiles as u64, skip as u64);
-            if let Some(apc_candidates_dir) = apc_candidates_dir {
+            if let Some(apc_candidates_dir) = &apc_candidates_dir {
                 powdr_config = powdr_config.with_apc_candidates_dir(apc_candidates_dir);
             }
+            let inputs = stdin_from(input);
             let execution_profile = powdr_openvm::execution_profile_from_guest(
                 &guest,
                 guest_opts.clone(),
-                stdin_from(input),
+                inputs.clone(),
             );
             let pgo_config = pgo_config(pgo, max_columns, execution_profile);
             let compile_and_prove = || {
diff --git a/openvm/src/customize_exe.rs b/openvm/src/customize_exe.rs
diff --git a/openvm/tests/common/mod.rs b/openvm/tests/common/mod.rs
