Make planner reusable

Cargo.toml

@@ -15,8 +15,14 @@ num-traits = "0.2.18"
 multiversion = "0.7"
 
 [dev-dependencies]
-utilities = { path = "utilities" }
+criterion = "0.5.1"
 fftw = "0.8.0"
+rand = "0.8.5"
+utilities = { path = "utilities" }
+
+[[bench]]
+name = "bench"
+harness = false
 
 [profile.release]
 codegen-units = 1

benches/bench.rs

@@ -0,0 +1,119 @@
+use criterion::{black_box, criterion_group, criterion_main, Criterion};
+use num_traits::Float;
+use phastft::{
+    fft_32_with_opts_and_plan, fft_64_with_opts_and_plan,
+    options::Options,
+    planner::{Direction, Planner32, Planner64},
+};
+use rand::{
+    distributions::{Distribution, Standard},
+    thread_rng, Rng,
+};
+
+const LENGTHS: &[usize] = &[
+    6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28,
+];
+
+fn generate_numbers<T: Float>(n: usize) -> (Vec<T>, Vec<T>)
+where
+    Standard: Distribution<T>,
+{
+    let mut rng = thread_rng();
+
+    let samples: Vec<T> = (&mut rng).sample_iter(Standard).take(2 * n).collect();
+
+    let mut reals = vec![T::zero(); n];
+    let mut imags = vec![T::zero(); n];
+
+    for ((z_re, z_im), rand_chunk) in reals
+        .iter_mut()
+        .zip(imags.iter_mut())
+        .zip(samples.chunks_exact(2))
+    {
+        *z_re = rand_chunk[0];
+        *z_im = rand_chunk[1];
+    }
+
+    (reals, imags)
+}
+
+fn benchmark_forward_f32(c: &mut Criterion) {
+    let options = Options::default();
+
+    for n in LENGTHS.iter() {
+        let len = 1 << n;
+        let id = format!("FFT Forward f32 {} elements", len);
+        c.bench_function(&id, |b| {
+            let (mut reals, mut imags) = generate_numbers(len);
+            let planner = Planner32::new(len, Direction::Forward);
+            b.iter(|| {
+                black_box(fft_32_with_opts_and_plan(
+                    &mut reals, &mut imags, &options, &planner,
+                ));
+            });
+        });
+    }
+}
+
+fn benchmark_inverse_f32(c: &mut Criterion) {
+    let options = Options::default();
+
+    for n in LENGTHS.iter() {
+        let len = 1 << n;
+        let id = format!("FFT Inverse f32 {} elements", len);
+        c.bench_function(&id, |b| {
+            let (mut reals, mut imags) = generate_numbers(len);
+            let planner = Planner32::new(len, Direction::Reverse);
+            b.iter(|| {
+                black_box(fft_32_with_opts_and_plan(
+                    &mut reals, &mut imags, &options, &planner,
+                ));
+            });
+        });
+    }
+}
+
+fn benchmark_forward_f64(c: &mut Criterion) {
+    let options = Options::default();
+
+    for n in LENGTHS.iter() {
+        let len = 1 << n;
+        let id = format!("FFT Forward f64 {} elements", len);
+        c.bench_function(&id, |b| {
+            let (mut reals, mut imags) = generate_numbers(len);
+            let planner = Planner64::new(len, Direction::Forward);
+            b.iter(|| {
+                black_box(fft_64_with_opts_and_plan(
+                    &mut reals, &mut imags, &options, &planner,
+                ));
+            });
+        });
+    }
+}
+
+fn benchmark_inverse_f64(c: &mut Criterion) {
+    let options = Options::default();
+
+    for n in LENGTHS.iter() {
+        let len = 1 << n;
+        let id = format!("FFT Inverse f64 {} elements", len);
+        c.bench_function(&id, |b| {
+            let (mut reals, mut imags) = generate_numbers(len);
+            let planner = Planner64::new(len, Direction::Reverse);
+            b.iter(|| {
+                black_box(fft_64_with_opts_and_plan(
+                    &mut reals, &mut imags, &options, &planner,
+                ));
+            });
+        });
+    }
+}
+
+criterion_group!(
+    benches,
+    benchmark_forward_f32,
+    benchmark_inverse_f32,
+    benchmark_forward_f64,
+    benchmark_inverse_f64
+);
+criterion_main!(benches);

