mpcs e2e batch in one proof

Cargo.toml

@@ -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" }

ceno_zkvm/benches/riscv_add.rs

@@ -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/src/circuit_builder.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/e2e.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/instructions/riscv/test.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/keygen.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;