Add tracing instrumentation for prover stages

Cargo.toml

@@ -9,6 +9,7 @@ rust-version = "1.88"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+tracing = "0.1"
 serde = { version = "1", features = ["derive"] }
 bincode = { version = "2", features = ["serde"] }
 p3-air = { git = "https://github.com/Plonky3/Plonky3", rev = "e9d75614dd6816f9b5dbb4413c69be63536efd64" }

src/lookup.rs

@@ -126,16 +126,20 @@ impl Lookup<Val> {
         // Flatten the references serially first so the parallel map operates
         // on an indexed slice and `collect` can write straight into the
         // output Vec without tree-reducing worker buffers.
+        let _g = tracing::info_span!("stark/lookup_messages").entered();
         let flat: Vec<&Self> = lookups.iter().flatten().flatten().collect();
         let messages: Vec<ExtVal> = flat
             .par_iter()
             .map(|lookup| lookup.compute_message(lookup_challenge, fingerprint_challenge))
             .collect();
+        drop(_g);
 
         // Compute the inverses of all messages in batch.
-        let messages_inverses = batch_multiplicative_inverse(&messages);
+        let messages_inverses = tracing::info_span!("stark/batch_inverse")
+            .in_scope(|| batch_multiplicative_inverse(&messages));
 
         // Compute and collect intermediate accumulators and traces.
+        let _g = tracing::info_span!("stark/lookup_traces").entered();
         let mut intermediate_accumulators = Vec::with_capacity(lookups.len());
         let mut traces = Vec::with_capacity(lookups.len());
         let mut offset = 0;
@@ -174,6 +178,7 @@ impl Lookup<Val> {
             intermediate_accumulators.push(accumulator);
             traces.push(trace);
         }
+        drop(_g);
         (traces, intermediate_accumulators)
     }
 

src/prover.rs

@@ -240,6 +240,7 @@ impl<A: BaseAir<Val> + for<'a> Air<ProverConstraintFolder<'a>>> System<A> {
     ///
     /// Each claim is a slice of field elements that is observed by the challenger
     /// before lookup challenges are sampled, binding the proof to the claimed values.
+    #[tracing::instrument(level = "info", skip_all, name = "stark/prove")]
     pub fn prove_multiple_claims(
         &self,
         fri_parameters: FriParameters,
@@ -254,6 +255,7 @@ impl<A: BaseAir<Val> + for<'a> Air<ProverConstraintFolder<'a>>> System<A> {
 
         // Cost: "Stage 1 commit" — coset LDE (FFT) of each trace from n_i to
         // n_i·B rows, then Merkle tree. FFT work: Σ w_i · n_i · B · log₂(n_i·B).
+        let _g = tracing::info_span!("stark/stage1_commit").entered();
         let mut log_degrees = vec![];
         let evaluations = witness.traces.into_iter().map(|trace| {
             let degree = trace.height();
@@ -265,6 +267,7 @@ impl<A: BaseAir<Val> + for<'a> Air<ProverConstraintFolder<'a>>> System<A> {
         });
         let (stage_1_trace_commit, stage_1_trace_data) =
             <Pcs as PcsTrait<ExtVal, Challenger>>::commit(pcs, evaluations);
+        drop(_g);
 
         if let Some(commit) = &self.preprocessed_commit {
             challenger.observe(commit);
@@ -299,14 +302,18 @@ impl<A: BaseAir<Val> + for<'a> Air<ProverConstraintFolder<'a>>> System<A> {
 
         // Cost: "Lookup trace construction" — fingerprint (Horner), batch
         // inversion, and accumulator update. Total: Σ n_i·L_i extension field ops.
+        let _g = tracing::info_span!("stark/lookup_construction").entered();
         let (stage_2_traces, intermediate_accumulators) = Lookup::stage_2_traces(
             &witness.lookups,
             lookup_argument_challenge,
             &fingerprint_challenge,
             acc,
         );
+        drop(_g);
+
         // Cost: "Stage 2 commit" — LDE + Merkle for flattened extension traces.
         // FFT work: Σ w2_i · D · n_i · B · log₂(n_i·B).
+        let _g = tracing::info_span!("stark/stage2_commit").entered();
         let evaluations = stage_2_traces.into_iter().map(|trace| {
             let degree = trace.height();
             let trace_domain =
@@ -315,6 +322,7 @@ impl<A: BaseAir<Val> + for<'a> Air<ProverConstraintFolder<'a>>> System<A> {
         });
         let (stage_2_trace_commit, stage_2_trace_data) =
             <Pcs as PcsTrait<ExtVal, Challenger>>::commit(pcs, evaluations);
+        drop(_g);
         challenger.observe(stage_2_trace_commit.clone());
 
         // generate constraint challenge
@@ -323,6 +331,7 @@ impl<A: BaseAir<Val> + for<'a> Air<ProverConstraintFolder<'a>>> System<A> {
         // Cost: "Quotient computation and commit" — constraint evaluation on the
         // quotient domain (Σ n_i·q_i·eval_cost(k_i)) plus LDE + Merkle of the
         // quotient sub-polynomials (Σ q_i·D·n_i·B·log₂(n_i·B)).
+        let _g = tracing::info_span!("stark/quotient").entered();
         debug_assert_eq!(intermediate_accumulators.len(), self.circuits.len());
         debug_assert_eq!(log_degrees.len(), self.circuits.len());
         let mut quotient_degrees = vec![];
@@ -407,6 +416,7 @@ impl<A: BaseAir<Val> + for<'a> Air<ProverConstraintFolder<'a>>> System<A> {
         let (quotient_commit, quotient_data) =
             <Pcs as PcsTrait<ExtVal, Challenger>>::commit(pcs, quotient_evaluations);
         challenger.observe(quotient_commit.clone());
+        drop(_g);
 
         // save the commitments
         let commitments = Commitments {
@@ -417,6 +427,7 @@ impl<A: BaseAir<Val> + for<'a> Air<ProverConstraintFolder<'a>>> System<A> {
 
         // Cost: "FRI opening" — barycentric interpolation (Σ n_i·B·W_i),
         // FRI folding (≈ H), and FRI queries (Q·R·log₂ H hash ops).
+        let _g = tracing::info_span!("stark/fri_open").entered();
         let zeta: ExtVal = challenger.sample_algebra_element();
         let mut round0_openings = vec![];
         let mut round1_openings = vec![];
@@ -446,6 +457,7 @@ impl<A: BaseAir<Val> + for<'a> Air<ProverConstraintFolder<'a>>> System<A> {
             rounds.push((key.preprocessed_data.as_ref().unwrap(), round0_openings));
         }
         let (opened_values, opening_proof) = pcs.open(rounds, &mut challenger);
+        drop(_g);
         let mut opened_values_iter = opened_values.into_iter();
         let stage_1_opened_values = opened_values_iter.next().unwrap();
         let stage_2_opened_values = opened_values_iter.next().unwrap();