feat: simplify QuotientCommitter trait for performance (#30)

* wip: simplify QuotientCommitter trait

* updated coordinator

* feat: simplify `QuotientCommitter` trait for performance

Just have one function for quotient poly eval and commit.
Importantly we update the CPU implementation to avoid additional memory allocation
when creating a **view** of the trace LDE matrices necessary for quotient poly
evaluation.

* chore: remove metric from doc

Author: jonathanpwang

Cargo.toml

@@ -1,5 +1,5 @@
 [workspace.package]
-version = "0.2.0-alpha"
+version = "0.3.0-alpha"
 edition = "2021"
 rust-version = "1.82"
 authors = ["OpenVM contributors"]

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

@@ -7,11 +7,8 @@ use p3_util::log2_strict_usize;
 use tracing::instrument;
 
 use super::{
-    hal::{ProverBackend, ProverDevice, QuotientCommitter},
-    types::{
-        DeviceMultiStarkProvingKey, HalProof, ProvingContext, RapSinglePhaseView, RapView,
-        SingleCommitPreimage,
-    },
+    hal::{ProverBackend, ProverDevice},
+    types::{DeviceMultiStarkProvingKey, HalProof, ProvingContext, SingleCommitPreimage},
     Prover,
 };
 use crate::{
@@ -176,6 +173,10 @@ where
             &mpk.per_air,
             pair_trace_view_per_air,
         );
+        // Challenger observes additional commitments if any exist:
+        for (commit, _) in &prover_data_after.committed_pcs_data_per_phase {
+            self.challenger.observe(commit.clone());
+        }
 
         // Collect exposed_values_per_air for the proof:
         // - transpose per_phase, per_air -> per_air, per_phase
@@ -205,47 +206,31 @@ where
             })
             .collect_vec();
 
-        let (commitments_after, pcs_data_after): (Vec<_>, Vec<_>) = prover_data_after
-            .committed_pcs_data_per_phase
-            .into_iter()
-            .unzip();
-        // Challenger observes additional commitments if any exist:
-        for commit in &commitments_after {
-            self.challenger.observe(commit.clone());
-        }
-
         // ==================== Quotient polynomial computation and commitment, if any ====================
         // Note[jpw]: Currently we always call this step, we could add a flag to skip it for protocols that
         // do not require quotient poly.
-        let extended_rap_views = metrics_span("quotient_extended_view_time_ms", || {
-            create_trace_view_per_air(
-                &self.device,
-                mpk,
-                &log_trace_height_per_air,
-                &cached_views_per_air,
-                &common_main_pcs_data,
-                &pvs_per_air,
-                &pcs_data_after,
-                prover_data_after.rap_views_per_phase,
-            )
-        });
-        let (constraints, quotient_degrees): (Vec<_>, Vec<_>) = mpk
-            .vk_view()
-            .per_air
-            .iter()
-            .map(|vk| (&vk.symbolic_constraints.constraints, vk.quotient_degree))
-            .unzip();
         let (quotient_commit, quotient_data) = self.device.eval_and_commit_quotient(
             &mut self.challenger,
-            &constraints,
-            extended_rap_views,
-            &quotient_degrees,
+            &mpk.per_air,
+            &pvs_per_air,
+            &cached_views_per_air,
+            &common_main_pcs_data,
+            &prover_data_after,
         );
         // Observe quotient commitment
         self.challenger.observe(quotient_commit.clone());
 