src/lib.rs

@@ -48,35 +48,15 @@ macro_rules! impl_fft_for {
                 imags.len()
             );
 
-            let mut planner = <$planner>::new(reals.len(), Direction::Forward);
+            let planner = <$planner>::new(reals.len(), direction);
             assert!(
                 planner.num_twiddles().is_power_of_two()
                     && planner.num_twiddles() == reals.len() / 2
             );
 
             let opts = Options::guess_options(reals.len());
 
-            match direction {
-                Direction::Reverse => {
-                    for z_im in imags.iter_mut() {
-                        *z_im = -*z_im;
-                    }
-                }
-                _ => (),
-            }
-
-            $opts_and_plan(reals, imags, &opts, &mut planner);
-
-            match direction {
-                Direction::Reverse => {
-                    let scaling_factor = (reals.len() as $precision).recip();
-                    for (z_re, z_im) in reals.iter_mut().zip(imags.iter_mut()) {
-                        *z_re *= scaling_factor;
-                        *z_im *= -scaling_factor;
-                    }
-                }
-                _ => (),
-            }
+            $opts_and_plan(reals, imags, &opts, &planner);
         }
     };
 }
@@ -112,30 +92,39 @@ macro_rules! impl_fft_with_opts_and_plan_for {
             reals: &mut [$precision],
             imags: &mut [$precision],
             opts: &Options,
-            planner: &mut $planner,
+            planner: &$planner,
         ) {
             assert!(reals.len() == imags.len() && reals.len().is_power_of_two());
             let n: usize = reals.len().ilog2() as usize;
 
-            let twiddles_re = &mut planner.twiddles_re;
-            let twiddles_im = &mut planner.twiddles_im;
+            let mut twiddles_re = planner.twiddles_re.clone();
+            let mut twiddles_im = planner.twiddles_im.clone();
 
             // We shouldn't be able to execute FFT if the # of twiddles isn't equal to the distance
             // between pairs
             assert!(twiddles_re.len() == reals.len() / 2 && twiddles_im.len() == imags.len() / 2);
 
+            match planner.direction {
+                Direction::Reverse => {
+                    for z_im in imags.iter_mut() {
+                        *z_im = -*z_im;
+                    }
+                }
+                _ => (),
+            }
+
             for t in (0..n).rev() {
                 let dist = 1 << t;
                 let chunk_size = dist << 1;
 
                 if chunk_size > 4 {
                     if t < n - 1 {
-                        filter_twiddles(twiddles_re, twiddles_im);
+                        (twiddles_re, twiddles_im) = filter_twiddles(&twiddles_re, &twiddles_im);
                     }
                     if chunk_size >= $lanes * 2 {
-                        $simd_butterfly_kernel(reals, imags, twiddles_re, twiddles_im, dist);
+                        $simd_butterfly_kernel(reals, imags, &twiddles_re, &twiddles_im, dist);
                     } else {
-                        fft_chunk_n(reals, imags, twiddles_re, twiddles_im, dist);
+                        fft_chunk_n(reals, imags, &twiddles_re, &twiddles_im, dist);
                     }
                 } else if chunk_size == 2 {
                     fft_chunk_2(reals, imags);
@@ -153,6 +142,17 @@ macro_rules! impl_fft_with_opts_and_plan_for {
                 cobra_apply(reals, n);
                 cobra_apply(imags, n);
             }
+
+            match planner.direction {
+                Direction::Reverse => {
+                    let scaling_factor = (reals.len() as $precision).recip();
+                    for (z_re, z_im) in reals.iter_mut().zip(imags.iter_mut()) {
+                        *z_re *= scaling_factor;
+                        *z_im *= -scaling_factor;
+                    }
+                }
+                _ => (),
+            }
         }
     };
 }
