tools: support whitelist shard id (#1173)

support `--shard-id=k` which only prove k-th shard and skip all others

Author: hero78119

ceno_cli/src/commands/common_args/ceno.rs

@@ -100,6 +100,11 @@ pub struct CenoOptions {
     /// Setting any value restricts logs to profiling information
     #[arg(long)]
     profiling: Option<usize>,
+
+    // for debug purpose
+    // only generate respective shard id and skip others
+    #[arg(long)]
+    shard_id: Option<u64>,
 }
 
 impl CenoOptions {
@@ -414,6 +419,7 @@ fn run_elf_inner<
         &public_io,
         options.max_steps,
         checkpoint,
+        options.shard_id.map(|v| v as usize),
     ))
 }
 

ceno_recursion/src/bin/e2e_aggregate.rs

@@ -261,6 +261,7 @@ fn main() {
             &public_io,
             max_steps,
             Checkpoint::Complete,
+            None,
         );
 
     let zkvm_proofs = result

ceno_zkvm/benches/fibonacci.rs

@@ -59,6 +59,7 @@ fn fibonacci_prove(c: &mut Criterion) {
             &[],
             max_steps,
             Checkpoint::Complete,
+            None,
         );
         let proof = result
             .proofs
@@ -100,6 +101,7 @@ fn fibonacci_prove(c: &mut Criterion) {
                             &[],
                             max_steps,
                             Checkpoint::PrepE2EProving,
+                            None,
                         );
                         let instant = std::time::Instant::now();
                         result.next_step();

ceno_zkvm/benches/fibonacci_witness.rs

@@ -71,6 +71,7 @@ fn fibonacci_witness(c: &mut Criterion) {
                         &[],
                         max_steps,
                         Checkpoint::PrepWitnessGen,
+                        None,
                     );
                     let instant = std::time::Instant::now();
                     result.next_step();

ceno_zkvm/benches/is_prime.rs

@@ -68,6 +68,7 @@ fn is_prime_1(c: &mut Criterion) {
                             &[],
                             max_steps,
                             Checkpoint::PrepE2EProving,
+                            None,
                         );
                         let instant = std::time::Instant::now();
                         result.next_step();

ceno_zkvm/benches/keccak.rs

@@ -56,6 +56,7 @@ fn keccak_prove(c: &mut Criterion) {
         &[],
         max_steps,
         Checkpoint::Complete,
+        None,
     );
     let proof = result
         .proofs
@@ -94,6 +95,7 @@ fn keccak_prove(c: &mut Criterion) {
                         &[],
                         max_steps,
                         Checkpoint::PrepE2EProving,
+                        None,
                     );
                     let instant = std::time::Instant::now();
                     result.next_step();

ceno_zkvm/benches/quadratic_sorting.rs

@@ -69,6 +69,7 @@ fn quadratic_sorting_1(c: &mut Criterion) {
                             &[],
                             max_steps,
                             Checkpoint::PrepE2EProving,
+                            None,
                         );
                         let instant = std::time::Instant::now();
                         result.next_step();

ceno_zkvm/src/bin/e2e.rs

