FFT/LDE benches (#188)

And expose FftStrategy so we can bench both
This commit is contained in:
Daniel Lubarov 2021-08-18 08:36:40 -07:00 committed by GitHub
parent f2ed563da9
commit c31c06d227
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 95 additions and 23 deletions

View File

@ -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

48
benches/ffts.rs Normal file
View File

@ -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);

View File

@ -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)
});
}

View File

@ -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 {
// expand coefficients by factor 2^r by filling with zeros
let zero_tail = coefficients.lde(r);
assert_eq!(fft(&zero_tail), fft_with_options(&zero_tail, Some(r), None));
for strategy in [FftStrategy::Classic, FftStrategy::Unrolled] {
for r in 0..4 {
// expand coefficients by factor 2^r by filling with zeros
let zero_tail = coefficients.lde(r);
assert_eq!(
fft(&zero_tail),
fft_with_options(&zero_tail, strategy, Some(r), None)
);
}
}
}

View File

@ -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(

View File

@ -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>