Replace computation method "InverseOrZero" by "QuotientOrZero". (#3457)

Replace the computation method "InverseOrZero" by "QuotientOrZero" which
is more flexible.

Author: chriseth

autoprecompiles/src/optimizer.rs

@@ -217,9 +217,10 @@ fn symbolic_machine_to_constraint_system<P: FieldElement>(
             .map(|(v, method)| {
                 let method = match method {
                     ComputationMethod::Constant(c) => ComputationMethod::Constant(*c),
-                    ComputationMethod::InverseOrZero(c) => {
-                        ComputationMethod::InverseOrZero(algebraic_to_grouped_expression(c))
-                    }
+                    ComputationMethod::QuotientOrZero(e1, e2) => ComputationMethod::QuotientOrZero(
+                        algebraic_to_grouped_expression(e1),
+                        algebraic_to_grouped_expression(e2),
+                    ),
                 };
                 DerivedVariable {
                     variable: v.clone(),
@@ -250,9 +251,10 @@ fn constraint_system_to_symbolic_machine<P: FieldElement>(
             .map(|derived_var| {
                 let method = match derived_var.computation_method {
                     ComputationMethod::Constant(c) => ComputationMethod::Constant(c),
-                    ComputationMethod::InverseOrZero(c) => {
-                        ComputationMethod::InverseOrZero(grouped_expression_to_algebraic(c))
-                    }
+                    ComputationMethod::QuotientOrZero(e1, e2) => ComputationMethod::QuotientOrZero(
+                        grouped_expression_to_algebraic(e1),
+                        grouped_expression_to_algebraic(e2),
+                    ),
                 };
                 (derived_var.variable, method)
             })

autoprecompiles/src/symbolic_machine_generator.rs

@@ -35,8 +35,9 @@ pub fn convert_machine_field_type<T, U>(
                     ComputationMethod::Constant(c) => {
                         ComputationMethod::Constant(convert_field_element(c))
                     }
-                    ComputationMethod::InverseOrZero(e) => ComputationMethod::InverseOrZero(
-                        convert_expression(e, convert_field_element),
+                    ComputationMethod::QuotientOrZero(e1, e2) => ComputationMethod::QuotientOrZero(
+                        convert_expression(e1, convert_field_element),
+                        convert_expression(e2, convert_field_element),
                     ),
                 };
                 (v, method)

constraint-solver/src/constraint_system.rs

@@ -94,15 +94,16 @@ pub struct DerivedVariable<T, V> {
 pub enum ComputationMethod<T, E> {
     /// A constant value.
     Constant(T),
-    /// The field inverse of an expression if it exists or zero otherwise.
-    InverseOrZero(E),
+    /// The quotiont (using inversion in the field) of the first argument
+    /// by the second argument, or zero if the latter is zero.
+    QuotientOrZero(E, E),
 }
 
 impl<T: Display, E: Display> Display for ComputationMethod<T, E> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             ComputationMethod::Constant(c) => write!(f, "{c}"),
-            ComputationMethod::InverseOrZero(e) => write!(f, "InverseOrZero({e})"),
+            ComputationMethod::QuotientOrZero(e1, e2) => write!(f, "QuotientOrZero({e1}, {e2})"),
         }
     }
 }
@@ -112,7 +113,10 @@ impl<T, F> ComputationMethod<T, GroupedExpression<T, F>> {
     pub fn referenced_unknown_variables(&self) -> Box<dyn Iterator<Item = &F> + '_> {
         match self {
             ComputationMethod::Constant(_) => Box::new(std::iter::empty()),
-            ComputationMethod::InverseOrZero(e) => e.referenced_unknown_variables(),
+            ComputationMethod::QuotientOrZero(e1, e2) => Box::new(
+                e1.referenced_unknown_variables()
+                    .chain(e2.referenced_unknown_variables()),
+            ),
         }
     }
 }
@@ -125,8 +129,9 @@ impl<T: RuntimeConstant + Substitutable<V>, V: Ord + Clone + Eq>
     pub fn substitute_by_known(&mut self, variable: &V, substitution: &T) {
         match self {
             ComputationMethod::Constant(_) => {}
-            ComputationMethod::InverseOrZero(e) => {
-                e.substitute_by_known(variable, substitution);
+            ComputationMethod::QuotientOrZero(e1, e2) => {
+                e1.substitute_by_known(variable, substitution);
+                e2.substitute_by_known(variable, substitution);
             }
         }
     }
@@ -138,8 +143,9 @@ impl<T: RuntimeConstant + Substitutable<V>, V: Ord + Clone + Eq>
     pub fn substitute_by_unknown(&mut self, variable: &V, substitution: &GroupedExpression<T, V>) {
         match self {
             ComputationMethod::Constant(_) => {}
-            ComputationMethod::InverseOrZero(e) => {
-                e.substitute_by_unknown(variable, substitution);
+            ComputationMethod::QuotientOrZero(e1, e2) => {
+                e1.substitute_by_unknown(variable, substitution);
+                e2.substitute_by_unknown(variable, substitution);
             }
         }
     }

constraint-solver/src/indexed_constraint_system.rs

@@ -829,30 +829,32 @@ mod tests {
             derived_variables: vec![
                 DerivedVariable {
                     variable: "d1",
-                    computation_method: ComputationMethod::InverseOrZero(
-                        GroupedExpression::from_unknown_variable("x"),
+                    computation_method: ComputationMethod::QuotientOrZero(
+                        GroupedExpression::from_unknown_variable("x1"),
+                        GroupedExpression::from_unknown_variable("x2"),
                     ),
                 },
                 DerivedVariable {
                     variable: "d2",
-                    computation_method: ComputationMethod::InverseOrZero(
-                        GroupedExpression::from_unknown_variable("y"),
+                    computation_method: ComputationMethod::QuotientOrZero(
+                        GroupedExpression::from_unknown_variable("y1"),
+                        GroupedExpression::from_unknown_variable("y2"),
                     ),
                 },
             ],
         }
         .into();
-        // We first substitute `y` by an expression that contains `x` such that when we
-        // substitute `x` in the next step, `d2` has to be updated again.
+        // We first substitute `y2` by an expression that contains `x1` such that when we
+        // substitute `x1` in the next step, `d2` has to be updated again.
         system.substitute_by_unknown(
-            &"y",
-            &(GroupedExpression::from_unknown_variable("x")
+            &"y2",
+            &(GroupedExpression::from_unknown_variable("x1")
                 + GroupedExpression::from_number(7.into())),
         );
-        system.substitute_by_known(&"x", &1.into());
+        system.substitute_by_known(&"x1", &1.into());
         assert_eq!(
             format!("{system}"),
-            "d1 := InverseOrZero(1)\nd2 := InverseOrZero(8)"
+            "d1 := QuotientOrZero(1, x2)\nd2 := QuotientOrZero(y1, 8)"
         );
     }
 }

openvm/src/powdr_extension/trace_generator/cpu/mod.rs

@@ -183,16 +183,19 @@ impl PowdrTraceGeneratorCpu {
                     let col_index = apc_poly_id_to_index[&column.id];
                     row_slice[col_index] = match computation_method {
                         ComputationMethod::Constant(c) => *c,
-                        ComputationMethod::InverseOrZero(expr) => {
+                        ComputationMethod::QuotientOrZero(e1, e2) => {
                             use powdr_number::ExpressionConvertible;
 
-                            let expr_val = expr.to_expression(&|n| *n, &|column_ref| {
+                            let divisor_val = e2.to_expression(&|n| *n, &|column_ref| {
                                 row_slice[apc_poly_id_to_index[&column_ref.id]]
                             });
-                            if expr_val.is_zero() {
+                            if divisor_val.is_zero() {
                                 BabyBear::ZERO
                             } else {
-                                expr_val.inverse()
+                                divisor_val.inverse()
+                                    * e1.to_expression(&|n| *n, &|column_ref| {
+                                        row_slice[apc_poly_id_to_index[&column_ref.id]]
+                                    })
                             }
                         }
                     };

openvm/src/powdr_extension/trace_generator/cuda/mod.rs

@@ -120,10 +120,12 @@ fn compile_derived_to_gpu(
                 bytecode.push(OpCode::PushConst as u32);
                 bytecode.push(c.as_canonical_u32());
             }
-            ComputationMethod::InverseOrZero(expr) => {
+            ComputationMethod::QuotientOrZero(e1, e2) => {
                 // Encode inner expression, then apply InvOrZero
-                emit_expr(&mut bytecode, expr, apc_poly_id_to_index, apc_height);
+                emit_expr(&mut bytecode, e2, apc_poly_id_to_index, apc_height);
                 bytecode.push(OpCode::InvOrZero as u32);
+                emit_expr(&mut bytecode, e1, apc_poly_id_to_index, apc_height);
+                bytecode.push(OpCode::Mul as u32);
             }
         }
         let len = (bytecode.len() as u32) - off;