Some notes while investigating '[wip] batch STARKs'

Cargo.toml

@@ -49,7 +49,16 @@ lto = "fat"
 lto = "thin"
 opt-level = 3
 
+[workspace.dependencies]
+plonky2 = { git = "https://github.com/0xmozak/plonky2.git", branch = "matthias/sai/add_batch_fri", default-features = false }
+plonky2_maybe_rayon = { git = "https://github.com/0xmozak/plonky2.git", branch = "matthias/sai/add_batch_fri", default-features = false }
+starky = { git = "https://github.com/0xmozak/plonky2.git", branch = "matthias/sai/add_batch_fri", default-features = false }
+
+plonky2_crypto = { git = "https://github.com/0xmozak/plonky2-crypto.git" }
+
 [patch.crates-io]
-plonky2 = { git = "https://github.com/0xmozak/plonky2.git" }
-plonky2_maybe_rayon = { git = "https://github.com/0xmozak/plonky2.git" }
-starky = { git = "https://github.com/0xmozak/plonky2.git" }
+unroll = { git = "https://github.com/0xmozak/unroll.git" }
+
+plonky2 = { git = "https://github.com/0xmozak/plonky2.git", branch = "matthias/sai/add_batch_fri" }
+plonky2_maybe_rayon = { git = "https://github.com/0xmozak/plonky2.git", branch = "matthias/sai/add_batch_fri" }
+starky = { git = "https://github.com/0xmozak/plonky2.git", branch = "matthias/sai/add_batch_fri" }

circuits/Cargo.toml

@@ -20,11 +20,11 @@ log = "0.4"
 mozak-circuits-derive = { path = "./derive" }
 mozak-runner = { path = "../runner" }
 mozak-sdk = { path = "../sdk" }
-plonky2 = { version = "0", default-features = false }
-plonky2_maybe_rayon = { version = "0", default-features = false }
+plonky2 = { workspace = true, default-features = false }
+plonky2_maybe_rayon = { workspace = true, default-features = false }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
-starky = { version = "0", default-features = false, features = ["std"] }
+starky = { workspace = true, default-features = false, features = ["std"] }
 thiserror = "1.0"
 tt-call = "1.0"
 

circuits/src/stark/proof.rs

@@ -325,6 +325,7 @@ pub struct AllProof<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, co
     pub mozak_memory_init_trace_cap: MerkleCap<F, C::Hasher>,
     pub public_inputs: PublicInputs<F>,
     pub public_sub_table_values: TableKindArray<Vec<PublicSubTableValues<F>>>,
+    // pub batch_fri_proof: FriProof<F, C::Hasher, D>,
 }
 
 pub(crate) struct AllProofChallenges<F: RichField + Extendable<D>, const D: usize> {

circuits/src/stark/prover.rs

@@ -1,5 +1,7 @@
 #![allow(clippy::too_many_lines)]
 
+use std::cmp::Reverse;
+use std::collections::BTreeSet;
 use std::fmt::Display;
 
 use anyhow::{ensure, Result};
@@ -12,6 +14,7 @@ use plonky2::field::extension::Extendable;
 use plonky2::field::packable::Packable;
 use plonky2::field::polynomial::PolynomialValues;
 use plonky2::field::types::Field;
+use plonky2::fri::batch_oracle::BatchFriOracle;
 use plonky2::fri::oracle::PolynomialBatch;
 use plonky2::hash::hash_types::RichField;
 use plonky2::iop::challenger::Challenger;
@@ -30,7 +33,7 @@ use crate::cross_table_lookup::ctl_utils::debug_ctl;
 use crate::cross_table_lookup::{cross_table_lookup_data, CtlData};
 use crate::generation::{debug_traces, generate_traces};
 use crate::public_sub_table::public_sub_table_data_and_values;
-use crate::stark::mozak_stark::{all_starks, PublicInputs};
+use crate::stark::mozak_stark::{all_kind, all_starks, PublicInputs};
 use crate::stark::permutation::challenge::GrandProductChallengeTrait;
 use crate::stark::poly::compute_quotient_polys;
 
@@ -86,6 +89,62 @@ where
     let rate_bits = config.fri_config.rate_bits;
     let cap_height = config.fri_config.cap_height;
 
+    // We cannot batch prove these tables because trace caps are needed as public
+    // inputs for the following tables.
+    let public_table_kinds = [
+        TableKind::Program,
+        TableKind::ElfMemoryInit,
+        TableKind::MozakMemoryInit,
+    ];
+
+    let _separate_trace_commitments = timed!(
+        timing,
+        "Compute trace commitments for separate tables",
+        all_kind!(|kind| public_table_kinds.contains(&kind).then(|| {
+            PolynomialBatch::<F, C, D>::from_values(
+                traces_poly_values[kind].clone(),
+                rate_bits,
+                false,
+                cap_height,
+                timing,
+                None,
+            )
+        }))
+    );
+
+    let batch_traces_poly_values = all_kind!(|kind| public_table_kinds
+        .contains(&kind)
+        .then_some(&traces_poly_values[kind]));
+    let _degree_logs: Vec<usize> = batch_traces_poly_values
+        .iter()
+        .filter_map(|p| p.map(Vec::len))
+        .collect::<BTreeSet<_>>()
+        .into_iter()
+        .rev()
+        .collect();
+
+    let batch_trace_polys: Vec<_> = batch_traces_poly_values
+        .iter()
+        .filter_map(|t| *t)
+        .flatten()
+        .cloned()
+        .sorted_by_key(|p| Reverse(p.len()))
+        .collect();
+    let bacth_trace_polys_len = batch_trace_polys.len();
+
+    let _batch_trace_commitments: BatchFriOracle<F, C, D> = timed!(
+        timing,
+        "Compute trace commitments for batch tables",
+        BatchFriOracle::from_values(
+            batch_trace_polys,
+            rate_bits,
+            false,
+            cap_height,
+            timing,
+            &vec![None; bacth_trace_polys_len],
+        )
+    );
+
     let trace_commitments = timed!(
         timing,
         "Compute trace commitments for each table",
@@ -108,6 +167,9 @@ where
             })
     );
 
