Merge branch 'reduce-padding-overhead' into main.

Mnual import of Plonky3 field, koala-bear, poseidon2 and symetric crates, and of whir-p3, for easier customization and R&D.

New optimization: remove padding (when there are zeros at the end of the polynomial committed in WHIR) overhead in DFT + in the hashing of the Merkle leaves (using a "right-to-left" t-sponge).

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

Author: TomWambsgans

Cargo.lock

@@ -8,7 +8,18 @@ version = "0.1.0"
 edition = "2024"
 
 [workspace]
-members = ["crates/*"]
+members = [
+    "crates/*",
+    "crates/backend/utils",
+    "crates/backend/field",
+    "crates/backend/koala-bear",
+    "crates/backend/poly",
+    "crates/backend/symetric",
+    "crates/backend/air",
+    "crates/backend/fiat-shamir",
+    "crates/backend/sumcheck",
+    "crates/whir",
+]
 
 [workspace.lints]
 rust.missing_debug_implementations = "warn"
@@ -37,6 +48,7 @@ should_panic_without_expect = "allow"
 similar_names = "allow"
 suboptimal_flops = "allow"
 cast_sign_loss = "allow"
+wildcard_imports = "allow"
 
 [workspace.dependencies]
 # Local
@@ -50,33 +62,21 @@ lean_prover = { path = "crates/lean_prover" }
 rec_aggregation = { path = "crates/rec_aggregation" }
 
 # External
-thiserror = "2.0"
-clap = { version = "4.5.52", features = ["derive"] }
+clap = { version = "4.5.59", features = ["derive"] }
 rand = "0.9.2"
-sha3 = "0.10.8"
+rayon = "1.11.0"
 pest = "2.7"
 pest_derive = "2.7"
 itertools = "0.14.0"
-colored = "3.0.0"
 tracing = "0.1.26"
 serde_json = "1.0.145"
-strum = { version = "0.27.2", features = ["derive"] }
 serde = { version = "1.0.228", features = ["derive"] }
-bincode = "1.3.3"
-lz4_flex = "0.12.0"
-num_enum = "0.7.5"
-owo-colors = "4.2.3"
 tracing-subscriber = { version = "0.3.19", features = ["std", "env-filter"] }
 tracing-forest = { version = "0.3.0", features = ["ansi", "smallvec"] }
-p3-koala-bear = { git = "https://github.com/TomWambsgans/Plonky3.git", branch = "lean-multisig" }
-p3-baby-bear = { git = "https://github.com/TomWambsgans/Plonky3.git", branch = "lean-multisig" }
-p3-poseidon2 = { git = "https://github.com/TomWambsgans/Plonky3.git", branch = "lean-multisig" }
-p3-symmetric = { git = "https://github.com/TomWambsgans/Plonky3.git", branch = "lean-multisig" }
-p3-goldilocks = { git = "https://github.com/TomWambsgans/Plonky3.git", branch = "lean-multisig" }
-p3-util = { git = "https://github.com/TomWambsgans/Plonky3.git", branch = "lean-multisig" }
-p3-monty-31 = { git = "https://github.com/TomWambsgans/Plonky3.git", branch = "lean-multisig" }
+postcard = { version = "1.1.3", features = ["alloc"] }
+lz4_flex = "0.12.0"
 
-multilinear-toolkit = { git = "https://github.com/leanEthereum/multilinear-toolkit.git" }
+backend = { path = "crates/backend" }
 
 [features]
 prox-gaps-conjecture = ["rec_aggregation/prox-gaps-conjecture"]
@@ -89,9 +89,9 @@ air.workspace = true
 rand.workspace = true
 sub_protocols.workspace = true
 utils.workspace = true
-p3-koala-bear.workspace = true
+
 lean_vm.workspace = true
-multilinear-toolkit.workspace = true
+backend.workspace = true
 
 
 [profile.release]

Cargo.toml

@@ -39,7 +39,6 @@ But we can get the bost of both worlds (suggested by Lev, TODO implement):
 - Double Check AIR constraints, logup overflows etc
 - Formal Verification
 - Padd with noop cycles to always ensure memory size >= bytecode size (liveness), and ensure this condition is checked by the verifier (soundness)
-- re-enable grinding at folding in whir (should be small, like 1 or 2 bits?)
 
 # Ideas
 

TODO.md

@@ -9,9 +9,9 @@ workspace = true
 [dependencies]
 tracing.workspace = true
 utils.workspace = true
-p3-util.workspace = true
-multilinear-toolkit.workspace = true
+
+backend.workspace = true
 
 [dev-dependencies]
-p3-koala-bear.workspace = true
+
 rand.workspace = true

crates/air/Cargo.toml

@@ -6,7 +6,7 @@ mod utils;
 mod validity_check;
 mod verify;
 
