fix: parallel polynomial multiplication iteration

0xWOLAND · Nov 19, 2023 · a352e30 · a352e30
1 parent fdcb457
commit a352e30
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 14 deletions.
diff --git a/README.md b/README.md
@@ -27,17 +27,34 @@ fast-ntt is a Rust package to compute polynomial multiplication in `O(nlog(n))`
 
 ## Benchmarks
 
+#### Number-Theoretic Transform
+
+| Polynomial Degree | NTT       |
+| ----------------- | --------- |
+| 64                | 201.11 µs |
+| 128               | 355.35 µs |
+| 256               | 845.66 µs |
+| 512               | 1.3013 ms |
+| 1024              | 2.1763 ms |
+| 2048              | 4.1270 ms |
+| 4096              | 7.9391 ms |
+| 8192              | 16.674 ms |
+| 16384             | 34.160 ms |
+| 32768             | 79.303 ms |
+
+#### Polynomial Multiplication
+
 ```
-| Polynomial Degree | FFT-Based | Brute-Force |
+| Polynomial Degree | NTT-Based | Brute-Force |
 |-------------------|-----------|-------------|
-| 64                | 7.4158 ms | 153.32 µs   |
-| 128               | 13.675 ms | 617.36 µs   |
-| 256               | 25.056 ms | 2.4731 ms   |
-| 512               | 54.627 ms | 9.9661 ms   |
-| 1024              | 91.212 ms | 39.673 ms   |
-| 2048              | 189.84 ms | 157.35 ms   |
-| 4096              | 407.08 ms | 631.73 ms   |
-| 8192              | 858.76 s  | 2.5609 s    |
-| 16384             | 1.7593 s  | 10.222 s    |
-| 32768             | 3.7149 s  | 41.119 s    |
+| 64                | 1.2677 ms | 50.389 µs   |
+| 128               | 2.3206 ms | 196.92 µs   |
+| 256               | 3.6952 ms | 777.25 µs   |
+| 512               | 6.9324 ms | 3.2495 ms   |
+| 1024              | 13.158 ms | 12.777 ms   |
+| 2048              | 26.159 ms | 51.868 ms   |
+| 4096              | 55.093 ms | 205.77 ms   |
+| 8192              | 115.62 ms | 812.68 ms   |
+| 16384             | 241.09 ms | 3.2130 s    |
+| 32768             | 502.79 ms | 12.959 s    |
 ```
diff --git a/src/polynomial.rs b/src/polynomial.rs
@@ -4,6 +4,9 @@ use std::{
 };
 
 use itertools::{EitherOrBoth::*, Itertools};
+use rayon::iter::{
+    IndexedParallelIterator, IntoParallelRefIterator, IntoParallelRefMutIterator, ParallelIterator,
+};
 
 use crate::{ntt::*, numbers::BigInt};
 
@@ -66,9 +69,9 @@ impl Polynomial {
         let b_forward = forward(v2, &c);
 
         let mut mul = vec![ZERO; n as usize];
-        for i in 0..n {
-            mul[i] = (a_forward[i] * b_forward[i]).rem(c.N)
-        }
+        mul.par_iter_mut()
+            .enumerate()
+            .for_each(|(i, x)| *x = (a_forward[i] * b_forward[i]).rem(c.N));
 
         let coef = inverse(mul, &c);
         let start = coef.iter().position(|&x| x != 0).unwrap();