+    // TODO: todo
+    // let trace_caps = batch_trace_commitments.field_merkle_tree.cap;
+
     let trace_caps = trace_commitments
         .each_ref()
         .map(|c| c.merkle_tree.cap.clone());
@@ -373,6 +435,167 @@ where
     }))
 }
 
+/// Given the traces generated from [`generate_traces`] along with their
+/// commitments, prove a [`MozakStark`].
+///
+/// # Errors
+/// Errors if proving fails.
+#[allow(clippy::too_many_arguments)]
+#[allow(clippy::similar_names)]
+pub fn batch_prove_with_commitments<F, C, const D: usize>(
+    mozak_stark: &MozakStark<F, D>,
+    config: &StarkConfig,
+    public_table_kinds: &[TableKind],
+    public_inputs: &PublicInputs<F>,
+    trace_commitments: &TableKindArray<PolynomialBatch<F, C, D>>,
+    traces_poly_values: &TableKindArray<Vec<PolynomialValues<F>>>,
+    _batch_trace_commitments: &BatchFriOracle<F, C, D>,
+    separate_trace_commitments: &TableKindArray<PolynomialBatch<F, C, D>>,
+    ctl_data_per_table: &TableKindArray<CtlData<F>>,
+    public_sub_data_per_table: &TableKindArray<CtlData<F>>,
+    challenger: &mut Challenger<F, C::Hasher>,
+    timing: &mut TimingTree,
+) -> Result<TableKindArray<StarkProof<F, C, D>>>
+where
+    F: RichField + Extendable<D>,
+    C: GenericConfig<D, F = F>, {
+    let rate_bits = config.fri_config.rate_bits;
+    let cap_height = config.fri_config.cap_height;
+
+    let cpu_stark = [public_inputs.entry_point];
+    let public_inputs = TableKindSetBuilder::<&[_]> {
+        cpu_stark: &cpu_stark,
+        ..Default::default()
+    }
+    .build();
+
+    let _separate_proofs = all_starks!(mozak_stark, |stark, kind| if public_table_kinds
+        .contains(&kind)
+    {
+        Some(prove_single_table(
+            stark,
+            config,
+            &traces_poly_values[kind],
+            &separate_trace_commitments[kind],
+            public_inputs[kind],
+            &ctl_data_per_table[kind],
+            &public_sub_data_per_table[kind],
+            challenger,
+            timing,
+        )?)
+    } else {
+        None
+    });
+
+    let batch_ctl_z_polys = all_kind!(|kind| {
+        if public_table_kinds.contains(&kind) {
+            None
+        } else {
+            let degree = traces_poly_values[kind][0].len();
+            let degree_bits = log2_strict(degree);
+            let fri_params = config.fri_params(degree_bits);
+            assert!(
+                fri_params.total_arities() <= degree_bits + rate_bits - cap_height,
+                "FRI total reduction arity is too large.",
+            );
+
+            let z_poly_public_sub_table = ctl_data_per_table[kind].z_polys();
+
+            let z_polys = vec![ctl_data_per_table[kind].z_polys(), z_poly_public_sub_table]
+                .into_iter()
+                .flatten()
+                .collect_vec();
+
+            // TODO(Matthias): make the code work with empty z_polys, too.
+            assert!(!z_polys.is_empty());
+
+            Some(z_polys)
+        }
+    });
+
+    // TODO: we can remove duplicates in the ctl polynomials.
+    let batch_ctl_zs_polys: Vec<_> = batch_ctl_z_polys
+        .iter()
+        .filter_map(|t| t.as_ref())
+        .flat_map(|v| v.iter().cloned())
+        .sorted_by_key(|b| Reverse(b.len()))
+        .collect();
+    let batch_ctl_zs_polys_len = batch_ctl_zs_polys.len();
+
+    let batch_ctl_zs_commitments: BatchFriOracle<F, C, D> = timed!(
+        timing,
+        "compute batch Zs commitment",
+        BatchFriOracle::from_values(
+            batch_ctl_zs_polys,
+            rate_bits,
+            false,
+            config.fri_config.cap_height,
+            timing,
+            &vec![None; batch_ctl_zs_polys_len],
+        )
+    );
+
+    let ctl_zs_commitments = all_starks!(mozak_stark, |stark, kind| timed!(
+        timing,
+        format!("{stark}: compute Zs commitment").as_str(),
+        batch_ctl_z_polys[kind]
+            .as_ref()
+            .map(|poly| PolynomialBatch::<F, C, D>::from_values(
+                poly.clone(),
+                rate_bits,
+                false,
+                config.fri_config.cap_height,
+                timing,
+                None,
+            ))
+    ));
+
+    let ctl_zs_cap = batch_ctl_zs_commitments.field_merkle_tree.cap.clone();
+    challenger.observe_cap(&ctl_zs_cap);
+
+    let alphas = challenger.get_n_challenges(config.num_challenges);
+
+    let _batch_quotient_polys = all_starks!(mozak_stark, |stark, kind| {
+        if let Some(ctl_zs_commitment) = ctl_zs_commitments[kind].as_ref() {
+            let degree = traces_poly_values[kind][0].len();
+            let degree_bits = log2_strict(degree);
+            timed!(
+                timing,
+                format!("{stark}: compute quotient polynomial").as_str(),
+                Some(
+                    compute_quotient_polys::<F, <F as Packable>::Packing, C, _, D>(
+                        stark,
+                        &trace_commitments[kind],
+                        ctl_zs_commitment,
+                        public_inputs[kind],
+                        &ctl_data_per_table[kind],
+                        &public_sub_data_per_table[kind],
+                        &alphas,
+                        degree_bits,
+                        config,
+                    )
+                )
+            )
+        } else {
+            None
+        }
+    });
+
+    Ok(all_starks!(mozak_stark, |stark, kind| {
+        prove_single_table(
+            stark,
+            config,
+            &traces_poly_values[kind],
+            &trace_commitments[kind],
+            public_inputs[kind],
+            &ctl_data_per_table[kind],
+            &public_sub_data_per_table[kind],
+            challenger,
+            timing,
+        )?
+    }))
+}
+
 #[cfg(test)]
 mod tests {
 

cli/Cargo.toml

@@ -30,11 +30,11 @@ env_logger = "0.11"
 itertools = "0.12"
 log = "0.4"
 mozak-examples = { path = "../examples-builder", features = ["mozak-sort"] }
-plonky2 = { version = "0", default-features = false }
+plonky2 = { workspace = true, default-features = false }
 rkyv = { version = "=0.8.0-alpha.1", default-features = false, features = ["pointer_width_32", "alloc"] }
 rkyv_derive = "=0.8.0-alpha.1"
 serde_json = "1.0"
-starky = { version = "0", default-features = false }
+starky = { workspace = true, default-features = false }
 tempfile = "3"
 
 [dev-dependencies]