par map

Author: VitaliiH

crates/prover/src/constraint_framework/component.rs

@@ -286,11 +286,22 @@ impl<E: FrameworkEval + Sync> ComponentProver<SimdBackend> for FrameworkComponen
         // Extend trace if necessary.
         // TODO: Don't extend when eval_size < committed_size. Instead, pick a good
         // subdomain. (For larger blowup factors).
-        nvtx_timed!("extend trace");
+        nvtx_timed!("need_to_extend");
+        #[cfg(not(feature = "parallel"))]
         let need_to_extend = component_evals
             .iter()
             .flatten()
             .any(|c| c.domain != eval_domain);
+        nvtx_timed_pop!();
+
+        #[cfg(feature = "parallel")]
+        let need_to_extend = component_evals
+            .par_iter()
+            .flat_map_iter(|v| v.iter())
+            .any(|c| c.domain != eval_domain);
+        nvtx_timed_pop!();
+
+        nvtx_timed!("SIMD extend trace");
         let trace: TreeVec<
             Vec<Cow<'_, CircleEvaluation<SimdBackend, BaseField, BitReversedOrder>>>,
         > = if need_to_extend {
@@ -578,13 +589,20 @@ impl<E: FrameworkEval + Sync> ComponentProver<IcicleBackend> for FrameworkCompon
         // Extend trace if necessary.
         // TODO: Don't extend when eval_size < committed_size. Instead, pick a good
         // subdomain. (For larger blowup factors).
-        nvtx_timed!("extend trace");
+        nvtx_timed!("need_to_extend");
+        #[cfg(not(feature = "parallel"))]
         let need_to_extend = component_evals
             .iter()
             .flatten()
             .any(|c| c.domain != eval_domain);
         nvtx_timed_pop!();
 
+        #[cfg(feature = "parallel")]
+        let need_to_extend = component_evals
+            .par_iter()
+            .flat_map_iter(|v| v.iter())
+            .any(|c| c.domain != eval_domain);
+
         // Denom inverses.
         nvtx_timed!("denom inverses");
         let log_expand = eval_domain.log_size() - trace_domain.log_size();

crates/prover/src/core/backend/cpu/mod.rs

@@ -48,7 +48,7 @@ pub fn bit_reverse<T>(v: &mut [T]) {
     }
 }
 
-impl<T: Debug + Clone + Default> ColumnOps<T> for CpuBackend {
+impl<T: Debug + Clone + Default + Sync> ColumnOps<T> for CpuBackend {
     type Column = Vec<T>;
 
     fn bit_reverse_column(column: &mut Self::Column) {
@@ -64,7 +64,7 @@ impl<F: Field> FieldOps<F> for CpuBackend {
     }
 }
 
-impl<T: Debug + Clone + Default> Column<T> for Vec<T> {
+impl<T: Debug + Clone + Default + Sync> Column<T> for Vec<T> {
     fn zeros(len: usize) -> Self {
         vec![T::default(); len]
     }

crates/prover/src/core/backend/mod.rs

@@ -48,7 +48,7 @@ pub trait ColumnOps<T> {
 pub type Col<B, T> = <B as ColumnOps<T>>::Column;
 
 // TODO(alont): Consider removing the generic parameter and only support BaseField.
-pub trait Column<T>: Clone + Debug + FromIterator<T> {
+pub trait Column<T>: Clone + Debug + FromIterator<T> + Sync {
     /// Creates a new column of zeros with the given length.
     fn zeros(len: usize) -> Self;
     /// Creates a new column of uninitialized values with the given length.

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

@@ -30,6 +30,7 @@ pub struct CommitmentSchemeProver<'a, B: BackendForChannel<MC>, MC: MerkleChanne
     pub config: PcsConfig,
     twiddles: &'a TwiddleTree<B>,
 }
+use crate::core::backend::ColumnOps;
 
 impl<'a, B: BackendForChannel<MC>, MC: MerkleChannel> CommitmentSchemeProver<'a, B, MC> {
     pub fn new(config: PcsConfig, twiddles: &'a TwiddleTree<B>) -> Self {
@@ -82,7 +83,6 @@ impl<'a, B: BackendForChannel<MC>, MC: MerkleChannel> CommitmentSchemeProver<'a,
         let evals = self.evaluations();
         Trace { polys, evals }
     }
-
     pub fn prove_values(
         self,
         sampled_points: TreeVec<ColumnVec<Vec<CirclePoint<SecureField>>>>,

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

@@ -67,7 +67,11 @@ impl<T> Default for TreeVec<T> {
     }
 }
 
+#[cfg(feature = "parallel")]
+    use rayon::prelude::*;
+
 impl<T> TreeVec<ColumnVec<T>> {
+    #[cfg(not(feature = "parallel"))]
     pub fn map_cols<U, F: FnMut(T) -> U>(self, mut f: F) -> TreeVec<ColumnVec<U>> {
         TreeVec(
             self.0
@@ -77,6 +81,28 @@ impl<T> TreeVec<ColumnVec<T>> {
         )
     }
 
+
+    #[cfg(feature = "parallel")]
+    pub fn map_cols<U, F>(self, f: F) -> TreeVec<ColumnVec<U>>
+    where
+        F: Fn(T) -> U + Send + Sync + Clone,
+        T: Send,
+        U: Send,
+    {
+        TreeVec(
+            self.0
+                .into_par_iter()
+                .map(|column| {
+                    let f = f.clone();
+                    column
+                        .into_par_iter()
+                        .map(f)
+                        .collect()
+                })
+                .collect(),
+        )
+    }
+    
     /// Zips two [`TreeVec<ColumVec<T>>`] with the same structure (number of columns in each tree).
     /// The resulting [`TreeVec<ColumVec<T>>`] has the same structure, with each value being a tuple
     /// of the corresponding values from the input [`TreeVec<ColumVec<T>>`].