Do not use bus field variables. (#3171)

Removes the special introduction and removal of bus variables because
the solver can now handle this on its own through the `LinearizedSolver`
component. It also has the benefit that we do not need to run the
optimizer twice any more (once with bus interaction variables and once
without).

Previously we would introduce a new variable type outside of the solver
and replace all bus interactions `bus_interaction(1, x + y * z)` by
`bus_interaction(1, bus_var_1)` and the algebraic constraint `bus_var_1
= x + y * z`.

Now the `LinearizedSolver` is a component inside the solver that
replaces all bus interaction fields (that are not already constants or
simple variables) and also replaces all non-linear components:

`bus_interaction(1, x + y * z)`

becomes

```
bus_interaction(1, v1)
v1 = x + v2
v2 = y * z
```

The reason behind also replacing non-linear parts is because range
constraints are only stored for variables, not for all sub-expressions.
Also if we find the same non-linear sub-expression multiple times, we
can now replace it directly by the same variable.

---------

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

autoprecompiles/src/bitwise_lookup_optimizer.rs

@@ -108,7 +108,9 @@ pub fn optimize_bitwise_lookup<T: FieldElement, V: Hash + Eq + Clone + Ord + Deb
             multiplicity: One::one(),
         });
     }
-    system.algebraic_constraints.extend(new_constraints);
+    system
+        .algebraic_constraints
+        .extend(new_constraints.into_iter().filter(|expr| !expr.is_zero()));
     system
 }
 

autoprecompiles/src/constraint_optimizer.rs

@@ -98,9 +98,53 @@ fn solver_based_optimization<T: FieldElement, V: Clone + Ord + Hash + Display>(
     // does not increase.
     assert!(assignments.iter().all(|(_, expr)| expr.is_affine()));
     constraint_system.apply_substitutions(assignments);
+
+    // Now try to replace bus interaction fields that the solver knows to be constant
+    let mut bus_interactions = vec![];
+    let mut new_algebraic_constraints = vec![];
+    // We remove all bus interactions because we do not want to change the order.
+    constraint_system.retain_bus_interactions(|bus_interaction| {
+        let mut modified = false;
+        let replacement = bus_interaction
+            .fields()
+            .map(|field| {
+                if let Some(n) = try_replace_by_number(field, solver) {
+                    modified = true;
+                    new_algebraic_constraints.push(&n - field);
+                    n
+                } else {
+                    field.clone()
+                }
+            })
+            .collect();
+        if modified {
+            log::trace!("Replacing bus interaction {bus_interaction} with {replacement}");
+        }
+        bus_interactions.push(replacement);
+        false
+    });
+    constraint_system.add_bus_interactions(bus_interactions);
+    constraint_system.add_algebraic_constraints(new_algebraic_constraints);
     Ok(constraint_system)
 }
 
+/// Tries to find a number that is equivalent to the expression and returns it
+/// as a GroupedExpression.
+/// Returns None if it was unsuccessful or if the expression already is a number.
+fn try_replace_by_number<T: FieldElement, V: Clone + Ord + Hash + Display>(
+    expr: &GroupedExpression<T, V>,
+    solver: &impl Solver<T, V>,
+) -> Option<GroupedExpression<T, V>> {
+    if expr.try_to_number().is_some() {
+        return None;
+    }
+    Some(GroupedExpression::from_number(
+        solver
+            .range_constraint_for_expression(expr)
+            .try_to_single_value()?,
+    ))
+}
+
 /// Removes free variables from the constraint system, under some conditions.
 ///
 /// Motivation: Suppose there is a constraint `2 * foo = bar` and `foo` only appears in this constraint.

autoprecompiles/src/low_degree_bus_interaction_optimizer.rs

@@ -1,4 +1,5 @@
 use itertools::Itertools;
+use num_traits::Zero;
 use powdr_constraint_solver::constraint_system::{
     BusInteraction, BusInteractionHandler, ConstraintSystem,
 };