+        let (commitments_after, pcs_data_after): (Vec<_>, Vec<_>) = prover_data_after
+            .committed_pcs_data_per_phase
+            .into_iter()
+            .unzip();
         // ==================== Polynomial Opening Proofs ====================
         let opening = metrics_span("pcs_opening_time_ms", || {
+            let quotient_degrees = mpk
+                .per_air
+                .iter()
+                .map(|pk| pk.vk.quotient_degree)
+                .collect_vec();
             let preprocessed = mpk
                 .per_air
                 .iter()
@@ -326,93 +311,3 @@ impl<'a, PB: ProverBackend> DeviceMultiStarkProvingKey<'a, PB> {
         MultiStarkVerifyingKeyView::new(self.per_air.iter().map(|pk| pk.vk).collect())
     }
 }
-
-/// Takes in views of pcs data and returns extended views of all matrices evaluated on quotient domains
-/// for quotient poly calculation.
-#[allow(clippy::too_many_arguments)]
-fn create_trace_view_per_air<PB: ProverBackend>(
-    device: &impl QuotientCommitter<PB>,
-    mpk: &DeviceMultiStarkProvingKey<PB>,
-    log_trace_height_per_air: &[u8],
-    cached_views_per_air: &[Vec<SingleCommitPreimage<&PB::Matrix, &PB::PcsData>>],
-    common_main_pcs_data: &PB::PcsData,
-    pvs_per_air: &[Vec<PB::Val>],
-    pcs_data_per_phase: &[PB::PcsData],
-    rap_views_per_phase: Vec<Vec<RapSinglePhaseView<usize, PB::Challenge>>>,
-) -> Vec<RapView<PB::Matrix, PB::Val, PB::Challenge>> {
-    let mut common_main_idx = 0;
-    izip!(
-        &mpk.per_air,
-        log_trace_height_per_air,
-        cached_views_per_air,
-        pvs_per_air
-    )
-    .enumerate()
-    .map(|(i, (pk, &log_trace_height, cached_views, pvs))| {
-        let quotient_degree = pk.vk.quotient_degree;
-        // The AIR will be treated as the full RAP with virtual columns after this
-        let preprocessed = pk.preprocessed_data.as_ref().map(|cv| {
-            device
-                .get_extended_matrix(&cv.data, cv.matrix_idx as usize, quotient_degree)
-                .unwrap()
-        });
-        let mut partitioned_main: Vec<_> = cached_views
-            .iter()
-            .map(|cv| {
-                device
-                    .get_extended_matrix(cv.data, cv.matrix_idx as usize, quotient_degree)
-                    .unwrap()
-            })
-            .collect();
-        if pk.vk.has_common_main() {
-            partitioned_main.push(
-                device
-                    .get_extended_matrix(common_main_pcs_data, common_main_idx, quotient_degree)
-                    .unwrap_or_else(|| {
-                        panic!("common main commitment could not get matrix_idx={common_main_idx}")
-                    }),
-            );
-            common_main_idx += 1;
-        }
-        let pair = PairView {
-            log_trace_height,
-            preprocessed,
-            partitioned_main,
-            public_values: pvs.to_vec(),
-        };
-        let mut per_phase = pcs_data_per_phase
-            .iter()
-            .zip_eq(&rap_views_per_phase)
-            .map(
-                |(pcs_data, rap_views)| -> Option<RapSinglePhaseView<PB::Matrix, PB::Challenge>> {
-                    let rap_view = rap_views.get(i)?;
-                    let matrix_idx = rap_view.inner?;
-                    let extended_matrix =
-                        device.get_extended_matrix(pcs_data, matrix_idx, quotient_degree);
-                    let extended_matrix = extended_matrix.unwrap_or_else(|| {
-                        panic!("could not get matrix_idx={matrix_idx} for rap {i}")
-                    });
-                    Some(RapSinglePhaseView {
-                        inner: Some(extended_matrix),
-                        challenges: rap_view.challenges.clone(),
-                        exposed_values: rap_view.exposed_values.clone(),
-                    })
-                },
-            )
-            .collect_vec();
-        while let Some(last) = per_phase.last() {
-            if last.is_none() {
-                per_phase.pop();
-            } else {
-                break;
-            }
-        }
-        let per_phase = per_phase
-            .into_iter()
-            .map(|v| v.unwrap_or_default())
-            .collect();
-
-        RapView { pair, per_phase }
-    })
-    .collect()
-}

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

@@ -1,7 +1,7 @@
-use std::{marker::PhantomData, ops::Deref, sync::Arc};
+use std::{iter::zip, marker::PhantomData, ops::Deref, sync::Arc};
 
 use derivative::Derivative;
-use itertools::{zip_eq, Itertools};
+use itertools::{izip, zip_eq, Itertools};
 use opener::OpeningProver;
 use p3_challenger::FieldChallenger;
 use p3_commit::{Pcs, PolynomialSpace};
@@ -18,7 +18,7 @@ use super::{
     },
 };
 use crate::{
-    air_builders::symbolic::{SymbolicConstraints, SymbolicExpressionDag},
+    air_builders::symbolic::SymbolicConstraints,
     config::{
         Com, PcsProof, PcsProverData, RapPhaseSeqPartialProof, RapPhaseSeqProvingKey,
         StarkGenericConfig, Val,
@@ -231,43 +231,112 @@ impl<SC: StarkGenericConfig> hal::RapPartialProver<CpuBackend<SC>> for CpuDevice
     }
 }
 
-type RapMatrixView<SC> =
-    RapView<Arc<RowMajorMatrix<Val<SC>>>, Val<SC>, <SC as StarkGenericConfig>::Challenge>;
-
 impl<SC: StarkGenericConfig> hal::QuotientCommitter<CpuBackend<SC>> for CpuDevice<'_, SC> {
