faster open_structured_columns in AIR

Author: TomWambsgans

Cargo.lock

@@ -195,16 +195,25 @@ fn open_structured_columns<EF: ExtensionField<PF<EF>> + ExtensionField<IF>, IF:
 
     let batched_column =
         multilinears_linear_combination(witness, &poly_eq_batching_scalars[..n_columns]);
-    let mut batched_column_mixed = column_up(&batched_column);
-    add_multilinears_inplace(
-        &mut batched_column_mixed,
-        &scale_poly(&column_down(&batched_column), alpha),
-    );
+
+    let batched_column_mixed = info_span!("mixing up / down").in_scope(|| {
+        let mut batched_column_mixed = column_down(&batched_column);
+        add_multilinears_inplace(
+            &mut batched_column_mixed,
+            &scale_poly(&column_up(&batched_column), alpha),
+        );
+        batched_column_mixed
+    });
+
     // TODO do not recompute this (we can deduce it from already computed values)
-    let sub_evals = fold_multilinear_chunks(
-        &batched_column_mixed,
-        &MultilinearPoint(outer_sumcheck_challenge[1..log_n_rows - univariate_skips + 1].to_vec()),
-    );
+    let sub_evals = info_span!("fold_multilinear_chunks").in_scope(|| {
+        fold_multilinear_chunks(
+            &batched_column_mixed,
+            &MultilinearPoint(
+                outer_sumcheck_challenge[1..log_n_rows - univariate_skips + 1].to_vec(),
+            ),
+        )
+    });
     prover_state.add_extension_scalars(&sub_evals);
 
     let epsilons = prover_state.sample_vec(univariate_skips);
@@ -216,29 +225,39 @@ fn open_structured_columns<EF: ExtensionField<PF<EF>> + ExtensionField<IF>, IF:
     .concat();
 
     // TODO do not recompute this (we can deduce it from already computed values)
-    let inner_sum = batched_column_mixed.evaluate(&MultilinearPoint(point.clone()));
-
-    let mut mat_up = matrix_up_folded(&point);
-    add_multilinears_inplace(&mut mat_up, &scale_poly(&matrix_down_folded(&point), alpha));
-    let inner_mle = MleGroupOwned::Extension(vec![mat_up, batched_column]);
-
-    let (inner_challenges, _, _) = sumcheck_prove::<EF, _, _, _>(
-        1,
-        inner_mle,
-        &ProductComputation,
-        &ProductComputation,
-        &[],
-        None,
-        false,
-        prover_state,
-        inner_sum,
-        None,
-    );
+    let inner_sum = info_span!("mixed column eval")
+        .in_scope(|| batched_column_mixed.evaluate(&MultilinearPoint(point.clone())));
+
+    let mut mat_up = matrix_up_folded(&point, alpha);
+    matrix_down_folded(&point, &mut mat_up);
+    let inner_mle = info_span!("packing").in_scope(|| {
+        MleGroupOwned::ExtensionPacked(vec![
+            pack_extension(&mat_up),
+            pack_extension(&batched_column),
+        ])
+    });
 
-    let evaluations_remaining_to_prove = witness
-        .iter()
-        .map(|col| col.evaluate(&inner_challenges))
-        .collect::<Vec<_>>();
+    let (inner_challenges, _, _) = info_span!("structured columns sumcheck").in_scope(|| {
+        sumcheck_prove::<EF, _, _, _>(
+            1,
+            inner_mle,
+            &ProductComputation,
+            &ProductComputation,
+            &[],
+            None,
+            false,
+            prover_state,
+            inner_sum,
+            None,
+        )
+    });
+
+    let evaluations_remaining_to_prove = info_span!("final evals").in_scope(|| {
+        witness
+            .iter()
+            .map(|col| col.evaluate(&inner_challenges))
+            .collect::<Vec<_>>()
+    });
     prover_state.add_extension_scalars(&evaluations_remaining_to_prove);
 
     (inner_challenges, evaluations_remaining_to_prove)

crates/air/src/prove.rs

@@ -1,31 +1,29 @@
 use multilinear_toolkit::prelude::*;
-use p3_field::Field;
+use tracing::instrument;
 
