skip the last 5 sumchecks of logup-GKR (send data in clear instead)

Author: TomWambsgans

crates/rec_aggregation/recursion.py

@@ -4,6 +4,9 @@
 
 N_TABLES = N_TABLES_PLACEHOLDER
 
+LOGUP_GKR_N_VARS_TO_SEND_COEFFS = LOGUP_GKR_N_VARS_TO_SEND_COEFFS_PLACEHOLDER
+LOGUP_GKR_N_COEFFS_SENT = 2**LOGUP_GKR_N_VARS_TO_SEND_COEFFS
+
 MIN_LOG_N_ROWS_PER_TABLE = MIN_LOG_N_ROWS_PER_TABLE_PLACEHOLDER
 MAX_LOG_N_ROWS_PER_TABLE = MAX_LOG_N_ROWS_PER_TABLE_PLACEHOLDER
 MIN_LOG_MEMORY_SIZE = MIN_LOG_MEMORY_SIZE_PLACEHOLDER
@@ -673,27 +676,33 @@ def fingerprint_bytecode(instr_evals, eval_on_pc, logup_alphas_eq_poly):
 
 
 def verify_gkr_quotient(fs: Mut, n_vars):
-    fs, nums = fs_receive_ef_inlined(fs, 2)
-    fs, denoms = fs_receive_ef_inlined(fs, 2)
-
-    q1 = div_extension_ret(nums, denoms)
-    q2 = div_extension_ret(nums + DIM, denoms + DIM)
-    quotient = add_extension_ret(q1, q2)
+    fs, nums = fs_receive_ef_inlined(fs, LOGUP_GKR_N_COEFFS_SENT)
+    fs, denoms = fs_receive_ef_inlined(fs, LOGUP_GKR_N_COEFFS_SENT)
+
+    initial_quotients = Array(LOGUP_GKR_N_COEFFS_SENT * DIM)
+    for k in unroll(0, LOGUP_GKR_N_COEFFS_SENT):
+        div_extension(nums + k * DIM, denoms + k * DIM, initial_quotients + k * DIM)
+    debug_assert(NUM_REPEATED_ONES <= LOGUP_GKR_N_COEFFS_SENT)
+    debug_assert(LOGUP_GKR_N_COEFFS_SENT % NUM_REPEATED_ONES == 0)
+    quotient: Mut = ZERO_VEC_PTR
+    for k in unroll(0, LOGUP_GKR_N_COEFFS_SENT / NUM_REPEATED_ONES):
+        quotient = add_extension_ret(quotient, sum_continuous_ef(initial_quotients + k * NUM_REPEATED_ONES * DIM, NUM_REPEATED_ONES))
 
     points = Array(n_vars)
     claims_num = Array(n_vars)
     claims_den = Array(n_vars)
 
-    fs, points[0] = fs_sample_ef(fs)
+    fs, initial_point = fs_sample_many_ef(fs, LOGUP_GKR_N_VARS_TO_SEND_COEFFS)
+    points[LOGUP_GKR_N_VARS_TO_SEND_COEFFS - 1] = initial_point
 
-    point_poly_eq = poly_eq_extension(points[0], 1)
+    point_poly_eq = poly_eq_extension(initial_point, LOGUP_GKR_N_VARS_TO_SEND_COEFFS)
 
-    first_claim_num = dot_product_ee_ret(nums, point_poly_eq, 2)
-    first_claim_den = dot_product_ee_ret(denoms, point_poly_eq, 2)
-    claims_num[0] = first_claim_num
-    claims_den[0] = first_claim_den
+    first_claim_num = dot_product_ee_ret(nums, point_poly_eq, LOGUP_GKR_N_COEFFS_SENT)
+    first_claim_den = dot_product_ee_ret(denoms, point_poly_eq, LOGUP_GKR_N_COEFFS_SENT)
+    claims_num[LOGUP_GKR_N_VARS_TO_SEND_COEFFS - 1] = first_claim_num
+    claims_den[LOGUP_GKR_N_VARS_TO_SEND_COEFFS - 1] = first_claim_den
 