@@ -126,6 +126,11 @@ struct Args {
     // => 2^30 * 16 / 4 / 2
     #[arg(long, default_value = "2147483648")]
     max_cell_per_shard: u64,
+
+    // for debug purpose
+    // only generate respective shard id and skip others
+    #[arg(long)]
+    shard_id: Option<u64>,
 }
 
 fn main() {
@@ -265,6 +270,8 @@ fn main() {
         args.max_cycle_per_shard,
     );
 
+    let target_shard_id = args.shard_id.map(|v| v as usize);
+
     match (args.pcs, args.field) {
         (PcsKind::Basefold, FieldType::Goldilocks) => {
             let backend = create_backend(args.max_num_variables, args.security_level);
@@ -280,6 +287,7 @@ fn main() {
                 args.proof_file,
                 args.vk_file,
                 Checkpoint::Complete,
+                target_shard_id,
             )
         }
         (PcsKind::Basefold, FieldType::BabyBear) => {
@@ -296,6 +304,7 @@ fn main() {
                 args.proof_file,
                 args.vk_file,
                 Checkpoint::Complete,
+                target_shard_id,
             )
         }
         (PcsKind::Whir, FieldType::Goldilocks) => {
@@ -312,6 +321,7 @@ fn main() {
                 args.proof_file,
                 args.vk_file,
                 Checkpoint::PrepVerify, // TODO: when whir and babybear is ready
+                target_shard_id,
             )
         }
         (PcsKind::Whir, FieldType::BabyBear) => {
@@ -328,6 +338,7 @@ fn main() {
                 args.proof_file,
                 args.vk_file,
                 Checkpoint::PrepVerify, // TODO: when whir and babybear is ready
+                target_shard_id,
             )
         }
     };
@@ -355,6 +366,7 @@ fn run_inner<
     proof_file: PathBuf,
     vk_file: PathBuf,
     checkpoint: Checkpoint,
+    target_shard_id: Option<usize>,
 ) {
     let result = run_e2e_with_checkpoint::<E, PCS, _, _>(
         pd,
@@ -365,6 +377,7 @@ fn run_inner<
         public_io,
         max_steps,
         checkpoint,
+        target_shard_id,
     );
 
     let zkvm_proofs = result
@@ -377,7 +390,7 @@ fn run_inner<
     let vk_bytes = bincode::serialize(&vk).unwrap();
     fs::write(&vk_file, vk_bytes).unwrap();
 
-    if checkpoint > Checkpoint::PrepVerify {
+    if checkpoint > Checkpoint::PrepVerify && target_shard_id.is_none() {
         let verifier = ZKVMVerifier::new(vk);
         verify(zkvm_proofs.clone(), &verifier).expect("Verification failed");
         soundness_test(zkvm_proofs.first().cloned().unwrap(), &verifier);

ceno_zkvm/src/e2e.rs

@@ -993,6 +993,8 @@ pub fn generate_witness<'a, E: ExtensionField>(
     system_config: &ConstraintSystemConfig<E>,
     mut emul_result: EmulationResult<'a>,
     program: &Program,
+    // this is for debug purpose, which only run target shard id and skip all others
+    target_shard_id: Option<usize>,
 ) -> impl Iterator<Item = (ZKVMWitnesses<E>, ShardContext<'a>, PublicValues)> {
     let mut shard_ctx_builder = std::mem::take(&mut emul_result.shard_ctx_builder);
     let all_records = std::mem::take(&mut emul_result.all_records);
@@ -1038,6 +1040,21 @@ pub fn generate_witness<'a, E: ExtensionField>(
         let current_shard_end_pc = all_records[cur_index..end].last().unwrap().pc().after.0;
 
         let mut zkvm_witness = ZKVMWitnesses::default();
+
+        if let Some(target_shard_id) = target_shard_id {
+            if shard_ctx.shard_id < target_shard_id {
+                tracing::debug!("{}th shard skipped", shard_ctx.shard_id);
+                // update next round start
+                // early stop and return empty zkvm witness, skipped all the potiential cost
+                cur_index = end;
+                return Some((zkvm_witness, shard_ctx, pi));
+            } else if shard_ctx.shard_id > target_shard_id {
+                tracing::debug!("{}th shard skipped", shard_ctx.shard_id);
+                return None;
+            }
+            // go ahead to generate witness
+        }
+
         let time = std::time::Instant::now();
         // assign opcode circuits
         let dummy_records = system_config
@@ -1353,6 +1370,8 @@ pub fn run_e2e_with_checkpoint<
     public_io: &[u32],
     max_steps: usize,
     checkpoint: Checkpoint,
+    // for debug purpose
+    target_shard_id: Option<usize>,
 ) -> E2ECheckpointResult<E, PCS> {
     let start = std::time::Instant::now();
     let ctx = setup_program::<E>(program, platform, multi_prover);
@@ -1377,7 +1396,13 @@ pub fn run_e2e_with_checkpoint<
             proofs: None,
             vk: Some(vk),
             next_step: Some(Box::new(move || {
-                _ = run_e2e_proof::<E, _, _, _>(&prover, &init_full_mem, max_steps, is_mock_proving)
+                _ = run_e2e_proof::<E, _, _, _>(
+                    &prover,
+                    &init_full_mem,
+                    max_steps,
+                    is_mock_proving,
+                    target_shard_id,
+                )
             })),
         };
     }
@@ -1408,12 +1433,22 @@ pub fn run_e2e_with_checkpoint<
                     &prover.pk.program_ctx.as_ref().unwrap().system_config,
                     emul_result,
                     &prover.pk.program_ctx.as_ref().unwrap().program,
+                    target_shard_id,
                 )
             })),
         };
     }
 
