speed up generate vk pk with multi-thread

axiom-crypto · Jul 15, 2022 · 17e3c4e · 17e3c4e
1 parent f586922
commit 17e3c4e
Showing 1 changed file with 19 additions and 13 deletions.
diff --git 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]