Padding free AIR-sumcheck (#196)

TomWambsgans · web-flow · commit e5cd331270cc · 2026-04-17T10:02:43.000+02:00
* padding free sumcheck wip

* w

* w

* clippy

---------

Co-authored-by: Tom Wambsgans &lt;TomWambsgans@users.noreply.github.com&gt;
diff --git a/crates/backend/poly/src/evals.rs b/crates/backend/poly/src/evals.rs
@@ -346,6 +346,25 @@ where
     }
 }
 
+/// evaluate the MLE of [0, 0, ..., 0, 1, 1, ..., 1] at `alphas` where there are t zeros
+pub fn evaluate_mle_of_zero_then_ones<EF: ExtensionField<PF<EF>>>(t: usize, alphas: &[EF]) -> EF {
+    let n = alphas.len();
+    if t == 0 {
+        return EF::ONE;
+    }
+    if t >= 1usize << n {
+        return EF::ZERO;
+    }
+    let half = 1usize << (n - 1);
+    let alpha = alphas[0];
+    let sub = &alphas[1..];
+    if t < half {
+        (EF::ONE - alpha) * evaluate_mle_of_zero_then_ones(t, sub) + alpha
+    } else {
+        alpha * evaluate_mle_of_zero_then_ones(t - half, sub)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::time::Instant;
diff --git a/crates/backend/poly/src/mle/mle_group.rs b/crates/backend/poly/src/mle/mle_group.rs
@@ -75,4 +75,10 @@ impl<'a, EF: ExtensionField<PF<EF>>> MleGroup<'a, EF> {
             Self::Ref(r) => r.clone_to_owned(),
         }
     }
+
+    pub fn unpack_if_needed(&mut self) {
+        if self.is_packed() && must_unpack_multilinears::<EF>(self.n_vars()) {
+            *self = self.by_ref().unpack().as_owned_or_clone().into();
+        }
+    }
 }
diff --git a/crates/backend/poly/src/mle/mle_group_ref.rs b/crates/backend/poly/src/mle/mle_group_ref.rs
@@ -118,6 +118,24 @@ impl<'a, EF: ExtensionField<PF<EF>>> MleGroupRef<'a, EF> {
         }
     }
 
+    pub fn fold_at_bit(&self, alpha: EF, bit: usize) -> MleGroupOwned<EF> {
+        match self {
+            Self::Base(pols) => {
+                MleGroupOwned::Extension(batch_fold_multilinears_at_bit(pols, alpha, bit, |a, b| b * a))
+            }
+            Self::Extension(pols) => {
+                MleGroupOwned::Extension(batch_fold_multilinears_at_bit(pols, alpha, bit, |a, b| b * a))
+            }
+            Self::BasePacked(pols) => {
+                let alpha_packed = EFPacking::<EF>::from(alpha);
+                MleGroupOwned::ExtensionPacked(batch_fold_multilinears_at_bit(pols, alpha_packed, bit, |a, b| b * a))
+            }
+            Self::ExtensionPacked(pols) => {
+                MleGroupOwned::ExtensionPacked(batch_fold_multilinears_at_bit(pols, alpha, bit, |a, b| a * b))
+            }
+        }
+    }
+
     pub fn merge(mles: &'a [&'a MleRef<'a, EF>]) -> Self {
         match &mles[0] {
             MleRef::Base(_) => Self::Base(mles.iter().map(|m| m.as_base().unwrap()).collect()),
diff --git a/crates/backend/poly/src/utils.rs b/crates/backend/poly/src/utils.rs
@@ -89,6 +89,45 @@ pub fn batch_fold_multilinears<
     }
 }
 
+pub fn fold_multilinear_at_bit<
+    EF: PrimeCharacteristicRing + Copy + Send + Sync,
+    IF: Copy + Sub<Output = IF> + Send + Sync,
+    OF: Copy + Add<IF, Output = OF> + Send + Sync,
+    Mul: Fn(IF, EF) -> OF + Sync + Send,
+>(
+    m: &[IF],
+    alpha: EF,
+    bit: usize,
+    mul_if_of: &Mul,
+) -> Vec<OF> {
+    let new_size = m.len() / 2;
+    assert!(m.len() >= 2 * (1 << bit), "bit out of range for slice length");
+    let stride = 1usize << bit;
+    let lo_mask = stride - 1;
+    let mut res = unsafe { uninitialized_vec(new_size) };
+
+    let compute = |new_j: usize| {
+        let i_hi = new_j >> bit;
+        let i_lo = new_j & lo_mask;
+        let i0 = (i_hi << (bit + 1)) | i_lo;
+        let i1 = i0 | stride;
+        mul_if_of(m[i1] - m[i0], alpha) + m[i0]
+    };
+
+    if new_size < PARALLEL_THRESHOLD {
+        for (new_j, res_v) in res.iter_mut().enumerate() {
+            *res_v = compute(new_j);
+        }
+    } else {
+        (0..new_size)
+            .into_par_iter()
+            .with_min_len(PARALLEL_THRESHOLD)
+            .map(compute)
+            .collect_into_vec(&mut res);
+    }
+    res
+}
+
 pub fn fold_multilinear<
     EF: PrimeCharacteristicRing + Copy + Send + Sync,
     IF: Copy + Sub<Output = IF> + Send + Sync,