@@ -179,7 +179,7 @@ mod tests {
 
     use utilities::rustfft::num_complex::Complex;
     use utilities::rustfft::FftPlanner;
-    use utilities::{assert_float_closeness, gen_random_signal};
+    use utilities::{assert_float_closeness, gen_random_signal, gen_random_signal_f32};
 
     use super::*;
 
@@ -308,4 +308,65 @@ mod tests {
             }
         }
     }
+
+    #[test]
+    fn fft_64_with_opts_and_plan_vs_fft_64() {
+        let num_points = 4096;
+
+        let mut reals = vec![0.0; num_points];
+        let mut imags = vec![0.0; num_points];
+        gen_random_signal(&mut reals, &mut imags);
+
+        let mut re = reals.clone();
+        let mut im = imags.clone();
+
+        let mut planner = Planner64::new(num_points, Direction::Forward);
+        let opts = Options::guess_options(reals.len());
+        fft_64_with_opts_and_plan(&mut reals, &mut imags, &opts, &mut planner);
+
+        fft_64(&mut re, &mut im, Direction::Forward);
+
+        reals
+            .iter()
+            .zip(imags.iter())
+            .zip(re.iter())
+            .zip(im.iter())
+            .for_each(|(((r, i), z_re), z_im)| {
+                assert_float_closeness(*r, *z_re, 1e-6);
+                assert_float_closeness(*i, *z_im, 1e-6);
+            });
+    }
+
+    #[test]
+    fn fft_32_with_opts_and_plan_vs_fft_64() {
+        let dirs = [Direction::Forward, Direction::Reverse];
+
+        for direction in dirs {
+            for n in 4..14 {
+                let num_points = 1 << n;
+                let mut reals = vec![0.0; num_points];
+                let mut imags = vec![0.0; num_points];
+                gen_random_signal_f32(&mut reals, &mut imags);
+
+                let mut re = reals.clone();
+                let mut im = imags.clone();
+
+                let mut planner = Planner32::new(num_points, direction);
+                let opts = Options::guess_options(reals.len());
+                fft_32_with_opts_and_plan(&mut reals, &mut imags, &opts, &mut planner);
+
+                fft_32(&mut re, &mut im, direction);
+
+                reals
+                    .iter()
+                    .zip(imags.iter())
+                    .zip(re.iter())
+                    .zip(im.iter())
+                    .for_each(|(((r, i), z_re), z_im)| {
+                        assert_float_closeness(*r, *z_re, 1e-6);
+                        assert_float_closeness(*i, *z_im, 1e-6);
+                    });
+            }
+        }
+    }
 }

src/planner.rs

@@ -25,6 +25,8 @@ macro_rules! impl_planner_for {
             pub twiddles_re: Vec<$precision>,
             /// The imaginary components of the twiddle factors
             pub twiddles_im: Vec<$precision>,
+            /// The direction of the FFT associated with this `Planner`
+            pub direction: Direction,
         }
         impl $struct_name {
             /// Create a `Planner` for an FFT of size `num_points`.
@@ -35,12 +37,18 @@ macro_rules! impl_planner_for {
             /// # Panics
             ///
             /// Panics if `num_points < 1` or if `num_points` is __not__ a power of 2.
-            pub fn new(num_points: usize, _direction: Direction) -> Self {
+            pub fn new(num_points: usize, direction: Direction) -> Self {
                 assert!(num_points > 0 && num_points.is_power_of_two());
+                let dir = match direction {
+                    Direction::Forward => Direction::Forward,
+                    Direction::Reverse => Direction::Reverse,
+                };
+
                 if num_points <= 4 {
                     return Self {
                         twiddles_re: vec![],
                         twiddles_im: vec![],
+                        direction: dir,
                     };
                 }
 
@@ -57,6 +65,7 @@ macro_rules! impl_planner_for {
                 Self {
                     twiddles_re,
                     twiddles_im,
+                    direction: dir,
                 }
             }