Splitting PR #44: ZKVM Profile and Cycle Reduction

src/basefold_verifier/query_phase.rs

@@ -324,7 +324,6 @@ pub(crate) fn batch_verifier_query_phase<C: Config>(
     builder: &mut Builder<C>,
     input: QueryPhaseVerifierInputVariable<C>,
 ) {
-    builder.cycle_tracker_start("Before checking opening proofs");
     let inv_2 = builder.constant(C::F::from_canonical_u32(0x3c000001));
     let two_adic_generators_inverses: Array<C, Felt<C::F>> = builder.dyn_array(28);
     for (index, val) in [
@@ -393,11 +392,11 @@ pub(crate) fn batch_verifier_query_phase<C: Config>(
         input.proof.commits.len() + Usize::from(1),
         input.fold_challenges.len(),
     );
-    builder.cycle_tracker_end("Before checking opening proofs");
 
-    builder.cycle_tracker_start("Build round context");
     let rounds_context: Array<C, RoundContextVariable<C>> = builder.dyn_array(input.rounds.len());
     let batch_coeffs_offset: Var<C::N> = builder.constant(C::N::ZERO);
+
+    builder.cycle_tracker_start("Construct round context");
     iter_zip!(builder, input.rounds, rounds_context).for_each(|ptr_vec, builder| {
         let round = builder.iter_ptr_get(&input.rounds, ptr_vec[0]);
         // This buffer is not initialized here in providing the context.
@@ -465,13 +464,10 @@ pub(crate) fn batch_verifier_query_phase<C: Config>(
         };
         builder.iter_ptr_set(&rounds_context, ptr_vec[1], round_context);
     });
-    builder.cycle_tracker_end("Build round context");
+    builder.cycle_tracker_end("Construct round context");
 
-    builder.cycle_tracker_start("Checking opening proofs");
     iter_zip!(builder, input.indices, input.proof.query_opening_proof).for_each(
         |ptr_vec, builder| {
-            builder.cycle_tracker_start("Checking opening proofs (per index)");
-            builder.cycle_tracker_start("Prepare");
             // TODO: change type of input.indices to be `Array<C, Array<C, Var<C::N>>>`
             let idx = builder.iter_ptr_get(&input.indices, ptr_vec[0]);
             let idx = builder.unsafe_cast_var_to_felt(idx);
@@ -496,9 +492,8 @@ pub(crate) fn batch_verifier_query_phase<C: Config>(
                 builder.set_value(&reduced_codeword_by_height, ptr_vec[0], zero_codeword);
             });
             let query = builder.iter_ptr_get(&input.proof.query_opening_proof, ptr_vec[1]);
-            builder.cycle_tracker_end("Prepare");
 
-            builder.cycle_tracker_start("MMCS Verify Loop Rounds");
+            builder.cycle_tracker_start("Batching and first FRI round");
             iter_zip!(builder, query.input_proofs, input.rounds, rounds_context).for_each(
                 |ptr_vec, builder| {
                     let batch_opening = builder.iter_ptr_get(&query.input_proofs, ptr_vec[0]);
@@ -584,9 +579,7 @@ pub(crate) fn batch_verifier_query_phase<C: Config>(
                     mmcs_verify_batch(builder, mmcs_verifier_input);
                 },
             );
-            builder.cycle_tracker_end("MMCS Verify Loop Rounds");
-
-            builder.cycle_tracker_start("Initial folding");
+            builder.cycle_tracker_end("Batching and first FRI round");
             let opening_ext = query.commit_phase_openings;
 
             // fold 1st codeword
@@ -614,7 +607,6 @@ pub(crate) fn batch_verifier_query_phase<C: Config>(
             // check commit phases
             let commits = &input.proof.commits;
             builder.assert_eq::<Var<C::N>>(commits.len(), opening_ext.len());
-            builder.cycle_tracker_end("Initial folding");
             builder.cycle_tracker_start("FRI rounds");
             let i: Var<C::N> = builder.constant(C::N::ZERO);
             iter_zip!(builder, commits, opening_ext).for_each(|ptr_vec, builder| {
@@ -697,8 +689,6 @@ pub(crate) fn batch_verifier_query_phase<C: Config>(
                 builder.assign(&i, i_plus_one);
             });
             builder.cycle_tracker_end("FRI rounds");
-
-            builder.cycle_tracker_start("Finalizing");
             // assert that final_value[i] = folded
             let final_idx: Var<C::N> = builder.constant(C::N::ZERO);
             builder
@@ -713,13 +703,8 @@ pub(crate) fn batch_verifier_query_phase<C: Config>(
                 });
             let final_value = builder.get(&final_codeword.values, final_idx);
             builder.assert_eq::<Ext<C::F, C::EF>>(final_value, folded);
-            builder.cycle_tracker_end("Finalizing");
-            builder.cycle_tracker_end("Checking opening proofs (per index)");
         },
     );
-    builder.cycle_tracker_end("Checking opening proofs");
-
-    builder.cycle_tracker_start("Checking sumcheck proofs");
     // 1. check initial claim match with first round sumcheck value
     let batch_coeffs_offset: Var<C::N> = builder.constant(C::N::ZERO);
     let expected_sum: Ext<C::F, C::EF> = builder.constant(C::EF::ZERO);
@@ -766,9 +751,7 @@ pub(crate) fn batch_verifier_query_phase<C: Config>(
             let right: Ext<C::F, C::EF> = builder.eval(eval0 + eval1);
             builder.assert_eq::<Ext<C::F, C::EF>>(left, right);
         });
-    builder.cycle_tracker_end("Checking sumcheck proofs");
 
-    builder.cycle_tracker_start("Checking final evaluations");
     // 3. check final evaluation are correct
     let final_evals = builder
         .get(&input.proof.sumcheck_proof, fold_len_minus_one.clone())
@@ -816,7 +799,6 @@ pub(crate) fn batch_verifier_query_phase<C: Config>(
     });
     builder.assert_eq::<Var<C::N>>(j, input.proof.final_message.len());
     builder.assert_eq::<Ext<C::F, C::EF>>(left, right);
-    builder.cycle_tracker_end("Checking final evaluations");
 }
 
 #[cfg(test)]

