metrics

Author: zlangley

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

@@ -11,15 +11,11 @@ use p3_field::{ExtensionField, Field};
 use p3_maybe_rayon::prelude::*;
 use thiserror::Error;
 
-use crate::{
-    gkr::types::{GkrArtifact, GkrBatchProof, GkrMask, Layer},
-    poly::{
-        multi::{hypercube_eq, Mle, MultivariatePolyOracle},
-        uni::{random_linear_combination, UnivariatePolynomial},
-    },
-    sumcheck,
-    sumcheck::SumcheckArtifacts,
-};
+use crate::{gkr::types::{GkrArtifact, GkrBatchProof, GkrMask, Layer}, poly::{
+    multi::{hypercube_eq, Mle, MultivariatePolyOracle},
+    uni::{random_linear_combination, UnivariatePolynomial},
+}, sumcheck, sumcheck::SumcheckArtifacts};
+use crate::utils::metrics_span;
 
 /// For a given `y`, stores evaluations of [hypercube_eq](x, y) on all 2^{n-1} boolean hypercube
 /// points of the form `x = (0, x_2, ..., x_n)`.
@@ -348,12 +344,14 @@ pub fn prove_batch<F: Field, EF: ExtensionField<F>>(
                 constant_poly_oracles,
                 ..
             },
-        ) = sumcheck::prove_batch(
-            sumcheck_claims,
-            sumcheck_oracles,
-            sumcheck_alpha,
-            challenger,
-        );
+        ) = metrics_span("sumcheck_prove_batch_ms", || {
+            sumcheck::prove_batch(
+                sumcheck_claims,
+                sumcheck_oracles,
+                sumcheck_alpha,
+                challenger,
+            )
+        });
 
         sumcheck_proofs.push(sumcheck_proof);
 

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

@@ -28,6 +28,7 @@ use crate::{
     },
     rap::PermutationAirBuilderWithExposedValues,
 };
+use crate::utils::metrics_span;
 
 pub struct GkrLogUpPhase<F, EF, Challenger> {
     // FIXME: USE THIS IN POW
@@ -132,22 +133,26 @@ where
             .map(|gkr_instance| gkr_instance.build_gkr_input_layer(alpha))
             .collect();
 
-        let (gkr_proof, gkr_artifact) = gkr::prove_batch(challenger, input_layers);
+        let (gkr_proof, gkr_artifact) = metrics_span("gkr_prove_batch_ms", || {
+            gkr::prove_batch(challenger, input_layers)
+        });
 
         // Generate after challenge trace, exposed values, and MLE claims for each instance.
         let GkrAuxData {
             after_challenge_trace_per_air,
             exposed_values_per_air,
             count_mle_claims_per_instance,
             sigma_mle_claims_per_instance,
-        } = Self::generate_aux_per_air(
-            challenger,
-            interactions_per_air,
-            trace_view_per_air,
-            alpha,
-            &gkr_instances,
-            &gkr_artifact,
-        );
+        } = metrics_span("generate_perm_trace_time_ms", || {
+            Self::generate_aux_per_air(
+                challenger,
+                interactions_per_air,
+                trace_view_per_air,
+                alpha,
+                &gkr_instances,
+                &gkr_artifact,
+            )
+        });
 
         let mut challenges = vec![beta, alpha];
         challenges.extend_from_slice(&gkr_artifact.ood_point);