@@ -116,6 +155,31 @@ pub fn fold_multilinear<
     res
 }
 
+pub fn batch_fold_multilinears_at_bit<
+    EF: PrimeCharacteristicRing + Copy + Send + Sync,
+    IF: Copy + Sub<Output = IF> + Send + Sync,
+    OF: Copy + Add<IF, Output = OF> + Send + Sync,
+    F: Fn(IF, EF) -> OF + Sync + Send,
+>(
+    polys: &[&[IF]],
+    alpha: EF,
+    bit: usize,
+    mul_if_of: F,
+) -> Vec<Vec<OF>> {
+    let total_size: usize = polys.iter().map(|p| p.len()).sum();
+    if total_size < PARALLEL_THRESHOLD {
+        polys
+            .iter()
+            .map(|poly| fold_multilinear_at_bit(poly, alpha, bit, &mul_if_of))
+            .collect()
+    } else {
+        polys
+            .par_iter()
+            .map(|poly| fold_multilinear_at_bit(poly, alpha, bit, &mul_if_of))
+            .collect()
+    }
+}
+
 /// Returns a vector of uninitialized elements of type `A` with the specified length.
 /// # Safety
 /// Entries should be overwritten before use.
diff --git a/crates/backend/sumcheck/src/prove.rs b/crates/backend/sumcheck/src/prove.rs
@@ -115,9 +115,7 @@ where
 
     let mut challenges = Vec::new();
     for _ in 0..n_rounds {
-        if multilinears.by_ref().is_packed() && must_unpack_multilinears::<EF>(n_vars) {
-            multilinears = multilinears.by_ref().unpack().as_owned_or_clone().into();
-        }
+        multilinears.unpack_if_needed();
 
         let ps = compute_and_send_polynomial(
             &mut multilinears,
diff --git a/crates/lean_prover/src/prove_execution.rs b/crates/lean_prover/src/prove_execution.rs
@@ -175,9 +175,11 @@ pub fn prove_execution(
         up_down.extend(shifted_rows[idx].iter().map(Vec::as_slice));
         let packed = MleGroupRef::<EF>::Base(up_down).pack();
 
+        let non_padded = traces[table].non_padded_n_rows;
+
         macro_rules! make_session {
             ($t:expr) => {{
-                let session = AirSumcheckSession::new(packed, eq_suffix, bus_final_value, *$t, extra_data);
+                let session = AirSumcheckSession::new(packed, eq_suffix, bus_final_value, *$t, extra_data, non_padded);
                 Box::new(session) as Box<dyn OuterSumcheckSession<EF> + '_>
             }};
         }
@@ -187,12 +189,14 @@ pub fn prove_execution(
     let sumcheck_air_point = info_span!("batched AIR sumcheck")
         .in_scope(|| prove_batched_air_sumcheck(&mut prover_state, &mut sessions, air_eta));
 
-    for (idx, (table, log_n_rows)) in tables_sorted.iter().enumerate() {
+    for (idx, (table, _)) in tables_sorted.iter().enumerate() {
         let col_evals = sessions[idx].final_column_evals();
         prover_state.add_extension_scalars(&col_evals);
 
+        let natural_ordering_point =
+            natural_ordering_point_for_session(&sumcheck_air_point.0, traces[table].log_n_rows);
         macro_rules! split {
-            ($t:expr) => {{ columns_evals_up_and_down($t, &col_evals, &sumcheck_air_point.0, *log_n_rows) }};
+            ($t:expr) => {{ columns_evals_up_and_down($t, &col_evals, &natural_ordering_point) }};
         }
         let claim = delegate_to_inner!(table => split);
         committed_statements.get_mut(table).unwrap().push(claim);
diff --git a/crates/lean_prover/src/verify_execution.rs b/crates/lean_prover/src/verify_execution.rs
@@ -149,11 +149,17 @@ pub fn verify_execution(
         let constraint_eval = delegate_to_inner!(&vd.table => eval_constraint);
 
         let bus_point = from_end(gkr_point, table_n_vars[&vd.table]);
-        my_air_final_value +=
-            back_loaded_table_contribution(bus_point, &sumcheck_air_point.0, constraint_eval, vd.eta_power);
+        let natural_ordering_point = natural_ordering_point_for_session(&sumcheck_air_point.0, table_n_vars[&vd.table]);
+        my_air_final_value += back_loaded_table_contribution(
+            bus_point,
+            &sumcheck_air_point.0,
+            &natural_ordering_point,
+            constraint_eval,
+            vd.eta_power,
+        );
 
         macro_rules! split {
-            ($t:expr) => {{ columns_evals_up_and_down($t, &col_evals, &sumcheck_air_point.0, table_n_vars[&vd.table]) }};
+            ($t:expr) => {{ columns_evals_up_and_down($t, &col_evals, &natural_ordering_point) }};
         }
         let claim = delegate_to_inner!(&vd.table => split);
 
@@ -218,3 +224,20 @@ pub fn verify_execution(
         verifier_state.into_raw_proof(),
     ))
 }
+
+fn back_loaded_table_contribution<EF: ExtensionField<PF<EF>>>(
+    bus_point: &[EF],
+    sumcheck_air_point: &[EF],
+    natural_ordering_point: &[EF],
+    constraint_eval: EF,
+    eta_power: EF,
+) -> EF {
+    let n_t = bus_point.len();
+    let n_max = sumcheck_air_point.len();
+    let suffix_start = n_max - n_t;
+    assert_eq!(natural_ordering_point.len(), n_t);
+    let eq_val =
+        MultilinearPoint(bus_point.to_vec()).eq_poly_outside(&MultilinearPoint(natural_ordering_point.to_vec()));
+    let k_t: EF = sumcheck_air_point[..suffix_start].iter().copied().product();
+    eta_power * k_t * eq_val * constraint_eval
+}
diff --git a/crates/rec_aggregation/recursion.py b/crates/rec_aggregation/recursion.py
@@ -432,8 +432,8 @@ def continue_recursion_ordered(
 
         bus_point = pcs_points[table_index][0]
         suffix_start = n_max - log_n_rows
-        challenge_suffix = all_challenges + suffix_start * DIM
-        eq_val = eq_mle_extension(bus_point, challenge_suffix, log_n_rows)
+        natural_ordering_point = natural_ordering_point_for_session(all_challenges, suffix_start, log_n_rows)
+        eq_val = eq_mle_extension(bus_point, natural_ordering_point, log_n_rows)
 
         k_t = product_first_n(all_challenges, suffix_start)
 
@@ -443,7 +443,7 @@ def continue_recursion_ordered(
         )
         check_sum = add_extension_ret(check_sum, contribution)
 
-        pcs_points[table_index].push(challenge_suffix)
+        pcs_points[table_index].push(natural_ordering_point)
         pcs_values[table_index].push(DynArray([]))
         pcs_values_down[table_index].push(DynArray([]))
         last_index = len(pcs_values[table_index]) - 1
diff --git a/crates/rec_aggregation/utils.py b/crates/rec_aggregation/utils.py
@@ -71,6 +71,21 @@ def product_first_n(values, n):
     return res
 
 
+def natural_ordering_point_for_session(all_challenges, suffix_start, log_n_rows):
+    res = Array(log_n_rows * DIM)
+    match_range(log_n_rows, range(1, 33), lambda m: _natural_ordering_point_const(all_challenges, suffix_start, res, m))
+    return res
+
+
+def _natural_ordering_point_const(all_challenges, suffix_start, dst, m: Const):
+    half = div_ceil(m, 2)
+    for t in unroll(0, half):
+        copy_5(all_challenges + (suffix_start + t) * DIM, dst + (half - 1 - t) * DIM)
+    for t in unroll(half, m):
+        copy_5(all_challenges + (suffix_start + t) * DIM, dst + t * DIM)
+    return
+
+
 @inline
 def product_first_n_const(values, n):
     debug_assert(n != 0)
diff --git a/crates/sub_protocols/src/air_sumcheck.rs b/crates/sub_protocols/src/air_sumcheck.rs
diff --git a/crates/whir/src/open.rs b/crates/whir/src/open.rs

Original file line number	Diff line number	Diff line change
`@@ -75,4 +75,10 @@ impl<'a, EF: ExtensionField<PF<EF>>> MleGroup<'a, EF> {`
`75`	`75`	`Self::Ref(r) => r.clone_to_owned(),`
`76`	`76`	`}`
`77`	`77`	`}`
	`78`	`+`
	`79`	`+ pub fn unpack_if_needed(&mut self) {`
	`80`	`+ if self.is_packed() && must_unpack_multilinears::<EF>(self.n_vars()) {`
	`81`	`+ *self = self.by_ref().unpack().as_owned_or_clone().into();`
	`82`	`+ }`
	`83`	`+ }`
`78`	`84`	`}`