prof

Author: zlangley

crates/stark-backend/src/air_builders/debug/check_constraints.rs

@@ -100,6 +100,7 @@ pub fn check_constraints<R, SC>(
     });
 }
 
+#[allow(dead_code)]
 fn check_gkr_log_up_adapter_constraints_for_row<SC: StarkGenericConfig>(
     after_challenge: &[RowMajorMatrixView<SC::Challenge>],
     challenges: &[Vec<SC::Challenge>],

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

@@ -11,11 +11,16 @@ 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::utils::metrics_span;
+use crate::utils::{metrics_span, metrics_span_increment};
+use crate::{
+    gkr::types::{GkrArtifact, GkrBatchProof, GkrMask, Layer},
+    poly::{
+        multi::{Mle, MultivariatePolyOracle},
+        uni::{random_linear_combination, UnivariatePolynomial},
+    },
+    sumcheck,
+    sumcheck::SumcheckArtifacts,
+};
 
 /// 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)`.
@@ -279,10 +284,12 @@ pub fn prove_batch<F: Field, EF: ExtensionField<F>>(
     let n_layers = *n_layers_by_instance.iter().max().unwrap();
 
     // Evaluate all instance circuits and collect the layer values.
-    let mut layers_by_instance: Vec<_> = input_layer_by_instance
-        .into_par_iter()
-        .map(|input_layer| gen_layers(input_layer).into_iter().rev())
-        .collect();
+    let mut layers_by_instance: Vec<_> = metrics_span("gkr_gen_layers_ms", || {
+        input_layer_by_instance
+            .into_par_iter()
+            .map(|input_layer| gen_layers(input_layer).into_iter().rev())
+            .collect()
+    });
 
     let mut output_claims_by_instance = vec![None; n_instances];
     let mut layer_masks_by_instance = vec![vec![]; n_instances];
@@ -297,6 +304,7 @@ pub fn prove_batch<F: Field, EF: ExtensionField<F>>(
         output_instances_by_layer[output_layer].push(instance);
     }
 
+    #[allow(clippy::needless_range_loop)]
     for layer in 0..n_layers {
         // Check all the instances for output layers.
         for &instance in &output_instances_by_layer[layer] {
@@ -344,7 +352,7 @@ pub fn prove_batch<F: Field, EF: ExtensionField<F>>(
                 constant_poly_oracles,
                 ..
             },
-        ) = metrics_span("sumcheck_prove_batch_ms", || {
+        ) = metrics_span_increment("sumcheck_prove_batch_ms", || {
             sumcheck::prove_batch(
                 sumcheck_claims,
                 sumcheck_oracles,

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

@@ -5,7 +5,6 @@ use openvm_stark_sdk::{
     config::baby_bear_blake3::default_engine, engine::StarkEngine, utils::create_seeded_rng,
 };
 use p3_baby_bear::BabyBear;
-use p3_field::FieldAlgebra;
 use rand::Rng;
 
 use crate::{

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

@@ -19,7 +19,7 @@ use crate::{
     },
     parizip,
     rap::PermutationAirBuilderWithExposedValues,
-    utils::metrics_span,
+    utils::{metrics_span, metrics_span_increment},
 };
 
 pub struct FriLogUpPhase<F, Challenge, Challenger> {
@@ -114,14 +114,15 @@ where
         let challenges: [Challenge; STARK_LU_NUM_CHALLENGES] =
             array::from_fn(|_| challenger.sample_ext_element::<Challenge>());
 
-        let after_challenge_trace_per_air = metrics_span("generate_perm_trace_time_ms", || {
-            Self::generate_after_challenge_traces_per_air(
-                &challenges,
-                interactions_per_air,
-                params_per_air,
-                trace_view_per_air,
-            )
-        });
+        let after_challenge_trace_per_air =
+            metrics_span_increment("generate_perm_trace_time_ms", || {
+                Self::generate_after_challenge_traces_per_air(
+                    &challenges,
+                    interactions_per_air,
+                    params_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)
@@ -248,7 +249,7 @@ where
         parizip!(interactions_per_air, trace_view_per_air, params_per_air)
             .map(|(interactions, trace_view, params)| {
                 Self::generate_after_challenge_trace(
-                    &interactions,
+                    interactions,
                     trace_view,
                     challenges,
                     &params.interaction_partitions,

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

@@ -13,11 +13,11 @@ use p3_field::{
 use p3_matrix::{dense::DenseMatrix, Matrix};
 use p3_maybe_rayon::prelude::*;
 use p3_util::{log2_ceil_usize, log2_strict_usize};
-use rayon::iter::IntoParallelRefIterator;
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
 
 use crate::interaction::SymbolicInteraction;
+use crate::utils::metrics_span;
 use crate::{
     air_builders::symbolic::SymbolicConstraints,
     gkr,
@@ -28,7 +28,6 @@ use crate::{
     },
     rap::PermutationAirBuilderWithExposedValues,
 };
-use crate::utils::metrics_span;
 
 pub struct GkrLogUpPhase<F, EF, Challenger> {
     // FIXME: USE THIS IN POW
@@ -113,7 +112,7 @@ where
         _params_per_air: &[&Self::PartialProvingKey],
         trace_view_per_air: &[PairTraceView<'_, F>],
     ) -> Option<(Self::PartialProof, RapPhaseProverData<EF>)> {
-        let all_interactions = interactions_per_air.iter().cloned().flatten().collect_vec();
+        let all_interactions = interactions_per_air.iter().flatten().collect_vec();
         if all_interactions.is_empty() {
             return None;
         }
@@ -143,16 +142,14 @@ where
             exposed_values_per_air,
             count_mle_claims_per_instance,
             sigma_mle_claims_per_instance,
-        } = 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,
-            )
-        });
+        } = 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);

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

@@ -12,13 +12,18 @@ use p3_util::log2_strict_usize;
 
 use crate::{
     air_builders::symbolic::symbolic_expression::{SymbolicEvaluator, SymbolicExpression},
-    gkr::{GkrArtifact, Layer, Layer::LogUpGeneric},
+    gkr::{
+        GkrArtifact,
+        Layer::{self, LogUpGeneric},
+    },
     interaction::{
         gkr_log_up::{num_interaction_dimensions, GkrAuxData, GkrLogUpPhase},
         trace::Evaluator,
         PairTraceView, SymbolicInteraction,
     },
+    parizip,
     poly::multi::{hypercube_eq_over_y, Mle},
+    utils::metrics_span,
 };
 
 /// A struct that holds the interaction-related data for a single GKR instance:
@@ -148,7 +153,7 @@ where
                 } else {
                     Some(GkrLogUpInstance::from_interactions(
                         trace_view,
-                        &interactions,
+                        interactions,
                         beta_pows,
                     ))
                 }
@@ -181,29 +186,30 @@ where
             })
             .collect_vec();
 
-        let results: Vec<_> = trace_view_per_air_filtered
-            .par_iter()
-            .zip(gkr_instances.par_iter())
-            .zip(gkr_artifact.n_variables_by_instance.par_iter())
-            .map(|((trace_view, gkr_instance), &n_vars)| {
-                let height = trace_view.partitioned_main[0].height();
-                let log_height = log2_strict_usize(height);
-
-                let instance_ood = &ood_point[ood_point.len() - n_vars..];
-                let (_, r) = instance_ood.split_at(n_vars - log_height);
-                debug_assert_eq!(r.len(), log_height);
-
-                let (after_challenge_trace, exposed_values, count_mle_claims, sigma_mle_claims) =
-                    Self::generate_aux(gkr_instance, &gamma_pows, r);
-
-                (
-                    after_challenge_trace,
-                    exposed_values,
-                    count_mle_claims,
-                    sigma_mle_claims,
-                )
-            })
-            .collect();
+        let results: Vec<_> = parizip!(
+            trace_view_per_air_filtered,
+            gkr_instances,
+            &gkr_artifact.n_variables_by_instance
+        )
+        .map(|(trace_view, gkr_instance, &n_vars)| {
+            let height = trace_view.partitioned_main[0].height();
+            let log_height = log2_strict_usize(height);
+
+            let instance_ood = &ood_point[ood_point.len() - n_vars..];
+            let (_, r) = instance_ood.split_at(n_vars - log_height);
+            debug_assert_eq!(r.len(), log_height);
+
+            let (after_challenge_trace, exposed_values, count_mle_claims, sigma_mle_claims) =
+                Self::generate_aux(gkr_instance, &gamma_pows, r);
+
+            (
+                after_challenge_trace,
+                exposed_values,
+                count_mle_claims,
+                sigma_mle_claims,
+            )
+        })
+        .collect();
 
         let mut results_iter = results.into_iter();
 
@@ -257,32 +263,38 @@ where
             .map(|row| Mle::eval_slice(row, r))
             .collect();
 
-        let s_at_rows: Vec<EF> = gkr_instance
-            .counts
-            .par_row_slices()
-            .zip(gkr_instance.sigmas.par_row_slices())
-            .map(|(count_row, sigma_row)| {
-                itertools::interleave(count_row, sigma_row)
-                    .zip(gamma_pows)
-                    .fold(EF::ZERO, |acc, (&val, &gamma_pow)| acc + gamma_pow * val)
-            })
-            .collect();
-
-        let eqs_at_r = hypercube_eq_over_y(r);
-
-        let mut partial_sum = EF::ZERO;
-        let mut after_challenge_trace_data = Vec::with_capacity(eqs_at_r.len() * 3);
-        for (eq_at_r, s_at_row) in eqs_at_r.iter().zip(s_at_rows.iter()) {
-            after_challenge_trace_data.push(*eq_at_r);
-            after_challenge_trace_data.push(*s_at_row);
-            after_challenge_trace_data.push(partial_sum);
+        let (after_challenge_trace, exposed_values) =
+            metrics_span("generate_perm_trace_time_ms", || {
+                let s_at_rows: Vec<EF> = gkr_instance
+                    .counts
+                    .par_row_slices()
+                    .zip(gkr_instance.sigmas.par_row_slices())
+                    .map(|(count_row, sigma_row)| {
+                        itertools::interleave(count_row, sigma_row)
+                            .zip(gamma_pows)
+                            .fold(EF::ZERO, |acc, (&val, &gamma_pow)| acc + gamma_pow * val)
+                    })
+                    .collect();
+
+                // TODO: Precompute these per height.
+                let eqs_at_r = hypercube_eq_over_y(r);
+
+                let mut partial_sum = EF::ZERO;
+                let mut after_challenge_trace_data = Vec::with_capacity(eqs_at_r.len() * 3);
+                for (eq_at_r, s_at_row) in eqs_at_r.iter().zip(s_at_rows.iter()) {
+                    after_challenge_trace_data.push(*eq_at_r);
+                    after_challenge_trace_data.push(*s_at_row);
+                    after_challenge_trace_data.push(partial_sum);
+
+                    partial_sum += *eq_at_r * *s_at_row;
+                }
+                let after_challenge_trace = RowMajorMatrix::new(after_challenge_trace_data, 3);
+                (after_challenge_trace, vec![partial_sum])
+            });
 
-            partial_sum += *eq_at_r * *s_at_row;
-        }
-        let after_challenge_trace = RowMajorMatrix::new(after_challenge_trace_data, 3);
         (
             after_challenge_trace,
-            vec![partial_sum],
+            exposed_values,
             count_mle_claims,
             sigma_mle_claims,
         )

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

@@ -1,16 +1,17 @@
 use p3_field::FieldAlgebra;
+use std::borrow::Borrow;
 
 use super::Interaction;
 
 /// Returns [beta^0, beta^1, ..., beta^{max_num_fields}]
 /// where max_num_fields is the maximum length of `fields` in any interaction.
 pub fn generate_betas<AF: FieldAlgebra, E>(
     beta: AF,
-    all_interactions: &[Interaction<E>],
+    all_interactions: &[impl Borrow<Interaction<E>>],
 ) -> Vec<AF> {
     let max_fields_len = all_interactions
         .iter()
-        .map(|interaction| interaction.message.len())
+        .map(|interaction| interaction.borrow().message.len())
         .max()
         .unwrap_or(0);
 

crates/stark-backend/src/poly/uni.rs

@@ -32,7 +32,7 @@ impl<F: Field> UnivariatePolynomial<F> {
     }
 
     pub fn evaluate_at_zero(&self) -> F {
-        self.coeffs.get(0).copied().unwrap_or(F::ZERO)
+        *self.coeffs.first().unwrap_or(&F::ZERO)
     }
 
     pub fn evaluate_at_one(&self) -> F {

crates/stark-backend/src/sumcheck.rs

@@ -120,6 +120,7 @@ where
             .zip(this_round_polys.par_iter())
             .for_each(|(claim, round_poly)| *claim = round_poly.evaluate(challenge));
 
+        // TODO: This can be optimized if we keep track of the active polynomials.
         polys.par_iter_mut().for_each(|multivariate_poly| {
             if n_remaining_rounds == multivariate_poly.arity() {
                 multivariate_poly.fix_first_in_place(challenge)

crates/stark-backend/src/utils.rs

@@ -59,6 +59,23 @@ pub fn metrics_span<R, F: FnOnce() -> R>(name: impl Into<Cow<'static, str>>, f:
     }
 }
 
+/// 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_increment<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()).increment(start.elapsed().as_millis() as f64);
+            res
+        } else {
+            f()
+        }
+    }
+}
+
 #[macro_export]
 #[cfg(feature = "parallel")]
 macro_rules! parizip {