Improve performance of vk & pk keygen and of default parallelize c…

…hunking size (#127)

* Squashed commit of the following:

commit 17e3c4e
Author: Mickey <[email protected]>
Date:   Fri Jul 15 11:10:32 2022 +0800

    speed up generate vk pk with multi-thread

* fix

* Improve performance of vk & pk keygen and of default `parallelize` chunking size.

Reduces proving time on large circuits consistently >=3%.
Builts upon [speed up generate vk pk with multi-thread](#88)
Fixes: #83

halo2_proofs/src/arithmetic.rs

@@ -373,7 +373,7 @@ pub fn parallelize<T: Send, F: Fn(&mut [T], usize) + Send + Sync + Clone>(v: &mu
     let num_threads = multicore::current_num_threads();
     let mut chunk = (n as usize) / num_threads;
     if chunk < num_threads {
-        chunk = n as usize;
+        chunk = 1;
     }
 
     multicore::scope(|scope| {

halo2_proofs/src/plonk/permutation/keygen.rs

@@ -3,7 +3,7 @@ use group::Curve;
 
 use super::{Argument, ProvingKey, VerifyingKey};
 use crate::{
-    arithmetic::{CurveAffine, FieldExt},
+    arithmetic::{parallelize, CurveAffine, FieldExt},
     plonk::{Any, Column, Error},
     poly::{
         commitment::{Blind, CommitmentScheme, Params},
@@ -109,49 +109,58 @@ impl Assembly {
         p: &Argument,
     ) -> VerifyingKey<C> {
         // Compute [omega^0, omega^1, ..., omega^{params.n - 1}]
-        let mut omega_powers = Vec::with_capacity(params.n() as usize);
+        let mut omega_powers = vec![C::Scalar::zero(); params.n() as usize];
         {
-            let mut cur = C::Scalar::one();
-            for _ in 0..params.n() {
-                omega_powers.push(cur);
-                cur *= &domain.get_omega();
-            }
+            let omega = domain.get_omega();
+            parallelize(&mut omega_powers, |o, start| {
+                let mut cur = omega.pow_vartime(&[start as u64]);
+                for v in o.iter_mut() {
+                    *v = cur;
+                    cur *= &omega;
+                }
+            })
         }
 
         // Compute [omega_powers * \delta^0, omega_powers * \delta^1, ..., omega_powers * \delta^m]
-        let mut deltaomega = Vec::with_capacity(p.columns.len());
+        let mut deltaomega = vec![omega_powers; p.columns.len()];
         {
-            let mut cur = C::Scalar::one();
-            for _ in 0..p.columns.len() {
-                let mut omega_powers = omega_powers.clone();
-                for o in &mut omega_powers {
-                    *o *= &cur;
+            parallelize(&mut deltaomega, |o, start| {
+                let mut cur = C::Scalar::DELTA.pow_vartime(&[start as u64]);
+                for omega_powers in o.iter_mut() {
+                    for v in omega_powers {
+                        *v *= &cur;
+                    }
+                    cur *= &C::Scalar::DELTA;
                 }
+            });
+        }
 
-                deltaomega.push(omega_powers);
-
-                cur *= &C::Scalar::DELTA;
-            }
+        // Computes the permutation polynomial based on the permutation
+        // description in the assembly.
+        let mut permutations = vec![domain.empty_lagrange(); p.columns.len()];
+        {
+            parallelize(&mut permutations, |o, start| {
+                for (x, permutation_poly) in o.iter_mut().enumerate() {
+                    let i = start + x;
+                    for (j, p) in permutation_poly.iter_mut().enumerate() {
+                        let (permuted_i, permuted_j) = self.mapping[i][j];
+                        *p = deltaomega[permuted_i][permuted_j];
+                    }
+                }
+            });
         }
 
         // Pre-compute commitments for the URS.
-        let mut commitments = vec![];
-        for i in 0..p.columns.len() {
-            // Computes the permutation polynomial based on the permutation
-            // description in the assembly.
-            let mut permutation_poly = domain.empty_lagrange();
-            for (j, p) in permutation_poly.iter_mut().enumerate() {
-                let (permuted_i, permuted_j) = self.mapping[i][j];
-                *p = deltaomega[permuted_i][permuted_j];
-            }
-
+        let mut commitments = Vec::with_capacity(p.columns.len());
+        for permutation in &permutations {
             // Compute commitment to permutation polynomial
             commitments.push(
                 params
-                    .commit_lagrange(&permutation_poly, Blind::default())
+                    .commit_lagrange(permutation, Blind::default())
                     .to_affine(),
             );
         }
+
         VerifyingKey { commitments }
     }
 
@@ -162,50 +171,68 @@ impl Assembly {
         p: &Argument,
     ) -> ProvingKey<C> {
         // Compute [omega^0, omega^1, ..., omega^{params.n - 1}]
-        let mut omega_powers = Vec::with_capacity(params.n() as usize);
+        let mut omega_powers = vec![C::Scalar::zero(); params.n() as usize];
         {
-            let mut cur = C::Scalar::one();
-            for _ in 0..params.n() {
-                omega_powers.push(cur);
-                cur *= &domain.get_omega();
-            }
+            let omega = domain.get_omega();
+            parallelize(&mut omega_powers, |o, start| {
+                let mut cur = omega.pow_vartime(&[start as u64]);
+                for v in o.iter_mut() {
+                    *v = cur;
+                    cur *= &omega;
+                }
+            })
         }
 
         // Compute [omega_powers * \delta^0, omega_powers * \delta^1, ..., omega_powers * \delta^m]
-        let mut deltaomega = Vec::with_capacity(p.columns.len());
+        let mut deltaomega = vec![omega_powers; p.columns.len()];
         {
-            let mut cur = C::Scalar::one();
-            for _ in 0..p.columns.len() {
-                let mut omega_powers = omega_powers.clone();
-                for o in &mut omega_powers {
-                    *o *= &cur;
+            parallelize(&mut deltaomega, |o, start| {
+                let mut cur = C::Scalar::DELTA.pow_vartime(&[start as u64]);
+                for omega_powers in o.iter_mut() {
+                    for v in omega_powers {
+                        *v *= &cur;
+                    }
+                    cur *= &C::Scalar::DELTA;
                 }
-
-                deltaomega.push(omega_powers);
-
-                cur *= &C::Scalar::DELTA;
-            }
+            });
         }
 
         // Compute permutation polynomials, convert to coset form.
-        let mut permutations = vec![];
-        let mut polys = vec![];
-        let mut cosets = vec![];
-        for i in 0..p.columns.len() {
-            // Computes the permutation polynomial based on the permutation
-            // description in the assembly.
-            let mut permutation_poly = domain.empty_lagrange();
-            for (j, p) in permutation_poly.iter_mut().enumerate() {
-                let (permuted_i, permuted_j) = self.mapping[i][j];
-                *p = deltaomega[permuted_i][permuted_j];
-            }
+        let mut permutations = vec![domain.empty_lagrange(); p.columns.len()];
+        {
+            parallelize(&mut permutations, |o, start| {
+                for (x, permutation_poly) in o.iter_mut().enumerate() {
+                    let i = start + x;
+                    for (j, p) in permutation_poly.iter_mut().enumerate() {
+                        let (permuted_i, permuted_j) = self.mapping[i][j];
+                        *p = deltaomega[permuted_i][permuted_j];
+                    }
+                }
+            });
+        }
 
-            // Store permutation polynomial and precompute its coset evaluation
-            permutations.push(permutation_poly.clone());
-            let poly = domain.lagrange_to_coeff(permutation_poly);
-            polys.push(poly.clone());
-            cosets.push(domain.coeff_to_extended(poly));
+        let mut polys = vec![domain.empty_coeff(); p.columns.len()];
+        {
+            parallelize(&mut polys, |o, start| {
+                for (x, poly) in o.iter_mut().enumerate() {
+                    let i = start + x;
+                    let permutation_poly = permutations[i].clone();
+                    *poly = domain.lagrange_to_coeff(permutation_poly);
+                }
+            });
         }
+
+        let mut cosets = vec![domain.empty_extended(); p.columns.len()];
+        {
+            parallelize(&mut cosets, |o, start| {
+                for (x, coset) in o.iter_mut().enumerate() {
+                    let i = start + x;
+                    let poly = polys[i].clone();
+                    *coset = domain.coeff_to_extended(poly);
+                }
+            });
+        }
+
         ProvingKey {
             permutations,
             polys,