put all AIR columns in the base field (columns in the extension fields are no more allowed) -> manually perform extension field operations via AIR constraints (trade 3% of perf for a nice simplification)

Co-authored-by: Thomas Coratger <thomas.coratger@gmail.com>

Author: TomWambsgans

crates/air/src/lib.rs

@@ -14,11 +14,9 @@ pub use verify::*;
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct AirClaims<EF: Field> {
     pub point: MultilinearPoint<EF>,
-    pub evals_f: Vec<EF>,
-    pub evals_ef: Vec<EF>,
+    pub evals: Vec<EF>,
 
     // only for columns with a "shift", in case univariate skip == 1
     pub down_point: Option<MultilinearPoint<EF>>,
-    pub evals_f_on_down_columns: Vec<EF>,
-    pub evals_ef_on_down_columns: Vec<EF>,
+    pub evals_on_down_columns: Vec<EF>,
 }

crates/air/src/prove.rs

@@ -16,23 +16,18 @@ pub fn prove_air<EF: ExtensionField<PF<EF>>, A: Air>(
     prover_state: &mut impl FSProver<EF>,
     air: &A,
     extra_data: A::ExtraData,
-    columns_f: &[impl AsRef<[PF<EF>]>],
-    columns_ef: &[impl AsRef<[EF]>],
+    columns: &[impl AsRef<[PF<EF>]>],
     virtual_column_statement: Option<Evaluation<EF>>, // point should be randomness generated after committing to the columns
     store_intermediate_foldings: bool,
 ) -> AirClaims<EF>
 where
     A::ExtraData: AlphaPowersMut<EF> + AlphaPowers<EF>,
 {
-    let columns_f: Vec<_> = columns_f.iter().map(|c| c.as_ref()).collect();
-    let columns_ef: Vec<_> = columns_ef.iter().map(|c| c.as_ref()).collect();
-    let n_rows = columns_f[0].len();
-    assert!(columns_f.iter().all(|col| col.len() == n_rows));
-    assert!(columns_ef.iter().all(|col| col.len() == n_rows));
+    let columns: Vec<_> = columns.iter().map(|c| c.as_ref()).collect();
+    let n_rows = columns[0].len();
+    assert!(columns.iter().all(|col| col.len() == n_rows));
     let log_n_rows = log2_strict_usize(n_rows);
 
-    // crate::check_air_validity(air, &extra_data, &columns_f, &columns_ef).unwrap();
-
     assert!(extra_data.alpha_powers().len() >= air.n_constraints() + virtual_column_statement.is_some() as usize);
 
     let zerocheck_challenges = virtual_column_statement
@@ -41,33 +36,22 @@ where
         .unwrap_or_else(|| prover_state.sample_vec(log_n_rows));
     assert_eq!(zerocheck_challenges.len(), log_n_rows);
 
-    let shifted_rows_f = air
-        .down_column_indexes_f()
-        .par_iter()
-        .map(|&col_index| column_shifted(columns_f[col_index]))
-        .collect::<Vec<_>>();
-    let shifted_rows_ef = air
-        .down_column_indexes_ef()
+    let shifted_rows = air
+        .down_column_indexes()
         .par_iter()
-        .map(|&col_index| column_shifted(columns_ef[col_index]))
+        .map(|&col_index| column_shifted(columns[col_index]))
         .collect::<Vec<_>>();
 
-    let mut columns_up_down_f = columns_f.to_vec(); // orginal columns, followed by shifted ones
-    columns_up_down_f.extend(shifted_rows_f.iter().map(Vec::as_slice));
+    let mut columns_up_down = columns.to_vec(); // orginal columns, followed by shifted ones
+    columns_up_down.extend(shifted_rows.iter().map(Vec::as_slice));
 
-    let mut columns_up_down_ef = columns_ef.to_vec(); // orginal columns, followed by shifted ones
-    columns_up_down_ef.extend(shifted_rows_ef.iter().map(Vec::as_slice));
+    let columns_up_down_group: MleGroupRef<'_, EF> = MleGroupRef::<'_, EF>::Base(columns_up_down);
 
-    let columns_up_down_group_f: MleGroupRef<'_, EF> = MleGroupRef::<'_, EF>::Base(columns_up_down_f);
-    let columns_up_down_group_ef: MleGroupRef<'_, EF> = MleGroupRef::<'_, EF>::Extension(columns_up_down_ef);
-
-    let columns_up_down_group_f_packed = columns_up_down_group_f.pack();
-    let columns_up_down_group_ef_packed = columns_up_down_group_ef.pack();
+    let columns_up_down_group_packed = columns_up_down_group.pack();
 
     let (outer_sumcheck_challenge, inner_sums, _) = info_span!("zerocheck").in_scope(|| {
         sumcheck_prove(
-            columns_up_down_group_f_packed,
-            Some(columns_up_down_group_ef_packed),
+            columns_up_down_group_packed,
             air,
             &extra_data,
             Some((zerocheck_challenges, None)),
@@ -86,10 +70,8 @@ where
     open_columns(
         prover_state,
         &inner_sums,
-        &air.down_column_indexes_f(),
-        &air.down_column_indexes_ef(),
-        &columns_f,
-        &columns_ef,
+        &air.down_column_indexes(),
+        &columns,
         &outer_sumcheck_challenge,
     )
 }
@@ -98,69 +80,43 @@ where
 fn open_columns<EF: ExtensionField<PF<EF>>>(
     prover_state: &mut impl FSProver<EF>,
     inner_evals: &[EF],
-    columns_with_shift_f: &[usize],
-    columns_with_shift_ef: &[usize],
-    columns_f: &[&[PF<EF>]],
-    columns_ef: &[&[EF]],
+    columns_with_shift: &[usize],
+    columns: &[&[PF<EF>]],
     outer_sumcheck_challenge: &[EF],
 ) -> AirClaims<EF> {
-    let n_columns_f_up = columns_f.len();
-    let n_columns_ef_up = columns_ef.len();
-    let n_columns_f_down = columns_with_shift_f.len();
-    let n_columns_ef_down = columns_with_shift_ef.len();
-    let n_down_columns = n_columns_f_down + n_columns_ef_down;
-    assert_eq!(inner_evals.len(), n_columns_f_up + n_columns_ef_up + n_down_columns);
-
-    let evals_up_f = inner_evals[..n_columns_f_up].to_vec();
-    let evals_down_f = &inner_evals[n_columns_f_up..][..n_columns_f_down];
-    let evals_up_ef = inner_evals[n_columns_f_up + n_columns_f_down..][..n_columns_ef_up].to_vec();
-    let evals_down_ef = &inner_evals[n_columns_f_up + n_columns_f_down + n_columns_ef_up..];
-
-    if n_down_columns == 0 {
+    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_f: evals_up_f,
-            evals_ef: evals_up_ef,
+            evals: evals_up,
             down_point: None,
-            evals_f_on_down_columns: vec![],
-            evals_ef_on_down_columns: vec![],
+            evals_on_down_columns: vec![],
         };
     }
 
     let batching_scalar = prover_state.sample();
-    let batching_scalar_powers = batching_scalar.powers().collect_n(n_down_columns);
-
-    let columns_shifted_f = &columns_with_shift_f.iter().map(|&i| columns_f[i]).collect::<Vec<_>>();
-    let columns_shifted_ef = &columns_with_shift_ef.iter().map(|&i| columns_ef[i]).collect::<Vec<_>>();
-
-    let mut batched_column_down =
-        multilinears_linear_combination(columns_shifted_f, &batching_scalar_powers[..n_columns_f_down]);
-
-    if n_columns_ef_down > 0 {
-        let batched_column_down_ef =
-            multilinears_linear_combination(columns_shifted_ef, &batching_scalar_powers[n_columns_f_down..]);
-        batched_column_down
-            .par_iter_mut()
-            .zip(&batched_column_down_ef)
-            .for_each(|(a, &b)| {
-                *a += b;
-            });
-    }
+    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_f.iter().chain(evals_down_ef).copied(),
-        batching_scalar_powers.iter().copied(),
-    );
+    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,
-            None,
             &ProductComputation {},
             &vec![],
             None,
@@ -171,27 +127,18 @@ fn open_columns<EF: ExtensionField<PF<EF>>>(
         )
     });
 
