Add new PolyOps trait for interpolating many columns. Add a default i…

…mplementation for it that parallelizes extend_evals with Rayon.

crates/prover/src/core/pcs/prover.rs

@@ -15,9 +15,9 @@ use super::quotients::{compute_fri_quotients, PointSample};
 use super::utils::TreeVec;
 use super::{PcsConfig, TreeSubspan};
 use crate::core::air::Trace;
-use crate::core::backend::BackendForChannel;
+use crate::core::backend::{BackendForChannel, ColumnOps};
 use crate::core::channel::{Channel, MerkleChannel};
-use crate::core::poly::circle::{CircleEvaluation, CirclePoly};
+use crate::core::poly::circle::{CircleEvaluation, CirclePoly, PolyOps};
 use crate::core::poly::twiddles::TwiddleTree;
 use crate::core::vcs::ops::MerkleHasher;
 use crate::core::vcs::prover::{MerkleDecommitment, MerkleProver};
@@ -168,13 +168,18 @@ impl<'a, 'b, B: BackendForChannel<MC>, MC: MerkleChannel> TreeBuilder<'a, 'b, B,
     pub fn extend_evals(
         &mut self,
         columns: ColumnVec<CircleEvaluation<B, BaseField, BitReversedOrder>>,
-    ) -> TreeSubspan {
+    ) -> TreeSubspan
+    where
+        CircleEvaluation<B, BaseField, BitReversedOrder>: Send + Sync,
+        <B as ColumnOps<BaseField>>::Column: Send + Sync,
+        <B as PolyOps>::Twiddles: Send + Sync,
+    {
         let span = span!(Level::INFO, "Interpolation for commitment").entered();
         let col_start = self.polys.len();
-        let polys = columns
-            .into_iter()
-            .map(|eval| eval.interpolate_with_twiddles(self.commitment_scheme.twiddles))
-            .collect_vec();
+
+        let polys: Vec<CirclePoly<B>> =
+            B::interpolate_columns(&columns, self.commitment_scheme.twiddles);
+
         span.exit();
         self.polys.extend(polys);
         TreeSubspan {

crates/prover/src/core/poly/circle/ops.rs

@@ -1,11 +1,15 @@
+#[cfg(feature = "parallel")]
+use rayon::prelude::*;
+
 use super::{CanonicCoset, CircleDomain, CircleEvaluation, CirclePoly};
-use crate::core::backend::Col;
+use crate::core::backend::{Col, ColumnOps};
 use crate::core::circle::{CirclePoint, Coset};
 use crate::core::fields::m31::BaseField;
 use crate::core::fields::qm31::SecureField;
 use crate::core::fields::FieldOps;
 use crate::core::poly::twiddles::TwiddleTree;
 use crate::core::poly::BitReversedOrder;
+use crate::core::ColumnVec;
 
 /// Operations on BaseField polynomials.
 pub trait PolyOps: FieldOps<BaseField> + Sized {
@@ -27,6 +31,28 @@ pub trait PolyOps: FieldOps<BaseField> + Sized {
         itwiddles: &TwiddleTree<Self>,
     ) -> CirclePoly<Self>;
 
+    fn interpolate_columns(
+        columns: &ColumnVec<CircleEvaluation<Self, BaseField, BitReversedOrder>>,
+        twiddles: &TwiddleTree<Self>,
+    ) -> Vec<CirclePoly<Self>>
+    where
+        CircleEvaluation<Self, BaseField, BitReversedOrder>: Send + Sync,
+        <Self as ColumnOps<BaseField>>::Column: Send + Sync,
+        <Self as PolyOps>::Twiddles: Send + Sync,
+    {
+        #[cfg(feature = "parallel")]
+        return columns
+            .into_par_iter()
+            .map(|eval| Self::interpolate(eval.clone(), twiddles))
+            .collect();
+
+        #[cfg(not(feature = "parallel"))]
+        return columns
+            .into_iter()
+            .map(|eval| Self::interpolate(eval.clone(), twiddles))
+            .collect();
+    }
+
     /// Evaluates the polynomial at a single point.
     /// Used by the [`CirclePoly::eval_at_point()`] function.
     fn eval_at_point(poly: &CirclePoly<Self>, point: CirclePoint<SecureField>) -> SecureField;

crates/prover/src/core/poly/circle/secure_poly.rs

@@ -64,7 +64,7 @@ pub struct SecureEvaluation<B: FieldOps<BaseField>, EvalOrder> {
     _eval_order: PhantomData<EvalOrder>,
 }
 
-impl<B: FieldOps<BaseField>, EvalOrder> SecureEvaluation<B, EvalOrder> {
+impl<B: FieldOps<BaseField> + PolyOps, EvalOrder> SecureEvaluation<B, EvalOrder> {
     pub fn new(domain: CircleDomain, values: SecureColumnByCoords<B>) -> Self {
         assert_eq!(domain.size(), values.len());
         Self {
@@ -76,7 +76,11 @@ impl<B: FieldOps<BaseField>, EvalOrder> SecureEvaluation<B, EvalOrder> {
 
     pub fn into_coordinate_evals(
         self,
-    ) -> [CircleEvaluation<B, BaseField, EvalOrder>; SECURE_EXTENSION_DEGREE] {
+    ) -> [CircleEvaluation<B, BaseField, EvalOrder>; SECURE_EXTENSION_DEGREE]
+    where
+        CircleEvaluation<B, BaseField, BitReversedOrder>: Send + Sync,
+        <B as PolyOps>::Twiddles: Send + Sync,
+    {
         let Self { domain, values, .. } = self;
         values.columns.map(|c| CircleEvaluation::new(domain, c))
     }

crates/prover/src/examples/poseidon/mod.rs

@@ -377,20 +377,25 @@ mod tests {
 
     use itertools::Itertools;
     use num_traits::One;
+    use tracing::{span, Level};
 
     use crate::constraint_framework::assert_constraints;
     use crate::constraint_framework::logup::{LogupAtRow, LookupElements};
     use crate::core::air::Component;
+    use crate::core::backend::simd::SimdBackend;
+    use crate::core::backend::CpuBackend;
     use crate::core::channel::Blake2sChannel;
     use crate::core::fields::m31::BaseField;
     use crate::core::fri::FriConfig;
-    use crate::core::pcs::{CommitmentSchemeVerifier, PcsConfig, TreeVec};
-    use crate::core::poly::circle::CanonicCoset;
+    use crate::core::pcs::{CommitmentSchemeProver, CommitmentSchemeVerifier, PcsConfig, TreeVec};
+    use crate::core::poly::circle::{CanonicCoset, CircleEvaluation, PolyOps};
+    use crate::core::poly::BitReversedOrder;
     use crate::core::prover::verify;
     use crate::core::vcs::blake2_merkle::Blake2sMerkleChannel;
+    use crate::core::ColumnVec;
     use crate::examples::poseidon::{
         apply_internal_round_matrix, apply_m4, eval_poseidon_constraints, gen_interaction_trace,
-        gen_trace, prove_poseidon, PoseidonElements,
+        gen_trace, prove_poseidon, PoseidonElements, LOG_EXPAND, N_COLUMNS
     };
     use crate::math::matrix::{RowMajorMatrix, SquareMatrix};
 
@@ -505,4 +510,49 @@ mod tests {
 
         verify(&[&component], channel, commitment_scheme, proof).unwrap();
     }
+
+    #[test_log::test]
+    fn test_optimize_interpolation() {
+        let blowup_factor = 2;
+        let domain =
+            CanonicCoset::new(N_COLUMNS.ilog2() + LOG_EXPAND + blowup_factor).circle_domain();
+        let config = PcsConfig {
+            pow_bits: 10,
+            fri_config: FriConfig::new(5, 1, 64),
+        };
+
+        let simd_span = span!(Level::INFO, "Test SIMD interpolation").entered();
+        let simd_columns: ColumnVec<CircleEvaluation<SimdBackend, BaseField, BitReversedOrder>> =
+            (0..N_COLUMNS)
+                .map( |_index| {
+                    CircleEvaluation::<SimdBackend, BaseField, BitReversedOrder>::new(
+                        domain,
+                        (0..domain.size()).map(BaseField::from).collect(),
+                    )
+                })
+                .collect();
+        let simd_twiddles = SimdBackend::precompute_twiddles(domain.half_coset);
+        let mut simd_commitment_scheme_prover: CommitmentSchemeProver<
+            '_, SimdBackend, Blake2sMerkleChannel,
+        > = CommitmentSchemeProver::new(config, &simd_twiddles);
+        simd_commitment_scheme_prover.tree_builder().extend_evals(simd_columns);
+        simd_span.exit();
+
+        let cpu_span = span!(Level::INFO, "Test CPU interpolation").entered();
+        let cpu: ColumnVec<CircleEvaluation<CpuBackend, BaseField, BitReversedOrder>> =
+            (0..N_COLUMNS)
+                .map( |_index| {
+                    CircleEvaluation::<CpuBackend, BaseField, BitReversedOrder>::new(
+                        domain,
+                        (0..domain.size()).map(BaseField::from).collect(),
+                    )
+                })
+                .collect();
+        let cpu_twiddles = CpuBackend::precompute_twiddles(domain.half_coset);
+        let mut simd_commitment_scheme_prover: CommitmentSchemeProver<
+            '_, CpuBackend, Blake2sMerkleChannel
+        > = CommitmentSchemeProver::new(config, &cpu_twiddles);
+        simd_commitment_scheme_prover.tree_builder().extend_evals(cpu);
+        cpu_span.exit();
+    }
 }