-    for i in range(1, n_vars):
+    for i in range(LOGUP_GKR_N_VARS_TO_SEND_COEFFS, n_vars):
         fs, points[i], claims_num[i], claims_den[i] = verify_gkr_quotient_step(fs, i, points[i - 1], claims_num[i - 1], claims_den[i - 1])
 
     return (

crates/rec_aggregation/src/compilation.rs

@@ -8,7 +8,7 @@ use lean_vm::*;
 use std::collections::{BTreeMap, HashMap};
 use std::path::Path;
 use std::sync::OnceLock;
-use sub_protocols::{min_stacked_n_vars, total_whir_statements};
+use sub_protocols::{N_VARS_TO_SEND_GKR_COEFFS, min_stacked_n_vars, total_whir_statements};
 use tracing::instrument;
 use utils::Counter;
 use xmss::{LOG_LIFETIME, MESSAGE_LEN_FE, RANDOMNESS_LEN_FE, TARGET_SUM, V, V_GRINDING, W};
@@ -183,6 +183,10 @@ fn build_replacements(
         "MAX_NUM_VARIABLES_TO_SEND_COEFFS_PLACEHOLDER".to_string(),
         MAX_NUM_VARIABLES_TO_SEND_COEFFS.to_string(),
     );
+    replacements.insert(
+        "LOGUP_GKR_N_VARS_TO_SEND_COEFFS_PLACEHOLDER".to_string(),
+        N_VARS_TO_SEND_GKR_COEFFS.to_string(),
+    );
     replacements.insert(
         "WHIR_INITIAL_FOLDING_FACTOR_PLACEHOLDER".to_string(),
         WHIR_INITIAL_FOLDING_FACTOR.to_string(),

crates/rec_aggregation/utils.py

@@ -304,9 +304,14 @@ def mul_base_extension_ret(a, b):
 @inline
 def div_extension_ret(n, d):
     quotient = Array(DIM)
-    dot_product_ee(d, quotient, n)
+    div_extension(n, d, quotient)
     return quotient
 
+@inline
+def div_extension(n, d, res):
+    dot_product_ee(d, res, n)
+    return
+
 
 @inline
 def sub_extension(a, b, c):

crates/sub_protocols/src/lib.rs

@@ -11,3 +11,4 @@ mod quotient_gkr;
 pub use quotient_gkr::*;
 
 pub(crate) const MIN_VARS_FOR_PACKING: usize = 8;
+pub const N_VARS_TO_SEND_GKR_COEFFS: usize = 5;

crates/sub_protocols/src/quotient_gkr.rs

@@ -2,7 +2,7 @@ use std::ops::Mul;
 
 use backend::*;
 
-use crate::MIN_VARS_FOR_PACKING;
+use crate::{MIN_VARS_FOR_PACKING, N_VARS_TO_SEND_GKR_COEFFS};
 
 /*
 GKR to compute sum of fractions.
@@ -13,6 +13,8 @@ pub fn prove_gkr_quotient<EF: ExtensionField<PF<EF>>>(
     numerators: &MleRef<'_, EF>,
     denominators: &MleRef<'_, EF>,
 ) -> (EF, MultilinearPoint<EF>, EF, EF) {
+    assert!(numerators.n_vars() == denominators.n_vars());
+    assert!(numerators.n_vars() > N_VARS_TO_SEND_GKR_COEFFS);
     assert!(numerators.is_packed() == denominators.is_packed());
     let mut layers: Vec<(Mle<'_, EF>, Mle<'_, EF>)> =
         vec![(numerators.soft_clone().into(), denominators.soft_clone().into())];
@@ -26,23 +28,23 @@ pub fn prove_gkr_quotient<EF: ExtensionField<PF<EF>>>(
                 prev_denominators.unpack().as_owned_or_clone().into(),
             )
         }
-        if prev_numerators.n_vars() == 1 {
+        if prev_numerators.n_vars() <= N_VARS_TO_SEND_GKR_COEFFS {
             break;
         }
         let (new_numerators, new_denominators) = sum_quotients(prev_numerators.by_ref(), prev_denominators.by_ref());
         layers.push((new_numerators.into(), new_denominators.into()));
     }
 
-    let (last_numerators, last_denominators) = layers.pop().unwrap();
-    let last_numerators = last_numerators.as_owned().unwrap();
-    let last_numerators = last_numerators.as_extension().unwrap();
-    let last_denominators = last_denominators.as_owned().unwrap();
-    let last_denominators = last_denominators.as_extension().unwrap();
+    let (last_numerators_mle, last_denominators_mle) = layers.pop().unwrap();
+    let last_numerators_owned = last_numerators_mle.unpack().as_owned_or_clone();
+    let last_denominators_owned = last_denominators_mle.unpack().as_owned_or_clone();
+    let last_numerators = last_numerators_owned.as_extension().unwrap();
+    let last_denominators = last_denominators_owned.as_extension().unwrap();
     prover_state.add_extension_scalars(last_numerators);
     prover_state.add_extension_scalars(last_denominators);
-    let quotient = last_numerators[0] / last_denominators[0] + last_numerators[1] / last_denominators[1];
+    let quotient = compute_quotient(last_numerators, last_denominators);
 
-    let mut point = MultilinearPoint(vec![prover_state.sample()]);
+    let mut point = MultilinearPoint(prover_state.sample_vec(N_VARS_TO_SEND_GKR_COEFFS));
     let mut claims = vec![last_numerators.evaluate(&point), last_denominators.evaluate(&point)];
 
     for (nums, denoms) in layers.iter().rev() {
@@ -208,13 +210,15 @@ pub fn verify_gkr_quotient<EF: ExtensionField<PF<EF>>>(
     verifier_state: &mut impl FSVerifier<EF>,
     n_vars: usize,
 ) -> Result<(EF, MultilinearPoint<EF>, EF, EF), ProofError> {
-    let last_nums = verifier_state.next_extension_scalars_vec(2)?;
-    let last_dens = verifier_state.next_extension_scalars_vec(2)?;
-    let quotient = last_nums[0] / last_dens[0] + last_nums[1] / last_dens[1];
-    let mut point = MultilinearPoint(vec![verifier_state.sample()]);
+    assert!(n_vars > N_VARS_TO_SEND_GKR_COEFFS);
+    let send_len = 1 << N_VARS_TO_SEND_GKR_COEFFS;
+    let last_nums = verifier_state.next_extension_scalars_vec(send_len)?;
+    let last_dens = verifier_state.next_extension_scalars_vec(send_len)?;
+    let quotient: EF = compute_quotient(&last_nums, &last_dens);
+    let mut point = MultilinearPoint(verifier_state.sample_vec(N_VARS_TO_SEND_GKR_COEFFS));
     let mut claims_num = last_nums.evaluate(&point);
     let mut claims_den = last_dens.evaluate(&point);
-    for i in 1..n_vars {
+    for i in N_VARS_TO_SEND_GKR_COEFFS..n_vars {
         (point, claims_num, claims_den) = verify_gkr_quotient_step(verifier_state, i, &point, claims_num, claims_den)?;
     }
     Ok((quotient, point, claims_num, claims_den))
@@ -291,6 +295,10 @@ where
     (new_numerators, new_denominators)
 }
 
+fn compute_quotient<EF: ExtensionField<PF<EF>>>(numerators: &[EF], denominators: &[EF]) -> EF {
+    numerators.iter().zip(denominators).map(|(&n, &d)| n / d).sum()
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;