fork transcript to support chip concurrent prove

Cargo.toml

@@ -27,16 +27,16 @@ version = "0.1.0"
 ceno_crypto_primitives = { git = "https://github.com/scroll-tech/ceno-patch.git", package = "ceno_crypto_primitives", branch = "main" }
 ceno_syscall = { git = "https://github.com/scroll-tech/ceno-patch.git", package = "ceno_syscall", branch = "main" }
 
-ff_ext = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "ff_ext", tag = "v1.0.0-alpha.21" }
-mpcs = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "mpcs", tag = "v1.0.0-alpha.21" }
-multilinear_extensions = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "multilinear_extensions", tag = "v1.0.0-alpha.21" }
-p3 = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "p3", tag = "v1.0.0-alpha.21" }
-poseidon = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "poseidon", tag = "v1.0.0-alpha.21" }
-sp1-curves = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "sp1-curves", tag = "v1.0.0-alpha.21" }
-sumcheck = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "sumcheck", tag = "v1.0.0-alpha.21" }
-transcript = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "transcript", tag = "v1.0.0-alpha.21" }
-whir = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "whir", tag = "v1.0.0-alpha.21" }
-witness = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "witness", tag = "v1.0.0-alpha.21" }
+ff_ext = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "ff_ext", tag = "v1.0.0-alpha.22" }
+mpcs = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "mpcs", tag = "v1.0.0-alpha.22" }
+multilinear_extensions = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "multilinear_extensions", tag = "v1.0.0-alpha.22" }
+p3 = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "p3", tag = "v1.0.0-alpha.22" }
+poseidon = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "poseidon", tag = "v1.0.0-alpha.22" }
+sp1-curves = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "sp1-curves", tag = "v1.0.0-alpha.22" }
+sumcheck = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "sumcheck", tag = "v1.0.0-alpha.22" }
+transcript = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "transcript", tag = "v1.0.0-alpha.22" }
+whir = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "whir", tag = "v1.0.0-alpha.22" }
+witness = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "witness", tag = "v1.0.0-alpha.22" }
 
 anyhow = { version = "1.0", default-features = false }
 bincode = "1"

Makefile.toml

@@ -63,6 +63,7 @@ args = [
   "nightly-features",
   "--path",
   "./ceno_cli",
+  "--locked",
 ]
 command = "cargo"
 env = { "JEMALLOC_SYS_WITH_MALLOC_CONF" = "retain:true,metadata_thp:always,thp:always,dirty_decay_ms:-1,muzzy_decay_ms:-1,abort_conf:true" }

ceno_recursion/src/transcript/mod.rs

@@ -31,3 +31,32 @@ pub fn transcript_check_pow_witness<C: Config>(
         builder.assert_eq::<Var<C::N>>(bit, Usize::from(0));
     });
 }
+
+pub fn clone_challenger_state<C: Config>(
+    builder: &mut Builder<C>,
+    src: &DuplexChallengerVariable<C>,
+) -> DuplexChallengerVariable<C> {
+    let dst = DuplexChallengerVariable::new(builder);
+    builder
+        .range(0, dst.sponge_state.len())
+        .for_each(|idx_vec, builder| {
+            let value = builder.get(&src.sponge_state, idx_vec[0]);
+            builder.set(&dst.sponge_state, idx_vec[0], value);
+        });
+
+    let input_offset = src.input_ptr - src.io_empty_ptr;
+    builder.assign(&dst.input_ptr, input_offset + dst.io_empty_ptr);
+
+    let output_offset = src.output_ptr - src.io_empty_ptr;
+    builder.assign(&dst.output_ptr, output_offset + dst.io_empty_ptr);
+    dst
+}
+
+pub fn challenger_add_forked_index<C: Config>(
+    builder: &mut Builder<C>,
+    challenger: &mut DuplexChallengerVariable<C>,
+    index: &Usize<C::N>,
+) {
+    let felt = builder.unsafe_cast_var_to_felt(index.get_var());
+    challenger.observe(builder, felt);
+}

ceno_recursion/src/zkvm_verifier/verifier.rs

