Parallelize OOD evaluation. (#1306)

Author: leo-starkware

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

@@ -77,6 +77,22 @@ impl<T> TreeVec<ColumnVec<T>> {
         )
     }
 
+    #[cfg(feature = "parallel")]
+    pub fn par_map_cols<U, F>(self, f: F) -> TreeVec<ColumnVec<U>>
+    where
+        T: Send,
+        U: Send,
+        F: Fn(T) -> U + Sync + Send,
+    {
+        use rayon::iter::{IntoParallelIterator, ParallelIterator};
+        TreeVec(
+            self.0
+                .into_par_iter()
+                .map(|column| column.into_par_iter().map(&f).collect::<Vec<_>>())
+                .collect(),
+        )
+    }
+
     /// Zips two [`TreeVec<ColumnVec<T>>`] with the same structure (number of columns in each tree).
     /// The resulting [`TreeVec<ColumnVec<T>>`] has the same structure, with each value being a
     /// tuple of the corresponding values from the input [`TreeVec<ColumnVec<T>>`].

crates/stwo/src/prover/pcs/mod.rs

@@ -150,18 +150,30 @@ impl<'a, B: BackendForChannel<MC>, MC: MerkleChannel> CommitmentSchemeProver<'a,
         } else {
             Some(self.build_weights_hash_map(&sampled_points))
         };
+
+        // Lambda that evaluates a polynomial on a collection of circle points and returns a vector
+        // of point samples.
+        let eval_at_points = |(poly, points): (&Poly<B>, &Vec<CirclePoint<SecureField>>)| {
+            points
+                .iter()
+                .map(|&point| PointSample {
+                    point,
+                    value: poly.eval_at_point(point, weights_hash_map.as_ref()),
+                })
+                .collect_vec()
+        };
+
+        #[cfg(not(feature = "parallel"))]
         let samples: TreeVec<Vec<Vec<PointSample>>> = self
             .polynomials()
             .zip_cols(&sampled_points)
-            .map_cols(|(poly, points)| {
-                points
-                    .iter()
-                    .map(|&point| PointSample {
-                        point,
-                        value: poly.eval_at_point(point, weights_hash_map.as_ref()),
-                    })
-                    .collect_vec()
-            });
+            .map_cols(eval_at_points);
+        #[cfg(feature = "parallel")]
+        let samples: TreeVec<Vec<Vec<PointSample>>> = self
+            .polynomials()
+            .zip_cols(&sampled_points)
+            .par_map_cols(eval_at_points);
+
         span.exit();
         let sampled_values = samples
             .as_cols_ref()