Add variant of fingerprint() that yields intermediate columns (powd…

…r-labs#2072)

Depends on powdr-labs#2151.

This PR solves all performance issues with `std::fingerprint`. It now
has two functions:
- `fingerprint()`: Is similar to a pre-powdr-labs#1985 version of `fingerprint()`.
It computes the fingerprint in a recursive way. It should *not* be used
to generate expressions, because they would be exponentially large. To
prevent that, I changed the type so that it can only be applied to `fe`.
- `fingerprint_inter()`: Is the analog for expressions. It stores
intermediate results in intermediate columns. Because of that, it needs
to be a `constr` function. It generates $O(n)$ intermediate columns of
constant size.

I also added a test for the `fingerprint_inter` function and changed all
protocols (bus, permutation, lookup) to use `fingerprint_inter` to
generate constraints.

pipeline/tests/powdr_std.rs

@@ -9,12 +9,20 @@ use powdr_pipeline::{
         evaluate_function, evaluate_integer_function, gen_estark_proof_with_backend_variant,
         gen_halo2_proof, make_simple_prepared_pipeline, regular_test,
         regular_test_without_small_field, std_analyzed, test_halo2_with_backend_variant,
-        test_mock_backend, test_plonky3_with_backend_variant, BackendVariant,
+        test_mock_backend, test_plonky3_pipeline, test_plonky3_with_backend_variant,
+        BackendVariant,
     },
     Pipeline,
 };
 use test_log::test;
 
+#[test]
+fn fingerprint_test() {
+    let f = "std/fingerprint_test.asm";
+    let pipeline = make_simple_prepared_pipeline::<GoldilocksField>(f);
+    test_plonky3_pipeline(pipeline);
+}
+
 #[test]
 #[ignore = "Too slow"]
 fn poseidon_bn254_test() {

std/array.asm

@@ -123,4 +123,10 @@ mod internal {
                     merge(l_short, right, lt) + [l_last]
                 }
         };
