Build optimistic precompiles

Author: georgwiese

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> {

autoprecompiles/src/empirical_constraints.rs

@@ -6,6 +6,13 @@ use std::hash::Hash;
 use itertools::Itertools;
 use serde::{Deserialize, Serialize};
 
+use crate::{
+    adapter::Adapter,
+    blocks::BasicBlock,
+    expression::{AlgebraicExpression, AlgebraicReference},
+    SymbolicConstraint,
+};
+
 /// "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)]
@@ -16,6 +23,9 @@ pub struct EmpiricalConstraints {
     /// 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)>>>,
+    /// Count of how many times each program counter was executed in the sampled executions.
+    /// This can be used to set a threshold for applying constraints only to frequently executed PCs.
+    pub pc_counts: BTreeMap<u32, u64>,
 }
 
 /// Debug information mapping AIR ids to program counters and column names.
@@ -61,6 +71,24 @@ impl EmpiricalConstraints {
                 })
                 .or_insert(classes);
         }
+
+        // Combine pc counts
+        for (pc, count) in other.pc_counts {
+            *self.pc_counts.entry(pc).or_insert(0) += count;
+        }
+    }
+
+    pub fn with_thresholded_pc_count(mut self, threshold: u64) -> Self {
+        self.column_ranges_by_pc
+            .retain(|pc, _| self.pc_counts.get(pc).cloned().unwrap_or(0) >= threshold);
+        self.equivalence_classes_by_block.retain(|block_id, _| {
+            self.pc_counts
+                .get(&(*block_id as u32))
+                .cloned()
+                .unwrap_or(0)
+                >= threshold
+        });
+        self
     }
 }
 
@@ -88,6 +116,99 @@ fn merge_maps<K: Ord, V: Eq + Debug>(map1: &mut BTreeMap<K, V>, map2: BTreeMap<K
     }
 }
 
+type ConstraintList<A> = Vec<SymbolicConstraint<<A as Adapter>::PowdrField>>;
+
+/// For any program line that was not executed at least this many times in the traces,
+/// discard any empirical constraints associated with it.
+const EXECUTION_COUNT_THRESHOLD: u64 = 100;
+
+pub fn add_empirical_constraints<A: Adapter>(
+    empirical_constraints: &EmpiricalConstraints,
+    subs: &[Vec<u64>],
+    block: &BasicBlock<A::Instruction>,
+    columns: impl Iterator<Item = AlgebraicReference>,
+) -> (ConstraintList<A>, ConstraintList<A>) {
+    // Apply execution count threshold to avoid overfitting on rarely executed code
+    let empirical_constraints = empirical_constraints
+        .clone()
+        .with_thresholded_pc_count(EXECUTION_COUNT_THRESHOLD);
+
+    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(), 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.
+                tracing::warn!(
+                    "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.
+                    tracing::warn!(
+                        "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)
+}
+
 /// 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.

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;
@@ -429,6 +430,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();
 
@@ -438,6 +440,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);
@@ -446,13 +456,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>(
-        machine,
+        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>(
+        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);
@@ -477,6 +507,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);

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(),

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(

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,12 @@ 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> {

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

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)
     }
 }