@@ -51,10 +52,13 @@ impl<
                 if let Some((replacement, range_constraints)) =
                     self.try_replace_bus_interaction(&bus_int)
                 {
-                    // If we found a replacement, add the polynomial constraints and replace
-                    // the bus interaction with interactions implementing the range constraints.
+                    // If we found a replacement, add the polynomial constraints (unless it is
+                    // trivially zero) and replace the bus interaction with interactions implementing
+                    // the range constraints.
                     // Note that many of these may be optimized away by the range constraint optimizer.
-                    new_constraints.push(replacement);
+                    if !replacement.is_zero() {
+                        new_constraints.push(replacement);
+                    }
 
                     self.bus_interaction_handler
                         .batch_make_range_constraints(range_constraints)

autoprecompiles/src/optimizer.rs

@@ -3,19 +3,14 @@ use std::hash::Hash;
 use std::{collections::BTreeMap, fmt::Display};
 
 use itertools::Itertools;
-use num_traits::Zero;
 use powdr_constraint_solver::constraint_system::BusInteractionHandler;
-use powdr_constraint_solver::indexed_constraint_system::{
-    apply_substitutions, IndexedConstraintSystem,
-};
+use powdr_constraint_solver::indexed_constraint_system::IndexedConstraintSystem;
 use powdr_constraint_solver::inliner::inline_everything_below_degree_bound;
-use powdr_constraint_solver::runtime_constant::RuntimeConstant;
 use powdr_constraint_solver::solver::{new_solver, Solver};
 use powdr_constraint_solver::{
     constraint_system::{BusInteraction, ConstraintSystem},
     grouped_expression::GroupedExpression,
     journaling_constraint_system::JournalingConstraintSystem,
-    runtime_constant::VarTransformable,
 };
 use powdr_number::FieldElement;
 
@@ -49,22 +44,6 @@ pub fn optimize<A: Adapter>(
     }
 
     let constraint_system = symbolic_machine_to_constraint_system(machine);
-    let constraint_system = introduce_bus_interaction_variables(constraint_system);
-
-    // Run the optimizer while avoiding inlining bus interaction field variables
-    let constraint_system =
-        run_optimization_loop_until_no_change::<_, _, _, A::MemoryBusInteraction<_>>(
-            constraint_system,
-            bus_interaction_handler.clone(),
-            only_inline_degree_one_and_no_bus_field_vars,
-            &mut stats_logger,
-            bus_map,
-            degree_bound,
-        )?;
-
-    // Now remove the bus interaction field variables and run the optimizer,
-    // allowing all inlining below the degree bound.
-    let constraint_system = remove_bus_interaction_variables(constraint_system);
     let constraint_system =
         run_optimization_loop_until_no_change::<_, _, _, A::MemoryBusInteraction<_>>(
             constraint_system,
@@ -117,24 +96,6 @@ pub fn optimize<A: Adapter>(
     Ok(constraint_system_to_symbolic_machine(constraint_system))
 }
 
-/// Inlining discriminator that prevents the inliner from inlining
-/// bus interaction field variables (unless they are replaced
-/// by a single other variable) and variables defined via bus interaction
-/// field variables. All other variables are inlined if the expression has degree at most one.
-fn only_inline_degree_one_and_no_bus_field_vars<T: RuntimeConstant, V: Ord + Clone + Hash + Eq>(
-    var: &Variable<V>,
-    expr: &GroupedExpression<T, Variable<V>>,
-    _constraint_system: &IndexedConstraintSystem<T, Variable<V>>,
-) -> bool {
-    let is_about_bus_field_var = (matches!(var, Variable::BusInteractionField(_, _))
-        && expr.try_to_simple_unknown().is_none())
-        || expr
-            .referenced_unknown_variables()
-            .any(|v| matches!(v, Variable::BusInteractionField(_, _)));
-    // and only inline if the degree is at most 1.
-    !is_about_bus_field_var && expr.degree() <= 1
-}
-
 fn run_optimization_loop_until_no_change<
     P: FieldElement,
     V: Ord + Clone + Eq + Hash + Debug + Display,
@@ -395,112 +356,3 @@ impl<V: Display> Display for Variable<V> {
         }
     }
 }
