diff --git a/Cargo.toml b/Cargo.toml index 73b486dbb..52abf8147 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,8 +19,13 @@ thiserror = "1.0" [dev-dependencies] assert_matches = "1.5.0" +criterion = "0.3" modinverse = "0.1.0" num-bigint = "0.4.0" +[[bench]] +name = "fft" +harness = false + [[example]] name = "sum" diff --git a/benches/fft.rs b/benches/fft.rs new file mode 100644 index 000000000..2a672e562 --- /dev/null +++ b/benches/fft.rs @@ -0,0 +1,57 @@ +// SPDX-License-Identifier: MPL-2.0 + +use criterion::{criterion_group, criterion_main, Criterion}; +use prio::fft; +use prio::finite_field::{Field, FieldElement}; +use prio::polynomial; + +pub fn fft(c: &mut Criterion) { + let test_sizes = [16, 256, 1024, 4096]; + for size in test_sizes.iter() { + let mut rng = rand::thread_rng(); + let mut inp = vec![Field::zero(); *size]; + let mut outp = vec![Field::zero(); *size]; + for i in 0..*size { + inp[i] = Field::rand(&mut rng); + } + + // Test recursive FFT, including auxiliary data computation. + c.bench_function(&format!("recursive/{}", *size), |b| { + b.iter(|| { + let mut mem = polynomial::PolyAuxMemory::new(*size / 2); + polynomial::poly_fft( + &mut outp, + &inp, + &mem.roots_2n, + *size, + false, + &mut mem.fft_memory, + ) + }) + }); + + // Test recursive FFT, but amortize auxiliary data computation across all of the + // invocations of the call. + let mut mem = polynomial::PolyAuxMemory::new(*size / 2); + c.bench_function(&format!("recursive/{} (amortized)", *size), |b| { + b.iter(|| { + polynomial::poly_fft( + &mut outp, + &inp, + &mem.roots_2n, + *size, + false, + &mut mem.fft_memory, + ) + }) + }); + + // Test iteratigve FFT. + c.bench_function(&format!("iterative/{}", *size), |b| { + b.iter(|| fft::discrete_fourier_transform::(&mut outp, &inp)) + }); + } +} + +criterion_group!(benches, fft); +criterion_main!(benches); diff --git a/src/lib.rs b/src/lib.rs index 2c4e35414..d02551413 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -14,7 +14,7 @@ pub mod encrypt; pub mod fft; pub mod finite_field; mod fp; -mod polynomial; +pub mod polynomial; mod prng; pub mod server; pub mod util;