-    let (evals_f_on_down_columns, evals_ef_on_down_columns) = info_span!("final evals").in_scope(|| {
-        (
-            columns_shifted_f
-                .par_iter()
-                .map(|col| col.evaluate(&inner_challenges))
-                .collect::<Vec<_>>(),
-            columns_shifted_ef
-                .par_iter()
-                .map(|col| col.evaluate(&inner_challenges))
-                .collect::<Vec<_>>(),
-        )
+    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_f_on_down_columns);
-    prover_state.add_extension_scalars(&evals_ef_on_down_columns);
+    prover_state.add_extension_scalars(&evals_on_down_columns);
 
     AirClaims {
         point: MultilinearPoint(outer_sumcheck_challenge.to_vec()),
-        evals_f: evals_up_f,
-        evals_ef: evals_up_ef,
+        evals: evals_up,
         down_point: Some(inner_challenges),
-        evals_f_on_down_columns,
-        evals_ef_on_down_columns,
+        evals_on_down_columns,
     }
 }

crates/air/src/validity_check.rs

@@ -3,10 +3,8 @@ use tracing::instrument;
 
 #[derive(Debug)]
 pub struct ConstraintChecker<EF: ExtensionField<PF<EF>>> {
-    pub up_f: Vec<PF<EF>>,
-    pub up_ef: Vec<EF>,
-    pub down_f: Vec<PF<EF>>,
-    pub down_ef: Vec<EF>,
+    pub up: Vec<PF<EF>>,
+    pub down: Vec<PF<EF>>,
     pub constraint_index: usize,
     pub errors: Vec<usize>,
 }
