feat: mul optimizations

0xWOLAND · Nov 6, 2023 · b2bf517 · b2bf517
1 parent 6091a62
commit b2bf517
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 21 deletions.
diff --git a/benches/benchmark.rs b/benches/benchmark.rs
@@ -6,13 +6,22 @@ use fast_ntt::{
 };
 use itertools::Itertools;
 
+const deg: usize = 16;
+
 fn bench_mul(x: usize, y: usize, c: &Constants) {
     let ONE = BigInt::from(1);
     let a = Polynomial::new(vec![0; x].iter().map(|_| ONE).collect_vec());
     let b = Polynomial::new(vec![0; y].iter().map(|_| ONE).collect_vec());
     let _ = a.mul(b, c);
 }
 
+fn bench_mul_brute(x: usize, y: usize) {
+    let ONE = BigInt::from(1);
+    let a = Polynomial::new(vec![0; x].iter().map(|_| ONE).collect_vec());
+    let b = Polynomial::new(vec![0; y].iter().map(|_| ONE).collect_vec());
+    let _ = a.mul_brute(b);
+}
+
 fn bench_forward(n: usize, c: &Constants) {
     let ONE = BigInt::from(1);
     let a = Polynomial::new(vec![0; n].iter().map(|_| ONE).collect_vec());
@@ -21,19 +30,17 @@ fn bench_forward(n: usize, c: &Constants) {
 
 fn criterion_forward(c: &mut Criterion) {
     let mut group = c.benchmark_group("bench_forward");
-    let deg = 16;
-    (1..deg).for_each(|x| {
+    (6..deg).for_each(|x| {
         group.bench_function(BenchmarkId::from_parameter(x), |b| {
             let c = working_modulus(BigInt::from(x), BigInt::from(2 * x + 1));
-            b.iter(|| bench_mul(black_box(1 << x), black_box(1 << x), black_box(&c)))
+            b.iter(|| bench_forward(black_box(1 << x), black_box(&c)))
         });
     });
 }
 
 fn criterion_mul(c: &mut Criterion) {
     let mut group = c.benchmark_group("bench_mul");
-    let deg = 16;
-    (1..deg).for_each(|x| {
+    (6..deg).for_each(|x| {
         group.bench_function(BenchmarkId::from_parameter(x), |b| {
             let N = BigInt::from((2 * x as usize).next_power_of_two());
             let M = N << 1 + 1;
@@ -44,5 +51,19 @@ fn criterion_mul(c: &mut Criterion) {
     group.finish();
 }
 
-criterion_group!(benches, criterion_forward);
+fn criterion_brute_mul(c: &mut Criterion) {
+    let mut group = c.benchmark_group("bench_brute_mul");
+    (6..deg).for_each(|x| {
+        group.bench_function(BenchmarkId::from_parameter(x), |b| {
+            b.iter(|| bench_mul_brute(black_box(1 << x), black_box(1 << x)))
+        });
+    });
+    group.finish();
+}
+
+criterion_group! {
+  name = benches;
+  config = Criterion::default().sample_size(10);
+  targets = criterion_forward, criterion_mul, criterion_brute_mul
+}
 criterion_main!(benches);
diff --git a/src/polynomial.rs b/src/polynomial.rs
@@ -30,37 +30,46 @@ impl Polynomial {
         Self { coef }
     }
 
+    pub fn mul_brute(self, rhs: Polynomial) -> Polynomial {
+        let a = self.len();
+        let b = rhs.len();
+        let ZERO = BigInt::from(0);
+
+        let mut out: Vec<BigInt> = vec![ZERO; a + b];
+
+        for i in 0..a {
+            for j in 0..b {
+                let e = i + j;
+                out[e] += self.coef[i] * rhs.coef[j];
+            }
+        }
+
+        Polynomial { coef: out }
+    }
+
     pub fn mul(self, rhs: Polynomial, c: &Constants) -> Polynomial {
         let v1_deg = self.degree();
         let v2_deg = rhs.degree();
         let n = (self.len() + rhs.len()).next_power_of_two();
+        let ZERO = BigInt::from(0);
 
-        let v1 = vec![BigInt::from(0); n - self.len()]
+        let v1 = vec![ZERO; n - self.len()]
             .into_iter()
             .chain(self.coef.into_iter())
             .collect();
-        let v2 = vec![BigInt::from(0); n - rhs.len()]
+        let v2 = vec![ZERO; n - rhs.len()]
             .into_iter()
             .chain(rhs.coef.into_iter())
             .collect();
 
         let a_forward = forward(v1, &c);
         let b_forward = forward(v2, &c);
 
-        let ZERO = BigInt::from(0);
-
         let mut mul = vec![ZERO; n as usize];
-        a_forward
-            .iter()
-            .rev()
-            .zip_longest(b_forward.iter().rev())
-            .enumerate()
-            .for_each(|(i, p)| match p {
-                Both(&a, &b) => mul[i] = (a * b).rem(c.N),
-                Left(_) => {}
-                Right(_) => {}
-            });
-        mul.reverse();
+        for i in 0..n {
+            mul[i] = (a_forward[i] * b_forward[i]).rem(c.N)
+        }
+
         let coef = inverse(mul, &c);
         let start = coef.iter().position(|&x| x != 0).unwrap();