From 9cce7ea4f98144fc073822bf581716fde7ec7c52 Mon Sep 17 00:00:00 2001 From: onurinanc Date: Sun, 30 Jun 2024 16:04:15 +0300 Subject: [PATCH] implement lookup arguments from logarithmic derivatives in pil (#1477) This PR solves #1374 There are 3 examples included in the PR: 1. `lookup_via_challenges_ext_simple.asm` is the most basic example which implements {a} in {b} using extension field without using selectors 2. `lookup_via_challenges_ext_no_selector.asm` implements {a1, a2} in {b1, b2} using extension field without using selectors 3. `lookup_via_challenges_ext.asm` is a more complex example than others which implements {a1, a2, a3} in {b1, b2, b3} using extension field (which also handles tuples using Reed-Solomon fingerprinting). It also use different lhs and rhs selectors. --------- Co-authored-by: Georg Wiese --- pipeline/tests/powdr_std.rs | 30 ++++ std/protocols/lookup.asm | 160 ++++++++++++++++++ std/protocols/mod.asm | 3 +- test_data/std/lookup_via_challenges.asm | 27 +++ test_data/std/lookup_via_challenges_ext.asm | 38 +++++ .../std/lookup_via_challenges_ext_simple.asm | 31 ++++ 6 files changed, 288 insertions(+), 1 deletion(-) create mode 100644 std/protocols/lookup.asm create mode 100644 test_data/std/lookup_via_challenges.asm create mode 100644 test_data/std/lookup_via_challenges_ext.asm create mode 100644 test_data/std/lookup_via_challenges_ext_simple.asm diff --git a/pipeline/tests/powdr_std.rs b/pipeline/tests/powdr_std.rs index 8121c9290..f7e54b8c4 100644 --- a/pipeline/tests/powdr_std.rs +++ b/pipeline/tests/powdr_std.rs @@ -115,6 +115,36 @@ fn permutation_via_challenges_ext() { .unwrap(); } +#[test] +fn lookup_via_challenges_bn() { + let f = "std/lookup_via_challenges.asm"; + test_halo2(f, Default::default()); +} + +#[test] +fn lookup_via_challenges_ext() { + let f = "std/lookup_via_challenges_ext.asm"; + test_halo2(f, Default::default()); + // Note that this does not actually run the second-phase witness generation, because no + // Goldilocks backend support challenges yet. + Pipeline::::default() + .from_file(resolve_test_file(f)) + .compute_witness() + .unwrap(); +} + +#[test] +fn lookup_via_challenges_ext_simple() { + let f = "std/lookup_via_challenges_ext_simple.asm"; + test_halo2(f, Default::default()); + // Note that this does not actually run the second-phase witness generation, because no + // Goldilocks backend support challenges yet. + Pipeline::::default() + .from_file(resolve_test_file(f)) + .compute_witness() + .unwrap(); +} + #[test] fn write_once_memory_test() { let f = "std/write_once_memory_test.asm"; diff --git a/std/protocols/lookup.asm b/std/protocols/lookup.asm new file mode 100644 index 000000000..68ca77404 --- /dev/null +++ b/std/protocols/lookup.asm @@ -0,0 +1,160 @@ +use std::prover::challenge; +use std::array::fold; +use std::utils::unwrap_or_else; +use std::array::len; +use std::array::map; +use std::check::assert; +use std::check::panic; +use std::field::known_field; +use std::field::KnownField; +use std::math::fp2::Fp2; +use std::math::fp2::add_ext; +use std::math::fp2::sub_ext; +use std::math::fp2::mul_ext; +use std::math::fp2::unpack_ext; +use std::math::fp2::next_ext; +use std::math::fp2::inv_ext; +use std::math::fp2::eval_ext; +use std::math::fp2::from_base; +use std::math::fp2::constrain_eq_ext; + + +let is_first: col = |i| if i == 0 { 1 } else { 0 }; + +// challenges to be used in polynomial evaluation and folding different columns +let alpha1: expr = challenge(0, 1); +let alpha2: expr = challenge(0, 2); + +let beta1: expr = challenge(0, 3); +let beta2: expr = challenge(0, 4); + +let unpack_lookup_constraint: Constr -> (expr, expr[], expr, expr[]) = |lookup_constraint| match lookup_constraint { + Constr::Lookup((lhs_selector, rhs_selector), values) => ( + unwrap_or_else(lhs_selector, || 1), + map(values, |(lhs, _)| lhs), + unwrap_or_else(rhs_selector, || 1), + map(values, |(_, rhs)| rhs) + ), + _ => panic("Expected lookup constraint") +}; + +/// Whether we need to operate on the F_{p^2} extension field (because the current field is too small). +let needs_extension: -> bool = || match known_field() { + Option::Some(KnownField::Goldilocks) => true, + Option::Some(KnownField::BN254) => false, + None => panic("The lookup argument is not implemented for the current field!") +}; + +//* Generic for both permutation and lookup arguments +/// Maps [x_1, x_2, ..., x_n] to alpha**(n - 1) * x_1 + alpha ** (n - 2) * x_2 + ... + x_n +let compress_expression_array: T[], Fp2 -> Fp2 = |expr_array, alpha| fold( + expr_array, + from_base(0), + |sum_acc, el| add_ext(mul_ext(alpha, sum_acc), from_base(el)) +); + +// Compute z' = z + 1/(beta-a_i) * lhs_selector - m_i/(beta-b_i) * rhs_selector, using extension field arithmetic +let compute_next_z: Fp2, Constr, expr -> fe[] = query |acc, lookup_constraint, multiplicities| { + let (lhs_selector, lhs, rhs_selector, rhs) = unpack_lookup_constraint(lookup_constraint); + let alpha = if len(lhs) > 1 { + Fp2::Fp2(alpha1, alpha2) + } else { + // The optimizer will have removed alpha, but the compression function + // still accesses it (to multiply by 0 in this case) + from_base(0) + }; + let beta = Fp2::Fp2(beta1, beta2); + + let lhs_denom = sub_ext(beta, compress_expression_array(lhs, alpha)); + let rhs_denom = sub_ext(beta, compress_expression_array(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)), + eval_ext(from_base(lhs_selector))), + mul_ext( + mul_ext(eval_ext(m_ext), inv_ext(eval_ext(rhs_denom))), + eval_ext(from_base(rhs_selector)) + ) + )); + match res { + Fp2::Fp2(a0_fe, a1_fe) => [a0_fe, a1_fe] + } +}; + +// Adds constraints that enforce that rhs is the lookup for lhs +// Arguments: +// - acc: A phase-2 witness column to be used as the accumulator. If 2 are provided, computations +// are done on the F_{p^2} extension field. +// - lookup_constraint: The lookup constraint +// - multiplicities: The multiplicities which shows how many times each RHS value appears in the LHS +let lookup: expr[], Constr, expr -> Constr[] = |acc, lookup_constraint, multiplicities| { + + let (lhs_selector, lhs, rhs_selector, rhs) = unpack_lookup_constraint(lookup_constraint); + + let _ = assert(len(lhs) == len(rhs), || "LHS and RHS should have equal length"); + + let with_extension = match len(acc) { + 1 => false, + 2 => true, + _ => panic("Expected 1 or 2 accumulator columns!") + }; + + let _ = if !with_extension { + assert(!needs_extension(), || "The Goldilocks field is too small and needs to move to the extension field. Pass two accumulators instead!") + } else { () }; + + // On the extension field, we'll need two field elements to represent the challenge. + // If we don't need an extension field, we can simply set the second component to 0, + // in which case the operations below effectively only operate on the first component. + let fp2_from_array = |arr| if with_extension { Fp2::Fp2(arr[0], arr[1]) } else { from_base(arr[0]) }; + let acc_ext = fp2_from_array(acc); + let alpha = fp2_from_array([alpha1, alpha2]); + let beta = fp2_from_array([beta1, beta2]); + + let lhs_denom = sub_ext(beta, compress_expression_array(lhs, alpha)); + let rhs_denom = sub_ext(beta, compress_expression_array(rhs, alpha)); + let m_ext = from_base(multiplicities); + + let next_acc = if with_extension { + next_ext(acc_ext) + } else { + // The second component is 0, but the next operator is not defined on it... + from_base(acc[0]') + }; + + // Update rule: + // acc' * (beta - A) * (beta - B) + m * rhs_selector * (beta - A) = acc * (beta - A) * (beta - B) + lhs_selector * (beta - B) + // => (acc' - acc) * lhs_denom * rhs_denom + // + m * rhs_selector * lhs_denom + // - lhs_selector * rhs_denom = 0 + + let update_expr = sub_ext( + add_ext( + mul_ext( + mul_ext(lhs_denom, rhs_denom), + sub_ext(next_acc, acc_ext) + ), + mul_ext( + mul_ext(m_ext, from_base(rhs_selector)), + lhs_denom + ) + ), + mul_ext( + from_base(lhs_selector), + rhs_denom + ) + ); + + let (acc_1, acc_2) = unpack_ext(acc_ext); + + [ + is_first * acc_1 = 0, + + is_first * acc_2 = 0 + ] + constrain_eq_ext(update_expr, from_base(0)) +}; \ No newline at end of file diff --git a/std/protocols/mod.asm b/std/protocols/mod.asm index 69b16e969..1ef638a24 100644 --- a/std/protocols/mod.asm +++ b/std/protocols/mod.asm @@ -1 +1,2 @@ -mod permutation; \ No newline at end of file +mod permutation; +mod lookup; \ No newline at end of file diff --git a/test_data/std/lookup_via_challenges.asm b/test_data/std/lookup_via_challenges.asm new file mode 100644 index 000000000..59ba20b32 --- /dev/null +++ b/test_data/std/lookup_via_challenges.asm @@ -0,0 +1,27 @@ +use std::prover::Query; +use std::convert::fe; +use std::protocols::lookup::lookup; + +machine Main with degree: 8 { + col fixed random_six = [1, 1, 1, 0, 1, 1, 1, 0]; + col fixed first_seven = [1, 1, 1, 1, 1, 1, 1, 0]; + + col fixed a1 = [1, 2, 4, 3, 1, 1, 4, 1]; + col fixed a2 = [1, 2, 4, 1, 1, 1, 4, 1]; + col fixed a3 = [1, 2, 4, 1, 1, 1, 4, 3]; + col witness b1(i) query Query::Hint(fe(i+1)); + col witness b2(i) query Query::Hint(fe(i+1)); + col witness b3(i) query Query::Hint(fe(i+1)); + col fixed m = [3, 1, 0, 2, 0, 0, 0, 0]; + + let lookup_constraint = Constr::Lookup( + (Option::Some(random_six), Option::Some(first_seven)), + [(a1, b1), (a2, b2), (a3, b3)] + ); + + // TODO: Functions currently cannot add witness columns at later stages, + // so we have to manually create it here and pass it to permutation(). + col witness stage(1) z; + lookup([z], lookup_constraint, m); + +} \ No newline at end of file diff --git a/test_data/std/lookup_via_challenges_ext.asm b/test_data/std/lookup_via_challenges_ext.asm new file mode 100644 index 000000000..01b24cfeb --- /dev/null +++ b/test_data/std/lookup_via_challenges_ext.asm @@ -0,0 +1,38 @@ +use std::prover::Query; +use std::convert::fe; +use std::protocols::lookup::lookup; +use std::protocols::lookup::compute_next_z; +use std::math::fp2::Fp2; + +machine Main with degree: 8 { + col fixed a_sel = [0, 1, 1, 1, 0, 1, 0, 0]; + col fixed b_sel = [1, 1, 0, 1, 1, 1, 1, 0]; + + col fixed a1 = [16, 20, 22, 17, 16, 20, 4, 1]; + col fixed a2 = [12, 5, 7, 2, 5, 5, 4, 1]; + col fixed a3 = [20, 36, 38, 33, 36, 36, 4, 3]; + col witness b1(i) query Query::Hint(fe(i+16)); + col witness b2(i) query Query::Hint(fe(i+1)); + col witness b3(i) query Query::Hint(fe(i+32)); + col fixed m = [0, 1, 0, 0, 2, 0, 1, 0]; + + let lookup_constraint = Constr::Lookup( + (Option::Some(a_sel), Option::Some(b_sel)), + [(a1, b1), (a2, b2), (a3, b3)] + ); + + // TODO: Functions currently cannot add witness columns at later stages, + // so we have to manually create it here and pass it to lookup(). + col witness stage(1) z1; + col witness stage(1) z2; + + lookup([z1, z2], lookup_constraint, m); + + // TODO: Helper columns, because we can't access the previous row in hints + let hint = query |i| Query::Hint(compute_next_z(Fp2::Fp2(z1, z2), lookup_constraint, m)[i]); + col witness stage(1) z1_next(i) query hint(0); + col witness stage(1) z2_next(i) query hint(1); + + z1' = z1_next; + z2' = z2_next; +} \ No newline at end of file diff --git a/test_data/std/lookup_via_challenges_ext_simple.asm b/test_data/std/lookup_via_challenges_ext_simple.asm new file mode 100644 index 000000000..033d62213 --- /dev/null +++ b/test_data/std/lookup_via_challenges_ext_simple.asm @@ -0,0 +1,31 @@ +use std::prover::Query; +use std::convert::fe; +use std::protocols::lookup::lookup; +use std::protocols::lookup::compute_next_z; +use std::math::fp2::Fp2; + +machine Main with degree: 8 { + col fixed a = [1, 1, 4, 1, 1, 2, 1, 1]; + col witness b(i) query Query::Hint(fe(i+1)); + col fixed m = [6, 1, 0, 1, 0, 0, 0, 0]; + + let lookup_constraint = Constr::Lookup( + (Option::None, Option::None), + [(a, b)] + ); + + // TODO: Functions currently cannot add witness columns at later stages, + // so we have to manually create it here and pass it to lookup(). + col witness stage(1) z1; + col witness stage(1) z2; + + lookup([z1, z2], lookup_constraint, m); + + // TODO: Helper columns, because we can't access the previous row in hints + let hint = query |i| Query::Hint(compute_next_z(Fp2::Fp2(z1, z2), lookup_constraint, m)[i]); + col witness stage(1) z1_next(i) query hint(0); + col witness stage(1) z2_next(i) query hint(1); + + z1' = z1_next; + z2' = z2_next; +} \ No newline at end of file