diff --git a/constraint-solver/src/rule_based_optimizer/environment.rs b/constraint-solver/src/rule_based_optimizer/environment.rs index 63c2974496..5510b1cc04 100644 --- a/constraint-solver/src/rule_based_optimizer/environment.rs +++ b/constraint-solver/src/rule_based_optimizer/environment.rs @@ -4,7 +4,7 @@ use std::{ hash::Hash, }; -use itertools::Itertools; +use itertools::{EitherOrBoth, Itertools}; use powdr_number::FieldElement; use crate::{ @@ -185,6 +185,64 @@ impl Environment { Some((self.insert(&l), self.insert(&r))) } + /// Returns Some(C) if `a - b = C' and both are affine. + pub fn constant_difference(&self, a: Expr, b: Expr) -> Option { + let db = self.expressions.borrow(); + let a = &db[a]; + let b = &db[b]; + (a.is_affine() + && b.is_affine() + && a.linear_components() + .zip(b.linear_components()) + .all(|(x, y)| x == y)) + .then(|| *a.constant_offset() - *b.constant_offset()) + } + + /// If this returns Some((v1, v2, factor)), then + /// a is obtained from b * factor by substituting v2 by v1. + pub fn differ_in_exactly_one_variable(&self, a_id: Expr, b_id: Expr) -> Option<(Var, Var, T)> { + let db = self.expressions.borrow(); + let a = &db[a_id]; + let b = &db[b_id]; + if !a.is_affine() + || !b.is_affine() + || a.linear_components().len() != b.linear_components().len() + || a.linear_components().len() < 2 + { + return None; + } + // First find the variables, ignoring the coefficients. + let (v1, v2) = a + .linear_components() + .merge_join_by(b.linear_components(), |(v1, _), (v2, _)| v1.cmp(v2)) + .filter(|either| !matches!(either, EitherOrBoth::Both(_, _))) + .collect_tuple()?; + let (left_var, right_var, factor) = match (v1, v2) { + (EitherOrBoth::Left((lv, lc)), EitherOrBoth::Right((rv, rc))) + | (EitherOrBoth::Right((rv, rc)), EitherOrBoth::Left((lv, lc))) => { + (*lv, *rv, *lc / *rc) + } + _ => return None, + }; + // Now verify that the other coefficients agree with the factor + if *a.constant_offset() != *b.constant_offset() * factor { + return None; + } + if !a + .linear_components() + .filter(|(v, _)| **v != left_var) + .map(|(_, c)| *c) + .eq(b + .linear_components() + .filter(|(v, _)| **v != right_var) + .map(|(_, bc)| *bc * factor)) + { + return None; + } + + Some((left_var, right_var, factor)) + } + /// Substitutes the variable `var` by the constant `value` in the expression `e` /// and returns the resulting expression. pub fn substitute_by_known(&self, e: Expr, var: Var, value: T) -> Expr { diff --git a/constraint-solver/src/rule_based_optimizer/rules.rs b/constraint-solver/src/rule_based_optimizer/rules.rs index f8668e1957..7443e3fdc2 100644 --- a/constraint-solver/src/rule_based_optimizer/rules.rs +++ b/constraint-solver/src/rule_based_optimizer/rules.rs @@ -75,6 +75,13 @@ crepe! { Expression(e), for v in env.on_expr(e, (), |e, _| e.referenced_unknown_variables().cloned().collect_vec()); + struct Product(Expr, Expr, Expr); + Product(e, l, r) <- + Expression(e), + Env(env), + let Some((l, r)) = env.try_as_single_product(e); + Product(e, r, l) <- Product(e, l, r); + // AffineExpression(e, coeff, var, offset) => e = coeff * var + offset struct AffineExpression(Expr, T, Var, T); AffineExpression(e, coeff, var, offset) <- @@ -113,6 +120,55 @@ crepe! { struct Equivalence(Var, Var); + //------- quadratic equivalence ----- + + // QuadraticEquivalenceCandidate(E, expr, offset) => + // E = ((expr) * (expr + offset) = 0) is a constraint and + // expr is affine with at least 2 variables. + struct QuadraticEquivalenceCandidate(Expr, Expr, T); + QuadraticEquivalenceCandidate(e, r, offset) <- + Env(env), + AlgebraicConstraint(e), + Product(e, l, r), // note that this will always produce two facts for (l, r) and (r, l) + ({env.affine_var_count(l).unwrap_or(0) > 1 && env.affine_var_count(r).unwrap_or(0) > 1}), + let Some(offset) = env.constant_difference(l, r); + + // QuadraticEquivalenceCandidatePair(expr1, expr2, offset1 / coeff, v1, v2) => + // (expr1) * (expr1 + offset1) = 0 and (expr2) * (expr2 + offset2) = 0 are constraints, + // expr1 is affine with at least 2 variables and is obtained from + // expr2 * factor by substituting v2 by v1 (factor != 0), + // offset1 == offset2 * factor and coeff is the coefficient of v1 in expr1. + // + // This means that v1 is always equal to (-expr1 / coeff) or equal to + // (-(expr1 + offset1) / coeff) = (-expr1 / coeff - offset1 / coeff). + // Because of the above, also v2 is equal to + // (-expr1 / coeff) or equal to (-(expr1 + offset1) / coeff) [Yes, expr1!]. + struct QuadraticEquivalenceCandidatePair(Expr, Expr, T, Var, Var); + QuadraticEquivalenceCandidatePair(expr1, expr2, offset1 / coeff, v1, v2) <- + Env(env), + QuadraticEquivalenceCandidate(_, expr1, offset1), + QuadraticEquivalenceCandidate(_, expr2, offset2), + (expr1 < expr2), + let Some((v1, v2,factor)) = env.differ_in_exactly_one_variable(expr1, expr2), + (offset1 == offset2 * factor), + let coeff = env.on_expr(expr1, (), |e, _| *e.coefficient_of_variable_in_affine_part(&v1).unwrap()); + + // QuadraticEquivalence(v1, v2) => v1 and v2 are equal in all satisfying assignments. + // Because of QuadraticEquivalenceCandidatePair, v1 is equal to X or X + offset, + // where X is some value that depends on other variables. Similarly, v2 is equal to X or X + offset. + // Because of the range constraints of v1 and v2, these two "or"s are exclusive ors. + // This means depending on the value of X, it is either X or X + offset. + // Since this "decision" only depens on X, both v1 and v2 are either X or X + offset at the same time + // and thus equal. + struct QuadraticEquivalence(Var, Var); + QuadraticEquivalence(v1, v2) <- + QuadraticEquivalenceCandidatePair(_, _, offset, v1, v2), + RangeConstraintOnVar(v1, rc), + RangeConstraintOnVar(v2, rc), + (rc.is_disjoint(&rc.combine_sum(&RangeConstraint::from_value(offset)))); + + Equivalence(v1, v2) <- QuadraticEquivalence(v1, v2); + ReplaceAlgebraicConstraintBy(e, env.substitute_by_known(e, v, val)) <- Env(env), Assignment(v, val), diff --git a/constraint-solver/src/rule_based_optimizer/tests.rs b/constraint-solver/src/rule_based_optimizer/tests.rs index d88bc303f5..a4333769c2 100644 --- a/constraint-solver/src/rule_based_optimizer/tests.rs +++ b/constraint-solver/src/rule_based_optimizer/tests.rs @@ -137,3 +137,87 @@ fn test_rule_based_optimization_simple_assignment() { ); expect!["(y) * (y - 1) - 3 = 0"].assert_eq(&optimized_system.0.to_string()); } + +#[test] +fn add_with_carry() { + // This tests a case of equivalent constraints that appear in the + // way "add with carry" is performed in openvm. + // X and Y end up being equivalent because they are both either + // A or A - 256, depending on whether the value of A is between + // 0 and 255 or not. + // A is the result of an addition with carry. + let mut system = IndexedConstraintSystem::default(); + system.add_algebraic_constraints([ + assert_zero( + (v("X") * c(7) - v("A") * c(7) + c(256) * c(7)) * (v("X") * c(7) - v("A") * c(7)), + ), + assert_zero((v("Y") - v("A") + c(256)) * (v("Y") - v("A"))), + ]); + system.add_bus_interactions([bit_constraint("X", 8), bit_constraint("Y", 8)]); + let optimized_system = rule_based_optimization( + system, + NoRangeConstraints, + TestBusInteractionHandler, + &mut new_var(), + None, + ); + // Y has been replaced by X + expect![[r#" + (7 * A - 7 * X - 1792) * (7 * A - 7 * X) = 0 + (A - X - 256) * (A - X) = 0 + BusInteraction { bus_id: 3, multiplicity: 1, payload: X, 8 } + BusInteraction { bus_id: 3, multiplicity: 1, payload: X, 8 }"#]] + .assert_eq(&optimized_system.0.to_string()); +} + +#[test] +fn test_rule_based_optimization_quadratic_equality() { + let mut system = IndexedConstraintSystem::default(); + system.add_algebraic_constraints([ + assert_zero( + (c(30720) * v("rs1_data__0_1") + c(7864320) * v("rs1_data__1_1") + - c(30720) * v("mem_ptr_limbs__0_1") + + c(737280)) + * (c(30720) * v("rs1_data__0_1") + c(7864320) * v("rs1_data__1_1") + - c(30720) * v("mem_ptr_limbs__0_1") + + c(737281)), + ), + assert_zero( + (c(30720) * v("rs1_data__0_1") + c(7864320) * v("rs1_data__1_1") + - c(30720) * v("mem_ptr_limbs__0_2") + + c(737280)) + * (c(30720) * v("rs1_data__0_1") + c(7864320) * v("rs1_data__1_1") + - c(30720) * v("mem_ptr_limbs__0_2") + + c(737281)), + ), + ]); + system.add_bus_interactions([ + bit_constraint("rs1_data__0_1", 8), + bit_constraint("rs1_data__1_1", 8), + BusInteraction { + bus_id: c(3), + multiplicity: c(1), + payload: vec![c(-503316480) * v("mem_ptr_limbs__0_1"), c(14)], + }, + BusInteraction { + bus_id: c(3), + multiplicity: c(1), + payload: vec![c(-503316480) * v("mem_ptr_limbs__0_2"), c(14)], + }, + ]); + let optimized_system = rule_based_optimization( + system, + NoRangeConstraints, + TestBusInteractionHandler, + &mut new_var(), + None, + ); + // Note that in the system below, mem_ptr_limbs__0_2 has been eliminated + expect![[r#" + (30720 * mem_ptr_limbs__0_1 - 30720 * rs1_data__0_1 - 7864320 * rs1_data__1_1 - 737280) * (30720 * mem_ptr_limbs__0_1 - 30720 * rs1_data__0_1 - 7864320 * rs1_data__1_1 - 737281) = 0 + (30720 * mem_ptr_limbs__0_1 - 30720 * rs1_data__0_1 - 7864320 * rs1_data__1_1 - 737280) * (30720 * mem_ptr_limbs__0_1 - 30720 * rs1_data__0_1 - 7864320 * rs1_data__1_1 - 737281) = 0 + BusInteraction { bus_id: 3, multiplicity: 1, payload: rs1_data__0_1, 8 } + BusInteraction { bus_id: 3, multiplicity: 1, payload: rs1_data__1_1, 8 } + BusInteraction { bus_id: 3, multiplicity: 1, payload: -(503316480 * mem_ptr_limbs__0_1), 14 } + BusInteraction { bus_id: 3, multiplicity: 1, payload: -(503316480 * mem_ptr_limbs__0_1), 14 }"#]].assert_eq(&optimized_system.0.to_string()); +} diff --git a/constraint-solver/src/solver.rs b/constraint-solver/src/solver.rs index 8514d40e5d..b74e6ad1ca 100644 --- a/constraint-solver/src/solver.rs +++ b/constraint-solver/src/solver.rs @@ -20,7 +20,6 @@ mod boolean_extractor; mod constraint_splitter; mod exhaustive_search; mod linearizer; -mod quadratic_equivalences; mod var_transformation; /// Solve a constraint system, i.e. derive assignments for variables in the system. diff --git a/constraint-solver/src/solver/base.rs b/constraint-solver/src/solver/base.rs index b81de11f77..a1e31f09ee 100644 --- a/constraint-solver/src/solver/base.rs +++ b/constraint-solver/src/solver/base.rs @@ -12,7 +12,7 @@ use crate::solver::boolean_extractor::BooleanExtractor; use crate::solver::constraint_splitter::try_split_constraint; use crate::solver::linearizer::Linearizer; use crate::solver::var_transformation::Variable; -use crate::solver::{exhaustive_search, quadratic_equivalences, Error, Solver, VariableAssignment}; +use crate::solver::{exhaustive_search, Error, Solver, VariableAssignment}; use crate::utils::possible_concrete_values; use std::collections::{BTreeSet, HashMap, HashSet}; @@ -326,8 +326,6 @@ where let mut progress = false; // Try solving constraints in isolation. progress |= self.solve_in_isolation()?; - // Try to find equivalent variables using quadratic constraints. - progress |= self.try_solve_quadratic_equivalences(); if !progress { // This might be expensive, so we only do it if we made no progress @@ -372,18 +370,6 @@ where Ok(progress) } - /// Tries to find equivalent variables using quadratic constraints. - fn try_solve_quadratic_equivalences(&mut self) -> bool { - let equivalences = quadratic_equivalences::find_quadratic_equalities( - self.constraint_system.system().algebraic_constraints(), - &*self, - ); - for (x, y) in &equivalences { - self.apply_assignment(y, &GroupedExpression::from_unknown_variable(x.clone())); - } - !equivalences.is_empty() - } - /// Find groups of variables with a small set of possible assignments. /// For each group, performs an exhaustive search in the possible assignments /// to deduce new range constraints (also on other variables). diff --git a/constraint-solver/src/solver/quadratic_equivalences.rs b/constraint-solver/src/solver/quadratic_equivalences.rs deleted file mode 100644 index 3e54470166..0000000000 --- a/constraint-solver/src/solver/quadratic_equivalences.rs +++ /dev/null @@ -1,145 +0,0 @@ -//! Module to help find certain pairs of equivalent variables in a system of quadratic constraints. - -use std::{collections::HashSet, fmt::Display, hash::Hash}; - -use itertools::Itertools; -use powdr_number::FieldElement; - -use crate::{ - constraint_system::AlgebraicConstraint, - grouped_expression::{GroupedExpression, RangeConstraintProvider}, - range_constraint::RangeConstraint, -}; - -/// Given a list of constraints, tries to determine pairs of equivalent variables. -pub fn find_quadratic_equalities( - constraints: &[AlgebraicConstraint>], - range_constraints: impl RangeConstraintProvider, -) -> Vec<(V, V)> { - let candidates = constraints - .iter() - .filter_map(QuadraticEqualityCandidate::try_from_constraint) - .filter(|c| c.variables.len() >= 2) - .collect::>(); - candidates - .iter() - .tuple_combinations() - .flat_map(|(c1, c2)| process_quadratic_equality_candidate_pair(c1, c2, &range_constraints)) - .collect() -} - -/// If we have two constraints of the form -/// `(X + expr) * (X + expr + offset) = 0` and -/// `(Y + expr) * (Y + expr + offset) = 0` -/// with the same range constraints on `X` and `Y` and -/// offset is a known value such that `X` and `Y` can be -/// determined by a (the same) conditional assignment from -/// `expr` or `expr + offset` (see [`QuadraticSymbolicExpression::solve_quadratic`]), -/// then `X` and `Y` must be equal and are returned. -fn process_quadratic_equality_candidate_pair< - T: FieldElement, - V: Ord + Clone + Hash + Eq + Display, ->( - c1: &QuadraticEqualityCandidate, - c2: &QuadraticEqualityCandidate, - range_constraints: &impl RangeConstraintProvider, -) -> Option<(V, V)> { - if c1.variables.len() != c2.variables.len() || c1.variables.len() < 2 { - return None; - } - let c1_var = c1.variables.difference(&c2.variables).exactly_one().ok()?; - let c2_var = c2.variables.difference(&c1.variables).exactly_one().ok()?; - // The expressions differ in exactly one variable. - let rc1 = range_constraints.get(c1_var); - let rc2 = range_constraints.get(c2_var); - - // And those variables have the same range constraint. - if rc1 != rc2 { - return None; - } - - let c1 = c1.normalized_for_var(c1_var); - let c2 = c2.normalized_for_var(c2_var); - - if c1.offset != c2.offset { - return None; - } - - // And the offset (the difference between the two alternatives) determines if - // we are inside the range constraint or not. - if !rc1.is_disjoint(&rc1.combine_sum(&RangeConstraint::from_value(c1.offset))) { - return None; - } - - // Now the only remaining check is to see if the affine expressions are the same. - // This could have been the first step, but it is rather expensive, so we do it last. - if c1.expr - GroupedExpression::from_unknown_variable(c1_var.clone()) - != c2.expr - GroupedExpression::from_unknown_variable(c2_var.clone()) - { - return None; - } - - // Now we have `(X + A) * (X + A + offset) = 0` and `(Y + A) * (Y + A + offset) = 0` - // Furthermore, the range constraints of `X` and `Y` are such that for both identities, - // the two alternatives can never be satisfied at the same time. Since both variables - // have the same range constraint, depending on the value of `A`, we either have - // - X = -A and Y = -A, or - // - X = -A - offset and Y = -A - offset - // Since `A` has to have some value, we can conclude `X = Y`. - - Some((c1_var.clone(), c2_var.clone())) -} - -/// This represents an identity `expr * (expr + offset) = 0`, -/// where `expr` is an affine expression and `offset` is a runtime constant. -/// -/// All unknown variables appearing in `expr` are stored in `variables`. -struct QuadraticEqualityCandidate { - expr: GroupedExpression, - offset: T, - /// All unknown variables in `expr`. - variables: HashSet, -} - -impl QuadraticEqualityCandidate { - fn try_from_constraint(constr: &AlgebraicConstraint>) -> Option { - let (left, right) = constr.expression.try_as_single_product()?; - if !left.is_affine() || !right.is_affine() { - return None; - } - // `constr = 0` is equivalent to `left * right = 0` - let offset = *(left - right).try_to_known()?; - // `offset + right = left` - // `constr = 0` is equivalent to `right * (right + offset) = 0` - let variables = right - .referenced_unknown_variables() - .cloned() - .collect::>(); - Some(Self { - expr: right.clone(), - offset, - variables, - }) - } - - /// Returns an equivalent candidate that is normalized - /// such that `var` has a coefficient of `1`. - fn normalized_for_var(&self, var: &V) -> Self { - let coefficient = *self - .expr - .coefficient_of_variable_in_affine_part(var) - .unwrap(); - - // self represents - // `(coeff * var + X) * (coeff * var + X + offset) = 0` - // Dividing by `coeff` twice results in - // `(var + X / coeff) * (var + X / coeff + offset / coeff) = 0` - let offset = self.offset / coefficient; - let expr = self.expr.clone() * (T::from(1) / coefficient); - Self { - expr, - offset, - variables: self.variables.clone(), - } - } -} diff --git a/constraint-solver/tests/solver.rs b/constraint-solver/tests/solver.rs index f426ea72e7..2903d48aaa 100644 --- a/constraint-solver/tests/solver.rs +++ b/constraint-solver/tests/solver.rs @@ -1,13 +1,11 @@ use std::collections::BTreeMap; -use itertools::Itertools; use num_traits::identities::{One, Zero}; use powdr_constraint_solver::{ constraint_system::{ BusInteraction, BusInteractionHandler, ConstraintSystem, DefaultBusInteractionHandler, }, grouped_expression::GroupedExpression, - indexed_constraint_system::apply_substitutions, range_constraint::RangeConstraint, solver::{solve_system, Error}, }; @@ -251,39 +249,6 @@ fn xor_invalid() { } } -#[test] -fn add_with_carry() { - // This tests a case of equivalent constraints that appear in the - // way "add with carry" is performed in openvm. - // X and Y end up being equivalent because they are both either - // A or A - 256, depending on whether the value of A is between - // 0 and 255 or not. - // A is the result of an addition with carry. - let constraint_system = ConstraintSystem::default() - .with_constraints(vec![ - (var("X") * constant(7) - var("A") * constant(7) + constant(256) * constant(7)) - * (var("X") * constant(7) - var("A") * constant(7)), - (var("Y") - var("A") + constant(256)) * (var("Y") - var("A")), - ]) - .with_bus_interactions(vec![ - // Byte range constraints on X and Y - send(BYTE_BUS_ID, vec![var("X")]), - send(BYTE_BUS_ID, vec![var("Y")]), - ]); - - let final_state = solve_system(constraint_system.clone(), TestBusInteractionHandler).unwrap(); - let final_state = apply_substitutions(constraint_system, final_state) - .algebraic_constraints - .iter() - .format("\n") - .to_string(); - assert_eq!( - final_state, - "(7 * A - 7 * X - 1792) * (7 * A - 7 * X) = 0 -(A - X - 256) * (A - X) = 0" - ); -} - #[test] fn one_hot_flags() { let constraint_system = ConstraintSystem::default().with_constraints(vec![ diff --git a/openvm/src/lib.rs b/openvm/src/lib.rs index bbcab05663..91ff41574c 100644 --- a/openvm/src/lib.rs +++ b/openvm/src/lib.rs @@ -1758,11 +1758,11 @@ mod tests { AirMetrics { widths: AirWidths { preprocessed: 0, - main: 17300, - log_up: 27896, + main: 17289, + log_up: 27884, }, - constraints: 8834, - bus_interactions: 11925, + constraints: 8823, + bus_interactions: 11919, } "#]], powdr_expected_machine_count: expect![[r#" @@ -1780,8 +1780,8 @@ mod tests { }, after: AirWidths { preprocessed: 0, - main: 17300, - log_up: 27896, + main: 17289, + log_up: 27884, }, } "#]]), @@ -1807,11 +1807,11 @@ mod tests { AirMetrics { widths: AirWidths { preprocessed: 0, - main: 19928, - log_up: 30924, + main: 19909, + log_up: 30904, }, - constraints: 11103, - bus_interactions: 13442, + constraints: 11084, + bus_interactions: 13432, } "#]], powdr_expected_machine_count: expect![[r#" @@ -1829,8 +1829,8 @@ mod tests { }, after: AirWidths { preprocessed: 0, - main: 19928, - log_up: 30924, + main: 19909, + log_up: 30904, }, } "#]]),