Add impl of FftField for Fp2 and Fp3 (#848)

* Add impl of FftField for Fp2 and Fp3 * Move to extensions * Testing out * Expand range of benchmarks and better parallelization handling for small FFTs * Format --------- Co-authored-by: Pratyush Mishra <[email protected]>
arkworks-rs · Aug 26, 2024 · 5a1db8a · 5a1db8a
1 parent 273bf21
commit 5a1db8a
Show file tree

Hide file tree

Showing 5 changed files with 75 additions and 18 deletions.
diff --git a/ff/src/fields/models/cubic_extension.rs b/ff/src/fields/models/cubic_extension.rs
@@ -1,6 +1,7 @@
 use crate::{
     fields::{Field, PrimeField},
-    AdditiveGroup, LegendreSymbol, One, SqrtPrecomputation, ToConstraintField, UniformRand, Zero,
+    AdditiveGroup, FftField, LegendreSymbol, One, SqrtPrecomputation, ToConstraintField,
+    UniformRand, Zero,
 };
 use ark_serialize::{
     CanonicalDeserialize, CanonicalDeserializeWithFlags, CanonicalSerialize,
@@ -770,3 +771,28 @@ mod cube_ext_tests {
         }
     }
 }
+
+impl<P: CubicExtConfig> FftField for CubicExtField<P>
+where
+    P::BaseField: FftField,
+{
+    const GENERATOR: Self = Self::new(
+        P::BaseField::GENERATOR,
+        P::BaseField::ZERO,
+        P::BaseField::ZERO,
+    );
+    const TWO_ADICITY: u32 = P::BaseField::TWO_ADICITY;
+    const TWO_ADIC_ROOT_OF_UNITY: Self = Self::new(
+        P::BaseField::TWO_ADIC_ROOT_OF_UNITY,
+        P::BaseField::ZERO,
+        P::BaseField::ZERO,
+    );
+    const SMALL_SUBGROUP_BASE: Option<u32> = P::BaseField::SMALL_SUBGROUP_BASE;
+    const SMALL_SUBGROUP_BASE_ADICITY: Option<u32> = P::BaseField::SMALL_SUBGROUP_BASE_ADICITY;
+    const LARGE_SUBGROUP_ROOT_OF_UNITY: Option<Self> =
+        if let Some(x) = P::BaseField::LARGE_SUBGROUP_ROOT_OF_UNITY {
+            Some(Self::new(x, P::BaseField::ZERO, P::BaseField::ZERO))
+        } else {
+            None
+        };
+}
diff --git a/ff/src/fields/models/fp3.rs b/ff/src/fields/models/fp3.rs
@@ -39,7 +39,6 @@ impl<P: Fp3Config> CubicExtConfig for Fp3ConfigWrapper<P> {
     type FrobCoeff = P::Fp;
 
     const DEGREE_OVER_BASE_PRIME_FIELD: usize = 3;
-
     const NONRESIDUE: Self::BaseField = P::NONRESIDUE;
 
     const SQRT_PRECOMP: Option<SqrtPrecomputation<CubicExtField<Self>>> =

diff --git a/ff/src/fields/models/quadratic_extension.rs b/ff/src/fields/models/quadratic_extension.rs
@@ -1,7 +1,7 @@
 use crate::{
     biginteger::BigInteger,
     fields::{Field, LegendreSymbol, PrimeField},
-    AdditiveGroup, One, SqrtPrecomputation, ToConstraintField, UniformRand, Zero,
+    AdditiveGroup, FftField, One, SqrtPrecomputation, ToConstraintField, UniformRand, Zero,
 };
 use ark_serialize::{
     CanonicalDeserialize, CanonicalDeserializeWithFlags, CanonicalSerialize,
@@ -825,3 +825,21 @@ mod quad_ext_tests {
         }
     }
 }
