Optimize memory (#2716)

Depends on #2794

---------

Co-authored-by: Leo <leo@powdrlabs.com>
Co-authored-by: Georg Wiese <georgwiese@gmail.com>

autoprecompiles/src/lib.rs

@@ -18,6 +18,7 @@ use powdr_number::{FieldElement, LargeInt};
 use powdr_pilopt::simplify_expression;
 
 pub mod constraint_optimizer;
+pub mod memory_optimizer;
 pub mod optimizer;
 pub mod powdr;
 pub mod register_optimizer;

autoprecompiles/src/memory_optimizer.rs

@@ -0,0 +1,336 @@
+use std::collections::{HashMap, HashSet};
+use std::fmt;
+use std::fmt::Display;
+use std::hash::Hash;
+
+use itertools::Itertools;
+use powdr_ast::analyzed::{
+    algebraic_expression_conversion, AlgebraicExpression, AlgebraicReference, Challenge,
+};
+use powdr_constraint_solver::boolean_extractor;
+use powdr_constraint_solver::constraint_system::{ConstraintRef, ConstraintSystem};
+use powdr_constraint_solver::indexed_constraint_system::IndexedConstraintSystem;
+use powdr_constraint_solver::quadratic_symbolic_expression::QuadraticSymbolicExpression;
+use powdr_constraint_solver::quadratic_symbolic_expression::RangeConstraintProvider;
+use powdr_constraint_solver::range_constraint::RangeConstraint;
+use powdr_constraint_solver::utils::possible_concrete_values;
+use powdr_number::FieldElement;
+
+use crate::{MemoryBusInteraction, MemoryOp, MemoryType, SymbolicConstraint, SymbolicMachine};
+
+/// Optimizes bus sends that correspond to general-purpose memory read and write operations.
+/// It works best if all read-write-operation addresses are fixed offsets relative to some
+/// symbolic base address. If stack and heap access operations are mixed, this is usually violated.
+pub fn optimize_memory<T: FieldElement>(mut machine: SymbolicMachine<T>) -> SymbolicMachine<T> {
+    let (to_remove, new_constraints) = redundant_memory_interactions_indices(&machine);
+    let to_remove = to_remove.into_iter().collect::<HashSet<_>>();
+    machine.bus_interactions = machine
+        .bus_interactions
+        .into_iter()
+        .enumerate()
+        .filter_map(|(i, bus)| (!to_remove.contains(&i)).then_some(bus))
+        .collect();
+    machine.constraints.extend(new_constraints);
+    machine
+}
+
+/// Tries to find indices of bus interactions that can be removed in the given machine
+/// and also returns a set of new constraints to be added.
+fn redundant_memory_interactions_indices<T: FieldElement>(
+    machine: &SymbolicMachine<T>,
+) -> (Vec<usize>, Vec<SymbolicConstraint<T>>) {
+    let address_comparator = MemoryAddressComparator::new(machine);
+    let mut new_constraints: Vec<SymbolicConstraint<T>> = Vec::new();
+
+    // Track memory contents while we go through bus interactions.
+    // This maps an address to the index of the previous send on that address and the
+    // data currently stored there.
+    let mut memory_contents: HashMap<_, (usize, Vec<AlgebraicExpression<_>>)> = Default::default();
+    let mut to_remove: Vec<usize> = Default::default();
+
+    // TODO we assume that memory interactions are sorted by timestamp.
+    for (index, bus_int) in machine.bus_interactions.iter().enumerate() {
+        let mem_int = match MemoryBusInteraction::try_from_symbolic_bus_interaction_with_memory_kind(
+            bus_int,
+            MemoryType::Memory,
+        ) {
+            Ok(Some(mem_int)) => mem_int,
+            Ok(None) => continue,
+            Err(_) => {
+                // This interaction might be going to memory, but we do not know
+                // the multiplicity. Delete all knowledge.
+                memory_contents.clear();
+                continue;
+            }
+        };
+        let addr = algebraic_to_quadratic_symbolic_expression(&mem_int.addr);
+
+        match mem_int.op {
+            MemoryOp::Receive => {
+                // If there is an unconsumed send to this address, consume it.
+                // 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) {
+                    for (existing, new) in existing_values.iter().zip_eq(mem_int.data.iter()) {
+                        new_constraints.push((existing.clone() - new.clone()).into());
+                    }
+                    to_remove.extend([index, previous_send]);
+                }
+            }
+            MemoryOp::Send => {
+                // We can only retain knowledge about addresses where we can prove
+                // that this send operation does not interfere with it, i.e.
+                // if we can prove that the two addresses differ by at least a word size.
+                memory_contents.retain(|other_addr, _| {
+                    address_comparator.are_addrs_known_to_be_different_by_word(&addr, other_addr)
+                });
+                memory_contents.insert(addr.clone(), (index, mem_int.data.clone()));
+            }
+        }
+    }
+
+    log::debug!(
+        "Removing {} memory interactions and adding {} new constraints",
+        to_remove.len(),
+        new_constraints.len()
+    );
+
+    (to_remove, new_constraints)
+}
+
+struct MemoryAddressComparator<T: FieldElement> {
+    /// For each address `a` contains a list of expressions `v` such that
+    /// `a = v` is true in the constraint system.
+    memory_addresses: HashMap<
+        QuadraticSymbolicExpression<T, Variable>,
+        Vec<QuadraticSymbolicExpression<T, Variable>>,
+    >,
+}
+
+impl<T: FieldElement> MemoryAddressComparator<T> {
+    fn new(machine: &SymbolicMachine<T>) -> Self {
+        let addresses = machine
+            .bus_interactions
+            .iter()
+            .flat_map(|bus| {
+                MemoryBusInteraction::try_from_symbolic_bus_interaction_with_memory_kind(
+                    bus,
+                    MemoryType::Memory,
+                )
+                .ok()
+                .flatten()
+            })
+            .map(|bus| algebraic_to_quadratic_symbolic_expression(&bus.addr));
+
+        let constraints = symbolic_to_simplified_constraints(&machine.constraints);
+        let constraint_system: IndexedConstraintSystem<_, _> = ConstraintSystem {
+            algebraic_constraints: constraints,
+            bus_interactions: vec![],
+        }
+        .into();
+
+        let memory_addresses = addresses
+            .map(|addr| {
+                (
+                    addr.clone(),
+                    find_equivalent_expressions(&addr, &constraint_system),
+                )
+            })
+            .collect();
+        Self { memory_addresses }
+    }
+
+    /// Returns true if we can prove that for two addresses `a` and `b`,
+    /// `a - b` never falls into the range `-3..=3`.
+    pub fn are_addrs_known_to_be_different_by_word(
+        &self,
+        a: &QuadraticSymbolicExpression<T, Variable>,
+        b: &QuadraticSymbolicExpression<T, Variable>,
+    ) -> bool {
+        let a_exprs = &self.memory_addresses[a];
+        let b_exprs = &self.memory_addresses[b];
+        a_exprs
+            .iter()
+            .cartesian_product(b_exprs)
+            .any(|(a_exprs, b_exprs)| {
+                is_value_known_to_be_different_by_word(
+                    a_exprs,
+                    b_exprs,
+                    &RangeConstraintsForBooleans,
+                )
+            })
+    }
+}
+
+/// Converts from SymbolicConstraint to QuadraticSymbolicExpression and
+/// simplifies constraints by introducing boolean variables.
+fn symbolic_to_simplified_constraints<T: FieldElement>(
+    constraints: &[SymbolicConstraint<T>],
+) -> Vec<QuadraticSymbolicExpression<T, Variable>> {
+    let mut counter = 0..;
+    let mut var_dispenser = || Variable::Boolean(counter.next().unwrap());
+
+    constraints
+        .iter()
+        .map(|constr| {
+            let constr = algebraic_to_quadratic_symbolic_expression(&constr.expr);
+            boolean_extractor::extract_boolean(&constr, &mut var_dispenser).unwrap_or(constr)
+        })
+        .collect_vec()
+}
+
+#[derive(Clone, PartialOrd, Ord, PartialEq, Eq, Hash, Debug)]
+pub enum Variable {
+    Reference(AlgebraicReference),
+    PublicReference(String),
+    Challenge(Challenge),
+    Boolean(usize),
+}
+
+impl Display for Variable {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Variable::Reference(r) => write!(f, "{r}"),
+            Variable::PublicReference(r) => write!(f, "{r}"),
+            Variable::Challenge(c) => write!(f, "{c}"),
+            Variable::Boolean(id) => write!(f, "boolean_{id}"),
+        }
+    }
+}
+
+#[derive(Default)]
+struct RangeConstraintsForBooleans;
+
+impl<T: FieldElement> RangeConstraintProvider<T, Variable> for RangeConstraintsForBooleans {
+    fn get(&self, variable: &Variable) -> RangeConstraint<T> {
+        match variable {
+            Variable::Boolean(_) => RangeConstraint::from_mask(1),
+            _ => Default::default(),
+        }
+    }
+}
+
+/// Tries to find equivalent expressions for the given expression
+/// according to the given constraint system.
+/// Returns at least one equivalent expression (in the worst case, the expression itself).
+fn find_equivalent_expressions<T: FieldElement, V: Clone + Ord + Hash + Eq + Display>(
+    expression: &QuadraticSymbolicExpression<T, V>,
+    constraints: &IndexedConstraintSystem<T, V>,
+) -> Vec<QuadraticSymbolicExpression<T, V>> {
+    if expression.is_quadratic() {
+        // This case is too complicated.
+        return vec![expression.clone()];
+    }
+
+    // Go through the constraints related to this address
+    // and try to solve for the expression
+    let mut exprs = constraints
+        .constraints_referencing_variables(expression.referenced_unknown_variables().cloned())
+        .filter_map(|constr| match constr {
+            ConstraintRef::AlgebraicConstraint(constr) => Some(constr),
+            ConstraintRef::BusInteraction(_) => None,
+        })
+        .flat_map(|constr| constr.try_solve_for_expr(expression))
+        .collect_vec();
+    if exprs.is_empty() {
+        // If we cannot solve for the expression, we just take the expression unmodified.
+        exprs.push(expression.clone());
+    }
+    exprs
+}
+
+/// Returns true if we can prove that `a - b` never falls into the range `-3..=3`.
+fn is_value_known_to_be_different_by_word<T: FieldElement, V: Clone + Ord + Hash + Eq + Display>(
+    a: &QuadraticSymbolicExpression<T, V>,
+    b: &QuadraticSymbolicExpression<T, V>,
+    range_constraints: &impl RangeConstraintProvider<T, V>,
+) -> bool {
+    let disallowed_range = RangeConstraint::from_range(-T::from(3), T::from(3));
+    possible_concrete_values(&(a - b), range_constraints, 20)
+        .is_some_and(|mut values| !values.any(|value| disallowed_range.allows_value(value)))
+}
+
+/// Turns an algebraic expression into a quadratic symbolic expression,
+/// assuming all [`AlgebraicReference`]s, public references and challenges
+/// are unknown variables.
+pub fn algebraic_to_quadratic_symbolic_expression<T: FieldElement>(
+    expr: &AlgebraicExpression<T>,
+) -> QuadraticSymbolicExpression<T, Variable> {
+    type Qse<T> = QuadraticSymbolicExpression<T, Variable>;
+
+    struct TerminalConverter;
+
+    impl<T: FieldElement> algebraic_expression_conversion::TerminalConverter<Qse<T>>
+        for TerminalConverter
+    {
+        fn convert_reference(&mut self, reference: &AlgebraicReference) -> Qse<T> {
+            Qse::from_unknown_variable(Variable::Reference(reference.clone()))
+        }
+        fn convert_public_reference(&mut self, reference: &str) -> Qse<T> {
+            Qse::from_unknown_variable(Variable::PublicReference(reference.to_string()))
+        }
+        fn convert_challenge(&mut self, challenge: &Challenge) -> Qse<T> {
+            Qse::from_unknown_variable(Variable::Challenge(*challenge))
+        }
+    }
+
+    algebraic_expression_conversion::convert(expr, &mut TerminalConverter)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    use powdr_constraint_solver::{
+        quadratic_symbolic_expression::NoRangeConstraints,
+        test_utils::{constant, var},
+    };
+
+    #[test]
+    fn difference_for_constants() {
+        assert!(!is_value_known_to_be_different_by_word(
+            &constant(7),
+            &constant(5),
+            &NoRangeConstraints
+        ));
+        assert!(!is_value_known_to_be_different_by_word(
+            &constant(5),
+            &constant(7),
+            &NoRangeConstraints
+        ));
+        assert!(is_value_known_to_be_different_by_word(
+            &constant(4),
+            &constant(0),
+            &NoRangeConstraints
+        ));
+        assert!(is_value_known_to_be_different_by_word(
+            &constant(0),
+            &constant(4),
+            &NoRangeConstraints
+        ));
+    }
+
+    #[test]
+    fn difference_for_vars() {
+        assert!(!is_value_known_to_be_different_by_word(
+            &(constant(7) + var("a")),
+            &(constant(5) + var("a")),
+            &NoRangeConstraints
+        ));
+        assert!(is_value_known_to_be_different_by_word(
+            &(constant(7) + var("a")),
+            &(constant(2) + var("a")),
+            &NoRangeConstraints
+        ));
+        assert!(!is_value_known_to_be_different_by_word(
+            &(constant(7) - var("a")),
+            &(constant(2) + var("a")),
+            &NoRangeConstraints
+        ));
+        assert!(!is_value_known_to_be_different_by_word(
+            &var("a"),
+            &var("b"),
+            &NoRangeConstraints
+        ));
+    }
+}

