perf: Concurrent chip proving

ceno_zkvm/benches/riscv_add.rs

@@ -3,7 +3,9 @@ use std::time::Duration;
 use ceno_zkvm::{
     self,
     instructions::{Instruction, riscv::arith::AddInstruction},
-    scheme::{create_backend, create_prover, hal::ProofInput, prover::ZKVMProver},
+    scheme::{
+        create_backend, create_prover, hal::ProofInput, prover::ZKVMProver, scheduler::ChipTask,
+    },
     structs::{ZKVMConstraintSystem, ZKVMFixedTraces},
 };
 mod alloc;
@@ -115,14 +117,21 @@ fn bench_add(c: &mut Criterion) {
                             num_instances: vec![num_instances],
                             has_ecc_ops: false,
                         };
+                        let task = ChipTask {
+                            task_id: 0,
+                            circuit_name: AddInstruction::<E>::name(),
+                            circuit_idx: 0,
+                            pk: circuit_pk,
+                            input,
+                            estimated_memory_bytes: 0,
+                            has_witness_or_fixed: true,
+                            challenges,
+                            witness_trace_idx: None,
+                            num_witin: 0,
+                            structural_rmm: None,
+                        };
                         let _ = prover
-                            .create_chip_proof(
-                                "ADD",
-                                circuit_pk,
-                                input,
-                                &mut transcript,
-                                &challenges,
-                            )
+                            .create_chip_proof(&task, &mut transcript)
                             .expect("create_proof failed");
                         let elapsed = instant.elapsed();
                         println!(

ceno_zkvm/src/scheme.rs

@@ -10,7 +10,7 @@ use std::{
     fmt::{self, Debug},
     iter,
     ops::Div,
-    rc::Rc,
+    sync::Arc,
 };
 use sumcheck::structs::IOPProverMessage;
 
@@ -31,6 +31,7 @@ pub mod cpu;
 pub mod gpu;
 pub mod hal;
 pub mod prover;
+pub mod scheduler;
 pub mod septic_curve;
 pub mod utils;
 pub mod verifier;
@@ -342,13 +343,16 @@ fn byte_to_mb(byte_size: u64) -> f64 {
 pub fn create_backend<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>>(
     max_num_variables: usize,
     security_level: mpcs::SecurityLevel,
-) -> Rc<gkr_iop::cpu::CpuBackend<E, PCS>> {
-    gkr_iop::cpu::CpuBackend::<E, PCS>::new(max_num_variables, security_level).into()
+) -> Arc<gkr_iop::cpu::CpuBackend<E, PCS>> {
+    Arc::new(gkr_iop::cpu::CpuBackend::<E, PCS>::new(
+        max_num_variables,
+        security_level,
+    ))
 }
 
 #[cfg(not(feature = "gpu"))]
 pub fn create_prover<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>>(
-    backend: Rc<gkr_iop::cpu::CpuBackend<E, PCS>>,
+    backend: Arc<gkr_iop::cpu::CpuBackend<E, PCS>>,
 ) -> gkr_iop::cpu::CpuProver<gkr_iop::cpu::CpuBackend<E, PCS>> {
     gkr_iop::cpu::CpuProver::new(backend)
 }
@@ -357,13 +361,16 @@ pub fn create_prover<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>>(
 pub fn create_backend<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>>(
     max_num_variables: usize,
     security_level: mpcs::SecurityLevel,
-) -> Rc<gkr_iop::gpu::GpuBackend<E, PCS>> {
-    gkr_iop::gpu::GpuBackend::<E, PCS>::new(max_num_variables, security_level).into()
+) -> Arc<gkr_iop::gpu::GpuBackend<E, PCS>> {
+    Arc::new(gkr_iop::gpu::GpuBackend::<E, PCS>::new(
+        max_num_variables,
+        security_level,
+    ))
 }
 
 #[cfg(feature = "gpu")]
 pub fn create_prover<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>>(
-    backend: Rc<gkr_iop::gpu::GpuBackend<E, PCS>>,
+    backend: Arc<gkr_iop::gpu::GpuBackend<E, PCS>>,
 ) -> gkr_iop::gpu::GpuProver<gkr_iop::gpu::GpuBackend<E, PCS>> {
     gkr_iop::gpu::GpuProver::new(backend)
 }

ceno_zkvm/src/scheme/cpu/mod.rs

@@ -1051,6 +1051,18 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> DeviceTransporter<Cp
 //     }
 // }
 
+impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>>
+    super::hal::ChipInputPreparer<CpuBackend<E, PCS>> for CpuProver<CpuBackend<E, PCS>>
+{
+    fn prepare_chip_input(
+        &self,
+        _task: &mut crate::scheme::scheduler::ChipTask<'_, CpuBackend<E, PCS>>,
+        _pcs_data: &<CpuBackend<E, PCS> as ProverBackend>::PcsData,
+    ) {
+        // No-op: CPU tasks are eagerly populated during build_chip_tasks
+    }
+}
+
 impl<E, PCS> ProverDevice<CpuBackend<E, PCS>> for CpuProver<CpuBackend<E, PCS>>
 where
     E: ExtensionField,