[experiment] more tracing span to identify mpcs latency breakdown on flamegraph

ceno_zkvm/examples/riscv_opcodes.rs

@@ -306,7 +306,7 @@ fn main() {
         println!(
             "riscv_opcodes::create_proof, instance_num_vars = {}, time = {}",
             instance_num_vars,
-            timer.elapsed().as_secs()
+            timer.elapsed()
         );
 
         let transcript = Transcript::new(b"riscv");

ceno_zkvm/src/instructions.rs

@@ -2,6 +2,7 @@ use std::mem::MaybeUninit;
 
 use ceno_emul::StepRecord;
 use ff_ext::ExtensionField;
+use multilinear_extensions::util::max_usable_threads;
 use rayon::{
     iter::{IndexedParallelIterator, ParallelIterator},
     slice::ParallelSlice,
@@ -47,8 +48,7 @@ pub trait Instruction<E: ExtensionField> {
         num_witin: usize,
         steps: Vec<StepRecord>,
     ) -> Result<(RowMajorMatrix<E::BaseField>, LkMultiplicity), ZKVMError> {
-        let nthreads =
-            std::env::var("RAYON_NUM_THREADS").map_or(8, |s| s.parse::<usize>().unwrap_or(8));
+        let nthreads = max_usable_threads();
         let num_instance_per_batch = if steps.len() > 256 {
             steps.len().div_ceil(nthreads)
         } else {

ceno_zkvm/src/scheme.rs

@@ -17,7 +17,7 @@ pub mod mock_prover;
 mod tests;
 
 #[derive(Clone)]
-pub struct ZKVMOpcodeProof<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> {
+pub struct ZKVMOpcodeProof<E: ExtensionField> {
     // TODO support >1 opcodes
     pub num_instances: usize,
 
@@ -39,8 +39,8 @@ pub struct ZKVMOpcodeProof<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>
     pub w_records_in_evals: Vec<E>,
     pub lk_records_in_evals: Vec<E>,
 
-    pub wits_commit: PCS::Commitment,
-    pub wits_opening_proof: PCS::Proof,
+    // pub wits_commit: PCS::Commitment,
+    // pub wits_opening_proof: PCS::Proof,
     pub wits_in_evals: Vec<E>,
 }
 
@@ -62,9 +62,9 @@ pub struct ZKVMTableProof<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>>
 
     pub fixed_in_evals: Vec<E>,
     pub fixed_opening_proof: Option<PCS::Proof>,
-    pub wits_commit: PCS::Commitment,
+    // pub wits_commit: PCS::Commitment,
     pub wits_in_evals: Vec<E>,
-    pub wits_opening_proof: PCS::Proof,
+    // pub wits_opening_proof: PCS::Proof,
 }
 
 /// each field will be interpret to (constant) polynomial
@@ -122,7 +122,7 @@ pub struct ZKVMProof<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> {
     pub raw_pi: Vec<Vec<E::BaseField>>,
     // the evaluation of raw_pi.
     pub pi_evals: Vec<E>,
-    opcode_proofs: BTreeMap<String, (usize, ZKVMOpcodeProof<E, PCS>)>,
+    opcode_proofs: BTreeMap<String, (usize, ZKVMOpcodeProof<E>)>,
     table_proofs: BTreeMap<String, (usize, ZKVMTableProof<E, PCS>)>,
 }
 

ceno_zkvm/src/scheme/prover.rs

@@ -6,7 +6,7 @@ use std::{
 
 use ff::Field;
 use itertools::{Itertools, izip};
-use mpcs::PolynomialCommitmentScheme;
+use mpcs::{CommitmentWithData, PolynomialCommitmentScheme};
 use multilinear_extensions::{
     mle::{IntoMLE, MultilinearExtension},
     util::ceil_log2,
@@ -99,8 +99,9 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
                     let witness = witness.into_mles();
                     commitments.insert(
                         circuit_name.clone(),
-                        PCS::batch_commit_and_write(&self.pk.pp, &witness, &mut transcript)
-                            .map_err(ZKVMError::PCSError)?,
+                        CommitmentWithData::<E, PCS>::default(),
+                        // PCS::batch_commit_and_write(&self.pk.pp, &witness, &mut transcript)
+                        //     .map_err(ZKVMError::PCSError)?,
                     );
                     witness
                 }
@@ -214,7 +215,7 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
         num_instances: usize,
         transcript: &mut Transcript<E>,
         challenges: &[E; 2],
-    ) -> Result<ZKVMOpcodeProof<E, PCS>, ZKVMError> {
+    ) -> Result<ZKVMOpcodeProof<E>, ZKVMError> {
         let cs = circuit_pk.get_cs();
         let next_pow2_instances = next_pow2_instance_padding(num_instances);
         let log2_num_instances = ceil_log2(next_pow2_instances);
@@ -604,22 +605,23 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
             witnesses.len(),
             input_open_point
         );
-        let wits_opening_proof = PCS::simple_batch_open(
-            pp,
-            &witnesses,
-            &wits_commit,
-            &input_open_point,
-            wits_in_evals.as_slice(),
-            transcript,
-        )
-        .map_err(ZKVMError::PCSError)?;
+        // let wits_opening_proof = PCS::simple_batch_open(
+        //     pp,
+        //     &witnesses,
+        //     &wits_commit,
+        //     &input_open_point,
+        //     wits_in_evals.as_slice(),
+        //     transcript,
+        // )
+        // .map_err(ZKVMError::PCSError)?;
+
         tracing::info!(
             "[opcode {}] build opening proof took {:?}",
             name,
             opening_dur.elapsed(),
         );
         exit_span!(pcs_open_span);
-        let wits_commit = PCS::get_pure_commitment(&wits_commit);
+        // let wits_commit = PCS::get_pure_commitment(&wits_commit);
 
         Ok(ZKVMOpcodeProof {
             num_instances,
@@ -634,8 +636,8 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
             r_records_in_evals,
             w_records_in_evals,
             lk_records_in_evals,
-            wits_commit,
-            wits_opening_proof,
+            // wits_commit,
+            // wits_opening_proof,
             wits_in_evals,
         })
     }
@@ -1038,22 +1040,24 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
 
         let pcs_opening = entered_span!("pcs_opening");
         let (fixed_opening_proof, fixed_commit) = if !fixed.is_empty() {
-            (
-                Some(
-                    PCS::simple_batch_open(
-                        pp,
-                        &fixed,
-                        circuit_pk.fixed_commit_wd.as_ref().unwrap(),
-                        &input_open_point,
-                        fixed_in_evals.as_slice(),
-                        transcript,
-                    )
-                    .map_err(ZKVMError::PCSError)?,
-                ),
-                Some(PCS::get_pure_commitment(
-                    circuit_pk.fixed_commit_wd.as_ref().unwrap(),
-                )),
+            let span = entered_span!("pcs_fixed_opening");
+            let fixed_opening_proof = PCS::simple_batch_open(
+                pp,
+                &fixed,
+                circuit_pk.fixed_commit_wd.as_ref().unwrap(),
+                &input_open_point,
+                fixed_in_evals.as_slice(),
+                transcript,
             )
+            .map_err(ZKVMError::PCSError)?;
+            exit_span!(span);
+
+            let span = entered_span!("pcs_fixed_commitment");
+            let fixed_commitment =
+                PCS::get_pure_commitment(circuit_pk.fixed_commit_wd.as_ref().unwrap());
+            exit_span!(span);
+
+            (Some(fixed_opening_proof), Some(fixed_commitment))
         } else {
             (None, None)
         };
@@ -1066,25 +1070,28 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
             fixed_in_evals,
             fixed_commit,
         );
-        let wits_opening_proof = PCS::simple_batch_open(
-            pp,
-            &witnesses,
-            &wits_commit,
-            &input_open_point,
-            wits_in_evals.as_slice(),
-            transcript,
-        )
-        .map_err(ZKVMError::PCSError)?;
+        let span = entered_span!("pcs_witin_opening");
+        // let wits_opening_proof = PCS::simple_batch_open(
+        //     pp,
+        //     &witnesses,
+        //     &wits_commit,
+        //     &input_open_point,
+        //     wits_in_evals.as_slice(),
+        //     transcript,
+        // )
+        // .map_err(ZKVMError::PCSError)?;
         exit_span!(pcs_opening);
-        let wits_commit = PCS::get_pure_commitment(&wits_commit);
-        tracing::debug!(
-            "[table {}] build opening proof for {} polys at {:?}: values = {:?}, commit = {:?}",
-            name,
-            witnesses.len(),
-            input_open_point,
-            wits_in_evals,
-            wits_commit,
-        );
+        let pcs_witin_commitment = entered_span!("pcs_witin_commitment");
+        // let wits_commit = PCS::get_pure_commitment(&wits_commit);
+        exit_span!(pcs_witin_commitment);
+        // tracing::debug!(
+        //     "[table {}] build opening proof for {} polys at {:?}: values = {:?}, commit = {:?}",
+        //     name,
+        //     witnesses.len(),
+        //     input_open_point,
+        //     wits_in_evals,
+        //     wits_commit,
+        // );
 
         Ok((
             ZKVMTableProof {
@@ -1099,8 +1106,8 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
                 fixed_opening_proof,
                 rw_hints_num_vars,
                 wits_in_evals,
-                wits_commit,
-                wits_opening_proof,
+                // wits_commit,
+                // wits_opening_proof,
             },
             pi_in_evals,
         ))

ceno_zkvm/src/scheme/verifier.rs

@@ -99,14 +99,14 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMVerifier<E, PCS>
             }
         }
 
-        for (_, (_, proof)) in vm_proof.opcode_proofs.iter() {
-            PCS::write_commitment(&proof.wits_commit, &mut transcript)
-                .map_err(ZKVMError::PCSError)?;
-        }
-        for (_, (_, proof)) in vm_proof.table_proofs.iter() {
-            PCS::write_commitment(&proof.wits_commit, &mut transcript)
-                .map_err(ZKVMError::PCSError)?;
-        }
+        // for (_, (_, proof)) in vm_proof.opcode_proofs.iter() {
+        //     PCS::write_commitment(&proof.wits_commit, &mut transcript)
+        //         .map_err(ZKVMError::PCSError)?;
+        // }
+        // for (_, (_, proof)) in vm_proof.table_proofs.iter() {
+        //     PCS::write_commitment(&proof.wits_commit, &mut transcript)
+        //         .map_err(ZKVMError::PCSError)?;
+        // }
 
         // alpha, beta
         let challenges = [
@@ -233,7 +233,7 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMVerifier<E, PCS>
         name: &str,
         vp: &PCS::VerifierParam,
         circuit_vk: &VerifyingKey<E, PCS>,
-        proof: &ZKVMOpcodeProof<E, PCS>,
+        proof: &ZKVMOpcodeProof<E>,
         pi: &[E],
         transcript: &mut Transcript<E>,
         num_product_fanin: usize,
@@ -463,15 +463,15 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMVerifier<E, PCS>
             proof.wits_in_evals.len(),
             input_opening_point
         );
-        PCS::simple_batch_verify(
-            vp,
-            &proof.wits_commit,
-            &input_opening_point,
-            &proof.wits_in_evals,
-            &proof.wits_opening_proof,
-            transcript,
-        )
-        .map_err(ZKVMError::PCSError)?;
+        // PCS::simple_batch_verify(
+        //     vp,
+        //     &proof.wits_commit,
+        //     &input_opening_point,
+        //     &proof.wits_in_evals,
+        //     &proof.wits_opening_proof,
+        //     transcript,
+        // )
+        // .map_err(ZKVMError::PCSError)?;
 
         Ok(input_opening_point)
     }