-    let zkvm_proofs = create_proofs_helper(emul_result, &prover, is_mock_proving);
+    let zkvm_proofs = create_proofs_helper(emul_result, &prover, is_mock_proving, target_shard_id);
+
+    if target_shard_id.is_some() {
+        // skip verify as the proof are in-completed
+        return E2ECheckpointResult {
+            proofs: Some(zkvm_proofs),
+            vk: Some(vk),
+            next_step: None,
+        };
+    }
 
     let verifier = ZKVMVerifier::new(vk.clone());
 
@@ -1451,6 +1486,8 @@ pub fn run_e2e_proof<
     init_full_mem: &InitMemState,
     max_steps: usize,
     is_mock_proving: bool,
+    // for debug purpose
+    target_shard_id: Option<usize>,
 ) -> Vec<ZKVMProof<E, PCS>> {
     let ctx = prover.pk.program_ctx.as_ref().unwrap();
     // Emulate program
@@ -1461,7 +1498,7 @@ pub fn run_e2e_proof<
         &ctx.platform,
         &ctx.multi_prover,
     );
-    create_proofs_helper(emul_result, prover, is_mock_proving)
+    create_proofs_helper(emul_result, prover, is_mock_proving, target_shard_id)
 }
 
 /// defines a lightweight CPU -> GPU pipeline for witness generation and proof creation.
@@ -1499,6 +1536,7 @@ fn create_proofs_helper<
     emulation_result: EmulationResult,
     prover: &ZKVMProver<E, PCS, PB, PD>,
     is_mock_proving: bool,
+    target_shard_id: Option<usize>,
 ) -> Vec<ZKVMProof<E, PCS>> {
     let ctx = prover.pk.program_ctx.as_ref().unwrap();
     #[cfg(feature = "gpu")]
@@ -1510,9 +1548,20 @@ fn create_proofs_helper<
             // cpu producer
             s.spawn({
                 move || {
-                    for proof_input in
-                        generate_witness(&ctx.system_config, emulation_result, &ctx.program)
-                    {
+                    let wit_iter = generate_witness(
+                        &ctx.system_config,
+                        emulation_result,
+                        &ctx.program,
+                        target_shard_id,
+                    );
+
+                    let wit_iter = if let Some(target_shard_id) = target_shard_id {
+                        Box::new(wit_iter.skip(target_shard_id)) as Box<dyn Iterator<Item = _>>
+                    } else {
+                        Box::new(wit_iter)
+                    };
+
+                    for proof_input in wit_iter {
                         tx.send(proof_input).unwrap()
                     }
                 }
@@ -1552,9 +1601,20 @@ fn create_proofs_helper<
     #[cfg(not(feature = "gpu"))]
     {
         // Generate witness
-        let zkvm_witness = generate_witness(&ctx.system_config, emulation_result, &ctx.program);
+        let wit_iter = generate_witness(
+            &ctx.system_config,
+            emulation_result,
+            &ctx.program,
+            target_shard_id,
+        );
 
-        zkvm_witness
+        let wit_iter = if let Some(target_shard_id) = target_shard_id {
+            Box::new(wit_iter.skip(target_shard_id)) as Box<dyn Iterator<Item = _>>
+        } else {
+            Box::new(wit_iter)
+        };
+
+        wit_iter
             .map(|(zkvm_witness, shard_ctx, pi)| {
                 if is_mock_proving {
                     MockProver::assert_satisfied_full(