fix: Upstream concurrent poly generation fixes

zcash · Feb 28, 2023 · a4cc9f6 · a4cc9f6
1 parent e49cf60
commit a4cc9f6
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 19 deletions.
diff --git a/halo2_proofs/src/plonk/vanishing/prover.rs b/halo2_proofs/src/plonk/vanishing/prover.rs
@@ -2,18 +2,15 @@ use std::iter;
 
 use ff::Field;
 use group::Curve;
+use rand_core::RngCore;
+
 #[cfg(feature = "multicore")]
-use maybe_rayon::{
-    current_num_threads,
-    prelude::{IntoParallelRefMutIterator, ParallelIterator},
-};
+use maybe_rayon::{current_num_threads, prelude::*};
 #[cfg(feature = "multicore")]
 use rand_chacha::ChaCha20Rng;
 #[cfg(feature = "multicore")]
 use rand_core::SeedableRng;
 
-use rand_core::RngCore;
-
 use super::Argument;
 use crate::{
     arithmetic::{eval_polynomial, CurveAffine},
@@ -55,9 +52,11 @@ impl<C: CurveAffine> Argument<C> {
         #[cfg(feature = "multicore")]
         let random_poly = {
             let n_threads = current_num_threads();
-            let needed_scalars = (1usize << domain.k as usize) / n_threads;
+            let n = 1usize << domain.k() as usize;
+            let n_chunks = n_threads + if n % n_threads != 0 { 1 } else { 0 };
+            let mut rand_vec = vec![C::Scalar::ZERO; n];
 
-            let mut thread_seeds: Vec<ChaCha20Rng> = (0..n_threads)
+            let mut thread_seeds: Vec<ChaCha20Rng> = (0..n_chunks)
                 .into_iter()
                 .map(|_| {
                     let mut seed = [0u8; 32];
@@ -66,17 +65,16 @@ impl<C: CurveAffine> Argument<C> {
                 })
                 .collect();
 
-            let rand_vec: Vec<C::Scalar> = thread_seeds
+            thread_seeds
                 .par_iter_mut()
-                .flat_map(|mut rng| {
-                    (0..needed_scalars)
-                        .into_iter()
-                        .map(|_| C::Scalar::random(&mut rng))
-                        .collect::<Vec<C::Scalar>>()
-                })
-                .collect();
-
-            Polynomial::<C::ScalarExt, Coeff>::from_evals(rand_vec)
+                .zip_eq(rand_vec.par_chunks_mut(n / n_threads))
+                .for_each(|(mut rng, chunk)| {
+                    chunk
+                        .iter_mut()
+                        .for_each(|v| *v = C::Scalar::random(&mut rng))
+                });
+
+            domain.coeff_from_vec(rand_vec)
         };
 
         #[cfg(not(feature = "multicore"))]

diff --git a/halo2_proofs/src/poly/domain.rs b/halo2_proofs/src/poly/domain.rs
@@ -19,7 +19,7 @@ use std::marker::PhantomData;
 #[derive(Clone, Debug)]
 pub struct EvaluationDomain<F: Field> {
     n: u64,
-    pub(crate) k: u32,
+    k: u32,
     extended_k: u32,
     omega: F,
     omega_inv: F,
@@ -142,6 +142,11 @@ impl<F: WithSmallOrderMulGroup<3>> EvaluationDomain<F> {
         }
     }
 
+    /// Returns the `k` or Degree associated to the domain.
+    pub(crate) fn k(&self) -> usize {
+        self.k as usize
+    }
+
     /// Obtains a polynomial in Lagrange form when given a vector of Lagrange
     /// coefficients of size `n`; panics if the provided vector is the wrong
     /// length.