@@ -16,23 +14,13 @@ impl<EF: ExtensionField<PF<EF>>> AirBuilder for ConstraintChecker<EF> {
     type EF = EF;
 
     #[inline]
-    fn up_f(&self) -> &[Self::F] {
-        &self.up_f
+    fn up(&self) -> &[Self::F] {
+        &self.up
     }
 
     #[inline]
-    fn up_ef(&self) -> &[Self::EF] {
-        &self.up_ef
-    }
-
-    #[inline]
-    fn down_f(&self) -> &[Self::F] {
-        &self.down_f
-    }
-
-    #[inline]
-    fn down_ef(&self) -> &[Self::EF] {
-        &self.down_ef
+    fn down(&self) -> &[Self::F] {
+        &self.down
     }
 
     #[inline]
@@ -60,13 +48,11 @@ impl<EF: ExtensionField<PF<EF>>> AirBuilder for ConstraintChecker<EF> {
 pub fn check_air_validity<A: Air, EF: ExtensionField<PF<EF>>>(
     air: &A,
     extra_data: &A::ExtraData,
-    columns_f: &[&[PF<EF>]],
-    columns_ef: &[&[EF]],
+    columns: &[&[PF<EF>]],
 ) -> Result<(), String> {
-    let n_rows = columns_f[0].len();
-    assert!(columns_f.iter().all(|col| col.len() == n_rows));
-    assert!(columns_ef.iter().all(|col| col.len() == n_rows));
-    if columns_f.len() != air.n_columns_f_air() || columns_ef.len() != air.n_columns_ef_air() {
+    let n_rows = columns[0].len();
+    assert!(columns.iter().all(|col| col.len() == n_rows));
+    if columns.len() != air.n_columns() {
         return Err("Invalid number of columns".to_string());
     }
     let handle_errors = |row: usize, constraint_checker: &ConstraintChecker<EF>| {
@@ -85,52 +71,29 @@ pub fn check_air_validity<A: Air, EF: ExtensionField<PF<EF>>>(
         Ok(())
     };
     for row in 0..n_rows - 1 {
-        let up_f = (0..air.n_columns_f_air())
-            .map(|j| columns_f[j][row])
-            .collect::<Vec<_>>();
-        let up_ef = (0..air.n_columns_ef_air())
-            .map(|j| columns_ef[j][row])
-            .collect::<Vec<_>>();
-        let down_f = air
-            .down_column_indexes_f()
-            .iter()
-            .map(|j| columns_f[*j][row + 1])
-            .collect::<Vec<_>>();
-        let down_ef = air
-            .down_column_indexes_ef()
+        let up = (0..air.n_columns()).map(|j| columns[j][row]).collect::<Vec<_>>();
+        let down = air
+            .down_column_indexes()
             .iter()
-            .map(|j| columns_ef[*j][row + 1])
+            .map(|j| columns[*j][row + 1])
             .collect::<Vec<_>>();
         let mut constraints_checker = ConstraintChecker {
-            up_f,
-            up_ef,
-            down_f,
-            down_ef,
+            up,
+            down,
             constraint_index: 0,
             errors: Vec::new(),
         };
         air.eval(&mut constraints_checker, extra_data);
         handle_errors(row, &constraints_checker)?;
     }
     // last transition:
-    let up_f = (0..air.n_columns_f_air())
-        .map(|j| columns_f[j][n_rows - 1])
-        .collect::<Vec<_>>();
-    let up_ef = (0..air.n_columns_ef_air())
-        .map(|j| columns_ef[j][n_rows - 1])
-        .collect::<Vec<_>>();
+    let up = (0..air.n_columns()).map(|j| columns[j][n_rows - 1]).collect::<Vec<_>>();
     let mut constraints_checker = ConstraintChecker {
-        up_f,
-        up_ef,
-        down_f: air
-            .down_column_indexes_f()
-            .iter()
-            .map(|j| columns_f[*j][n_rows - 1])
-            .collect::<Vec<_>>(),
-        down_ef: air
-            .down_column_indexes_ef()
+        up,
+        down: air
+            .down_column_indexes()
             .iter()
-            .map(|j| columns_ef[*j][n_rows - 1])
+            .map(|j| columns[*j][n_rows - 1])
             .collect::<Vec<_>>(),
         constraint_index: 0,
         errors: Vec::new(),