perf: cpu TraceCommitter should reserve capacity for LDE (#68)

Author: jonathanpwang

crates/stark-backend/src/engine.rs

@@ -13,7 +13,7 @@ use crate::{
     },
     proof::Proof,
     prover::{
-        cpu::{CpuBackend, CpuDevice, PcsData},
+        cpu::PcsData,
         hal::{DeviceDataTransporter, TraceCommitter},
         types::{
             AirProofInput, AirProvingContext, ProofInput, ProvingContext, SingleCommitPreimage,
@@ -57,14 +57,7 @@ pub trait StarkEngine<SC: StarkGenericConfig> {
 
     fn prover<'a>(&'a self) -> MultiTraceStarkProver<'a, SC>
     where
-        Self: 'a,
-    {
-        MultiTraceStarkProver::new(
-            CpuBackend::<SC>::default(),
-            CpuDevice::new(self.config()),
-            self.new_challenger(),
-        )
-    }
+        Self: 'a;
 
     fn verifier(&self) -> MultiTraceStarkVerifier<SC> {
         MultiTraceStarkVerifier::new(self.config())

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

@@ -35,12 +35,6 @@ pub mod opener;
 /// Computation of DEEP quotient polynomial and commitment
 pub mod quotient;
 
-/// Proves multiple chips with interactions together.
-/// This prover implementation is specialized for Interactive AIRs.
-pub struct MultiTraceStarkProver<'c, SC: StarkGenericConfig> {
-    pub config: &'c SC,
-}
-
 /// CPU backend using Plonky3 traits.
 #[derive(Derivative)]
 #[derivative(Clone(bound = ""), Copy(bound = ""), Default(bound = ""))]
@@ -52,6 +46,9 @@ pub struct CpuBackend<SC> {
 #[derivative(Clone(bound = ""), Copy(bound = ""))]
 pub struct CpuDevice<'a, SC> {
     config: &'a SC,
+    /// When committing a matrix, the matrix is cloned into newly allocated memory.
+    /// The size of the newly allocated memory will be `matrix.size() << log_blowup_factor`.
+    log_blowup_factor: usize,
 }
 
 impl<SC: StarkGenericConfig> ProverBackend for CpuBackend<SC> {
@@ -103,6 +100,7 @@ impl<SC: StarkGenericConfig> ProverDevice<CpuBackend<SC>> for CpuDevice<'_, SC>
 
 impl<SC: StarkGenericConfig> TraceCommitter<CpuBackend<SC>> for CpuDevice<'_, SC> {
     fn commit(&self, traces: &[Arc<RowMajorMatrix<Val<SC>>>]) -> (Com<SC>, PcsData<SC>) {
+        let log_blowup_factor = self.log_blowup_factor;
         let pcs = self.pcs();
         let (log_trace_heights, traces_with_domains): (Vec<_>, Vec<_>) = traces
             .iter()
@@ -111,7 +109,33 @@ impl<SC: StarkGenericConfig> TraceCommitter<CpuBackend<SC>> for CpuDevice<'_, SC
                 let log_height: u8 = log2_strict_usize(height).try_into().unwrap();
                 // Recomputing the domain is lightweight
                 let domain = pcs.natural_domain_for_degree(height);
-                (log_height, (domain, matrix.as_ref().clone()))
+                // pcs.commit takes the trace matrix and in the case of FRI, does in-place cosetDFT
+                // which requires resizing to a larger buffer size. Since we are cloning anyways,
+                // we should just allocate the larger size to avoid memory-reallocation
+                // ref: https://github.com/Plonky3/Plonky3/blob/8c8bbb4c17bd2b7ef2404338ab8f9036d5f08337/dft/src/traits.rs#L116
+                let trace_slice = &matrix.as_ref().values;
+                let new_buffer_size = trace_slice
+                    .len()
+                    .checked_shl(log_blowup_factor.try_into().unwrap())
+                    .unwrap();
+                let mut new_buffer = Vec::with_capacity(new_buffer_size);
+                // SAFETY:
+                // - `trace_slice` is allocated for `trace_slice.len() * size_of::<F>` bytes, obviously
+                // - we just allocated `new_buffer` for at least `trace_slice.len() * size_of::<F>` bytes above (more if there's blowup)
+                // - both are slices of &[F] so alignment is guaranteed
+                // - `new_buffer` is newly allocated so non-overlapping with `trace_slice`
+                unsafe {
+                    std::ptr::copy_nonoverlapping(
+                        trace_slice.as_ptr(),
+                        new_buffer.as_mut_ptr(),
+                        trace_slice.len(),
+                    );
+                    new_buffer.set_len(trace_slice.len());
+                }
+                (
+                    log_height,
+                    (domain, RowMajorMatrix::new(new_buffer, matrix.width)),
+                )
             })
             .unzip();
         let (commit, data) = pcs.commit(traces_with_domains);

crates/stark-sdk/src/config/baby_bear_bytehash.rs

@@ -4,6 +4,10 @@ use openvm_stark_backend::{
     p3_challenger::{HashChallenger, SerializingChallenger32},
     p3_commit::ExtensionMmcs,
     p3_field::extension::BinomialExtensionField,
+    prover::{
+        cpu::{CpuBackend, CpuDevice},
+        MultiTraceStarkProver,
+    },
 };
 use p3_baby_bear::BabyBear;
 use p3_dft::Radix2DitParallel;
@@ -56,6 +60,17 @@ where
         &self.config
     }
 
+    fn prover<'a>(&'a self) -> MultiTraceStarkProver<'a, BabyBearByteHashConfig<H>>
+    where
+        Self: 'a,
+    {
+        MultiTraceStarkProver::new(
+            CpuBackend::default(),
+            CpuDevice::new(self.config(), self.fri_params.log_blowup),
+            self.new_challenger(),
+        )
+    }
+
     fn max_constraint_degree(&self) -> Option<usize> {
         Some(self.max_constraint_degree)
     }

crates/stark-sdk/src/config/baby_bear_poseidon2.rs

@@ -6,6 +6,10 @@ use openvm_stark_backend::{
     p3_challenger::DuplexChallenger,
     p3_commit::ExtensionMmcs,
     p3_field::{extension::BinomialExtensionField, Field, FieldAlgebra},
+    prover::{
+        cpu::{CpuBackend, CpuDevice},
+        MultiTraceStarkProver,
+    },
 };
 use p3_baby_bear::{BabyBear, Poseidon2BabyBear};
 use p3_dft::Radix2DitParallel;
@@ -81,6 +85,17 @@ where
         &self.config
     }
 