-
-/// Transform the variable type of a constraint system by introducing
-/// new variables for bus interaction fields.
-fn introduce_bus_interaction_variables<T: FieldElement, V: Clone + Ord>(
-    constraint_system: ConstraintSystem<T, V>,
-) -> ConstraintSystem<T, Variable<V>> {
-    let mut new_constraints = Vec::new();
-    let mut bus_interaction_vars = BTreeMap::new();
-    let bus_interactions = constraint_system
-        .bus_interactions
-        .iter()
-        .enumerate()
-        .map(|(bus_interaction_index, bus_interaction)| {
-            BusInteraction::from_iter(bus_interaction.fields().enumerate().map(
-                |(field_index, expr)| {
-                    let transformed_expr =
-                        expr.transform_var_type(&mut |v| Variable::Variable(v.clone()));
-                    if transformed_expr.is_affine()
-                        && transformed_expr.referenced_unknown_variables().count() <= 1
-                    {
-                        transformed_expr
-                    } else {
-                        let v = Variable::BusInteractionField(bus_interaction_index, field_index);
-                        new_constraints.push(
-                            transformed_expr - GroupedExpression::from_unknown_variable(v.clone()),
-                        );
-                        bus_interaction_vars.insert(v.clone(), expr.clone());
-                        GroupedExpression::from_unknown_variable(v)
-                    }
-                },
-            ))
-        })
-        .collect();
-    ConstraintSystem {
-        algebraic_constraints: constraint_system
-            .algebraic_constraints
-            .iter()
-            .map(|expr| expr.transform_var_type(&mut |v| Variable::Variable(v.clone())))
-            .chain(new_constraints)
-            .collect(),
-        bus_interactions,
-    }
-}
-
-/// Reverses the effect of `introduce_bus_interaction_variables`, by inlining bus interaction
-/// field variables. This might fail in some cases.
-fn remove_bus_interaction_variables<T: FieldElement, V: Clone + Ord + Hash + Eq + Display>(
-    constraint_system: ConstraintSystem<T, Variable<V>>,
-) -> ConstraintSystem<T, V> {
-    let bus_interaction_var_definitions = constraint_system
-        .algebraic_constraints
-        .iter()
-        .flat_map(|expr| try_solve_for_single_bus_interaction_variable(expr))
-        .into_grouping_map()
-        .min_by_key(|_, expr| expr.degree());
-    let substituted_system = apply_substitutions(
-        constraint_system,
-        bus_interaction_var_definitions
-            .into_iter()
-            .sorted_by_key(|(var, _)| var.clone()),
-    );
-    ConstraintSystem {
-        algebraic_constraints: substituted_system
-            .algebraic_constraints
-            .into_iter()
-            .map(|expr| expr.transform_var_type(&mut transform_to_original_variable))
-            .filter(|expr| expr != &Zero::zero())
-            .collect(),
-        bus_interactions: substituted_system
-            .bus_interactions
-            .into_iter()
-            .map(|bi| {
-                BusInteraction::from_iter(
-                    bi.fields()
-                        .map(|expr| expr.transform_var_type(&mut transform_to_original_variable)),
-                )
-            })
-            .collect(),
-    }
-}
-
-fn transform_to_original_variable<V: Clone>(v: &Variable<V>) -> V {
-    match v {
-        Variable::Variable(v) => v.clone(),
-        Variable::BusInteractionField(_, _) => {
-            panic!("Unexpected bus interaction field in transformation to original variable")
-        }
-    }
-}
-
-/// Returns `Some(var, expr)` if `constraint` is equivalent to `var = expr`
-/// and `var` is the only bus interaction variable in `constraint`.
-#[allow(clippy::type_complexity)]
-fn try_solve_for_single_bus_interaction_variable<
-    T: FieldElement,
-    V: Clone + Ord + Hash + Eq + Display,
->(
-    constraint: &GroupedExpression<T, Variable<V>>,
-) -> Option<(Variable<V>, GroupedExpression<T, Variable<V>>)> {
-    let var = constraint
-        .referenced_unknown_variables()
-        .filter(|var| matches!(var, Variable::BusInteractionField(_, _)))
-        .unique()
-        .exactly_one()
-        .ok()?
-        .clone();
-    let solution = constraint.try_solve_for(&var)?;
-    Some((var, solution))
-}