@@ -761,23 +761,23 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMVerifier<E, PCS>
             circuit_vk.fixed_commit,
         );
 
-        PCS::simple_batch_verify(
-            vp,
-            &proof.wits_commit,
-            &input_opening_point,
-            &proof.wits_in_evals,
-            &proof.wits_opening_proof,
-            transcript,
-        )
-        .map_err(ZKVMError::PCSError)?;
-        tracing::debug!(
-            "[table {}] verified opening proof for {} polys at {:?}: values = {:?}, commit = {:?}",
-            name,
-            proof.wits_in_evals.len(),
-            input_opening_point,
-            proof.wits_in_evals,
-            proof.wits_commit
-        );
+        // PCS::simple_batch_verify(
+        //     vp,
+        //     &proof.wits_commit,
+        //     &input_opening_point,
+        //     &proof.wits_in_evals,
+        //     &proof.wits_opening_proof,
+        //     transcript,
+        // )
+        // .map_err(ZKVMError::PCSError)?;
+        // tracing::debug!(
+        //     "[table {}] verified opening proof for {} polys at {:?}: values = {:?}, commit = {:?}",
+        //     name,
+        //     proof.wits_in_evals.len(),
+        //     input_opening_point,
+        //     proof.wits_in_evals,
+        //     proof.wits_commit
+        // );
 
         Ok(input_opening_point)
     }

