From 892868f901cfb727f401e521fb62056cf05d931a Mon Sep 17 00:00:00 2001
From: Saveliy Yusufov <sy2685@columbia.edu>
Date: Tue, 18 Jun 2024 16:58:19 -0400
Subject: [PATCH] Generate twiddles once for untangling & Planner

---
 examples/profile.rs | 13 ++++++++++++-
 src/fft.rs          | 36 +++++++++++++++++++++++-------------
 2 files changed, 35 insertions(+), 14 deletions(-)

diff --git a/examples/profile.rs b/examples/profile.rs
index 6725046..941f420 100644
--- a/examples/profile.rs
+++ b/examples/profile.rs
@@ -1,9 +1,11 @@
 use std::env;
 use std::str::FromStr;
 
+use phastft::fft::r2c_fft_f64;
 use phastft::fft_64;
 use phastft::planner::Direction;
 
+#[allow(dead_code)]
 fn benchmark_fft(num_qubits: usize) {
     let n = 1 << num_qubits;
     let mut reals: Vec<f64> = (1..=n).map(|i| i as f64).collect();
@@ -11,10 +13,19 @@ fn benchmark_fft(num_qubits: usize) {
     fft_64(&mut reals, &mut imags, Direction::Forward);
 }
 
+fn benchmark_r2c_fft(n: usize) {
+    let big_n = 1 << n;
+    let reals: Vec<f64> = (1..=big_n).map(|i| i as f64).collect();
+    let mut output_re = vec![0.0; big_n];
+    let mut output_im = vec![0.0; big_n];
+    r2c_fft_f64(&reals, &mut output_re, &mut output_im);
+}
+
 fn main() {
     let args: Vec<String> = env::args().collect();
     assert_eq!(args.len(), 2, "Usage {} <n>", args[0]);
 
     let n = usize::from_str(&args[1]).unwrap();
-    benchmark_fft(n);
+    // benchmark_fft(n);
+    benchmark_r2c_fft(n);
 }
diff --git a/src/fft.rs b/src/fft.rs
index ea322f0..3acacdf 100644
--- a/src/fft.rs
+++ b/src/fft.rs
@@ -1,11 +1,15 @@
 //! Implementation of Real valued FFT
 use std::simd::prelude::f64x8;
 
-use crate::twiddles::{generate_twiddles_simd_32, generate_twiddles_simd_64};
-use crate::{fft_32, fft_64, twiddles::generate_twiddles, Direction};
+use crate::planner::{Planner32, Planner64};
+use crate::twiddles::filter_twiddles;
+use crate::{
+    fft_32_with_opts_and_plan, fft_64, fft_64_with_opts_and_plan, twiddles::generate_twiddles,
+    Direction, Options,
+};
 
 macro_rules! impl_r2c_fft {
-    ($func_name:ident, $precision:ty, $fft_func:ident, $gen_twiddles_simd:ident) => {
+    ($func_name:ident, $precision:ty, $planner:ident, $fft_w_opts_and_plan:ident) => {
         /// Performs a Real-Valued Fast Fourier Transform (FFT)
         ///
         /// This function computes the FFT of a real-valued input signal and produces
@@ -57,9 +61,19 @@ macro_rules! impl_r2c_fft {
             let (mut z_even, mut z_odd): (Vec<_>, Vec<_>) =
                 input_re.chunks_exact(2).map(|c| (c[0], c[1])).unzip();
 
-            // Z = np.fft.fft(z)
-            $fft_func(&mut z_even, &mut z_odd, Direction::Forward);
+            let mut planner = <$planner>::new(big_n, Direction::Forward);
+
+            // save these for the untanngling step
+            let twiddle_re = planner.twiddles_re.clone();
+            let twiddle_im = planner.twiddles_im.clone();
+
+            // We only need (N / 2) / 2 twiddle factors for the actual FFT call, so we filter
+            filter_twiddles(&mut planner.twiddles_re, &mut planner.twiddles_im);
+
+            let opts = Options::guess_options(z_even.len());
+            $fft_w_opts_and_plan(&mut z_even, &mut z_odd, &opts, &mut planner);
 
+            // Z = np.fft.fft(z)
             let mut z_x_re = vec![0.0; big_n / 2];
             let mut z_x_im = vec![0.0; big_n / 2];
             let mut z_y_re = vec![0.0; big_n / 2];
@@ -111,12 +125,6 @@ macro_rules! impl_r2c_fft {
             z_y_re[0] = w;
             z_y_im[0] = v;
 
-            let (twiddle_re, twiddle_im): (Vec<$precision>, Vec<$precision>) = if big_n / 2 >= 16 {
-                $gen_twiddles_simd(big_n / 2, Direction::Forward)
-            } else {
-                generate_twiddles(big_n / 2, Direction::Forward)
-            };
-
             // Zall = np.concatenate([Zx + W*Zy, Zx - W*Zy])
             let (output_re_first_half, output_re_second_half) = output_re.split_at_mut(big_n / 2);
             let (output_im_first_half, output_im_second_half) = output_im.split_at_mut(big_n / 2);
@@ -156,8 +164,8 @@ macro_rules! impl_r2c_fft {
     };
 }
 
-impl_r2c_fft!(r2c_fft_f32, f32, fft_32, generate_twiddles_simd_32);
-impl_r2c_fft!(r2c_fft_f64, f64, fft_64, generate_twiddles_simd_64);
+impl_r2c_fft!(r2c_fft_f32, f32, Planner32, fft_32_with_opts_and_plan);
+impl_r2c_fft!(r2c_fft_f64, f64, Planner64, fft_64_with_opts_and_plan);
 
 /// Performs a Real-Valued, Inverse, Fast Fourier Transform (FFT)
 ///
@@ -245,6 +253,8 @@ pub fn r2c_ifft_f64(reals: &mut [f64], imags: &mut [f64], output: &mut [f64]) {
 mod tests {
     use utilities::assert_float_closeness;
 
+    use crate::fft_32;
+
     use super::*;
 
     macro_rules! impl_r2c_vs_c2c_test {