use coeffs form (instead of evals) for the final poly in whir (less verifier work)

Author: TomWambsgans

crates/backend/poly/src/evals.rs

@@ -41,6 +41,46 @@ impl<F: Field, EL: Borrow<[F]>> EvaluationsList<F> for EL {
     }
 }
 
+pub fn evals_to_coeffs<F: PrimeCharacteristicRing + Copy>(data: &mut [F]) {
+    let n = data.len();
+    let mut half = 1;
+    while half < n {
+        for i in (0..n).step_by(2 * half) {
+            for j in 0..half {
+                data[i + j + half] -= data[i + j];
+            }
+        }
+        half <<= 1;
+    }
+    bit_reverse_permutation(data);
+}
+
+pub fn bit_reverse_permutation<T>(data: &mut [T]) {
+    let n = data.len();
+    let log_n = n.ilog2() as usize;
+    for i in 0..n {
+        let j = i.reverse_bits() >> (usize::BITS as usize - log_n);
+        if i < j {
+            data.swap(i, j);
+        }
+    }
+}
+
+pub fn eval_multilinear_coeffs<F, EF>(coeffs: &[F], point: &[EF]) -> EF
+where
+    F: Field,
+    EF: ExtensionField<F>,
+{
+    debug_assert_eq!(coeffs.len(), 1 << point.len());
+    match point {
+        [] => EF::from(coeffs[0]),
+        [x, tail @ ..] => {
+            let (c0, c1) = coeffs.split_at(coeffs.len() / 2);
+            eval_multilinear_coeffs(c0, tail) + eval_multilinear_coeffs(c1, tail) * *x
+        }
+    }
+}
+
 /// Multiply the polynomial by a scalar factor.
 #[must_use]
 pub fn scale_poly<F: Field, EF: ExtensionField<F>>(poly: &[F], factor: EF) -> Vec<EF> {

crates/rec_aggregation/utils.py

@@ -76,20 +76,6 @@ def poly_eq_extension(point, n: Const):
             )
     return res + (2**n - 1) * DIM
 
-
-def poly_eq_base(point, n: Const):
-    # Example: for n = 2: eq(x, y) = [(1 - x)(1 - y), (1 - x)y, x(1 - y), xy]
-
-    res = Array((2 ** (n + 1) - 1))
-    res[0] = 1
-    for s in unroll(0, n):
-        p = point[n - 1 - s]
-        for i in unroll(0, 2**s):
-            res[2 ** (s + 1) - 1 + 2**s + i] = p * res[2**s - 1 + i]
-            res[2 ** (s + 1) - 1 + i] = res[2**s - 1 + i] - res[2 ** (s + 1) - 1 + 2**s + i]
-    return res + (2**n - 1)
-
-
 def eq_mle_extension(a, b, n):
     debug_assert(n < 30)
     debug_assert(0 < n)
@@ -207,6 +193,31 @@ def expand_from_univariate_ext(alpha, n):
     return res
 
 
+def univariate_eval_on_base(coeffs, alpha, n: Const):
+    # coeffs= univariate poly of degree 2^n
+    # alpha: base field element
+    # -> evaluates it at (1, alpha, alpha^2, alpha^4, ..., alpha^(2^(n-1)))
+    alpha_powers = Array(2**n)
+    alpha_powers[0] = 1
+    for i in unroll(0, 2**n - 1):
+        alpha_powers[i + 1] = alpha_powers[i] * alpha
+    result = Array(DIM)
+    dot_product(alpha_powers, coeffs, result, 2**n, BE)
+    return result
+
+
+def eval_multilinear_coeffs_rev(coeffs, point, n: Const):
+    # Evaluate multilinear polynomial in coefficient form (bit-reversed) at point.
+    basis = Array(2**n * DIM)
+    set_to_one(basis)
+    for k in unroll(0, n):
+        for j in unroll(0, 2**k):
+            mul_extension(basis + j * DIM, point + k * DIM, basis + (j + 2**k) * DIM)
+    result = Array(DIM)
+    dot_product(coeffs, basis, result, 2**n, EE)
+    return result
+
+
 def dot_product_be_dynamic(a, b, res, n):
     debug_assert(n <= 256)
     match_range(n, range(1, 257), lambda i: dot_product(a, b, res, i, BE))

crates/rec_aggregation/whir.py

@@ -24,7 +24,6 @@ def whir_open(
     claimed_sum: Mut,
 ):
     n_rounds, n_final_vars, num_queries, num_oods, query_grinding_bits, folding_grinding = get_whir_params(n_vars, initial_log_inv_rate)
-    n_final_coeffs = powers_of_two(n_final_vars)
     folding_factors = Array(n_rounds + 1)
     folding_factors[0] = WHIR_INITIAL_FOLDING_FACTOR
     for i in range(1, n_rounds + 1):
