feat(fft): add fast mul for fft friendly fields (#997)

* feat(fft): add fast multiplication

* chore: add fast_mul benchmark

* impl suggestion

* chore: impl suggestion

Author: 0xLucqs

crates/math/benches/polynomials/polynomial.rs

@@ -1,8 +1,10 @@
-use super::utils::{rand_field_elements, rand_poly, FE};
+use super::utils::{rand_complex_mersenne_poly, rand_field_elements, rand_poly, FE};
 use const_random::const_random;
 use core::hint::black_box;
 use criterion::Criterion;
-use lambdaworks_math::polynomial::Polynomial;
+use lambdaworks_math::{
+    field::fields::mersenne31::extensions::Degree2ExtensionField, polynomial::Polynomial,
+};
 
 pub fn polynomial_benchmarks(c: &mut Criterion) {
     let mut group = c.benchmark_group("Polynomial");
@@ -43,6 +45,20 @@ pub fn polynomial_benchmarks(c: &mut Criterion) {
         bench.iter(|| black_box(&x_poly) * black_box(&y_poly));
     });
 
+    let big_order = 9;
+    let x_poly = rand_complex_mersenne_poly(big_order);
+    let y_poly = rand_complex_mersenne_poly(big_order);
+    group.bench_function("fast_mul big poly", |bench| {
+        bench.iter(|| {
+            black_box(&x_poly)
+                .fast_fft_multiplication::<Degree2ExtensionField>(black_box(&y_poly))
+                .unwrap()
+        });
+    });
+    group.bench_function("slow mul big poly", |bench| {
+        bench.iter(|| black_box(&x_poly) * black_box(&y_poly));
+    });
+
     group.bench_function("div", |bench| {
         let x_poly = rand_poly(order);
         let y_poly = rand_poly(order);

crates/math/benches/polynomials/utils.rs

@@ -1,6 +1,12 @@
 use const_random::const_random;
 use lambdaworks_math::{
-    field::fields::u64_prime_field::{U64FieldElement, U64PrimeField},
+    field::{
+        element::FieldElement,
+        fields::{
+            mersenne31::{extensions::Degree2ExtensionField, field::Mersenne31Field},
+            u64_prime_field::{U64FieldElement, U64PrimeField},
+        },
+    },
     polynomial::{
         dense_multilinear_poly::DenseMultilinearPolynomial,
         sparse_multilinear_poly::SparseMultilinearPolynomial, Polynomial,
@@ -36,6 +42,28 @@ pub fn rand_field_elements(order: u64) -> Vec<FE> {
 pub fn rand_poly(order: u64) -> Polynomial<FE> {
     Polynomial::new(&rand_field_elements(order))
 }
+#[allow(dead_code)]
+#[inline(never)]
+#[export_name = "u64_utils::rand_complex_mersenne_field_elements"]
+pub fn rand_complex_mersenne_field_elements(
+    order: u32,
+) -> Vec<FieldElement<Degree2ExtensionField>> {
+    let mut result = Vec::with_capacity(1 << order);
+    for _ in 0..result.capacity() {
+        result.push(FieldElement::<Degree2ExtensionField>::new([
+            FieldElement::<Mersenne31Field>::new(random()),
+            FieldElement::<Mersenne31Field>::new(random()),
+        ]));
+    }
+    result
+}
+
+#[allow(dead_code)]
+#[inline(never)]
+#[export_name = "u64_utils::rand_complex_mersenne_poly"]
+pub fn rand_complex_mersenne_poly(order: u32) -> Polynomial<FieldElement<Degree2ExtensionField>> {
+    Polynomial::new(&rand_complex_mersenne_field_elements(order))
+}
 
 #[allow(dead_code)]
 #[inline(never)]

crates/math/src/fft/errors.rs

@@ -10,6 +10,7 @@ pub enum FFTError {
     RootOfUnityError(u64),
     InputError(usize),
     OrderError(u64),
+    DomainSizeError(usize),
     #[cfg(feature = "cuda")]
     CudaError(CudaError),
 }
@@ -24,6 +25,9 @@ impl Display for FFTError {
             FFTError::OrderError(v) => {
                 write!(f, "Order should be less than or equal to 63, but is {v}")
             }
+            FFTError::DomainSizeError(_) => {
+                write!(f, "Domain size exceeds two adicity of the field")
+            }
             #[cfg(feature = "cuda")]
             FFTError::CudaError(_) => {
                 write!(f, "A CUDA related error has ocurred")

crates/math/src/fft/polynomial.rs

@@ -27,7 +27,9 @@ impl<E: IsField> Polynomial<FieldElement<E>> {
     ) -> Result<Vec<FieldElement<E>>, FFTError> {
         let domain_size = domain_size.unwrap_or(0);
         let len = core::cmp::max(poly.coeff_len(), domain_size).next_power_of_two() * blowup_factor;
-
+        if len.trailing_zeros() as u64 > F::TWO_ADICITY {
+            return Err(FFTError::DomainSizeError(len.trailing_zeros() as usize));
+        }
         if poly.coefficients().is_empty() {
             return Ok(vec![FieldElement::zero(); len]);
         }
@@ -97,6 +99,22 @@ impl<E: IsField> Polynomial<FieldElement<E>> {
         let scaled = Polynomial::interpolate_fft::<F>(fft_evals)?;
         Ok(scaled.scale(&offset.inv().unwrap()))
     }
+
+    /// Multiplies two polynomials using FFT.
+    /// It's faster than naive multiplication when the degree of the polynomials is large enough (>=2**6).
+    /// This works best with polynomials whose highest degree is equal to a power of 2 - 1.
+    /// Will return an error if the degree of the resulting polynomial is greater than 2**63.
+    pub fn fast_fft_multiplication<F: IsFFTField + IsSubFieldOf<E>>(
+        &self,
+        other: &Self,
+    ) -> Result<Self, FFTError> {
+        let domain_size = self.degree() + other.degree() + 1;
+        let p = Polynomial::evaluate_fft::<F>(self, 1, Some(domain_size))?;
+        let q = Polynomial::evaluate_fft::<F>(other, 1, Some(domain_size))?;
+        let r = p.into_iter().zip(q).map(|(a, b)| a * b).collect::<Vec<_>>();
+
+        Polynomial::interpolate_fft::<F>(&r)
+    }
 }
 
 pub fn compose_fft<F, E>(
@@ -313,6 +331,11 @@ mod tests {
 
                 prop_assert_eq!(poly, new_poly);
             }
+
+            #[test]
+            fn test_fft_multiplication_works(poly in poly(7), other in poly(7)) {
+                prop_assert_eq!(poly.fast_fft_multiplication::<F>(&other).unwrap(), poly * other);
+            }
         }
 
         #[test]
@@ -408,6 +431,11 @@ mod tests {
                 let (poly, new_poly) = gen_fft_interpolate_and_evaluate(poly);
                 prop_assert_eq!(poly, new_poly);
             }
+
+            #[test]
+            fn test_fft_multiplication_works(poly in poly(7), other in poly(7)) {
+                prop_assert_eq!(poly.fast_fft_multiplication::<F>(&other).unwrap(), poly * other);
+            }
         }
     }
 

crates/math/src/field/fields/mersenne31/extensions.rs

@@ -2,7 +2,7 @@ use super::field::Mersenne31Field;
 use crate::field::{
     element::FieldElement,
     errors::FieldError,
-    traits::{IsField, IsSubFieldOf},
+    traits::{IsFFTField, IsField, IsSubFieldOf},
 };
 #[cfg(feature = "alloc")]
 use alloc::vec::Vec;
@@ -93,6 +93,14 @@ impl IsField for Degree2ExtensionField {
     }
 }
 
+impl IsFFTField for Degree2ExtensionField {
+    // Values taken from stwo
+    // https://github.com/starkware-libs/stwo/blob/dev/crates/prover/src/core/circle.rs#L203-L209
+    const TWO_ADICITY: u64 = 31;
+    const TWO_ADIC_PRIMITVE_ROOT_OF_UNITY: Self::BaseType =
+        [FpE::const_from_raw(2), FpE::const_from_raw(1268011823)];
+}
+
 impl IsSubFieldOf<Degree2ExtensionField> for Mersenne31Field {
     fn add(
         a: &Self::BaseType,