replace the dedicated sumcheck that evaluates AIR shifted (down) columns, by "next" weights in WHIR (-> removes 2 sumchecks overall)

Author: TomWambsgans

Cargo.lock

@@ -8,10 +8,7 @@ workspace = true
 
 [dependencies]
 tracing.workspace = true
-utils.workspace = true
-
 backend.workspace = true
 
 [dev-dependencies]
-
 rand.workspace = true

crates/air/Cargo.toml

@@ -1,8 +1,6 @@
 #![cfg_attr(not(test), warn(unused_crate_dependencies))]
 
 mod prove;
-mod uni_skip_utils;
-mod utils;
 mod validity_check;
 mod verify;
 
@@ -15,8 +13,5 @@ pub use verify::*;
 pub struct AirClaims<EF: Field> {
     pub point: MultilinearPoint<EF>,
     pub evals: Vec<EF>,
-
-    // only for columns with a "shift", in case univariate skip == 1
-    pub down_point: Option<MultilinearPoint<EF>>,
-    pub evals_on_down_columns: Vec<EF>,
+    pub evals_down: Vec<EF>, // columns 'shifted down' by one row
 }

crates/air/src/lib.rs

@@ -1,8 +1,7 @@
 use backend::*;
 use tracing::{info_span, instrument};
-use utils::multilinears_linear_combination;
 
