Add parallel coset shift method (speeds up coset_fft), add fft correctness tests

CHANGELOG.md

@@ -35,6 +35,7 @@
 - #115 (ark-poly) Add parallel implementation to operations on `Evaluations`.
 - #115 (ark-ff) Add parallel implementation of `batch_inversion`.
 - #122 (ark-poly) Add infrastructure for benchmarking `FFT`s.
+- #125 (ark-poly) Add parallelization to applying coset shifts within `coset_fft`.
 
 ### Bug fixes
 - #36 (ark-ec) In Short-Weierstrass curves, include an infinity bit in `ToConstraintField`.

poly-benches/benches/fft.rs

@@ -49,7 +49,7 @@ fn setup_bench(c: &mut Criterion, name: &str, bench_fn: fn(&mut Bencher, &usize)
 fn fft_common_setup<F: FftField, D: EvaluationDomain<F>>(degree: usize) -> (D, Vec<F>) {
     let mut rng = &mut rand::thread_rng();
     let domain = D::new(degree).unwrap();
-    let a = DensePolynomial::<F>::rand(degree, &mut rng)
+    let a = DensePolynomial::<F>::rand(degree - 1, &mut rng)
         .coeffs()
         .to_vec();
     (domain, a)

poly/src/domain/mod.rs

@@ -10,6 +10,9 @@
 use ark_ff::FftField;
 use ark_std::{fmt, hash, vec::Vec};
 use rand::Rng;
+
+#[cfg(feature = "parallel")]
+use ark_std::cmp::max;
 #[cfg(feature = "parallel")]
 use rayon::prelude::*;
 
@@ -80,6 +83,7 @@ pub trait EvaluationDomain<F: FftField>:
     fn ifft_in_place<T: DomainCoeff<F>>(&self, evals: &mut Vec<T>);
 
     /// Multiply the `i`-th element of `coeffs` with the `i`-th power of `g`.
+    #[cfg(not(feature = "parallel"))]
     fn distribute_powers<T: DomainCoeff<F>>(coeffs: &mut [T], g: F) {
         let mut pow = F::one();
         coeffs.iter_mut().for_each(|c| {
@@ -88,6 +92,28 @@ pub trait EvaluationDomain<F: FftField>:
         })
     }
 
+    /// Multiply the `i`-th element of `coeffs` with the `i`-th power of `g`.
+    #[cfg(feature = "parallel")]
+    fn distribute_powers<T: DomainCoeff<F>>(coeffs: &mut [T], g: F) {
+        // compute the number of threads we will be using.
+        let num_cpus_available = rayon::current_num_threads();
+        let min_elements_per_thread = 256;
+        let num_elem_per_thread = max(coeffs.len() / num_cpus_available, min_elements_per_thread);
+
+        // Split up the coeffs to coset-shift across each thread evenly,
+        // and then apply coset shift.
+        coeffs
+            .par_chunks_mut(num_elem_per_thread)
+            .enumerate()
+            .for_each(|(i, chunk)| {
+                let mut pow = g.pow(&[(i * num_elem_per_thread) as u64]);
+                chunk.iter_mut().for_each(|c| {
+                    *c *= pow;
+                    pow *= &g
+                });
+            });
+    }
+
     /// Compute a FFT over a coset of the domain.
     #[inline]
     fn coset_fft<T: DomainCoeff<F>>(&self, coeffs: &[T]) -> Vec<T> {

poly/src/domain/radix2.rs

@@ -10,6 +10,7 @@ use crate::domain::{
 };
 use ark_ff::{FftField, FftParameters};
 use ark_std::{convert::TryFrom, fmt, vec::Vec};
+
 #[cfg(feature = "parallel")]
 use rayon::prelude::*;
 
@@ -256,9 +257,9 @@ pub(crate) fn serial_radix2_fft<T: DomainCoeff<F>, F: FftField>(a: &mut [T], ome
 #[cfg(test)]
 mod tests {
     use crate::domain::Vec;
-    use crate::polynomial::Polynomial;
+    use crate::polynomial::{univariate::*, Polynomial, UVPolynomial};
     use crate::{EvaluationDomain, Radix2EvaluationDomain};
-    use ark_ff::{test_rng, Field, One, UniformRand, Zero};
+    use ark_ff::{test_rng, FftField, Field, One, UniformRand, Zero};
     use ark_test_curves::bls12_381::Fr;
     use rand::Rng;
 
@@ -313,9 +314,6 @@ mod tests {
     /// Test that lagrange interpolation for a random polynomial at a random point works.
     #[test]
     fn non_systematic_lagrange_coefficients_test() {
-        use crate::polynomial::univariate::*;
-        use crate::polynomial::Polynomial;
-        use crate::polynomial::UVPolynomial;
         for domain_dim in 1..10 {
             let domain_size = 1 << domain_dim;
             let domain = Radix2EvaluationDomain::<Fr>::new(domain_size).unwrap();
@@ -361,6 +359,48 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_fft_correctness() {
+        // Tests that the ffts output the correct result.
+        // This assumes a correct polynomial evaluation at point procedure.
+        // It tests consistency of FFT/IFFT, and coset_fft/coset_ifft,
+        // along with testing that each individual evaluation is correct.
+
+        // Runs in time O(degree^2)
+        let log_degree = 5;
+        let degree = 1 << log_degree;
+        let rand_poly = DensePolynomial::<Fr>::rand(degree - 1, &mut test_rng());
+
+        for log_domain_size in log_degree..(log_degree + 2) {
+            let domain_size = 1 << log_domain_size;
+            let domain = Radix2EvaluationDomain::<Fr>::new(domain_size).unwrap();
+            let poly_evals = domain.fft(&rand_poly.coeffs);
+            let poly_coset_evals = domain.coset_fft(&rand_poly.coeffs);
+            for (i, x) in domain.elements().enumerate() {
+                let coset_x = Fr::multiplicative_generator() * x;
+
+                assert_eq!(poly_evals[i], rand_poly.evaluate(&x));
+                assert_eq!(poly_coset_evals[i], rand_poly.evaluate(&coset_x));
+            }
+
+            let rand_poly_from_subgroup =
+                DensePolynomial::from_coefficients_vec(domain.ifft(&poly_evals));
+            let rand_poly_from_coset =
+                DensePolynomial::from_coefficients_vec(domain.coset_ifft(&poly_coset_evals));
+
+            assert_eq!(
+                rand_poly, rand_poly_from_subgroup,
+                "degree = {}, domain size = {}",
+                degree, domain_size
+            );
+            assert_eq!(
+                rand_poly, rand_poly_from_coset,
+                "degree = {}, domain size = {}",
+                degree, domain_size
+            );
+        }
+    }
+
     #[test]
     #[cfg(feature = "parallel")]
     fn parallel_fft_consistency() {

poly/src/polynomial/univariate/dense.rs

@@ -506,13 +506,16 @@ mod tests {
         for degree in 0..70 {
             let poly = DensePolynomial::<Fr>::rand(degree, &mut rng);
             let neg = -poly.clone();
-            assert!((poly + neg).is_zero());
+            let result = poly + neg;
+            assert!(result.is_zero());
+            assert_eq!(result.degree(), 0);
 
             // Test with SubAssign trait
             let poly = DensePolynomial::<Fr>::rand(degree, &mut rng);
             let mut result = poly.clone();
             result -= &poly;
             assert!(result.is_zero());
+            assert_eq!(result.degree(), 0);
         }
     }