src/basefold_verifier/verifier.rs

@@ -392,7 +392,8 @@ pub mod tests {
 
         let system_config = SystemConfig::default()
             .with_public_values(4)
-            .with_max_segment_len((1 << 25) - 100);
+            .with_max_segment_len((1 << 25) - 100)
+            .with_profiling();
         let config = NativeConfig::new(system_config, Native);
 
         let executor = VmExecutor::<BabyBear, NativeConfig>::new(config);

src/e2e/mod.rs

@@ -1,5 +1,5 @@
 use crate::basefold_verifier::basefold::BasefoldCommitment;
-use crate::tower_verifier::binding::IOPProverMessage;
+use crate::tower_verifier::binding::IOPProverMessageVec;
 use crate::zkvm_verifier::binding::{TowerProofInput, ZKVMChipProofInput, ZKVMProofInput, E, F};
 use crate::zkvm_verifier::verifier::verify_zkvm_proof;
 
@@ -81,10 +81,11 @@ pub fn parse_zkvm_proof_import(
 
         // Tower proof
         let mut tower_proof = TowerProofInput::default();
-        let mut proofs: Vec<Vec<IOPProverMessage>> = vec![];
+        let mut proofs: Vec<IOPProverMessageVec> = vec![];
 
         for proof in &chip_proof.tower_proof.proofs {
-            let mut proof_messages: Vec<IOPProverMessage> = vec![];
+            let mut proof_messages: Vec<E> = vec![];
+            let mut prover_message_size = None;
             for m in proof {
                 let mut evaluations_vec: Vec<E> = vec![];
 
@@ -93,11 +94,17 @@ pub fn parse_zkvm_proof_import(
                         serde_json::from_value(serde_json::to_value(v.clone()).unwrap()).unwrap();
                     evaluations_vec.push(v_e);
                 }
-                proof_messages.push(IOPProverMessage {
-                    evaluations: evaluations_vec,
-                });
+                if let Some(size) = prover_message_size {
+                    assert_eq!(size, evaluations_vec.len());
+                } else {
+                    prover_message_size = Some(evaluations_vec.len());
+                }
+                proof_messages.extend_from_slice(&evaluations_vec);
             }
-            proofs.push(proof_messages);
+            proofs.push(IOPProverMessageVec {
+                prover_message_size: prover_message_size.unwrap(),
+                data: proof_messages,
+            });
         }
         tower_proof.num_proofs = proofs.len();
         tower_proof.proofs = proofs;
@@ -118,7 +125,7 @@ pub fn parse_zkvm_proof_import(
             prod_specs_eval.push(inner_v);
         }
         tower_proof.num_prod_specs = prod_specs_eval.len();
-        tower_proof.prod_specs_eval = prod_specs_eval;
+        tower_proof.prod_specs_eval = prod_specs_eval.into();
 
         let mut logup_specs_eval: Vec<Vec<Vec<E>>> = vec![];
         for inner_val in &chip_proof.tower_proof.logup_specs_eval {
@@ -136,23 +143,36 @@ pub fn parse_zkvm_proof_import(
             logup_specs_eval.push(inner_v);
         }
         tower_proof.num_logup_specs = logup_specs_eval.len();
-        tower_proof.logup_specs_eval = logup_specs_eval;
+        tower_proof.logup_specs_eval = logup_specs_eval.into();
 
         // main constraint and select sumcheck proof
-        let mut main_sumcheck_proofs: Vec<IOPProverMessage> = vec![];
-        if chip_proof.main_sumcheck_proofs.is_some() {
+        let main_sumcheck_proofs = if chip_proof.main_sumcheck_proofs.is_some() {
+            let mut main_sumcheck_proofs: Vec<E> = vec![];
+            let mut prover_message_size = None;
             for m in chip_proof.main_sumcheck_proofs.as_ref().unwrap() {
                 let mut evaluations_vec: Vec<E> = vec![];
                 for v in &m.evaluations {
                     let v_e: E =
                         serde_json::from_value(serde_json::to_value(v.clone()).unwrap()).unwrap();
                     evaluations_vec.push(v_e);
                 }
-                main_sumcheck_proofs.push(IOPProverMessage {
-                    evaluations: evaluations_vec,
-                });
+                main_sumcheck_proofs.extend_from_slice(&evaluations_vec);
+                if let Some(size) = prover_message_size {
+                    assert_eq!(size, evaluations_vec.len());
+                } else {
+                    prover_message_size = Some(evaluations_vec.len());
+                }
             }
-        }
+            IOPProverMessageVec {
+                prover_message_size: prover_message_size.unwrap(),
+                data: main_sumcheck_proofs,
+            }
+        } else {
+            IOPProverMessageVec {
+                prover_message_size: 0,
+                data: vec![],
+            }
+        };
 
         let mut wits_in_evals: Vec<E> = vec![];
         for v in &chip_proof.wits_in_evals {