-use multilinear_toolkit::prelude::{Field, MultilinearPoint};
+use backend::{Field, MultilinearPoint};
 pub use prove::*;
 pub use validity_check::*;
 pub use verify::*;

crates/air/src/lib.rs

@@ -1,7 +1,6 @@
-use multilinear_toolkit::prelude::*;
-use p3_util::{log2_ceil_usize, log2_strict_usize};
+use backend::*;
 use tracing::{info_span, instrument};
-use utils::{fold_multilinear_chunks, multilinears_linear_combination};
+use utils::multilinears_linear_combination;
 
 use crate::{AirClaims, uni_skip_utils::matrix_next_mle_folded, utils::column_shifted};
 
@@ -17,7 +16,6 @@ pub fn prove_air<EF: ExtensionField<PF<EF>>, A: Air>(
     prover_state: &mut impl FSProver<EF>,
     air: &A,
     extra_data: A::ExtraData,
-    univariate_skips: usize,
     columns_f: &[impl AsRef<[PF<EF>]>],
     columns_ef: &[impl AsRef<[EF]>],
     virtual_column_statement: Option<Evaluation<EF>>, // point should be randomness generated after committing to the columns
@@ -32,22 +30,16 @@ where
     assert!(columns_f.iter().all(|col| col.len() == n_rows));
     assert!(columns_ef.iter().all(|col| col.len() == n_rows));
     let log_n_rows = log2_strict_usize(n_rows);
-    assert!(
-        univariate_skips < log_n_rows,
-        "TODO handle the case UNIVARIATE_SKIPS >= log_length"
-    );
 
     // 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 n_sc_rounds = log_n_rows + 1 - univariate_skips;
-
     let zerocheck_challenges = virtual_column_statement
         .as_ref()
         .map(|st| st.point.0.clone())
-        .unwrap_or_else(|| prover_state.sample_vec(n_sc_rounds));
-    assert_eq!(zerocheck_challenges.len(), n_sc_rounds);
+        .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()
@@ -74,7 +66,6 @@ where
 
     let (outer_sumcheck_challenge, inner_sums, _) = info_span!("zerocheck").in_scope(|| {
         sumcheck_prove(
-            univariate_skips,
             columns_up_down_group_f_packed,
             Some(columns_up_down_group_ef_packed),
             air,
@@ -92,121 +83,19 @@ where
 
     prover_state.add_extension_scalars(&inner_sums);
 
-    if univariate_skips == 1 {
-        open_columns_no_skip(
-            prover_state,
-            &inner_sums,
-            &air.down_column_indexes_f(),
-            &air.down_column_indexes_ef(),
-            &columns_f,
-            &columns_ef,
-            &outer_sumcheck_challenge,
-        )
-    } else if shifted_rows_f.is_empty() && shifted_rows_ef.is_empty() {
-        // usefull for poseidon2 benchmark
-        open_flat_columns(
-            prover_state,
-            univariate_skips,
-            &columns_f,
-            &columns_ef,
-            &outer_sumcheck_challenge,
-        )
-    } else {
-        panic!(
-            "Currently unsupported for simplicty (checkout c7944152a4325b1e1913446e6684112099db5d78 for a version that supported this case)"
-        );
-    }
-}
-
-#[instrument(skip_all)]
-fn open_flat_columns<EF: ExtensionField<PF<EF>>>(
-    prover_state: &mut impl FSProver<EF>,
-    univariate_skips: usize,
-    columns_f: &[&[PF<EF>]],
-    columns_ef: &[&[EF]],
-    outer_sumcheck_challenge: &[EF],
-) -> AirClaims<EF> {
-    let n_up_down_columns = columns_f.len() + columns_ef.len();
-    let batching_scalars = prover_state.sample_vec(log2_ceil_usize(n_up_down_columns));
-
-    let eval_eq_batching_scalars = eval_eq(&batching_scalars)[..n_up_down_columns].to_vec();
-
-    let mut batched_column_up =
-        multilinears_linear_combination(columns_f, &eval_eq_batching_scalars[..columns_f.len()]);
-    if !columns_ef.is_empty() {
-        let batched_column_up_ef = multilinears_linear_combination(
-            columns_ef,
-            &eval_eq_batching_scalars[columns_f.len()..][..columns_ef.len()],
-        );
-        batched_column_up
-            .par_iter_mut()
-            .zip(&batched_column_up_ef)
-            .for_each(|(a, &b)| {
-                *a += b;
-            });
-    }
-
-    let sub_evals = fold_multilinear_chunks(
-        &batched_column_up,
-        &MultilinearPoint(outer_sumcheck_challenge[1..].to_vec()),
-    );
-    prover_state.add_extension_scalars(&sub_evals);
-
-    let epsilons = prover_state.sample_vec(univariate_skips);
-
-    let inner_sum = sub_evals.evaluate(&MultilinearPoint(epsilons.clone()));
-
-    let point = [epsilons, outer_sumcheck_challenge[1..].to_vec()].concat();
-
-    // TODO opti in case of flat AIR (no need of `matrix_next_mle_folded`)
-    let matrix_up = eval_eq(&point);
-    let inner_mle = info_span!("packing").in_scope(|| {
-        MleGroupOwned::ExtensionPacked(vec![pack_extension(&matrix_up), pack_extension(&batched_column_up)])
-    });
-
-    let (inner_challenges, _, _) = info_span!("structured columns sumcheck").in_scope(|| {
-        sumcheck_prove::<EF, _, _>(
-            1,
-            inner_mle,
-            None,
-            &ProductComputation {},
-            &vec![],
-            None,
-            false,
-            prover_state,
-            inner_sum,
-            false,
-        )
-    });
-
-    let (evaluations_remaining_to_prove_f, evaluations_remaining_to_prove_ef) =
-        info_span!("final evals").in_scope(|| {
-            (
-                columns_f
-                    .par_iter()
-                    .map(|col| col.evaluate(&inner_challenges))
-                    .collect::<Vec<_>>(),
-                columns_ef
-                    .par_iter()
-                    .map(|col| col.evaluate(&inner_challenges))
-                    .collect::<Vec<_>>(),
-            )
-        });
-    prover_state.add_extension_scalars(&evaluations_remaining_to_prove_f);
-    prover_state.add_extension_scalars(&evaluations_remaining_to_prove_ef);
-
-    AirClaims {
-        point: inner_challenges,
-        evals_f: evaluations_remaining_to_prove_f,
-        evals_ef: evaluations_remaining_to_prove_ef,
-        down_point: None,
-        evals_f_on_down_columns: vec![],
-        evals_ef_on_down_columns: vec![],
-    }
+    open_columns(
+        prover_state,
+        &inner_sums,
+        &air.down_column_indexes_f(),
+        &air.down_column_indexes_ef(),
+        &columns_f,
+        &columns_ef,
+        &outer_sumcheck_challenge,
+    )
 }
 
 #[instrument(skip_all)]
-fn open_columns_no_skip<EF: ExtensionField<PF<EF>>>(
+fn open_columns<EF: ExtensionField<PF<EF>>>(
     prover_state: &mut impl FSProver<EF>,
     inner_evals: &[EF],
     columns_with_shift_f: &[usize],
@@ -270,7 +159,6 @@ fn open_columns_no_skip<EF: ExtensionField<PF<EF>>>(
 
     let (inner_challenges, _, _) = info_span!("structured columns sumcheck").in_scope(|| {
         sumcheck_prove::<EF, _, _>(
-            1,
             inner_mle,
             None,
             &ProductComputation {},

crates/air/src/prove.rs

@@ -1,4 +1,4 @@
-use multilinear_toolkit::prelude::*;
+use backend::*;
 use tracing::instrument;
 
 #[instrument(skip_all)]
@@ -27,7 +27,7 @@ mod tests {
     use crate::utils::next_mle;
 
     use super::*;
-    type F = p3_koala_bear::KoalaBear;
+    type F = KoalaBear;
 
     #[test]
     fn test_matrix_down_folded() {

crates/air/src/uni_skip_utils.rs

@@ -1,4 +1,4 @@
-use multilinear_toolkit::prelude::*;
+use backend::*;
 
 /// Returns a multilinear polynomial in 2n variables that evaluates to 1
 /// if and only if the second n-bit vector is equal to the first vector plus one (viewed as big-endian integers).

crates/air/src/utils.rs

@@ -1,6 +1,60 @@
-use multilinear_toolkit::prelude::*;
+use backend::*;
 use tracing::instrument;
-use utils::ConstraintChecker;
+
+#[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 constraint_index: usize,
+    pub errors: Vec<usize>,
+}
+
+impl<EF: ExtensionField<PF<EF>>> AirBuilder for ConstraintChecker<EF> {
+    type F = PF<EF>;
+    type EF = EF;
+
+    #[inline]
+    fn up_f(&self) -> &[Self::F] {
+        &self.up_f
+    }
+
+    #[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
+    }
+
+    #[inline]
+    fn assert_zero(&mut self, x: Self::F) {
+        if !x.is_zero() {
+            self.errors.push(self.constraint_index);
+        }
+        self.constraint_index += 1;
+    }
+
+    #[inline]
+    fn assert_zero_ef(&mut self, x: Self::EF) {
+        if !x.is_zero() {
+            self.errors.push(self.constraint_index);
+        }
+        self.constraint_index += 1;
+    }
+
+    fn eval_virtual_column(&mut self, _: Self::EF) {
+        // do nothing
+    }
+}
 
 #[instrument(name = "Check trace validity", skip_all)]
 pub fn check_air_validity<A: Air, EF: ExtensionField<PF<EF>>>(