+    fn prover<'a>(&'a self) -> MultiTraceStarkProver<'a, BabyBearPermutationConfig<P>>
+    where
+        Self: 'a,
+    {
+        MultiTraceStarkProver::new(
+            CpuBackend::default(),
+            CpuDevice::new(self.config(), self.fri_params.log_blowup),
+            self.new_challenger(),
+        )
+    }
+
     fn max_constraint_degree(&self) -> Option<usize> {
         Some(self.max_constraint_degree)
     }

crates/stark-sdk/src/config/baby_bear_poseidon2_root.rs

@@ -1,8 +1,14 @@
 use ff::PrimeField;
 use openvm_stark_backend::{
-    config::StarkConfig, interaction::fri_log_up::FriLogUpPhase,
-    p3_challenger::MultiField32Challenger, p3_commit::ExtensionMmcs,
+    config::StarkConfig,
+    interaction::fri_log_up::FriLogUpPhase,
+    p3_challenger::MultiField32Challenger,
+    p3_commit::ExtensionMmcs,
     p3_field::extension::BinomialExtensionField,
+    prover::{
+        cpu::{CpuBackend, CpuDevice},
+        MultiTraceStarkProver,
+    },
 };
 use p3_baby_bear::BabyBear;
 use p3_bn254_fr::{Bn254Fr, FFBn254Fr, Poseidon2Bn254};
@@ -69,6 +75,17 @@ where
         &self.config
     }
 
+    fn prover<'a>(&'a self) -> MultiTraceStarkProver<'a, BabyBearPermutationRootConfig<P>>
+    where
+        Self: 'a,
+    {
+        MultiTraceStarkProver::new(
+            CpuBackend::default(),
+            CpuDevice::new(self.config(), self.fri_params.log_blowup),
+            self.new_challenger(),
+        )
+    }
+
     fn max_constraint_degree(&self) -> Option<usize> {
         Some(self.max_constraint_degree)
     }

crates/stark-sdk/src/config/goldilocks_poseidon.rs

@@ -6,6 +6,10 @@ use openvm_stark_backend::{
     p3_challenger::DuplexChallenger,
     p3_commit::ExtensionMmcs,
     p3_field::{extension::BinomialExtensionField, Field},
+    prover::{
+        cpu::{CpuBackend, CpuDevice},
+        MultiTraceStarkProver,
+    },
 };
 use p3_dft::Radix2DitParallel;
 use p3_fri::{FriConfig, TwoAdicFriPcs};
@@ -73,6 +77,17 @@ where
         &self.config
     }
 
+    fn prover<'a>(&'a self) -> MultiTraceStarkProver<'a, GoldilocksPermutationConfig<P>>
+    where
+        Self: 'a,
+    {
+        MultiTraceStarkProver::new(
+            CpuBackend::default(),
+            CpuDevice::new(self.config(), self.security_params.fri_params.log_blowup),
+            self.new_challenger(),
+        )
+    }
+
     fn max_constraint_degree(&self) -> Option<usize> {
         Some(self.max_constraint_degree)
     }

crates/stark-sdk/src/dummy_airs/interaction/dummy_interaction_air.rs

@@ -162,7 +162,7 @@ where
     ) -> Self {
         let air = DummyInteractionAir::new(field_width, is_send, bus_index).partition();
         Self {
-            device: Some(CpuDevice::new(config)),
+            device: Some(CpuDevice::new(config, 0)),
             data: None,
             air,
         }

crates/stark-sdk/src/dummy_airs/interaction/mod.rs

@@ -53,7 +53,7 @@ pub fn verify_interactions(
         .collect();
 
     let challenger = config::baby_bear_poseidon2::Challenger::new(perm.clone());
-    let mut prover = MultiTraceStarkProver::new(backend, CpuDevice::new(&config), challenger);
+    let mut prover = MultiTraceStarkProver::new(backend, CpuDevice::new(&config, 1), challenger);
     let proof = prover.prove(&pk, ProvingContext::new(per_air));
 
     // Verify the proof: