mpcs e2e batch in one proof (#894)

Build on top of #901.

### Design rationale
zkvm + mpcs api: use `batch_commit/batch_opening` to commit/open **ALL**
opcode together. In more detail, we commit `fixed_commit` and
`witin_commit` separately with just one opening.
basefold: batch opcodes with rmm with different "height", in other
words, different "num_vars".

For basefold to batch opcodes with different variables, for sumcheck
part we apply suffix alignment techniques based on
#870. For FRI part, we also apply suffix alignment techniques. So the
smaller codeword will "involved" into the folding process when current
length match.

During implementation, there is key principle in mind: verifier only
rely on one untrusted information "num_instances" from prover, and other
information should derived from verifier key.

### Working items
- [x] mpcs batch api + batch basefold prover
- [x] batch basefold verifier

### Fibonacci Benchmark Results

| Size | Branch | Proof Size | E2E Latency Change | Prover Speed (kHz) |

|----------|----------------------|------------|--------------------|---------------------|
| 2²⁰ | Master | 14.68 MB | - | 363 |
| 2²⁰ | Batched + conjecture | 1.16 MB | -14.195% | 422 |
| 2²⁰ | Batched | 2.02 MB | -11.599% | 411 |
| 2²¹ | Master | 15.53 MB | - | 411 |
| 2²¹ | Batched + conjecture | 1.24 MB | -11.178% | 451 |
| 2²¹ | Batched | 2.16 MB | -6.4525% | 428 |
| 2²² | Master | 16.42 MB | - | 406 |
| 2²² | Batched + conjecture | 1.31 MB | -11.178% | 433 |
| 2²² | Batched | 2.31 MB | -13.691% | 448 |

Author: hero78119

Cargo.lock

@@ -45,6 +45,7 @@ p3-challenger = { git = "https://github.com/scroll-tech/plonky3", rev = "20c3cb9
 p3-commit = { git = "https://github.com/scroll-tech/plonky3", rev = "20c3cb9" }
 p3-dft = { git = "https://github.com/scroll-tech/plonky3", rev = "20c3cb9" }
 p3-field = { git = "https://github.com/scroll-tech/plonky3", rev = "20c3cb9" }
+p3-fri = { git = "https://github.com/scroll-tech/plonky3", rev = "20c3cb9" }
 p3-goldilocks = { git = "https://github.com/scroll-tech/plonky3", rev = "20c3cb9" }
 p3-matrix = { git = "https://github.com/scroll-tech/plonky3", rev = "20c3cb9" }
 p3-maybe-rayon = { git = "https://github.com/scroll-tech/plonky3", rev = "20c3cb9" }

Cargo.toml

@@ -1,4 +1,4 @@
-use std::time::Duration;
+use std::{collections::BTreeMap, time::Duration};
 
 use ceno_zkvm::{
     self,
@@ -73,14 +73,16 @@ fn bench_add(c: &mut Criterion) {
                     for _ in 0..iters {
                         // generate mock witness
                         let num_instances = 1 << instance_num_vars;
-                        let rmm =
-                            RowMajorMatrix::rand(&mut OsRng, num_instances, num_witin as usize);
+                        let rmms = BTreeMap::from([(
+                            0,
+                            RowMajorMatrix::rand(&mut OsRng, num_instances, num_witin as usize),
+                        )]);
 
                         let instant = std::time::Instant::now();
                         let num_instances = 1 << instance_num_vars;
                         let mut transcript = BasicTranscript::new(b"riscv");
                         let commit =
-                            Pcs::batch_commit_and_write(&prover.pk.pp, rmm, &mut transcript)
+                            Pcs::batch_commit_and_write(&prover.pk.pp, rmms, &mut transcript)
                                 .unwrap();
                         let polys = Pcs::get_arc_mle_witness_from_commitment(&commit);
                         let challenges = [
@@ -91,10 +93,8 @@ fn bench_add(c: &mut Criterion) {
                         let _ = prover
                             .create_opcode_proof(
                                 "ADD",
-                                &prover.pk.pp,
                                 circuit_pk,
                                 polys,
-                                commit,
                                 &[],
                                 num_instances,
                                 &mut transcript,

ceno_zkvm/benches/riscv_add.rs

@@ -3,7 +3,6 @@ use serde::de::DeserializeOwned;
 use std::{collections::HashMap, iter::once, marker::PhantomData};
 
 use ff_ext::ExtensionField;
-use mpcs::PolynomialCommitmentScheme;
 
 use crate::{
     ROMType,
@@ -13,7 +12,6 @@ use crate::{
     structs::{ProgramParams, ProvingKey, RAMType, VerifyingKey, WitnessId},
 };
 use p3::field::PrimeCharacteristicRing;
-use witness::RowMajorMatrix;
 
 /// namespace used for annotation, preserve meta info during circuit construction
 #[derive(Clone, Debug, Default, serde::Serialize, serde::Deserialize)]
@@ -178,24 +176,9 @@ impl<E: ExtensionField> ConstraintSystem<E> {
         }
     }
 
-    pub fn key_gen<PCS: PolynomialCommitmentScheme<E>>(
-        self,
-        pp: &PCS::ProverParam,
-        fixed_traces: Option<RowMajorMatrix<E::BaseField>>,
-    ) -> ProvingKey<E, PCS> {
-        // transpose from row-major to column-major
-        let fixed_traces_polys = fixed_traces.as_ref().map(|rmm| rmm.to_mles());
-
-        let fixed_commit_wd = fixed_traces.map(|traces| PCS::batch_commit(pp, traces).unwrap());
-        let fixed_commit = fixed_commit_wd.as_ref().map(PCS::get_pure_commitment);
-
+    pub fn key_gen(self) -> ProvingKey<E> {
         ProvingKey {
-            fixed_traces: fixed_traces_polys,
-            fixed_commit_wd,
-            vk: VerifyingKey {
-                cs: self,
-                fixed_commit,
-            },
+            vk: VerifyingKey { cs: self },
         }
     }
 

ceno_zkvm/src/circuit_builder.rs

@@ -675,7 +675,11 @@ pub fn verify(
     // print verification statistics like proof size and hash count
     let stat_recorder = StatisticRecorder::default();
     let transcript = BasicTranscriptWithStat::new(&stat_recorder, b"riscv");
-    verifier.verify_proof_halt(zkvm_proof.clone(), transcript, zkvm_proof.has_halt())?;
+    verifier.verify_proof_halt(
+        zkvm_proof.clone(),
+        transcript,
+        zkvm_proof.has_halt(&verifier.vk),
+    )?;
     info!("e2e proof stat: {}", zkvm_proof);
     info!(
         "hashes count = {}",

ceno_zkvm/src/e2e.rs

@@ -23,6 +23,6 @@ fn test_multiple_opcode() {
         |cs| SubInstruction::construct_circuit(&mut CircuitBuilder::<E>::new(cs)),
     );
     let param = Pcs::setup(1 << 10).unwrap();
-    let (pp, _) = Pcs::trim(param, 1 << 10).unwrap();
-    cs.key_gen::<Pcs>(&pp, None);
+    let (_, _) = Pcs::trim(param, 1 << 10).unwrap();
+    cs.key_gen();
 }

ceno_zkvm/src/instructions/riscv/test.rs

@@ -1,3 +1,5 @@
+use std::collections::BTreeMap;
+
 use crate::{
     error::ZKVMError,
     structs::{ZKVMConstraintSystem, ZKVMFixedTraces, ZKVMProvingKey},
@@ -12,24 +14,27 @@ impl<E: ExtensionField> ZKVMConstraintSystem<E> {
         vp: PCS::VerifierParam,
         mut vm_fixed_traces: ZKVMFixedTraces<E>,
     ) -> Result<ZKVMProvingKey<E, PCS>, ZKVMError> {
-        let mut vm_pk = ZKVMProvingKey::new(pp, vp);
+        let mut vm_pk = ZKVMProvingKey::new(pp.clone(), vp);
+        let mut fixed_traces = BTreeMap::new();
 
-        for (c_name, cs) in self.circuit_css {
+        for (circuit_index, (c_name, cs)) in self.circuit_css.into_iter().enumerate() {
             // fixed_traces is optional
             // verifier will check it existent if cs.num_fixed > 0
-            let fixed_traces = if cs.num_fixed > 0 {
-                vm_fixed_traces
+            if cs.num_fixed > 0 {
+                let fixed_trace_rmm = vm_fixed_traces
                     .circuit_fixed_traces
                     .remove(&c_name)
-                    .ok_or(ZKVMError::FixedTraceNotFound(c_name.clone()))?
-            } else {
-                None
+                    .flatten()
+                    .ok_or(ZKVMError::FixedTraceNotFound(c_name.clone()))?;
+                fixed_traces.insert(circuit_index, fixed_trace_rmm);
             };
 
-            let circuit_pk = cs.key_gen(&vm_pk.pp, fixed_traces);
+            let circuit_pk = cs.key_gen();
             assert!(vm_pk.circuit_pks.insert(c_name, circuit_pk).is_none());
         }
 
+        vm_pk.commit_fixed(fixed_traces)?;
+
         vm_pk.initial_global_state_expr = self.initial_global_state_expr;
         vm_pk.finalize_global_state_expr = self.finalize_global_state_expr;