+
+impl<P: QuadExtConfig> FftField for QuadExtField<P>
+where
+    P::BaseField: FftField,
+{
+    const GENERATOR: Self = Self::new(P::BaseField::GENERATOR, P::BaseField::ZERO);
+    const TWO_ADICITY: u32 = P::BaseField::TWO_ADICITY;
+    const TWO_ADIC_ROOT_OF_UNITY: Self =
+        Self::new(P::BaseField::TWO_ADIC_ROOT_OF_UNITY, P::BaseField::ZERO);
+    const SMALL_SUBGROUP_BASE: Option<u32> = P::BaseField::SMALL_SUBGROUP_BASE;
+    const SMALL_SUBGROUP_BASE_ADICITY: Option<u32> = P::BaseField::SMALL_SUBGROUP_BASE_ADICITY;
+    const LARGE_SUBGROUP_ROOT_OF_UNITY: Option<Self> =
+        if let Some(x) = P::BaseField::LARGE_SUBGROUP_ROOT_OF_UNITY {
+            Some(Self::new(x, P::BaseField::ZERO))
+        } else {
+            None
+        };
+}
diff --git a/poly/benches/fft.rs b/poly/benches/fft.rs
@@ -12,7 +12,7 @@ use criterion::{criterion_group, criterion_main, Bencher, BenchmarkId, Criterion
 
 // degree bounds to benchmark on
 // e.g. degree bound of 2^{15}, means we do an FFT for a degree (2^{15} - 1) polynomial
-const BENCHMARK_MIN_DEGREE: usize = 1 << 15;
+const BENCHMARK_MIN_DEGREE: usize = 1 << 4;
 const BENCHMARK_MAX_DEGREE_BLS12_381: usize = 1 << 22;
 const BENCHMARK_MAX_DEGREE_MNT6_753: usize = 1 << 17;
 const BENCHMARK_LOG_INTERVAL_DEGREE: usize = 1;

diff --git a/poly/src/domain/radix2/fft.rs b/poly/src/domain/radix2/fft.rs
@@ -224,23 +224,33 @@ impl<F: FftField> Radix2EvaluationDomain<F> {
         max_threads: usize,
         gap: usize,
     ) {
-        cfg_chunks_mut!(xi, chunk_size).for_each(|cxi| {
-            let (lo, hi) = cxi.split_at_mut(gap);
-            // If the chunk is sufficiently big that parallelism helps,
-            // we parallelize the butterfly operation within the chunk.
-
-            if gap > MIN_GAP_SIZE_FOR_PARALLELISATION && num_chunks < max_threads {
-                cfg_iter_mut!(lo)
-                    .zip(hi)
-                    .zip(cfg_iter!(roots).step_by(step))
-                    .for_each(g);
-            } else {
+        if xi.len() <= MIN_INPUT_SIZE_FOR_PARALLELIZATION {
+            xi.chunks_mut(chunk_size).for_each(|cxi| {
+                let (lo, hi) = cxi.split_at_mut(gap);
                 lo.iter_mut()
                     .zip(hi)
                     .zip(roots.iter().step_by(step))
                     .for_each(g);
-            }
-        });
+            });
+        } else {
+            cfg_chunks_mut!(xi, chunk_size).for_each(|cxi| {
+                let (lo, hi) = cxi.split_at_mut(gap);
+                // If the chunk is sufficiently big that parallelism helps,
+                // we parallelize the butterfly operation within the chunk.
+
+                if gap > MIN_GAP_SIZE_FOR_PARALLELIZATION && num_chunks < max_threads {
+                    cfg_iter_mut!(lo)
+                        .zip(hi)
+                        .zip(cfg_iter!(roots).step_by(step))
+                        .for_each(g);
+                } else {
+                    lo.iter_mut()
+                        .zip(hi)
+                        .zip(roots.iter().step_by(step))
+                        .for_each(g);
+                }
+            });
+        }
     }
 
     fn io_helper<T: DomainCoeff<F>>(&self, xi: &mut [T], root: F) {
@@ -349,7 +359,11 @@ const MIN_NUM_CHUNKS_FOR_COMPACTION: usize = 1 << 7;
 
 /// The minimum size of a chunk at which parallelization of `butterfly`s is
 /// beneficial. This value was chosen empirically.
-const MIN_GAP_SIZE_FOR_PARALLELISATION: usize = 1 << 10;
+const MIN_GAP_SIZE_FOR_PARALLELIZATION: usize = 1 << 10;
+
+/// The minimum size of a chunk at which parallelization of `butterfly`s is
+/// beneficial. This value was chosen empirically.
+const MIN_INPUT_SIZE_FOR_PARALLELIZATION: usize = 1 << 10;
 
 // minimum size at which to parallelize.
 #[cfg(feature = "parallel")]