From c31c06d2278102593b0807c0fc5d3ad568c86846 Mon Sep 17 00:00:00 2001 From: Daniel Lubarov Date: Wed, 18 Aug 2021 08:36:40 -0700 Subject: [PATCH] FFT/LDE benches (#188) And expose FftStrategy so we can bench both --- Cargo.toml | 8 ++++++ benches/ffts.rs | 48 ++++++++++++++++++++++++++++++++++++ benches/field_arithmetic.rs | 7 +++--- src/field/fft.rs | 33 ++++++++++++++++--------- src/fri/commitment.rs | 8 +++++- src/polynomial/polynomial.rs | 14 ++++++----- 6 files changed, 95 insertions(+), 23 deletions(-) create mode 100644 benches/ffts.rs diff --git a/Cargo.toml b/Cargo.toml index 4d40b44e..420aafc8 100644 --- a/Cargo.toml +++ b/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 diff --git a/benches/ffts.rs b/benches/ffts.rs new file mode 100644 index 00000000..2808d951 --- /dev/null +++ b/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(c: &mut Criterion, strategy: FftStrategy) { + let mut group = c.benchmark_group(&format!("fft-{:?}<{}>", strategy, type_name::())); + + 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(c: &mut Criterion, strategy: FftStrategy) { + const RATE_BITS: usize = 3; + + let mut group = c.benchmark_group(&format!("lde-{:?}<{}>", strategy, type_name::())); + + 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::(c, FftStrategy::Classic); + bench_ffts::(c, FftStrategy::Unrolled); + bench_ldes::(c, FftStrategy::Classic); + bench_ldes::(c, FftStrategy::Unrolled); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/benches/field_arithmetic.rs b/benches/field_arithmetic.rs index e5ec13ab..47c11b45 100644 --- a/benches/field_arithmetic.rs +++ b/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(c: &mut Criterion) { - c.bench_function(&format!("{} mul", type_name::()), |b| { + c.bench_function(&format!("mul<{}>", type_name::()), |b| { b.iter_batched( || (F::rand(), F::rand()), |(x, y)| x * y, @@ -14,7 +13,7 @@ pub(crate) fn bench_field(c: &mut Criterion) { ) }); - c.bench_function(&format!("{} try_inverse", type_name::()), |b| { + c.bench_function(&format!("try_inverse<{}>", type_name::()), |b| { b.iter_batched(|| F::rand(), |x| x.try_inverse(), BatchSize::SmallInput) }); } diff --git a/src/field/fft.rs b/src/field/fft.rs index 1566db84..4a60f888 100644 --- a/src/field/fft.rs +++ b/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 = Vec>; @@ -68,11 +69,12 @@ fn fft_unrolled_root_table(n: usize) -> FftRootTable { #[inline] fn fft_dispatch( input: &[F], + strategy: FftStrategy, zero_factor: Option, root_table: Option>, ) -> Vec { 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( #[inline] pub fn fft(poly: &PolynomialCoeffs) -> PolynomialValues { - fft_with_options(poly, None, None) + fft_with_options(poly, DEFAULT_STRATEGY, None, None) } #[inline] pub fn fft_with_options( poly: &PolynomialCoeffs, + strategy: FftStrategy, zero_factor: Option, root_table: Option>, ) -> PolynomialValues { 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(poly: &PolynomialValues) -> PolynomialCoeffs { - ifft_with_options(poly, None, None) + ifft_with_options(poly, DEFAULT_STRATEGY, None, None) } pub fn ifft_with_options( poly: &PolynomialValues, + strategy: FftStrategy, zero_factor: Option, root_table: Option>, ) -> PolynomialCoeffs { @@ -118,7 +122,7 @@ pub fn ifft_with_options( 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(input: &[F], r_orig: usize, root_table: FftRootTable PolynomialBatchCommitment { .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( diff --git a/src/polynomial/polynomial.rs b/src/polynomial/polynomial.rs index 386d49ae..055c0813 100644 --- a/src/polynomial/polynomial.rs +++ b/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 PolynomialValues { 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 PolynomialCoeffs { 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 PolynomialCoeffs { pub fn fft_with_options( &self, + strategy: FftStrategy, zero_factor: Option, root_table: Option>, ) -> PolynomialValues { - 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 { - 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, root_table: Option>, ) -> PolynomialValues { @@ -235,7 +237,7 @@ impl PolynomialCoeffs { .map(|(r, &c)| r * c) .collect::>() .into(); - modified_poly.fft_with_options(zero_factor, root_table) + modified_poly.fft_with_options(strategy, zero_factor, root_table) } pub fn to_extension(&self) -> PolynomialCoeffs