-pub fn matrix_up_folded<F: Field>(outer_challenges: &[F]) -> Vec<F> {
+#[instrument(skip_all)]
+pub fn matrix_up_folded<F: ExtensionField<PF<F>>>(outer_challenges: &[F], alpha: F) -> Vec<F> {
     let n = outer_challenges.len();
-    let mut folded = eval_eq(outer_challenges);
+    let mut folded = eval_eq_scaled(outer_challenges, alpha);
     let outer_challenges_prod: F = outer_challenges.iter().copied().product();
-    folded[(1 << n) - 1] -= outer_challenges_prod;
-    folded[(1 << n) - 2] += outer_challenges_prod;
+    folded[(1 << n) - 1] -= outer_challenges_prod * alpha;
+    folded[(1 << n) - 2] += outer_challenges_prod * alpha;
     folded
 }
 
-pub fn matrix_down_folded<F: Field>(outer_challenges: &[F]) -> Vec<F> {
+#[instrument(skip_all)]
+pub fn matrix_down_folded<F: ExtensionField<PF<F>>>(outer_challenges: &[F], dest: &mut [F]) {
     let n = outer_challenges.len();
-    let mut folded = vec![F::ZERO; 1 << n];
     for k in 0..n {
         let outer_challenges_prod = (F::ONE - outer_challenges[n - k - 1])
             * outer_challenges[n - k..].iter().copied().product::<F>();
-        let mut eq_mle = eval_eq(&outer_challenges[0..n - k - 1]);
-        eq_mle = scale_poly(&eq_mle, outer_challenges_prod);
+        let mut eq_mle = eval_eq_scaled(&outer_challenges[0..n - k - 1], outer_challenges_prod);
         for (mut i, v) in eq_mle.iter_mut().enumerate() {
             i <<= k + 1;
             i += 1 << k;
-            folded[i] += *v;
+            dest[i] += *v;
         }
     }
     // bottom left corner:
-    folded[(1 << n) - 1] += outer_challenges.iter().copied().product::<F>();
-
-    folded
+    dest[(1 << n) - 1] += outer_challenges.iter().copied().product::<F>();
 }

crates/air/src/uni_skip_utils.rs

@@ -145,13 +145,14 @@ pub(crate) fn columns_up_and_down<F: Field>(columns: &[&[F]]) -> Vec<Vec<F>> {
 }
 
 pub(crate) fn column_up<F: Field>(column: &[F]) -> Vec<F> {
-    let mut up = column.to_vec();
+    let mut up = parallel_clone_vec(column);
     up[column.len() - 1] = up[column.len() - 2];
     up
 }
 
 pub(crate) fn column_down<F: Field>(column: &[F]) -> Vec<F> {
-    let mut down = column[1..].to_vec();
-    down.push(*down.last().unwrap());
+    let mut down = unsafe { uninitialized_vec(column.len()) };
+    parallel_clone(&column[1..], &mut down[..column.len() - 1]);
+    down[column.len() - 1] = down[column.len() - 2];
     down
 }

crates/air/src/utils.rs

@@ -177,10 +177,11 @@ fn verify_structured_columns<EF: ExtensionField<PF<EF>>>(
     ) != dot_product::<EF, _, _>(
         all_witness_up.iter().copied(),
         poly_eq_batching_scalars.iter().copied(),
-    ) + dot_product::<EF, _, _>(
-        all_witness_down.iter().copied(),
-        poly_eq_batching_scalars.iter().copied(),
     ) * alpha
+        + dot_product::<EF, _, _>(
+            all_witness_down.iter().copied(),
+            poly_eq_batching_scalars.iter().copied(),
+        )
     {
         return Err(ProofError::InvalidProof);
     }
@@ -202,7 +203,7 @@ fn verify_structured_columns<EF: ExtensionField<PF<EF>>>(
     let up = matrix_up_lde(&matrix_lde_point);
     let down = matrix_down_lde(&matrix_lde_point);
 
-    let final_value = inner_sumcheck_stement.value / (up + alpha * down);
+    let final_value = inner_sumcheck_stement.value / (up * alpha + down);
 
     let evaluations_remaining_to_verify = verifier_state.next_extension_scalars_vec(n_columns)?;
 

crates/air/src/verify.rs

@@ -17,6 +17,7 @@ tracing-forest.workspace = true
 p3-poseidon2.workspace = true
 p3-poseidon2-air.workspace = true
 tracing-subscriber.workspace = true
+tracing.workspace = true
 p3-util.workspace = true
 multilinear-toolkit.workspace = true
 

crates/utils/Cargo.toml

@@ -5,7 +5,9 @@ use p3_field::{ExtensionField, Field, dot_product};
 use p3_util::log2_strict_usize;
 
 use multilinear_toolkit::prelude::*;
+use tracing::instrument;
 
+#[instrument(skip_all)]
 pub fn multilinears_linear_combination<
     F: Field,
     EF: ExtensionField<F>,