constraint-solver/src/indexed_constraint_system.rs

@@ -217,7 +217,7 @@ fn retain<V, Item>(
     occurrences.retain(|_, occurrences| !occurrences.is_empty());
 }
 
-impl<T: RuntimeConstant, V: Clone + Ord + Hash> IndexedConstraintSystem<T, V> {
+impl<T: RuntimeConstant, V: Clone + Eq + Hash> IndexedConstraintSystem<T, V> {
     /// Adds new algebraic constraints to the system.
     pub fn add_algebraic_constraints(
         &mut self,

constraint-solver/src/journaling_constraint_system.rs

@@ -68,6 +68,24 @@ impl<T: RuntimeConstant + Substitutable<V>, V: Ord + Clone + Eq + Hash + Display
     }
 }
 
+impl<T: RuntimeConstant, V: Clone + Eq + Hash> JournalingConstraintSystem<T, V> {
+    /// Adds new algebraic constraints to the system.
+    pub fn add_algebraic_constraints(
+        &mut self,
+        constraints: impl IntoIterator<Item = GroupedExpression<T, V>>,
+    ) {
+        self.system.add_algebraic_constraints(constraints);
+    }
+
+    /// Adds new bus interactions to the system.
+    pub fn add_bus_interactions(
+        &mut self,
+        bus_interactions: impl IntoIterator<Item = BusInteraction<GroupedExpression<T, V>>>,
+    ) {
+        self.system.add_bus_interactions(bus_interactions);
+    }
+}
+
 impl<T: RuntimeConstant, V: Clone + Eq> JournalingConstraintSystem<T, V> {
     /// Removes all algebraic constraints that do not fulfill the predicate.
     pub fn retain_algebraic_constraints(

constraint-solver/src/solver/base.rs

@@ -1,4 +1,5 @@
 use itertools::Itertools;
+use num_traits::Zero;
 use powdr_number::{ExpressionConvertible, FieldElement};
 
 use crate::constraint_system::{BusInteraction, BusInteractionHandler, ConstraintRef};
@@ -84,7 +85,7 @@ where
     ) {
         self.equivalent_expressions_cache.clear();
         self.constraint_system
-            .add_algebraic_constraints(constraints);
+            .add_algebraic_constraints(constraints.into_iter().filter(|c| !c.is_zero()));
     }
 
     fn add_bus_interactions(

openvm/src/lib.rs

@@ -1670,7 +1670,7 @@ mod tests {
                             main: 16416,
                             log_up: 26292,
                         },
-                        constraints: 9752,
+                        constraints: 9132,
                         bus_interactions: 11135,
                     }
                 "#]],

openvm/tests/apc_builder_outputs/single_storeb.txt

@@ -74,5 +74,5 @@ flags__3_0 * ((flags__3_0 - 1) * (flags__3_0 - 2)) = 0
 (30720 * mem_ptr_limbs__0_0 - (30720 * rs1_data__0_0 + 7864320 * rs1_data__1_0 + 92160 * is_valid)) * (30720 * mem_ptr_limbs__0_0 - (30720 * rs1_data__0_0 + 7864320 * rs1_data__1_0 + 92161)) = 0
 (943718400 * rs1_data__0_0 + 30720 * mem_ptr_limbs__1_0 + 817889279 * is_valid - (120 * rs1_data__1_0 + 30720 * rs1_data__2_0 + 7864320 * rs1_data__3_0 + 943718400 * mem_ptr_limbs__0_0)) * (943718400 * rs1_data__0_0 + 30720 * mem_ptr_limbs__1_0 + 817889278 - (120 * rs1_data__1_0 + 30720 * rs1_data__2_0 + 7864320 * rs1_data__3_0 + 943718400 * mem_ptr_limbs__0_0)) = 0
 flags__1_0 * (flags__1_0 - 1) + flags__2_0 * (flags__2_0 - 1) + 4 * flags__0_0 * flags__1_0 + 4 * flags__0_0 * flags__2_0 + 5 * flags__0_0 * flags__3_0 + 5 * flags__1_0 * flags__2_0 + 5 * flags__1_0 * flags__3_0 + 5 * flags__2_0 * flags__3_0 - (1006632960 * flags__3_0 * (flags__3_0 - 1) + flags__0_0 * (flags__0_0 + flags__1_0 + flags__2_0 + flags__3_0 - 2) + flags__1_0 * (flags__0_0 + flags__1_0 + flags__2_0 + flags__3_0 - 2) + flags__2_0 * (flags__0_0 + flags__1_0 + flags__2_0 + flags__3_0 - 2) + 3 * flags__3_0 * (flags__0_0 + flags__1_0 + flags__2_0 + flags__3_0 - 2) + 5 * is_valid) = 0
-flags__0_0 * (flags__0_0 + flags__1_0 + flags__2_0 + flags__3_0 - 2) + 2 * flags__1_0 * (flags__0_0 + flags__1_0 + flags__2_0 + flags__3_0 - 2) + 3 * flags__2_0 * (flags__0_0 + flags__1_0 + flags__2_0 + flags__3_0 - 2) - flags__2_0 * (flags__2_0 - 1) = 0
+flags__2_0 * (flags__2_0 - 1) - (flags__0_0 * (flags__0_0 + flags__1_0 + flags__2_0 + flags__3_0 - 2) + 2 * flags__1_0 * (flags__0_0 + flags__1_0 + flags__2_0 + flags__3_0 - 2) + 3 * flags__2_0 * (flags__0_0 + flags__1_0 + flags__2_0 + flags__3_0 - 2)) = 0
 is_valid * (is_valid - 1) = 0

openvm/tests/apc_builder_outputs/single_storeh.txt

@@ -70,5 +70,5 @@ write_data__3_0 - (flags__2_0 * read_data__1_0 + (flags__1_0 * flags__2_0 + flag
 (30720 * mem_ptr_limbs__0_0 - (30720 * rs1_data__0_0 + 7864320 * rs1_data__1_0 + 184320 * is_valid)) * (30720 * mem_ptr_limbs__0_0 - (30720 * rs1_data__0_0 + 7864320 * rs1_data__1_0 + 184321)) = 0
 (943718400 * rs1_data__0_0 + 30720 * mem_ptr_limbs__1_0 - (120 * rs1_data__1_0 + 30720 * rs1_data__2_0 + 7864320 * rs1_data__3_0 + 943718400 * mem_ptr_limbs__0_0 + 377456642 * is_valid)) * (943718400 * rs1_data__0_0 + 30720 * mem_ptr_limbs__1_0 - (120 * rs1_data__1_0 + 30720 * rs1_data__2_0 + 7864320 * rs1_data__3_0 + 943718400 * mem_ptr_limbs__0_0 + 377456643)) = 0
 flags__1_0 * (flags__1_0 - 1) + flags__2_0 * (flags__2_0 - 1) + 5 * flags__1_0 * flags__2_0 + 3 * flags__1_0 + 3 * flags__2_0 - (flags__1_0 * (flags__1_0 + flags__2_0 - 1) + flags__2_0 * (flags__1_0 + flags__2_0 - 1) + 3 * is_valid) = 0
-2 * flags__1_0 * (flags__1_0 + flags__2_0 - 1) + 3 * flags__2_0 * (flags__1_0 + flags__2_0 - 1) + flags__1_0 + flags__2_0 - (flags__2_0 * (flags__2_0 - 1) + 1 * is_valid) = 0
+flags__2_0 * (flags__2_0 - 1) + 1 * is_valid - (2 * flags__1_0 * (flags__1_0 + flags__2_0 - 1) + 3 * flags__2_0 * (flags__1_0 + flags__2_0 - 1) + flags__1_0 + flags__2_0) = 0
 is_valid * (is_valid - 1) = 0