Collect empirical constraints

Author: georgwiese

autoprecompiles/src/empirical_constraints.rs

@@ -0,0 +1,117 @@
+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 {
+    /// 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,
+}
+
+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)
+}

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;

cli-openvm/src/main.rs

@@ -146,8 +146,14 @@ fn run_command(command: Commands) {
                 stdin_from(input),
             );
             let pgo_config = pgo_config(pgo, max_columns, execution_profile);
-            let program =
-                powdr_openvm::compile_guest(&guest, guest_opts, powdr_config, pgo_config).unwrap();
+            let program = powdr_openvm::compile_guest(
+                &guest,
+                guest_opts,
+                powdr_config,
+                pgo_config,
+                stdin_from(input),
+            )
+            .unwrap();
             write_program_to_file(program, &format!("{guest}_compiled.cbor")).unwrap();
         }
 
@@ -172,9 +178,14 @@ fn run_command(command: Commands) {
             );
             let pgo_config = pgo_config(pgo, max_columns, execution_profile);
             let compile_and_exec = || {
-                let program =
-                    powdr_openvm::compile_guest(&guest, guest_opts, powdr_config, pgo_config)
-                        .unwrap();
+                let program = powdr_openvm::compile_guest(
+                    &guest,
+                    guest_opts,
+                    powdr_config,
+                    pgo_config,
+                    stdin_from(input),
+                )
+                .unwrap();
                 powdr_openvm::execute(program, stdin_from(input)).unwrap();
             };
             if let Some(metrics_path) = metrics {
@@ -210,9 +221,14 @@ fn run_command(command: Commands) {
             );
             let pgo_config = pgo_config(pgo, max_columns, execution_profile);
             let compile_and_prove = || {
-                let program =
-                    powdr_openvm::compile_guest(&guest, guest_opts, powdr_config, pgo_config)
-                        .unwrap();
+                let program = powdr_openvm::compile_guest(
+                    &guest,
+                    guest_opts,
+                    powdr_config,
+                    pgo_config,
+                    stdin_from(input),
+                )
+                .unwrap();
                 powdr_openvm::prove(&program, mock, recursion, stdin_from(input), None).unwrap()
             };
             if let Some(metrics_path) = metrics {

openvm/Cargo.toml

@@ -61,6 +61,7 @@ derive_more = { version = "2.0.1", default-features = false, features = [
   "from",
 ] }
 itertools = "0.14.0"
+serde_json = "1.0.140"
 
 tracing = "0.1.40"
 tracing-subscriber = { version = "0.3.17", features = ["std", "env-filter"] }

openvm/src/customize_exe.rs

@@ -13,12 +13,13 @@ use crate::memory_bus_interaction::OpenVmMemoryBusInteraction;
 use crate::opcode::branch_opcodes_bigint_set;
 use crate::powdr_extension::chip::PowdrAir;
 use crate::utils::UnsupportedOpenVmReferenceError;
-use crate::OriginalCompiledProgram;
+use crate::{empirical_constraints::empirical_constraints, OriginalCompiledProgram};
 use crate::{CompiledProgram, SpecializedConfig};
 use itertools::Itertools;
 use openvm_instructions::instruction::Instruction as OpenVmInstruction;
 use openvm_instructions::program::{Program as OpenVmProgram, DEFAULT_PC_STEP};
 use openvm_instructions::VmOpcode;
+use openvm_sdk::StdIn;
 use openvm_stark_backend::{
     interaction::SymbolicInteraction,
     p3_field::{FieldAlgebra, PrimeField32},
@@ -147,6 +148,7 @@ pub fn customize<'a, P: PgoAdapter<Adapter = BabyBearOpenVmApcAdapter<'a>>>(
     }: OriginalCompiledProgram,
     config: PowdrConfig,
     pgo: P,
+    inputs: StdIn,
 ) -> CompiledProgram {
     let labels = elf.text_labels();
     let debug_info = elf.debug_info();
@@ -217,6 +219,17 @@ pub fn customize<'a, P: PgoAdapter<Adapter = BabyBearOpenVmApcAdapter<'a>>>(
         })
         .collect();
 
+    let _empirical_constraints = empirical_constraints(
+        exe.clone(),
+        SpecializedConfig::new(
+            original_config.clone(),
+            vec![],
+            config.degree_bound.identities,
+        ),
+        inputs,
+        &blocks,
+    );
+
     let start = std::time::Instant::now();
     let apcs = pgo.filter_blocks_and_create_apcs_with_pgo(blocks, &config, vm_config, labels);
     metrics::gauge!("total_apc_gen_time_ms").set(start.elapsed().as_millis() as f64);