-}
+}
+
+/// Applies the next operator to all elements of the array.
+let next: expr[] -> expr[] = query |arr| new(len(arr), |i| arr[i]');
+
+/// Evaluates an array of expressions
+let eval: expr[] -> fe[] = query |arr| new(len(arr), |i| std::prover::eval(arr[i]));

std/protocols/bus.asm

@@ -13,6 +13,7 @@ use std::math::fp2::needs_extension;
 use std::math::fp2::fp2_from_array;
 use std::math::fp2::constrain_eq_ext;
 use std::protocols::fingerprint::fingerprint_with_id;
+use std::protocols::fingerprint::fingerprint_with_id_inter;
 use std::math::fp2::required_extension_size;
 use std::prover::eval;
 
@@ -35,7 +36,7 @@ let bus_interaction: expr, expr[], expr -> () = constr |id, tuple, multiplicity|
     let beta = fp2_from_array(array::new(required_extension_size(), |i| challenge(0, i + 3)));
 
     // Implemented as: folded = (beta - fingerprint(id, tuple...));
-    let folded = sub_ext(beta, fingerprint_with_id(id, tuple, alpha));
+    let folded = sub_ext(beta, fingerprint_with_id_inter(id, tuple, alpha));
     let folded_next = next_ext(folded);
 
     let m_ext = from_base(multiplicity);
@@ -83,8 +84,7 @@ let bus_interaction: expr, expr[], expr -> () = constr |id, tuple, multiplicity|
 let compute_next_z: expr, expr, expr[], expr, Fp2<expr>, Fp2<expr>, Fp2<expr> -> fe[] = query |is_first, id, tuple, multiplicity, acc, alpha, beta| {
     // Implemented as: folded = (beta - fingerprint(id, tuple...));
     // `multiplicity / (beta - fingerprint(id, tuple...))` to `acc`
-    let folded = sub_ext(beta, fingerprint_with_id(id, tuple, alpha));
-    let folded_next = next_ext(folded);
+    let folded_next = sub_ext(eval_ext(beta), fingerprint_with_id(eval(id'), array::eval(array::next(tuple)), alpha));
 
     let m_ext = from_base(multiplicity);
     let m_ext_next = next_ext(m_ext);
@@ -95,7 +95,7 @@ let compute_next_z: expr, expr, expr[], expr, Fp2<expr>, Fp2<expr>, Fp2<expr> ->
     // acc' = current_acc + multiplicity' / folded'
     let res = add_ext(
         current_acc,
-        mul_ext(eval_ext(m_ext_next), inv_ext(eval_ext(folded_next)))
+        mul_ext(eval_ext(m_ext_next), inv_ext(folded_next))
     );
 
     unpack_ext_array(res)

std/protocols/fingerprint.asm

@@ -1,24 +1,51 @@
+use std::array;
 use std::array::len;
-use std::utils::fold;
 use std::math::fp2::Fp2;
 use std::math::fp2::add_ext;
 use std::math::fp2::mul_ext;
 use std::math::fp2::pow_ext;
 use std::math::fp2::from_base;
+use std::math::fp2::eval_ext;
+use std::check::assert;
 
 /// Maps [x_1, x_2, ..., x_n] to its Read-Solomon fingerprint, using a challenge alpha: $\sum_{i=1}^n alpha**{(n - i)} * x_i$
-let<T: Add + Mul + FromLiteral> fingerprint: T[], Fp2<T> -> Fp2<T> = |expr_array, alpha| {
-    let n = len(expr_array);
-    fold(
-        n,
-        |i| mul_ext(pow_ext(alpha, n - i - 1), from_base(expr_array[i])),
-        from_base(0),
-        |sum_acc, el| add_ext(sum_acc, el)
-    )
+/// To generate an expression that computes the fingerprint, use `fingerprint_inter` instead.
+/// Note that alpha is passed as an expressions, so that it is only evaluated if needed (i.e., if len(expr_array) > 1).
+let fingerprint: fe[], Fp2<expr> -> Fp2<fe> = query |expr_array, alpha| if len(expr_array) == 1 {
+    // Base case
+    from_base(expr_array[0])
+} else {
+    assert(len(expr_array) > 1, || "fingerprint requires at least one element");
+
+    // Recursively compute the fingerprint as fingerprint(expr_array[:-1], alpha) * alpha + expr_array[-1]
+    let intermediate_fingerprint = fingerprint(array::sub_array(expr_array, 0, len(expr_array) - 1), alpha);
+    add_ext(mul_ext(eval_ext(alpha), intermediate_fingerprint), from_base(expr_array[len(expr_array) - 1]))
+};
+
+/// Like `fingerprint`, but "materializes" the intermediate results as intermediate columns.
+/// Inlining them would lead to an exponentially-sized expression.
+let fingerprint_inter: expr[], Fp2<expr> -> Fp2<expr> = |expr_array, alpha| if len(expr_array) == 1 {
+    // Base case
+    from_base(expr_array[0])
+} else {
+    assert(len(expr_array) > 1, || "fingerprint requires at least one element");
+
+    // Recursively compute the fingerprint as fingerprint(expr_array[:-1], alpha) * alpha + expr_array[-1]
+    let intermediate_fingerprint = match fingerprint_inter(array::sub_array(expr_array, 0, len(expr_array) - 1), alpha) {
+        Fp2::Fp2(a0, a1) => {
+            let intermediate_fingerprint_0: inter = a0;
+            let intermediate_fingerprint_1: inter = a1;
+            Fp2::Fp2(intermediate_fingerprint_0, intermediate_fingerprint_1)
+        }
+    };
+    add_ext(mul_ext(alpha, intermediate_fingerprint), from_base(expr_array[len(expr_array) - 1]))
 };
 
 /// Maps [id, x_1, x_2, ..., x_n] to its Read-Solomon fingerprint, using a challenge alpha: $\sum_{i=1}^n alpha**{(n - i)} * x_i$
-let<T: Add + Mul + FromLiteral> fingerprint_with_id: T, T[], Fp2<T> -> Fp2<T> = |id, expr_array, alpha| fingerprint([id] + expr_array, alpha);
+let fingerprint_with_id: fe, fe[], Fp2<expr> -> Fp2<fe> = query |id, expr_array, alpha| fingerprint([id] + expr_array, alpha);
+
+/// Maps [id, x_1, x_2, ..., x_n] to its Read-Solomon fingerprint, using a challenge alpha: $\sum_{i=1}^n alpha**{(n - i)} * x_i$
+let fingerprint_with_id_inter: expr, expr[], Fp2<expr> -> Fp2<expr> = |id, expr_array, alpha| fingerprint_inter([id] + expr_array, alpha);
 
 mod test {
     use super::fingerprint;
@@ -27,8 +54,7 @@ mod test {
     use std::math::fp2::from_base;
 
     /// Helper function to assert that the fingerprint of a tuple is equal to the expected value.
-    /// We are working on integers here, wrapping them as Fp2 elements.
-    let assert_fingerprint_equal: int[], int, int -> () = |tuple, challenge, expected| {
+    let assert_fingerprint_equal: fe[], expr, fe -> () = query |tuple, challenge, expected| {
         let result = fingerprint(tuple, from_base(challenge));
         match result {
             Fp2::Fp2(actual, should_be_zero) => {
@@ -38,7 +64,7 @@ mod test {
         }
     };
 
-    let test_fingerprint = || {
+    let test_fingerprint = query || {
         // A tuple t of size n with challenge x should be mapped to:
         // t[0] * x**(n-1) + t[1] * x**(n-2) + ... + t[n-2] * x + t[n-1]
 

std/protocols/lookup.asm

@@ -20,6 +20,7 @@ use std::math::fp2::constrain_eq_ext;
 use std::math::fp2::required_extension_size;
 use std::math::fp2::needs_extension;
 use std::protocols::fingerprint::fingerprint;
+use std::protocols::fingerprint::fingerprint_inter;
 use std::utils::unwrap_or_else;
 
 let unpack_lookup_constraint: Constr -> (expr, expr[], expr, expr[]) = |lookup_constraint| match lookup_constraint {
@@ -36,19 +37,19 @@ let unpack_lookup_constraint: Constr -> (expr, expr[], expr, expr[]) = |lookup_c
 let compute_next_z: Fp2<expr>, Fp2<expr>, Fp2<expr>, Constr, expr -> fe[] = query |acc, alpha, beta, lookup_constraint, multiplicities| {
     let (lhs_selector, lhs, rhs_selector, rhs) = unpack_lookup_constraint(lookup_constraint);
 
-    let lhs_denom = sub_ext(beta, fingerprint(lhs, alpha));
-    let rhs_denom = sub_ext(beta, fingerprint(rhs, alpha));
+    let lhs_denom = sub_ext(eval_ext(beta), fingerprint(array::eval(lhs), alpha));
+    let rhs_denom = sub_ext(eval_ext(beta), fingerprint(array::eval(rhs), alpha));
     let m_ext = from_base(multiplicities);
 
     // acc' = acc + 1/(beta-a_i) * lhs_selector - m_i/(beta-b_i) * rhs_selector
     let res = add_ext(
         eval_ext(acc),
         sub_ext(
             mul_ext(
-                inv_ext(eval_ext(lhs_denom)), 
+                inv_ext(lhs_denom), 
                 eval_ext(from_base(lhs_selector))),
             mul_ext(
-                mul_ext(eval_ext(m_ext), inv_ext(eval_ext(rhs_denom))),
+                mul_ext(eval_ext(m_ext), inv_ext(rhs_denom)),
                 eval_ext(from_base(rhs_selector))
         )
     ));
@@ -67,8 +68,8 @@ let lookup: Constr -> () = constr |lookup_constraint| {
 
     let (lhs_selector, lhs, rhs_selector, rhs) = unpack_lookup_constraint(lookup_constraint);
 
-    let lhs_denom = sub_ext(beta, fingerprint(lhs, alpha));
-    let rhs_denom = sub_ext(beta, fingerprint(rhs, alpha));
+    let lhs_denom = sub_ext(beta, fingerprint_inter(lhs, alpha));
+    let rhs_denom = sub_ext(beta, fingerprint_inter(rhs, alpha));
     let multiplicities;
     let m_ext = from_base(multiplicities);
 

std/protocols/permutation.asm

@@ -16,6 +16,9 @@ use std::math::fp2::constrain_eq_ext;
 use std::math::fp2::required_extension_size;
 use std::math::fp2::needs_extension;
 use std::protocols::fingerprint::fingerprint;
+use std::protocols::fingerprint::fingerprint_inter;
+use std::prover::eval;
+use std::array;
 use std::utils::unwrap_or_else;
 use std::constraints::to_phantom_permutation;
 
@@ -42,13 +45,13 @@ let compute_next_z: Fp2<expr>, Fp2<expr>, Fp2<expr>, Constr -> fe[] = query |acc
 
     let (lhs_selector, lhs, rhs_selector, rhs) = unpack_permutation_constraint(permutation_constraint);
 
-    let lhs_folded = selected_or_one(lhs_selector, sub_ext(beta, fingerprint(lhs, alpha)));
-    let rhs_folded = selected_or_one(rhs_selector, sub_ext(beta, fingerprint(rhs, alpha)));
+    let lhs_folded = selected_or_one(eval(lhs_selector), sub_ext(eval_ext(beta), fingerprint(array::eval(lhs), alpha)));
+    let rhs_folded = selected_or_one(eval(rhs_selector), sub_ext(eval_ext(beta), fingerprint(array::eval(rhs), alpha)));
 
     // acc' = acc * lhs_folded / rhs_folded
     let res = mul_ext(
-        eval_ext(mul_ext(acc, lhs_folded)),
-        inv_ext(eval_ext(rhs_folded))
+        mul_ext(eval_ext(acc), lhs_folded),
+        inv_ext(rhs_folded)
     );
 
     unpack_ext_array(res)
@@ -83,8 +86,8 @@ let permutation: Constr -> () = constr |permutation_constraint| {
     // If the selector is 1, contribute a factor of `beta - fingerprint(lhs)` to accumulator.
     // If the selector is 0, contribute a factor of 1 to the accumulator.
     // Implemented as: folded = selector * (beta - fingerprint(values) - 1) + 1;
-    let lhs_folded = selected_or_one(lhs_selector, sub_ext(beta, fingerprint(lhs, alpha)));
-    let rhs_folded = selected_or_one(rhs_selector, sub_ext(beta, fingerprint(rhs, alpha)));
+    let lhs_folded = selected_or_one(lhs_selector, sub_ext(beta, fingerprint_inter(lhs, alpha)));
+    let rhs_folded = selected_or_one(rhs_selector, sub_ext(beta, fingerprint_inter(rhs, alpha)));
 
     let acc = std::array::new(required_extension_size(), |i| std::prover::new_witness_col_at_stage("acc", 1));
     let acc_ext = fp2_from_array(acc);

test_data/std/fingerprint_test.asm

@@ -0,0 +1,41 @@
+use std::math::fp2::from_base;
+use std::math::fp2::Fp2;
+use std::math::fp2::eval_ext;
+use std::math::fp2::unpack_ext_array;
+use std::math::fp2::constrain_eq_ext;
+use std::prover::challenge;
+use std::protocols::fingerprint::fingerprint;
+use std::protocols::fingerprint::fingerprint_inter;
+use std::array;
+use std::convert::expr;
+use std::prover::eval;
+
+machine Main with degree: 2048 {
+
+    col witness x(i) query Query::Hint(42);
+
+    // Fold tuple [x, x + 1, ..., x + n - 1]
+    // Note that, by setting a fairly large `n`, we test that performance is not exponential in `n`.
+    let n = 100;
+    let tuple = array::new(n, |i| x + 1);
+
+    // Add `fingerprint_value` witness columns and constrain them using `fingerprint_inter`
+    col witness stage(1) fingerprint_value0, fingerprint_value1;
+    let fingerprint_value = Fp2::Fp2(fingerprint_value0, fingerprint_value1);
+    let alpha = Fp2::Fp2(challenge(0, 0), challenge(0, 1));
+    constrain_eq_ext(fingerprint_inter(tuple, alpha), fingerprint_value);
+
+    // Add `fingerprint_value_hint` witness columns and compute the fingerprint in a hint using `fingerprint`
+    let fingerprint_hint: -> fe[] = query || {
+        let tuple_eval = array::new(array::len(tuple), |i| eval(tuple[i]));
+        unpack_ext_array(fingerprint(tuple_eval, alpha))
+    };
+
+    col witness stage(1) fingerprint_value0_hint(i) query Query::Hint(fingerprint_hint()[0]);
+    col witness stage(1) fingerprint_value1_hint(i) query Query::Hint(fingerprint_hint()[1]);
+
+    // Assert consistency between `fingerprint` and `fingerprint_inter`
+    fingerprint_value0 = fingerprint_value0_hint;
+    fingerprint_value1 = fingerprint_value1_hint;
+
+}