From 3813a4a1b94667b3d3a440a529c46901d2915cd8 Mon Sep 17 00:00:00 2001 From: Georg Wiese Date: Wed, 27 Nov 2024 12:43:50 +0100 Subject: [PATCH] Add variant of `fingerprint()` that yields intermediate columns (#2072) Depends on #2151. This PR solves all performance issues with `std::fingerprint`. It now has two functions: - `fingerprint()`: Is similar to a pre-#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 | 10 +++++- std/array.asm | 8 ++++- std/protocols/bus.asm | 8 ++--- std/protocols/fingerprint.asm | 52 ++++++++++++++++++++++-------- std/protocols/lookup.asm | 13 ++++---- std/protocols/permutation.asm | 15 +++++---- test_data/std/fingerprint_test.asm | 41 +++++++++++++++++++++++ 7 files changed, 116 insertions(+), 31 deletions(-) create mode 100644 test_data/std/fingerprint_test.asm diff --git a/pipeline/tests/powdr_std.rs b/pipeline/tests/powdr_std.rs index 87909b7d54..db54738c3b 100644 --- a/pipeline/tests/powdr_std.rs +++ b/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::(f); + test_plonky3_pipeline(pipeline); +} + #[test] #[ignore = "Too slow"] fn poseidon_bn254_test() { diff --git a/std/array.asm b/std/array.asm index d2eefdfdc4..886463384c 100644 --- a/std/array.asm +++ b/std/array.asm @@ -123,4 +123,10 @@ mod internal { merge(l_short, right, lt) + [l_last] } }; -} \ No newline at end of file +} + +/// 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])); \ No newline at end of file diff --git a/std/protocols/bus.asm b/std/protocols/bus.asm index a7fb29805a..c0e1ccb1ac 100644 --- a/std/protocols/bus.asm +++ b/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, Fp2, Fp2 -> 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, Fp2, Fp2 -> // 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) diff --git a/std/protocols/fingerprint.asm b/std/protocols/fingerprint.asm index 3fb586e618..c19bb6d634 100644 --- a/std/protocols/fingerprint.asm +++ b/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 fingerprint: T[], Fp2 -> Fp2 = |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 -> Fp2 = 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 -> Fp2 = |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 fingerprint_with_id: T, T[], Fp2 -> Fp2 = |id, expr_array, alpha| fingerprint([id] + expr_array, alpha); +let fingerprint_with_id: fe, fe[], Fp2 -> Fp2 = 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 -> Fp2 = |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] diff --git a/std/protocols/lookup.asm b/std/protocols/lookup.asm index 85b91a3ada..7bfdf5a087 100644 --- a/std/protocols/lookup.asm +++ b/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,8 +37,8 @@ let unpack_lookup_constraint: Constr -> (expr, expr[], expr, expr[]) = |lookup_c let compute_next_z: Fp2, Fp2, Fp2, 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 @@ -45,10 +46,10 @@ let compute_next_z: Fp2, Fp2, Fp2, Constr, expr -> fe[] = quer 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); diff --git a/std/protocols/permutation.asm b/std/protocols/permutation.asm index 983618ea2a..37eaa40213 100644 --- a/std/protocols/permutation.asm +++ b/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, Fp2, Fp2, 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); diff --git a/test_data/std/fingerprint_test.asm b/test_data/std/fingerprint_test.asm new file mode 100644 index 0000000000..bb15305d44 --- /dev/null +++ b/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; + +} \ No newline at end of file