Unrolled ntt_multiply.

cryspen · Aug 9, 2023 · a74bc9d · a74bc9d
1 parent f532cd2
commit a74bc9d
Showing 1 changed file with 16 additions and 9 deletions.
diff --git a/src/kem/kyber768/ntt.rs b/src/kem/kyber768/ntt.rs
@@ -6,6 +6,7 @@ pub(crate) mod kyber_polynomial_ring_element_mod {
     use crate::kem::kyber768::parameters::{
         self, KyberPolynomialRingElement, COEFFICIENTS_IN_RING_ELEMENT,
     };
+    use crate::kem::kyber768::field_element::KyberFieldElement;
 
     const ZETAS: [u16; 128] = [
         1, 1729, 2580, 3289, 2642, 630, 1897, 848, 1062, 1919, 193, 797, 2786, 3260, 569, 1746,
@@ -74,21 +75,27 @@ pub(crate) mod kyber_polynomial_ring_element_mod {
         out
     }
 
+    fn ntt_multiply_binomials((a0, a1): (KyberFieldElement, KyberFieldElement),
+                              (b0, b1): (KyberFieldElement, KyberFieldElement),
+                              zeta: u16) -> (KyberFieldElement, KyberFieldElement) {
+            ((a0 * b0) + ((a1 * b1) * zeta),
+             (a0 * b1) + (a1 * b0))
+    }
+
     pub fn ntt_multiply(
         left: &KyberPolynomialRingElement,
-        other: &KyberPolynomialRingElement,
+        right: &KyberPolynomialRingElement,
     ) -> KyberPolynomialRingElement {
         let mut out = KyberPolynomialRingElement::ZERO;
 
-        for i in (0..COEFFICIENTS_IN_RING_ELEMENT).step_by(2) {
-            let a0_times_b0 = left[i] * other[i];
-            let a1_times_b1 = left[i + 1] * other[i + 1];
-
-            let a0_times_b1 = left[i + 1] * other[i];
-            let a1_times_b0 = left[i] * other[i + 1];
+        for i in (0..out.coefficients.len()).step_by(4) {
+            let product = ntt_multiply_binomials((left[i], left[i+1]), (right[i], right[i + 1]), MOD_ROOTS[i / 2]);
+            out[i] = product.0;
+            out[i + 1] = product.1;
 
-            out[i] = a0_times_b0 + (a1_times_b1 * MOD_ROOTS[i / 2]);
-            out[i + 1] = a0_times_b1 + a1_times_b0;
+            let product = ntt_multiply_binomials((left[i + 2], left[i + 3]), (right[i + 2], right[i + 3]), MOD_ROOTS[(i + 2) / 2]);
+            out[i + 2] = product.0;
+            out[i + 3] = product.1;
         }
         out
     }