@@ -34,6 +34,7 @@ use crate::{
 use ceno_zkvm::structs::{ComposedConstrainSystem, VerifyingKey, ZKVMVerifyingKey};
 use ff_ext::BabyBearExt4;
 
+use crate::transcript::{challenger_add_forked_index, clone_challenger_state};
 use gkr_iop::{
     evaluation::EvalExpression,
     gkr::{
@@ -234,7 +235,7 @@ pub fn verify_zkvm_proof<C: Config<F = F>>(
     // not each chip has witness or fixed opening
     // therefore we need to truncate these two opening arrays
     let witin_openings: Array<C, RoundOpeningVariable<C>> = builder.dyn_array(proofs_len.clone());
-    let fixed_openings: Array<C, RoundOpeningVariable<C>> = builder.dyn_array(proofs_len);
+    let fixed_openings: Array<C, RoundOpeningVariable<C>> = builder.dyn_array(proofs_len.clone());
 
     let shard_ec_sum = SepticPointVariable {
         x: SepticExtensionVariable {
@@ -260,6 +261,10 @@ pub fn verify_zkvm_proof<C: Config<F = F>>(
             builder.set(&chip_indices, i, chip_idx);
         });
 
+    // collect fork sampling result
+    let forked_samples: Array<C, Ext<C::F, C::EF>> = builder.dyn_array(proofs_len.get_var());
+    let forked_sample_index: Usize<C::N> = builder.eval(C::N::ZERO);
+
     for (i, (circuit_name, chip_vk)) in vk.circuit_vks.iter().enumerate() {
         let circuit_vk = &vk.circuit_vks[circuit_name];
         let chip_id: Var<C::N> = builder.get(&chip_indices, num_chips_verified.get_var());
@@ -270,6 +275,9 @@ pub fn verify_zkvm_proof<C: Config<F = F>>(
 
             iter_zip!(builder, chip_proofs).for_each(|ptr_vec, builder| {
                 let chip_proof = builder.iter_ptr_get(&chip_proofs, ptr_vec[0]);
+                // fork transcript to support chip concurrently proved
+                let mut chip_challenger = clone_challenger_state(builder, &challenger);
+                challenger_add_forked_index(builder, &mut chip_challenger, &forked_sample_index);
                 builder.assert_usize_eq(
                     chip_proof.wits_in_evals.len(),
                     Usize::from(circuit_vk.get_cs().num_witin()),
@@ -302,7 +310,7 @@ pub fn verify_zkvm_proof<C: Config<F = F>>(
                     builder.assign(&chip_logup_sum, chip_logup_sum + p1 * q1.inverse());
                     builder.assign(&chip_logup_sum, chip_logup_sum + p2 * q2.inverse());
                 });
-                challenger.observe(builder, chip_proof.idx_felt);
+                chip_challenger.observe(builder, chip_proof.idx_felt);
 
                 if circuit_vk.get_cs().is_with_lk_table() {
                     builder.assign(&logup_sum, logup_sum - chip_logup_sum);
@@ -344,7 +352,7 @@ pub fn verify_zkvm_proof<C: Config<F = F>>(
                 let (input_opening_point, chip_shard_ec_sum) = verify_chip_proof(
                     circuit_name,
                     builder,
-                    &mut challenger,
+                    &mut chip_challenger,
                     &chip_proof,
                     &zkvm_proof_input.pi_evals,
                     &zkvm_proof_input.raw_pi,
@@ -396,6 +404,10 @@ pub fn verify_zkvm_proof<C: Config<F = F>>(
                     .then(|builder| {
                         add_septic_points_in_place(builder, &shard_ec_sum, &chip_shard_ec_sum);
                     });
+
+                let chip_sample = chip_challenger.sample_ext(builder);
+                builder.set(&forked_samples, forked_sample_index.get_var(), chip_sample);
+                builder.inc(&forked_sample_index);
             });
             builder.inc(&num_chips_verified);
         });
@@ -414,6 +426,13 @@ pub fn verify_zkvm_proof<C: Config<F = F>>(
         logup_sum - dummy_table_item_multiplicity * dummy_table_item.inverse(),
     );
 
+    // merge forked transcripts into transcript
+    iter_zip!(builder, forked_samples).for_each(|ptr_vec, builder| {
+        let sample = builder.iter_ptr_get(&forked_samples, ptr_vec[0]);
+        let sample_felts = builder.ext2felt(sample);
+        challenger.observe_slice(builder, sample_felts);
+    });
+
     let rounds: Array<C, RoundVariable<C>> = if num_fixed_opening > 0 {
         builder.dyn_array(2)
     } else {

ceno_zkvm/src/scheme/prover.rs

@@ -25,7 +25,7 @@ use sumcheck::{
     structs::IOPProverMessage,
 };
 use tracing::info_span;
-use transcript::Transcript;
+use transcript::{ForkableTranscript, Transcript};
 
 use super::{PublicValues, ZKVMChipProof, ZKVMProof, hal::ProverDevice};
 use crate::{
@@ -128,7 +128,7 @@ impl<
         shard_ctx: &ShardContext,
         witnesses: ZKVMWitnesses<E>,
         pi: PublicValues,
-        mut transcript: impl Transcript<E> + 'static,
+        mut transcript: impl ForkableTranscript<E> + 'static,
     ) -> Result<ZKVMProof<E, PCS>, ZKVMError> {
         info_span!(
             "[ceno] create_proof_of_shard",
@@ -253,6 +253,21 @@ impl<
             let mut witness_iter = self
                 .device
                 .extract_witness_mles(&mut witness_mles, &witness_data);
+
+            let num_proofs = name_and_instances
+                .iter()
+                .filter(|(circuit_name, num_instances)| {
+                    let pk = self.pk.circuit_pks.get(circuit_name).unwrap();
+                    let cs = pk.get_cs();
+                    let has_instances = !num_instances.is_empty();
+                    let skip_omc_init = !shard_ctx.is_first_shard() && cs.with_omc_init_only();
+                    has_instances && !skip_omc_init
+                })
+                .count();
+
+            // fork transcript to support chip concurrently proved
+            let mut forked_transcripts = transcript.fork(num_proofs);
+            let mut proof_index = 0;
             for ((circuit_name, num_instances), structural_rmm) in name_and_instances
                 .into_iter()
                 .zip_eq(structural_rmms.into_iter())
@@ -277,6 +292,8 @@ impl<
                     }
                     continue;
                 }
+                let transcript = &mut forked_transcripts[proof_index];
+                proof_index += 1;
                 transcript
                     .append_field_element(&E::BaseField::from_canonical_u64(circuit_idx as u64));
 
@@ -320,7 +337,7 @@ impl<
                                 circuit_name.as_str(),
                                 pk,
                                 input,
-                                &mut transcript,
+                                transcript,
                                 &challenges,
                             )
                         },
@@ -352,6 +369,15 @@ impl<
             drop(witness_iter);
             exit_span!(main_proofs_span);
 
+            // merge forked transcript into transcript
+            let forked_sampling = forked_transcripts
+                .into_iter()
+                .map(|mut fork_transcript| fork_transcript.sample_vec(1)[0])
+                .collect_vec();
+            for sample in forked_sampling {
+                transcript.append_field_element_ext(&sample);
+            }
+
             // batch opening pcs
             // generate static info from prover key for expected num variable
             let pcs_opening = entered_span!("pcs_opening", profiling_1 = true);

ceno_zkvm/src/scheme/verifier.rs

@@ -259,6 +259,7 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMVerifier<E, PCS>
         let mut shard_ec_sum = SepticPoint::<E::BaseField>::default();
 
         // check num proofs
+        let mut num_proofs = 0;
         for (index, proofs) in &vm_proof.chip_proofs {
             let circuit_name = &self.vk.circuit_index_to_name[index];
             let circuit_vk = &self.vk.circuit_vks[circuit_name];
@@ -274,12 +275,16 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMVerifier<E, PCS>
                         .into(),
                 ));
             }
+            num_proofs += proofs.len();
         }
 
-        for (index, proof) in vm_proof
+        // fork transcript to support chip concurrently proved
+        let mut forked_transcripts = transcript.fork(num_proofs);
+        for ((index, proof), transcript) in vm_proof
             .chip_proofs
             .iter()
             .flat_map(|(index, proofs)| iter::repeat_n(index, proofs.len()).zip(proofs))
+            .zip_eq(forked_transcripts.iter_mut())
         {
             let num_instance: usize = proof.num_instances.iter().sum();
             assert!(num_instance > 0);
@@ -359,7 +364,7 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMVerifier<E, PCS>
                 proof,
                 pi_evals,
                 &vm_proof.raw_pi,
-                &mut transcript,
+                transcript,
                 NUM_FANIN,
                 &point_eval,
                 &challenges,
@@ -396,6 +401,15 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMVerifier<E, PCS>
             );
         }
 
+        // merge forked transcript into transcript
+        let forked_samples = forked_transcripts
+            .into_iter()
+            .map(|mut fork_transcript| fork_transcript.sample_vec(1)[0])
+            .collect_vec();
+        for sample in forked_samples {
+            transcript.append_field_element_ext(&sample);
+        }
+
         // verify mpcs
         let mut rounds = vec![(vm_proof.witin_commit.clone(), witin_openings)];