feat: add granular prove metrics and docs (#17)

* feat: add granular prove metrics and docs

* chore: turn off feature

* chore: make serialize_metric_snapshot public

crates/stark-backend/src/engine.rs

@@ -145,15 +145,7 @@ pub trait StarkEngine<SC: StarkGenericConfig> {
         };
 
         let mut challenger = self.new_challenger();
-
-        #[cfg(feature = "bench-metrics")]
-        let prove_start = std::time::Instant::now();
-        let _proof = prover.prove(&mut challenger, pk, proof_input);
-        #[cfg(feature = "bench-metrics")]
-        metrics::gauge!("stark_prove_excluding_trace_time_ms")
-            .set(prove_start.elapsed().as_millis() as f64);
-
-        _proof
+        prover.prove(&mut challenger, pk, proof_input)
     }
 
     fn verify(

crates/stark-backend/src/interaction/stark_log_up.rs

@@ -23,6 +23,7 @@ use crate::{
     parizip,
     prover::PairTraceView,
     rap::PermutationAirBuilderWithExposedValues,
+    utils::metrics_span,
 };
 
 #[derive(Default)]
@@ -105,15 +106,14 @@ where
         let challenges: [Challenge; STARK_LU_NUM_CHALLENGES] =
             array::from_fn(|_| challenger.sample_ext_element::<Challenge>());
 
-        let after_challenge_trace_per_air = tracing::info_span!("generate permutation traces")
-            .in_scope(|| {
-                Self::generate_after_challenge_traces_per_air(
-                    &challenges,
-                    constraints_per_air,
-                    rap_pk_per_air,
-                    trace_view_per_air,
-                )
-            });
+        let after_challenge_trace_per_air = metrics_span("generate_perm_trace_time_ms", || {
+            Self::generate_after_challenge_traces_per_air(
+                &challenges,
+                constraints_per_air,
+                rap_pk_per_air,
+                trace_view_per_air,
+            )
+        });
         let cumulative_sum_per_air = Self::extract_cumulative_sums(&after_challenge_trace_per_air);
 
         // Challenger needs to observe what is exposed (cumulative_sums)

crates/stark-backend/src/prover/mod.rs

@@ -27,6 +27,7 @@ use crate::{
         types::{AirProofData, Commitments, Proof, ProofInput},
     },
     rap::AnyRap,
+    utils::metrics_span,
 };
 
 pub mod helper;