ceno_zkvm/src/tables/mod.rs

@@ -4,6 +4,7 @@ use crate::{
 };
 use ff::Field;
 use ff_ext::ExtensionField;
+use multilinear_extensions::util::max_usable_threads;
 use rayon::iter::{IndexedParallelIterator, ParallelIterator};
 use std::{collections::HashMap, mem::MaybeUninit};
 mod range;
@@ -49,8 +50,7 @@ pub trait TableCircuit<E: ExtensionField> {
         // Fill the padding with zeros, if any.
         let num_padding_instances = table.num_padding_instances();
         if num_padding_instances > 0 {
-            let nthreads =
-                std::env::var("RAYON_NUM_THREADS").map_or(8, |s| s.parse::<usize>().unwrap_or(8));
+            let nthreads = max_usable_threads();
             let padding_instance = vec![MaybeUninit::new(E::BaseField::ZERO); num_witin];
             let num_padding_instance_per_batch = if num_padding_instances > 256 {
                 num_padding_instances.div_ceil(nthreads)

mpcs/Cargo.toml

@@ -22,6 +22,9 @@ plonky2.workspace = true
 poseidon.workspace = true
 rand.workspace = true
 rand_chacha.workspace = true
+tracing.workspace = true
+tracing-flame.workspace = true
+tracing-subscriber.workspace = true
 rayon = { workspace = true, optional = true }
 serde.workspace = true
 transcript = { path = "../transcript" }