feat: higher FRI arity

fri/Cargo.toml

@@ -17,6 +17,7 @@ itertools.workspace = true
 rand.workspace = true
 tracing.workspace = true
 serde = { workspace = true, features = ["derive", "alloc"] }
+tracing-subscriber.workspace = true
 
 [dev-dependencies]
 p3-baby-bear.workspace = true

fri/src/config.rs

@@ -8,9 +8,11 @@ use p3_matrix::Matrix;
 pub struct FriConfig<M> {
     pub log_blowup: usize,
     // TODO: This parameter and FRI early stopping are not yet implemented in `CirclePcs`.
+    /// log of the number of coefficients in the final polynomial
     pub log_final_poly_len: usize,
     pub num_queries: usize,
     pub proof_of_work_bits: usize,
+    pub arity_bits: usize,
     pub mmcs: M,
 }
 
@@ -23,6 +25,10 @@ impl<M> FriConfig<M> {
         1 << self.log_final_poly_len
     }
 
+    pub const fn arity(&self) -> usize {
+        1 << self.arity_bits
+    }
+
     /// Returns the soundness bits of this FRI instance based on the
     /// [ethSTARK](https://eprint.iacr.org/2021/582) conjecture.
     ///
@@ -66,6 +72,7 @@ pub fn create_test_fri_config<Mmcs>(mmcs: Mmcs) -> FriConfig<Mmcs> {
         log_final_poly_len: 0,
         num_queries: 2,
         proof_of_work_bits: 1,
+        arity_bits: 1,
         mmcs,
     }
 }
@@ -78,6 +85,7 @@ pub fn create_benchmark_fri_config<Mmcs>(mmcs: Mmcs) -> FriConfig<Mmcs> {
         log_final_poly_len: 0,
         num_queries: 100,
         proof_of_work_bits: 16,
+        arity_bits: 1,
         mmcs,
     }
 }

fri/src/proof.rs

@@ -13,6 +13,7 @@ pub struct FriProof<F: Field, M: Mmcs<F>, Witness, InputProof> {
     pub commit_phase_commits: Vec<M::Commitment>,
     pub query_proofs: Vec<QueryProof<F, M, InputProof>>,
     pub final_poly: Vec<F>,
+    pub log_max_height: usize,
     pub pow_witness: Witness,
 }
 
@@ -31,10 +32,10 @@ pub struct QueryProof<F: Field, M: Mmcs<F>, InputProof> {
 #[derive(Debug, Serialize, Deserialize, Clone)]
 #[serde(bound = "")]
 pub struct CommitPhaseProofStep<F: Field, M: Mmcs<F>> {
-    /// The opening of the commit phase codeword at the sibling location.
-    // This may change to Vec<FC::Challenge> if the library is generalized to support other FRI
-    // folding arities besides 2, meaning that there can be multiple siblings.
-    pub sibling_value: F,
+    /// The opened rows of the commit phase. The first element is the evaluation of the folded
+    /// polynomials so far, and the other elements are evaluations of polynomials of smaller size
+    /// that enter before the next commitment round. See prover::commit_phase for more details
+    pub opened_rows: Vec<Vec<F>>,
 
     pub opening_proof: M::Proof,
 }

fri/src/prover.rs

@@ -7,7 +7,7 @@ use p3_challenger::{CanObserve, FieldChallenger, GrindingChallenger};
 use p3_commit::Mmcs;
 use p3_dft::{Radix2Dit, TwoAdicSubgroupDft};
 use p3_field::{ExtensionField, Field, TwoAdicField};
-use p3_matrix::dense::RowMajorMatrix;
+use p3_matrix::dense::{DenseMatrix, RowMajorMatrix};
 use p3_util::{log2_strict_usize, reverse_slice_index_bits};
 use tracing::{debug_span, info_span, instrument};
 
@@ -65,6 +65,7 @@ where
         commit_phase_commits: commit_phase_result.commits,
         query_proofs,
         final_poly: commit_phase_result.final_poly,
+        log_max_height,
         pow_witness,
     }
 }
