Add speed test for iterative vs. recursive FFT

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"

benches/fft.rs

@@ -0,0 +1,33 @@
+// SPDX-License-Identifier: MPL-2.0
+
+use criterion::{criterion_group, criterion_main, Criterion};
+
+use prio::benchmarked::{benchmarked_iterative_fft, benchmarked_recursive_fft};
+use prio::finite_field::{Field, FieldElement};
+
+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);
+        }
+
+        c.bench_function(&format!("iterative/{}", *size), |b| {
+            b.iter(|| {
+                benchmarked_iterative_fft(&mut outp, &inp);
+            })
+        });
+
+        c.bench_function(&format!("recursive/{}", *size), |b| {
+            b.iter(|| {
+                benchmarked_recursive_fft(&mut outp, &inp);
+            })
+        });
+    }
+}
+
+criterion_group!(benches, fft);
+criterion_main!(benches);

src/benchmarked.rs

@@ -0,0 +1,26 @@
+// SPDX-License-Identifier: MPL-2.0
+
+//! This packages provides wrappers around internal components of this crate that we wnat to
+//! benchmark, but which we don't want to expose in the public API.
+
+use crate::fft::discrete_fourier_transform;
+use crate::finite_field::{Field, FieldElement};
+use crate::polynomial::{poly_fft, PolyAuxMemory};
+
+/// Sets `outp` to the Discrete Fourier Transform (DFT) using an iterative FFT algorithm.
+pub fn benchmarked_iterative_fft<F: FieldElement>(outp: &mut [F], inp: &[F]) {
+    discrete_fourier_transform(outp, inp).expect("encountered unexpected error");
+}
+
+/// Sets `outp` to the Discrete Fourier Transform (DFT) using a recursive FFT algorithm.
+pub fn benchmarked_recursive_fft(outp: &mut [Field], inp: &[Field]) {
+    let mut mem = PolyAuxMemory::new(inp.len() / 2);
+    poly_fft(
+        outp,
+        inp,
+        &mem.roots_2n,
+        inp.len(),
+        false,
+        &mut mem.fft_memory,
+    )
+}

src/fft.rs

@@ -67,6 +67,7 @@ pub fn discrete_fourier_transform<F: FieldElement>(
 }
 
 /// Sets `outp` to the inverse of the DFT of `inp`.
+#[allow(dead_code)]
 pub fn discrete_fourier_transform_inv<F: FieldElement>(
     outp: &mut [F],
     inp: &[F],

src/lib.rs

@@ -9,9 +9,10 @@
 //!
 //! For now we only support 0 / 1 vectors.
 
+pub mod benchmarked;
 pub mod client;
 pub mod encrypt;
-pub mod fft;
+mod fft;
 pub mod finite_field;
 mod fp;
 mod polynomial;