autoprecompiles/src/optimizer.rs

@@ -18,6 +18,7 @@ use powdr_pilopt::{
 
 use crate::{
     constraint_optimizer::{optimize_constraints, IsBusStateful},
+    memory_optimizer::optimize_memory,
     powdr::{self},
     register_optimizer::{check_register_operation_consistency, optimize_register_operations},
     stats_logger::StatsLogger,
@@ -74,9 +75,11 @@ fn optimization_loop_iteration<T: FieldElement>(
     let machine =
         optimize_register_operations(constraint_system_to_symbolic_machine(constraint_system));
     assert!(check_register_operation_consistency(&machine));
-    let constraint_system = symbolic_machine_to_constraint_system(machine);
-    stats_logger.log("register optimization", &constraint_system);
-    Ok(constraint_system)
+    stats_logger.log("register optimization", &machine);
+    let machine = optimize_memory(machine);
+    stats_logger.log("memory optimization", &machine);
+
+    Ok(symbolic_machine_to_constraint_system(machine))
 }
 
 fn system_size<T: FieldElement>(constraint_system: &ConstraintSystem<T, Variable>) -> [usize; 3] {

constraint-solver/src/lib.rs

@@ -10,6 +10,5 @@ pub mod solver;
 pub mod symbolic_expression;
 pub mod symbolic_to_quadratic;
 pub mod test_utils;
+pub mod utils;
 pub mod variable_update;
-
-mod utils;

constraint-solver/src/solver/exhaustive_search.rs

@@ -85,7 +85,7 @@ fn find_unique_assignment_for_set<T: FieldElement, V: Clone + Hash + Ord + Eq +
     rc: impl RangeConstraintProvider<T, V> + Clone,
     bus_interaction_handler: &impl BusInteractionHandler<T>,
 ) -> Result<Option<BTreeMap<V, T>>, Error> {
-    match get_all_possible_assignments(variables.iter().cloned(), rc.clone())
+    match get_all_possible_assignments(variables.iter().cloned(), &rc)
         .filter(|assignments| {
             !constraint_system.is_assignment_conflicting(assignments, &rc, bus_interaction_handler)
         })
@@ -122,7 +122,7 @@ fn get_brute_force_candidates<'a, T: FieldElement, V: Clone + Hash + Ord>(
         .unique()
         .filter(|variables| !variables.is_empty())
         .filter(move |variables| {
-            count_possible_assignments(variables.iter().cloned(), rc.clone())
+            count_possible_assignments(variables.iter().cloned(), &rc)
                 .is_some_and(|count| count <= MAX_SEARCH_WIDTH)
         })
 }

constraint-solver/src/utils.rs

@@ -6,6 +6,7 @@ use itertools::Itertools;
 use powdr_number::{FieldElement, LargeInt};
 
 use crate::quadratic_symbolic_expression::{QuadraticSymbolicExpression, RangeConstraintProvider};
+use crate::symbolic_expression::SymbolicExpression;
 
 /// Returns the set of all known variables in a list of algebraic expressions.
 /// Panics if a variable appears as both known and unknown.
@@ -47,7 +48,7 @@ pub fn known_variables<'a, T: FieldElement, V: Clone + Hash + Ord + Eq + Display
 /// Returns `None` if this number would not fit a `u64`.
 pub fn count_possible_assignments<T: FieldElement, V: Clone + Ord>(
     variables: impl Iterator<Item = V>,
-    rc: impl RangeConstraintProvider<T, V>,
+    rc: &impl RangeConstraintProvider<T, V>,
 ) -> Option<u64> {
     variables
         .map(|v| rc.get(&v))
@@ -63,7 +64,7 @@ pub fn count_possible_assignments<T: FieldElement, V: Clone + Ord>(
 /// the function `has_few_possible_assignments`.
 pub fn get_all_possible_assignments<T: FieldElement, V: Clone + Ord>(
     variables: impl IntoIterator<Item = V>,
-    rc: impl RangeConstraintProvider<T, V>,
+    rc: &impl RangeConstraintProvider<T, V>,
 ) -> impl Iterator<Item = BTreeMap<V, T>> {
     variables
         .into_iter()
@@ -77,3 +78,30 @@ pub fn get_all_possible_assignments<T: FieldElement, V: Clone + Ord>(
         .multi_cartesian_product()
         .map(|assignment| assignment.into_iter().collect::<BTreeMap<_, _>>())
 }
+
+/// Returns all possible concrete values for `expr` using exhaustive search.
+/// Returns None if the number of possible assignments exceeds `max_elements`.
+pub fn possible_concrete_values<'a, T: FieldElement, V: Clone + Ord + Hash>(
+    expr: &'a QuadraticSymbolicExpression<T, V>,
+    rc: &'a impl RangeConstraintProvider<T, V>,
+    max_elements: u64,
+) -> Option<impl Iterator<Item = T> + 'a> {
+    let variables = expr.referenced_unknown_variables().cloned().collect_vec();
+    if count_possible_assignments(variables.iter().cloned(), rc)
+        .is_none_or(|count| count > max_elements)
+    {
+        // If there are too many possible assignments, we do not try to perform exhaustive search.
+        None
+    } else {
+        Some(
+            get_all_possible_assignments(variables, rc).map(|assignment| {
+                let mut expr = expr.clone();
+                for (variable, value) in assignment.iter() {
+                    expr.substitute_by_known(variable, &SymbolicExpression::Concrete(*value));
+                }
+                // We substitute all variables, so this has to be a number.
+                expr.try_to_number().unwrap()
+            }),
+        )
+    }
+}

openvm/src/lib.rs

@@ -940,9 +940,9 @@ mod tests {
             .powdr_airs_metrics();
         assert_eq!(machines.len(), 1);
         let m = &machines[0];
-        assert_eq!(m.width, 3195);
+        assert_eq!(m.width, 2175);
         assert_eq!(m.constraints, 161);
-        assert_eq!(m.bus_interactions, 2861);
+        assert_eq!(m.bus_interactions, 2181);
     }
 
     #[test]

openvm/tests/optimizer.rs

@@ -45,6 +45,7 @@ fn test_optimize() {
         machine.bus_interactions.len(),
         machine.constraints.len()
     );
+
     // This cbor file above has the `is_valid` column removed, this is why the number below
     // might be one less than in other tests.
     assert_eq!(
@@ -53,7 +54,7 @@ fn test_optimize() {
             machine.bus_interactions.len(),
             machine.constraints.len()
         ],
-        [3194, 2861, 160]
+        [2174, 2181, 160]
     );
 }