From 17e3c4ed34727267a2e336e735e7d16aff2b1a9e Mon Sep 17 00:00:00 2001
From: Mickey <hedgefund996@gmail.com>
Date: Fri, 15 Jul 2022 11:10:32 +0800
Subject: [PATCH] speed up generate vk pk with multi-thread

---
 halo2_proofs/src/plonk/permutation/keygen.rs | 32 ++++++++++++--------
 1 file changed, 19 insertions(+), 13 deletions(-)

diff --git a/halo2_proofs/src/plonk/permutation/keygen.rs b/halo2_proofs/src/plonk/permutation/keygen.rs
index 7c57b74fd8..4d5a23213f 100644
--- a/halo2_proofs/src/plonk/permutation/keygen.rs
+++ b/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, Params},
@@ -104,13 +104,16 @@ 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]
@@ -153,13 +156,16 @@ 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]