@@ -81,6 +82,8 @@ impl<'c, SC: StarkGenericConfig> MultiTraceStarkProver<'c, SC> {
         mpk: &'a MultiStarkProvingKey<SC>,
         proof_input: ProofInput<SC>,
     ) -> Proof<SC> {
+        #[cfg(feature = "bench-metrics")]
+        let start = std::time::Instant::now();
         assert!(mpk.validate(&proof_input), "Invalid proof input");
         let pcs = self.config.pcs();
         let rap_phase_seq = self.config.rap_phase_seq();
@@ -117,16 +120,17 @@ impl<'c, SC: StarkGenericConfig> MultiTraceStarkProver<'c, SC> {
         challenger.observe_slice(&preprocessed_commits);
 
         // Commit all common main traces in a commitment. Traces inside are ordered by AIR id.
-        let (common_main_trace_views, common_main_prover_data) = {
-            let committer = TraceCommitter::<SC>::new(pcs);
-            let (trace_views, traces): (Vec<_>, Vec<_>) = common_main_per_air
-                .iter()
-                .filter_map(|cm: &Option<RowMajorMatrix<_>>| cm.as_ref())
-                .map(|m| (m.as_view(), m.clone()))
-                .unzip();
-
-            (trace_views, committer.commit(traces))
-        };
+        let (common_main_trace_views, common_main_prover_data) =
+            metrics_span("main_trace_commit_time_ms", || {
+                let committer = TraceCommitter::<SC>::new(pcs);
+                let (trace_views, traces): (Vec<_>, Vec<_>) = common_main_per_air
+                    .iter()
+                    .filter_map(|cm: &Option<RowMajorMatrix<_>>| cm.as_ref())
+                    .map(|m| (m.as_view(), m.clone()))
+                    .unzip();
+
+                (trace_views, committer.commit(traces))
+            });
 
         // Commitments order:
         // - for each air:
@@ -232,8 +236,9 @@ impl<'c, SC: StarkGenericConfig> MultiTraceStarkProver<'c, SC> {
 
         // Commit to permutation traces: this means only 1 challenge round right now
         // One shared commit for all permutation traces
-        let perm_prover_data = tracing::info_span!("commit to permutation traces")
-            .in_scope(|| commit_perm_traces::<SC>(pcs, perm_trace_per_air, &domain_per_air));
+        let perm_prover_data = metrics_span("perm_trace_commit_time_ms", || {
+            commit_perm_traces::<SC>(pcs, perm_trace_per_air, &domain_per_air)
+        });
 
         // Challenger observes commitment if exists
         if let Some(data) = &perm_prover_data {
@@ -262,18 +267,24 @@ impl<'c, SC: StarkGenericConfig> MultiTraceStarkProver<'c, SC> {
             .flatten()
             .chain(iter::once(common_main_prover_data))
             .collect();
-        prove_raps_with_committed_traces(
-            pcs,
-            challenger,
-            mpk,
-            &main_prover_data,
-            perm_prover_data,
-            exposed_values_after_challenge,
-            quotient_data,
-            domain_per_air,
-            pvs_per_air,
-            rap_phase_seq_proof,
-        )
+        let proof = metrics_span("pcs_opening_time_ms", || {
+            prove_raps_with_committed_traces(
+                pcs,
+                challenger,
+                mpk,
+                &main_prover_data,
+                perm_prover_data,
+                exposed_values_after_challenge,
+                quotient_data,
+                domain_per_air,
+                pvs_per_air,
+                rap_phase_seq_proof,
+            )
+        });
+        #[cfg(feature = "bench-metrics")]
+        ::metrics::gauge!("stark_prove_excluding_trace_time_ms")
+            .set(start.elapsed().as_millis() as f64);
+        proof
     }
 }
 

crates/stark-backend/src/prover/opener.rs

@@ -23,7 +23,7 @@ impl<'pcs, SC: StarkGenericConfig> OpeningProver<'pcs, SC> {
     /// - main trace matrices can have multiple commitments
     /// - for each after_challenge phase, all matrices in the phase share a commitment
     /// - quotient poly chunks are all committed together
-    #[instrument(name = "FRI opening proofs", skip_all)]
+    #[instrument(name = "PCS opening proofs", skip_all)]
     pub fn open(
         &self,
         challenger: &mut SC::Challenger,

crates/stark-backend/src/prover/trace.rs

@@ -16,6 +16,7 @@ use crate::{
     keygen::view::MultiStarkProvingKeyView,
     prover::quotient::{helper::QuotientVkDataHelper, ProverQuotientData, QuotientCommitter},
     rap::AnyRap,
+    utils::metrics_span,
 };
 
 #[allow(clippy::too_many_arguments)]
@@ -46,10 +47,13 @@ pub(super) fn commit_quotient_traces<'a, SC: StarkGenericConfig>(
         .iter()
         .map(|pk| pk.get_quotient_vk_data())
         .collect_vec();
-    let quotient_values =
-        quotient_committer.quotient_values(raps, &qvks, &trace_views, public_values_per_air);
-    // Commit to quotient polynomias. One shared commit for all quotient polynomials
-    quotient_committer.commit(quotient_values)
+    let quotient_values = metrics_span("quotient_poly_compute_time_ms", || {
+        quotient_committer.quotient_values(raps, &qvks, &trace_views, public_values_per_air)
+    });
+    // Commit to quotient polynomials. One shared commit for all quotient polynomials
+    metrics_span("quotient_poly_commit_time_ms", || {
+        quotient_committer.commit(quotient_values)
+    })
 }
 
 fn create_trace_view_per_air<'a, SC: StarkGenericConfig>(

crates/stark-backend/src/utils.rs

@@ -1,3 +1,6 @@
+use std::borrow::Cow;
+
+use cfg_if::cfg_if;
 use p3_field::Field;
 use tracing::instrument;
 
@@ -39,6 +42,23 @@ pub fn disable_debug_builder() {
     });
 }
 
+/// A span that will run the given closure `f`,
+/// and emit a metric with the given `name` using [`gauge`](metrics::gauge)
+/// when the feature `"bench-metrics"` is enabled.
+#[allow(unused_variables)]
+pub fn metrics_span<R, F: FnOnce() -> R>(name: impl Into<Cow<'static, str>>, f: F) -> R {
+    cfg_if! {
+        if #[cfg(feature = "bench-metrics")] {
+            let start = std::time::Instant::now();
+            let res = f();
+            metrics::gauge!(name.into()).set(start.elapsed().as_millis() as f64);
+            res
+        } else {
+            f()
+        }
+    }
+}
+
 #[macro_export]
 #[cfg(feature = "parallel")]
 macro_rules! parizip {

crates/stark-sdk/src/bench/mod.rs

@@ -93,7 +93,7 @@ fn serialize_metric(ckey: CompositeKey, value: DebugValue) -> serde_json::Value
 ///   ...
 /// }
 ///
-fn serialize_metric_snapshot(snapshot: Snapshot) -> serde_json::Value {
+pub fn serialize_metric_snapshot(snapshot: Snapshot) -> serde_json::Value {
     let mut ret = BTreeMap::<_, Vec<serde_json::Value>>::new();
     for (ckey, _, _, value) in snapshot.into_vec() {
         match ckey.kind() {

docs/README.md

@@ -2,3 +2,4 @@
 
 - [STARK Backend](./stark-backend.md)
   - [AIR Interactions](./interactions.md)
+  - [Metrics](./metrics.md): Guide to metrics collected by the prover.

docs/metrics.md

@@ -0,0 +1,16 @@
+# Metrics
+
+We use the [`metrics`](https://docs.rs/metrics/latest/metrics/) crate to collect metrics for the STARK prover. We refer to [reth docs](https://github.com/paradigmxyz/reth/blob/main/docs/design/metrics.md) for more guidelines on how to use metrics.
+
+Each invocation of `MultiTraceStarkProver::prove` will collect the following metrics, all as `Gauge`s. We use gauge instead of histogram because these metrics are not frequently sampled and we care about the exact value. Any application that uses this backend is responsible for adding additional namespace labels if they wish to distinguish between different proof invocations.
+
+- `stark_prove_excluding_trace_time_ms`: The total elapsed time in milliseconds of `prove`. This excludes the main trace generation because that is not done by `stark-backend`.
+
+The following metrics comprise the main breakdown of the components of `prove`. They are disjoint and _expected_ to sum up to almost `stark_prove_excluding_trace_time_ms` (if it does not, an issue should be opened as it means there is an unexpected source of slowdown).
+
+- `main_trace_commit_time_ms`: The time to commit the main trace matrices, depending on the PCS.
+- `generate_perm_trace_time_ms`: When FRI is used for the log up argument, this is the time to generate the permutation trace.
+- `perm_trace_commit_time_ms`: When FRI is used for the log up argument, this is the time to commit the permutation trace.
+- `quotient_poly_compute_time_ms`: The time to compute the quotient polynomials from the trace matrices according to AIR constraints.
+- `quotient_poly_commit_time_ms`: The time to commit the quotient polynomials.
+- `pcs_opening_time_ms`: The time to compute all polynomial commitment scheme (PCS) opening proofs necessary for the proof. Currently the PCS is FRI over a base field with high `2`-adicity.