-use crate::{AirClaims, uni_skip_utils::matrix_next_mle_folded, utils::column_shifted};
+use crate::AirClaims;
 
 /*
 
@@ -66,77 +65,16 @@ where
 
     prover_state.add_extension_scalars(&inner_sums);
 
-    open_columns(
-        prover_state,
-        &inner_sums,
-        &air.down_column_indexes(),
-        &columns,
-        &outer_sumcheck_challenge,
-    )
-}
-
-#[instrument(skip_all)]
-fn open_columns<EF: ExtensionField<PF<EF>>>(
-    prover_state: &mut impl FSProver<EF>,
-    inner_evals: &[EF],
-    columns_with_shift: &[usize],
-    columns: &[&[PF<EF>]],
-    outer_sumcheck_challenge: &[EF],
-) -> AirClaims<EF> {
-    let n_columns_up = columns.len();
-    let n_columns_down = columns_with_shift.len();
-    assert_eq!(inner_evals.len(), n_columns_up + n_columns_down);
-
-    let evals_up = inner_evals[..n_columns_up].to_vec();
-    let evals_down = &inner_evals[n_columns_up..];
-
-    if n_columns_down == 0 {
-        return AirClaims {
-            point: MultilinearPoint(outer_sumcheck_challenge.to_vec()),
-            evals: evals_up,
-            down_point: None,
-            evals_on_down_columns: vec![],
-        };
-    }
-
-    let batching_scalar = prover_state.sample();
-    let batching_scalar_powers = batching_scalar.powers().collect_n(n_columns_down);
-
-    let columns_shifted = &columns_with_shift.iter().map(|&i| columns[i]).collect::<Vec<_>>();
-
-    let batched_column_down = multilinears_linear_combination(columns_shifted, &batching_scalar_powers);
-
-    let matrix_down = matrix_next_mle_folded(outer_sumcheck_challenge);
-    let inner_mle = info_span!("packing").in_scope(|| {
-        MleGroupOwned::ExtensionPacked(vec![pack_extension(&matrix_down), pack_extension(&batched_column_down)])
-    });
-
-    let inner_sum = dot_product(evals_down.iter().copied(), batching_scalar_powers.iter().copied());
-
-    let (inner_challenges, _, _) = info_span!("structured columns sumcheck").in_scope(|| {
-        sumcheck_prove::<EF, _, _>(
-            inner_mle,
-            &ProductComputation {},
-            &vec![],
-            None,
-            prover_state,
-            inner_sum,
-            false,
-        )
-    });
-
-    let evals_on_down_columns = info_span!("final evals").in_scope(|| {
-        columns_shifted
-            .par_iter()
-            .map(|col| col.evaluate(&inner_challenges))
-            .collect::<Vec<_>>()
-    });
-    prover_state.add_extension_scalars(&evals_on_down_columns);
-
     AirClaims {
         point: MultilinearPoint(outer_sumcheck_challenge.to_vec()),
-        evals: evals_up,
-        down_point: Some(inner_challenges),
-        evals_on_down_columns,
+        evals: inner_sums[..columns.len()].to_vec(),
+        evals_down: inner_sums[columns.len()..].to_vec(),
     }
 }
+
+pub(crate) fn column_shifted<F: Field>(column: &[F]) -> Vec<F> {
+    let mut down = unsafe { uninitialized_vec(column.len()) };
+    parallel_clone(&column[1..], &mut down[..column.len() - 1]);
+    down[column.len() - 1] = column[column.len() - 1];
+    down
+}

crates/air/src/prove.rs

@@ -1,6 +1,6 @@
 use backend::*;
 
-use crate::{AirClaims, utils::next_mle};
+use crate::AirClaims;
 
 #[allow(clippy::type_complexity)]
 #[allow(clippy::too_many_arguments)]
@@ -42,55 +42,13 @@ where
         return Err(ProofError::InvalidProof);
     }
 
-    open_columns(verifier_state, air, log_n_rows, &inner_evals, &outer_statement.point)
-}
-
-fn open_columns<A: Air, EF: ExtensionField<PF<EF>>>(
-    verifier_state: &mut impl FSVerifier<EF>,
-    air: &A,
-    log_n_rows: usize,
-    inner_evals: &[EF],
-    outer_sumcheck_challenge: &[EF],
-) -> ProofResult<AirClaims<EF>> {
     let n_columns_up = air.n_columns();
     let n_columns_down = air.n_down_columns();
     assert_eq!(inner_evals.len(), n_columns_up + n_columns_down);
 
-    let evals_up = inner_evals[..n_columns_up].to_vec();
-    let evals_down = inner_evals[n_columns_up..].to_vec();
-
-    if n_columns_down == 0 {
-        return Ok(AirClaims {
-            point: MultilinearPoint(outer_sumcheck_challenge.to_vec()),
-            evals: evals_up,
-            down_point: None,
-            evals_on_down_columns: vec![],
-        });
-    }
-
-    let batching_scalar = verifier_state.sample();
-    let batching_scalar_powers = batching_scalar.powers().collect_n(n_columns_down);
-
-    let inner_sum: EF = dot_product(evals_down.into_iter(), batching_scalar_powers.iter().copied());
-
-    let inner_sumcheck_stement = sumcheck_verify(verifier_state, log_n_rows, 2, inner_sum, None)?;
-
-    let matrix_down_sc_eval = next_mle(outer_sumcheck_challenge, &inner_sumcheck_stement.point);
-
-    let evals_on_down_columns = verifier_state.next_extension_scalars_vec(n_columns_down)?;
-    let batched_col_down_sc_eval = dot_product::<EF, _, _>(
-        batching_scalar_powers.iter().copied(),
-        evals_on_down_columns.iter().copied(),
-    );
-
-    if inner_sumcheck_stement.value != matrix_down_sc_eval * batched_col_down_sc_eval {
-        return Err(ProofError::InvalidProof);
-    }
-
     Ok(AirClaims {
-        point: MultilinearPoint(outer_sumcheck_challenge.to_vec()),
-        evals: evals_up,
-        down_point: Some(inner_sumcheck_stement.point.clone()),
-        evals_on_down_columns,
+        point: outer_statement.point,
+        evals: inner_evals[..n_columns_up].to_vec(),
+        evals_down: inner_evals[n_columns_up..].to_vec(),
     })
 }

crates/air/src/uni_skip_utils.rs

@@ -15,6 +15,9 @@ pub use utils::*;
 mod eq_mle;
 pub use eq_mle::*;
 
+mod next_mle;
+pub use next_mle::*;
+
 mod evals;
 pub use evals::*;