-    fn get_extended_matrix(
-        &self,
-        view: &PcsData<SC>,
-        matrix_idx: usize,
-        quotient_degree: u8,
-    ) -> Option<Arc<RowMajorMatrix<Val<SC>>>> {
-        let pcs = self.pcs();
-        let log_trace_height = *view.log_trace_heights.get(matrix_idx)?;
-        let trace_domain = pcs.natural_domain_for_degree(1usize << log_trace_height);
-        let quotient_domain =
-            trace_domain.create_disjoint_domain(trace_domain.size() * quotient_degree as usize);
-        // NOTE[jpw]: (perf) this clones under the hood!
-        let lde_matrix = self
-            .pcs()
-            .get_evaluations_on_domain(&view.data, matrix_idx, quotient_domain)
-            .to_row_major_matrix();
-        Some(Arc::new(lde_matrix))
-    }
-
     fn eval_and_commit_quotient(
         &self,
         challenger: &mut SC::Challenger,
-        constraints: &[&SymbolicExpressionDag<Val<SC>>],
-        extended_views: Vec<RapMatrixView<SC>>,
-        quotient_degrees: &[u8],
+        pk_views: &[DeviceStarkProvingKey<CpuBackend<SC>>],
+        public_values: &[Vec<Val<SC>>],
+        cached_views_per_air: &[Vec<
+            SingleCommitPreimage<&Arc<RowMajorMatrix<Val<SC>>>, &PcsData<SC>>,
+        >],
+        common_main_pcs_data: &PcsData<SC>,
+        prover_data_after: &ProverDataAfterRapPhases<CpuBackend<SC>>,
     ) -> (Com<SC>, PcsData<SC>) {
+        let pcs = self.pcs();
         // Generate `alpha` challenge
         let alpha: SC::Challenge = challenger.sample_ext_element();
         tracing::debug!("alpha: {alpha:?}");
+        // Prepare extended views:
+        let mut common_main_idx = 0;
+        let extended_views = izip!(pk_views, cached_views_per_air, public_values)
+            .enumerate()
+            .map(|(i, (pk, cached_views, pvs))| {
+                let quotient_degree = pk.vk.quotient_degree;
+                let log_trace_height = if pk.vk.has_common_main() {
+                    common_main_pcs_data.log_trace_heights[common_main_idx]
+                } else {
+                    log2_strict_usize(cached_views[0].trace.height()) as u8
+                };
+                let trace_domain = pcs.natural_domain_for_degree(1usize << log_trace_height);
+                let quotient_domain = trace_domain
+                    .create_disjoint_domain(trace_domain.size() * quotient_degree as usize);
+                // **IMPORTANT**: the return type of `get_evaluations_on_domain` is a matrix view. DO NOT call to_row_major_matrix as this will allocate new memory
+                let preprocessed = pk.preprocessed_data.as_ref().map(|cv| {
+                    pcs.get_evaluations_on_domain(
+                        &cv.data.data,
+                        cv.matrix_idx as usize,
+                        quotient_domain,
+                    )
+                });
+                let mut partitioned_main: Vec<_> = cached_views
+                    .iter()
+                    .map(|cv| {
+                        pcs.get_evaluations_on_domain(
+                            &cv.data.data,
+                            cv.matrix_idx as usize,
+                            quotient_domain,
+                        )
+                    })
+                    .collect();
+                if pk.vk.has_common_main() {
+                    partitioned_main.push(pcs.get_evaluations_on_domain(
+                        &common_main_pcs_data.data,
+                        common_main_idx,
+                        quotient_domain,
+                    ));
+                    common_main_idx += 1;
+                }
+                let pair = PairView {
+                    log_trace_height,
+                    preprocessed,
+                    partitioned_main,
+                    public_values: pvs.to_vec(),
+                };
+                let mut per_phase = zip(
+                    &prover_data_after.committed_pcs_data_per_phase,
+                    &prover_data_after.rap_views_per_phase,
+                )
+                .map(|((_, pcs_data), rap_views)| -> Option<_> {
+                    let rap_view = rap_views.get(i)?;
+                    let matrix_idx = rap_view.inner?;
+                    let extended_matrix =
+                        pcs.get_evaluations_on_domain(&pcs_data.data, matrix_idx, quotient_domain);
+                    Some(RapSinglePhaseView {
+                        inner: Some(extended_matrix),
+                        challenges: rap_view.challenges.clone(),
+                        exposed_values: rap_view.exposed_values.clone(),
+                    })
+                })
+                .collect_vec();
+                while let Some(last) = per_phase.last() {
+                    if last.is_none() {
+                        per_phase.pop();
+                    } else {
+                        break;
+                    }
+                }
+                let per_phase = per_phase
+                    .into_iter()
+                    .map(|v| v.unwrap_or_default())
+                    .collect();
+
+                RapView { pair, per_phase }
+            })
+            .collect_vec();
 
+        let (constraints, quotient_degrees): (Vec<_>, Vec<_>) = pk_views
+            .iter()
+            .map(|pk| {
+                (
+                    &pk.vk.symbolic_constraints.constraints,
+                    pk.vk.quotient_degree,
+                )
+            })
+            .unzip();
         let qc = QuotientCommitter::new(self.pcs(), alpha);
         let quotient_values = metrics_span("quotient_poly_compute_time_ms", || {
-            qc.quotient_values(constraints, extended_views, quotient_degrees)
+            qc.quotient_values(&constraints, extended_views, &quotient_degrees)
         });
 
         // Commit to quotient polynomials. One shared commit for all quotient polynomials