FFT/LDE benches (#188)

And expose FftStrategy so we can bench both

Cargo.toml

@@ -26,12 +26,20 @@ serde_cbor = "0.11.1"
 
 [dev-dependencies]
 criterion = "0.3.5"
+tynm = "0.1.6"
 
 [[bench]]
 name = "field_arithmetic"
 harness = false
 
+[[bench]]
+name = "ffts"
+harness = false
+
 [profile.release]
 opt-level = 3
 #lto = "fat"
 #codegen-units = 1
+
+[profile.bench]
+opt-level = 3

benches/ffts.rs

@@ -0,0 +1,48 @@
+use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
+use plonky2::field::crandall_field::CrandallField;
+use plonky2::field::extension_field::quartic::QuarticCrandallField;
+use plonky2::field::fft::FftStrategy;
+use plonky2::field::field_types::Field;
+use plonky2::polynomial::polynomial::PolynomialCoeffs;
+use tynm::type_name;
+
+pub(crate) fn bench_ffts<F: Field>(c: &mut Criterion, strategy: FftStrategy) {
+    let mut group = c.benchmark_group(&format!("fft-{:?}<{}>", strategy, type_name::<F>()));
+
+    for size_log in [13, 14, 15, 16] {
+        let size = 1 << size_log;
+        group.bench_with_input(BenchmarkId::from_parameter(size), &size, |b, _| {
+            let coeffs = PolynomialCoeffs::new(F::rand_vec(size));
+            b.iter(|| coeffs.fft_with_options(strategy, None, None));
+        });
+    }
+}
+
+pub(crate) fn bench_ldes<F: Field>(c: &mut Criterion, strategy: FftStrategy) {
+    const RATE_BITS: usize = 3;
+
+    let mut group = c.benchmark_group(&format!("lde-{:?}<{}>", strategy, type_name::<F>()));
+
+    for size_log in [16] {
+        let orig_size = 1 << (size_log - RATE_BITS);
+        let lde_size = 1 << size_log;
+
+        group.bench_with_input(BenchmarkId::from_parameter(lde_size), &lde_size, |b, _| {
+            let coeffs = PolynomialCoeffs::new(F::rand_vec(orig_size));
+            b.iter(|| {
+                let padded_coeffs = coeffs.lde(RATE_BITS);
+                padded_coeffs.fft_with_options(strategy, Some(RATE_BITS), None)
+            });
+        });
+    }
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    bench_ffts::<CrandallField>(c, FftStrategy::Classic);
+    bench_ffts::<CrandallField>(c, FftStrategy::Unrolled);
+    bench_ldes::<CrandallField>(c, FftStrategy::Classic);
+    bench_ldes::<CrandallField>(c, FftStrategy::Unrolled);
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

benches/field_arithmetic.rs

@@ -1,12 +1,11 @@
-use std::any::type_name;
-
 use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
 use plonky2::field::crandall_field::CrandallField;
 use plonky2::field::extension_field::quartic::QuarticCrandallField;
 use plonky2::field::field_types::Field;
+use tynm::type_name;
 
 pub(crate) fn bench_field<F: Field>(c: &mut Criterion) {
-    c.bench_function(&format!("{} mul", type_name::<F>()), |b| {
+    c.bench_function(&format!("mul<{}>", type_name::<F>()), |b| {
         b.iter_batched(
             || (F::rand(), F::rand()),
             |(x, y)| x * y,
@@ -14,7 +13,7 @@ pub(crate) fn bench_field<F: Field>(c: &mut Criterion) {
         )
     });
 
-    c.bench_function(&format!("{} try_inverse", type_name::<F>()), |b| {
+    c.bench_function(&format!("try_inverse<{}>", type_name::<F>()), |b| {
         b.iter_batched(|| F::rand(), |x| x.try_inverse(), BatchSize::SmallInput)
     });
 }

src/field/fft.rs

@@ -6,12 +6,13 @@ use crate::util::{log2_strict, reverse_index_bits};
 
 // TODO: Should really do some "dynamic" dispatch to handle the
 // different FFT algos rather than C-style enum dispatch.
-enum FftStrategy {
+#[derive(Copy, Clone, Debug)]
+pub enum FftStrategy {
     Classic,
     Unrolled,
 }
 
-const FFT_STRATEGY: FftStrategy = FftStrategy::Classic;
+pub(crate) const DEFAULT_STRATEGY: FftStrategy = FftStrategy::Classic;
 
 pub(crate) type FftRootTable<F> = Vec<Vec<F>>;
 
@@ -68,11 +69,12 @@ fn fft_unrolled_root_table<F: Field>(n: usize) -> FftRootTable<F> {
 #[inline]
 fn fft_dispatch<F: Field>(
     input: &[F],
+    strategy: FftStrategy,
     zero_factor: Option<usize>,
     root_table: Option<FftRootTable<F>>,
 ) -> Vec<F> {
     let n = input.len();
-    match FFT_STRATEGY {
+    match strategy {
         FftStrategy::Classic => fft_classic(
             input,
             zero_factor.unwrap_or(0),
@@ -88,28 +90,30 @@ fn fft_dispatch<F: Field>(
 
 #[inline]
 pub fn fft<F: Field>(poly: &PolynomialCoeffs<F>) -> PolynomialValues<F> {
-    fft_with_options(poly, None, None)
+    fft_with_options(poly, DEFAULT_STRATEGY, None, None)
 }
 
 #[inline]
 pub fn fft_with_options<F: Field>(
     poly: &PolynomialCoeffs<F>,
+    strategy: FftStrategy,
     zero_factor: Option<usize>,
     root_table: Option<FftRootTable<F>>,
 ) -> PolynomialValues<F> {
     let PolynomialCoeffs { coeffs } = poly;
     PolynomialValues {
-        values: fft_dispatch(coeffs, zero_factor, root_table),
+        values: fft_dispatch(coeffs, strategy, zero_factor, root_table),
     }
 }
 
 #[inline]
 pub fn ifft<F: Field>(poly: &PolynomialValues<F>) -> PolynomialCoeffs<F> {
-    ifft_with_options(poly, None, None)
+    ifft_with_options(poly, DEFAULT_STRATEGY, None, None)
 }
 
 pub fn ifft_with_options<F: Field>(
     poly: &PolynomialValues<F>,
+    strategy: FftStrategy,
     zero_factor: Option<usize>,
     root_table: Option<FftRootTable<F>>,
 ) -> PolynomialCoeffs<F> {
@@ -118,7 +122,7 @@ pub fn ifft_with_options<F: Field>(
     let n_inv = F::inverse_2exp(lg_n);
 
     let PolynomialValues { values } = poly;
-    let mut coeffs = fft_dispatch(values, zero_factor, root_table);
+    let mut coeffs = fft_dispatch(values, strategy, zero_factor, root_table);
 
     // We reverse all values except the first, and divide each by n.
     coeffs[0] *= n_inv;
@@ -305,7 +309,7 @@ fn fft_unrolled<F: Field>(input: &[F], r_orig: usize, root_table: FftRootTable<F
 #[cfg(test)]
 mod tests {
     use crate::field::crandall_field::CrandallField;
-    use crate::field::fft::{fft, fft_with_options, ifft};
+    use crate::field::fft::{fft, fft_with_options, ifft, FftStrategy};
     use crate::field::field_types::Field;
     use crate::polynomial::polynomial::{PolynomialCoeffs, PolynomialValues};
     use crate::util::{log2_ceil, log2_strict};
@@ -332,10 +336,15 @@ mod tests {
             assert_eq!(interpolated_coefficients.coeffs[i], F::ZERO);
         }
 
-        for r in 0..4 {
+        for strategy in [FftStrategy::Classic, FftStrategy::Unrolled] {
-            // expand coefficients by factor 2^r by filling with zeros
+            for r in 0..4 {
-            let zero_tail = coefficients.lde(r);
+                // expand coefficients by factor 2^r by filling with zeros
-            assert_eq!(fft(&zero_tail), fft_with_options(&zero_tail, Some(r), None));
+                let zero_tail = coefficients.lde(r);
+                assert_eq!(
+                    fft(&zero_tail),
+                    fft_with_options(&zero_tail, strategy, Some(r), None)
+                );
+            }
         }
     }
 

src/fri/commitment.rs

@@ -1,6 +1,7 @@
 use rayon::prelude::*;
 
 use crate::field::extension_field::Extendable;
+use crate::field::fft::DEFAULT_STRATEGY;
 use crate::field::field_types::Field;
 use crate::fri::proof::FriProof;
 use crate::fri::prover::fri_proof;
@@ -92,7 +93,12 @@ impl<F: Field> PolynomialBatchCommitment<F> {
             .map(|p| {
                 assert_eq!(p.len(), degree, "Polynomial degrees inconsistent");
                 p.lde(rate_bits)
-                    .coset_fft_with_options(F::coset_shift(), Some(rate_bits), None)
+                    .coset_fft_with_options(
+                        F::coset_shift(),
+                        DEFAULT_STRATEGY,
+                        Some(rate_bits),
+                        None,
+                    )
                     .values
             })
             .chain(

src/polynomial/polynomial.rs

@@ -6,8 +6,8 @@ use anyhow::{ensure, Result};
 use serde::{Deserialize, Serialize};
 
 use crate::field::extension_field::{Extendable, FieldExtension};
-use crate::field::fft::FftRootTable;
 use crate::field::fft::{fft, fft_with_options, ifft};
+use crate::field::fft::{FftRootTable, FftStrategy, DEFAULT_STRATEGY};
 use crate::field::field_types::Field;
 use crate::util::log2_strict;
 
@@ -57,7 +57,7 @@ impl<F: Field> PolynomialValues<F> {
 
     pub fn lde(&self, rate_bits: usize) -> Self {
         let coeffs = ifft(self).lde(rate_bits);
-        fft_with_options(&coeffs, Some(rate_bits), None)
+        fft_with_options(&coeffs, DEFAULT_STRATEGY, Some(rate_bits), None)
     }
 
     pub fn degree(&self) -> usize {
@@ -151,7 +151,7 @@ impl<F: Field> PolynomialCoeffs<F> {
         polys.into_iter().map(|p| p.lde(rate_bits)).collect()
     }
 
-    pub(crate) fn lde(&self, rate_bits: usize) -> Self {
+    pub fn lde(&self, rate_bits: usize) -> Self {
         self.padded(self.len() << rate_bits)
     }
 
@@ -211,21 +211,23 @@ impl<F: Field> PolynomialCoeffs<F> {
 
     pub fn fft_with_options(
         &self,
+        strategy: FftStrategy,
         zero_factor: Option<usize>,
         root_table: Option<FftRootTable<F>>,
     ) -> PolynomialValues<F> {
-        fft_with_options(self, zero_factor, root_table)
+        fft_with_options(self, strategy, zero_factor, root_table)
     }
 
     /// Returns the evaluation of the polynomial on the coset `shift*H`.
     pub fn coset_fft(&self, shift: F) -> PolynomialValues<F> {
-        self.coset_fft_with_options(shift, None, None)
+        self.coset_fft_with_options(shift, DEFAULT_STRATEGY, None, None)
     }
 
     /// Returns the evaluation of the polynomial on the coset `shift*H`.
     pub fn coset_fft_with_options(
         &self,
         shift: F,
+        strategy: FftStrategy,
         zero_factor: Option<usize>,
         root_table: Option<FftRootTable<F>>,
     ) -> PolynomialValues<F> {
@@ -235,7 +237,7 @@ impl<F: Field> PolynomialCoeffs<F> {
             .map(|(r, &c)| r * c)
             .collect::<Vec<_>>()
             .into();
-        modified_poly.fft_with_options(zero_factor, root_table)
+        modified_poly.fft_with_options(strategy, zero_factor, root_table)
     }
 
     pub fn to_extension<const D: usize>(&self) -> PolynomialCoeffs<F::Extension>