@@ -89,26 +90,85 @@ where
     Challenger: FieldChallenger<Val> + CanObserve<M::Commitment>,
     G: FriGenericConfig<Challenge>,
 {
+    // To illustrate the folding logic with arity > 2, let's go through an example.
+    // Suppose `inputs` consists of three polynomials with degrees 16, 8, and 4, and
+    // suppose that arity = 4 and final_poly_len = 1.
+    // There will be two FRI commitment layers: one at height 16 and one at height 4.
+
+    // The first commitment layer will consist of two matrices:
+    // - one of dimensions 4x4 corresponding to the first polynomial's evaluations
+    // - one of dimensions 4x2 corresponding to the second polynomial's evaluations
+
+    // The polynomial folding happens incrementally as follows: the first polynomial is folded
+    // once so its number of evaluations is halved, the second polynomial's evaluations are added
+    // to that, and then the sum is folded by 2 further to reduce the number of evaluations to 4.
+
+    // At that point, the third polynomial's evaluations are added to the running sum, and that sum
+    // is committed to form the second FRI commitment layer. The only matrix in that layer is of
+    // dimensions 4x1. The final polynomial's evaluation can then be computed through those 4
+    // evaluations.
+
     let mut inputs_iter = inputs.into_iter().peekable();
     let mut folded = inputs_iter.next().unwrap();
     let mut commits = vec![];
     let mut data = vec![];
 
+    let arity = config.arity();
+
     while folded.len() > config.blowup() * config.final_poly_len() {
-        let leaves = RowMajorMatrix::new(folded, 2);
-        let (commit, prover_data) = config.mmcs.commit_matrix(leaves);
+        let cur_arity = arity.min(folded.len());
+
+        let next_folded_len = folded.len() / cur_arity;
+
+        // First, we collect the polynomial evaluations that will be committed this round.
+        // Those are `folded` and polynomials in `inputs` not consumed yet with number of
+        // evaluations more than `next_folded_len`
+        let mut polys_before_next_round = vec![];
+
+        let mut cur_folded_len = folded.len();
+        while cur_folded_len > next_folded_len {
+            if let Some(poly_eval) = inputs_iter.next_if(|v| v.len() == cur_folded_len) {
+                let width = poly_eval.len() / next_folded_len;
+                let poly_eval_matrix = RowMajorMatrix::new(poly_eval, width);
+                polys_before_next_round.push(poly_eval_matrix);
+            }
+
+            cur_folded_len /= 2;
+        }
+
+        let folded_matrix = RowMajorMatrix::new(folded.clone(), cur_arity);
+        let matrices_to_commit: Vec<DenseMatrix<Challenge>> = iter::once(folded_matrix)
+            .chain(polys_before_next_round)
+            .collect();
+
+        let (commit, prover_data) = config.mmcs.commit(matrices_to_commit);
         challenger.observe(commit.clone());
 
-        let beta: Challenge = challenger.sample_ext_element();
-        // We passed ownership of `current` to the MMCS, so get a reference to it
-        let leaves = config.mmcs.get_matrices(&prover_data).pop().unwrap();
-        folded = g.fold_matrix(beta, leaves.as_view());
+        // Next, we fold `folded` and `polys_before_next_round` to prepare for the next round
+        let mut beta: Challenge = challenger.sample_ext_element();
+
+        // Get a reference to the committed matrices
+        let leaves = config.mmcs.get_matrices(&prover_data);
+        let mut leaves_iter = leaves.into_iter().peekable();
+        // Skip `folded`
+        leaves_iter.next();
+
+        while folded.len() > next_folded_len {
+            let matrix_to_fold = RowMajorMatrix::new(folded, 2);
+            folded = g.fold_matrix(beta, matrix_to_fold);
+            beta = beta.square();
+
+            if let Some(poly_eval) = leaves_iter.next_if(|v| v.values.len() == folded.len()) {
+                izip!(&mut folded, &poly_eval.values).for_each(|(f, v)| *f += *v);
+            }
+        }
 
         commits.push(commit);
         data.push(prover_data);
 
-        if let Some(v) = inputs_iter.next_if(|v| v.len() == folded.len()) {
-            izip!(&mut folded, v).for_each(|(c, x)| *c += x);
+        // We directly add the next polynomial's evaluations into `folded` in case their lengths match
+        if let Some(poly_eval) = inputs_iter.next_if(|v| v.len() == folded.len()) {
+            izip!(&mut folded, poly_eval).for_each(|(c, x)| *c += x);
         }
     }
 
@@ -129,7 +189,7 @@ where
     debug_assert!(
         final_poly
             .iter()
-            .skip(1 << config.log_final_poly_len)
+            .skip(config.final_poly_len())
             .all(|x| x.is_zero()),
         "All coefficients beyond final_poly_len must be zero"
     );
@@ -159,18 +219,12 @@ where
         .iter()
         .enumerate()
         .map(|(i, commit)| {
-            let index_i = index >> i;
-            let index_i_sibling = index_i ^ 1;
-            let index_pair = index_i >> 1;
+            let index_row = index >> ((i + 1) * config.arity_bits);
 
-            let (mut opened_rows, opening_proof) = config.mmcs.open_batch(index_pair, commit);
-            assert_eq!(opened_rows.len(), 1);
-            let opened_row = opened_rows.pop().unwrap();
-            assert_eq!(opened_row.len(), 2, "Committed data should be in pairs");
-            let sibling_value = opened_row[index_i_sibling % 2];
+            let (opened_rows, opening_proof) = config.mmcs.open_batch(index_row, commit);
 
             CommitPhaseProofStep {
-                sibling_value,
+                opened_rows,
                 opening_proof,
             }
         })

fri/src/two_adic_pcs.rs

@@ -79,7 +79,7 @@ impl<F: TwoAdicField, InputProof, InputError: Debug> FriGenericConfig<F>
         let log_arity = 1;
         let (e0, e1) = evals
             .collect_tuple()
-            .expect("TwoAdicFriFolder only supports arity=2");
+            .expect("TwoAdicFriGenericConfig only supports folding rows of size 2");
         // If performance critical, make this API stateful to avoid this
         // This is a bit more math than is necessary, but leaving it here
         // in case we want higher arity in the future
@@ -122,7 +122,9 @@ impl<F: TwoAdicField, InputProof, InputError: Debug> FriGenericConfig<F>
         m.par_rows()
             .zip(powers)
             .map(|(mut row, power)| {
-                let (lo, hi) = row.next_tuple().unwrap();
+                let (lo, hi) = row
+                    .next_tuple()
+                    .expect("TwoAdicFriGenericConfig only supports folding rows of size 2");
                 (one_half + power) * lo + (one_half - power) * hi
             })
             .collect()