@@ -88,15 +87,8 @@ def whir_open(
 
     final_circle_values = all_circle_values[n_rounds]
     for i in range(0, num_queries[n_rounds]):
-        powers_of_2_rev = expand_from_univariate_base(final_circle_values[i], n_final_vars)
-        poly_eq = match_range(n_final_vars, range(MAX_NUM_VARIABLES_TO_SEND_COEFFS - WHIR_SUBSEQUENT_FOLDING_FACTOR, MAX_NUM_VARIABLES_TO_SEND_COEFFS + 1), lambda n: poly_eq_base(powers_of_2_rev, n))
-        final_pol_evaluated_on_circle = Array(DIM)
-        dot_product_be_dynamic(
-            poly_eq,
-            final_coeffcients,
-            final_pol_evaluated_on_circle,
-            n_final_coeffs,
-        )
+        alpha = final_circle_values[i]
+        final_pol_evaluated_on_circle = match_range(n_final_vars, range(MAX_NUM_VARIABLES_TO_SEND_COEFFS - WHIR_SUBSEQUENT_FOLDING_FACTOR, MAX_NUM_VARIABLES_TO_SEND_COEFFS + 1), lambda n: univariate_eval_on_base(final_coeffcients, alpha, n))
         copy_5(final_pol_evaluated_on_circle, final_folds + i * DIM)
 
     fs, all_folding_randomness[n_rounds + 1], end_sum = sumcheck_verify(fs, n_final_vars, claimed_sum, 2)
@@ -158,9 +150,7 @@ def whir_open(
         )
         s = add_extension_ret(s, s7)
         s = add_extension_ret(summed_ood, s)
-    poly_eq_final = poly_eq_extension_dynamic(all_folding_randomness[n_rounds + 1], n_final_vars)
-    final_value = Array(DIM)
-    dot_product_ee_dynamic(poly_eq_final, final_coeffcients, final_value, n_final_coeffs)
+    final_value = match_range(n_final_vars, range(MAX_NUM_VARIABLES_TO_SEND_COEFFS - WHIR_SUBSEQUENT_FOLDING_FACTOR, MAX_NUM_VARIABLES_TO_SEND_COEFFS + 1), lambda n: eval_multilinear_coeffs_rev(final_coeffcients, all_folding_randomness[n_rounds + 1], n))
     # copy_5(mul_extension_ret(s, final_value), end_sum);
 
     return fs, folding_randomness_global, s, final_value, end_sum

crates/whir/src/open.rs

@@ -184,13 +184,14 @@ where
         prover_state: &mut impl FSProver<EF>,
         round_state: &mut RoundState<EF>,
     ) -> ProofResult<()> {
-        // Directly send coefficients of the polynomial to the verifier.
-
-        prover_state.add_extension_scalars(&match &round_state.sumcheck_prover.evals {
+        // Convert evaluations to coefficient form and send to the verifier.
+        let mut coeffs = match &round_state.sumcheck_prover.evals {
             MleOwned::Extension(evals) => evals.clone(),
             MleOwned::ExtensionPacked(evals) => unpack_extension::<EF>(evals),
             _ => unreachable!(),
-        });
+        };
+        evals_to_coeffs(&mut coeffs);
+        prover_state.add_extension_scalars(&coeffs);
 
         prover_state.pow_grinding(self.final_query_pow_bits);
 

crates/whir/src/verify.rs

@@ -161,7 +161,7 @@ where
 
         // In the final round we receive the full polynomial instead of a commitment.
         let n_final_coeffs = 1 << self.n_vars_of_final_polynomial();
-        let final_evaluations = verifier_state.next_extension_scalars_vec(n_final_coeffs)?;
+        let final_coefficients = verifier_state.next_extension_scalars_vec(n_final_coeffs)?;
 
         // Verify in-domain challenges on the previous commitment.
         let stir_constraints = self.verify_stir_challenges(
@@ -175,7 +175,7 @@ where
         // Verify stir constraints directly on final polynomial
         stir_constraints
             .iter()
-            .all(|c| verify_constraint(c, &final_evaluations))
+            .all(|c| verify_constraint_coeffs(c, &final_coefficients))
             .then_some(())
             .ok_or(ProofError::InvalidProof)
             .unwrap();
@@ -194,8 +194,10 @@ where
 
         let evaluation_of_weights = self.eval_constraints_poly(&round_constraints, folding_randomness.clone());
 
-        // Check the final sumcheck evaluation
-        let final_value = final_evaluations.evaluate(&final_sumcheck_randomness);
+        // Check the final sumcheck evaluation (coefficient form, reversed point)
+        let mut reversed_point = final_sumcheck_randomness.0.clone();
+        reversed_point.reverse();
+        let final_value = eval_multilinear_coeffs(&final_coefficients, &reversed_point);
         if claimed_sum != evaluation_of_weights * final_value {
             panic!();
         }
@@ -417,12 +419,19 @@ where
     }
 }
 
-fn verify_constraint<EF: Field>(constraint: &SparseStatement<EF>, poly: &[EF]) -> bool {
-    // poly.evaluate(&constraint.point) == constraint.value
-    constraint
-        .values
-        .iter()
-        .all(|e| poly.evaluate_sparse(e.selector, &constraint.point) == e.value)
+fn verify_constraint_coeffs<EF: Field>(constraint: &SparseStatement<EF>, coeffs: &[EF]) -> bool {
+    assert_eq!(constraint.selector_num_variables(), 0);
+    let alpha = constraint.point[0];
+    // Verify the point is expand_from_univariate(alpha, n): [alpha, alpha^2, alpha^4, ...]
+    assert!(
+        constraint
+            .point
+            .iter()
+            .zip(constraint.point.iter().skip(1))
+            .all(|(&a, &b)| a * a == b)
+    );
+    let univariate_eval = coeffs.iter().rfold(EF::ZERO, |acc, &c| acc * alpha + c);
+    constraint.values.iter().all(|e| univariate_eval == e.value)
 }
 
 /// The full vector of folding randomness values, in reverse round order.