[AArch64 / NEON] Vectorize matrix transpose and fuse twiddle multiplication in MixedRadixSmall

benches/bench_rustfft.rs

@@ -219,6 +219,30 @@ fn bench_good_thomas(b: &mut Bencher, width: usize, height: usize) {
 #[bench] fn good_thomas_2048_3(b: &mut Bencher) { bench_good_thomas(b,  2048, 3); }
 #[bench] fn good_thomas_2048_2187(b: &mut Bencher) { bench_good_thomas(b,  2048, 2187); }
 
+/// Times just the FFT execution (not allocation and pre-calculation)
+/// for a given length, specific to the Good-Thomas algorithm
+fn bench_good_thomas_f64(b: &mut Bencher, width: usize, height: usize) {
+
+    let mut planner = rustfft::FftPlanner::new();
+    let width_fft = planner.plan_fft_forward(width);
+    let height_fft = planner.plan_fft_forward(height);
+
+    let fft : Arc<dyn Fft<f64>> = Arc::new(GoodThomasAlgorithm::new(width_fft, height_fft));
+
+    let mut buffer = vec![Complex::zero(); width * height];
+    let mut scratch = vec![Complex::zero(); fft.get_inplace_scratch_len()];
+    b.iter(|| {fft.process_with_scratch(&mut buffer, &mut scratch);} );
+}
+
+#[bench] fn good_thomas_64_0002_3(b: &mut Bencher) { bench_good_thomas_f64(b,  2, 3); }
+#[bench] fn good_thomas_64_0003_4(b: &mut Bencher) { bench_good_thomas_f64(b,  3, 4); }
+#[bench] fn good_thomas_64_0004_5(b: &mut Bencher) { bench_good_thomas_f64(b,  4, 5); }
+#[bench] fn good_thomas_64_0007_32(b: &mut Bencher) { bench_good_thomas_f64(b, 7, 32); }
+#[bench] fn good_thomas_64_0032_27(b: &mut Bencher) { bench_good_thomas_f64(b,  32, 27); }
+#[bench] fn good_thomas_64_0256_243(b: &mut Bencher) { bench_good_thomas_f64(b,  256, 243); }
+#[bench] fn good_thomas_64_2048_3(b: &mut Bencher) { bench_good_thomas_f64(b,  2048, 3); }
+#[bench] fn good_thomas_64_2048_2187(b: &mut Bencher) { bench_good_thomas_f64(b,  2048, 2187); }
+
 /// Times just the FFT setup (not execution)
 /// for a given length, specific to the Good-Thomas algorithm
 fn bench_good_thomas_setup(b: &mut Bencher, width: usize, height: usize) {
@@ -310,6 +334,30 @@ fn bench_mixed_radix(b: &mut Bencher, width: usize, height: usize) {
 #[bench] fn mixed_radix_2048_3(b: &mut Bencher) { bench_mixed_radix(b,  2048, 3); }
 #[bench] fn mixed_radix_2048_2187(b: &mut Bencher) { bench_mixed_radix(b,  2048, 2187); }
 
+/// Times just the FFT execution (not allocation and pre-calculation)
+/// for a given length, specific to the Mixed-Radix algorithm
+fn bench_mixed_radix_f64(b: &mut Bencher, width: usize, height: usize) {
+
+    let mut planner = rustfft::FftPlanner::new();
+    let width_fft = planner.plan_fft_forward(width);
+    let height_fft = planner.plan_fft_forward(height);
+
+    let fft : Arc<dyn Fft<_>> = Arc::new(MixedRadix::new(width_fft, height_fft));
+
+    let mut buffer = vec![Complex{re: 0_f64, im: 0_f64}; fft.len()];
+    let mut scratch = vec![Complex{re: 0_f64, im: 0_f64}; fft.get_inplace_scratch_len()];
+    b.iter(|| {fft.process_with_scratch(&mut buffer, &mut scratch);} );
+}
+
+#[bench] fn mixed_radix_64_0002_3(b: &mut Bencher) { bench_mixed_radix_f64(b,  2, 3); }
+#[bench] fn mixed_radix_64_0003_4(b: &mut Bencher) { bench_mixed_radix_f64(b,  3, 4); }
+#[bench] fn mixed_radix_64_0004_5(b: &mut Bencher) { bench_mixed_radix_f64(b,  4, 5); }
+#[bench] fn mixed_radix_64_0007_32(b: &mut Bencher) { bench_mixed_radix_f64(b, 7, 32); }
+#[bench] fn mixed_radix_64_0032_27(b: &mut Bencher) { bench_mixed_radix_f64(b,  32, 27); }
+#[bench] fn mixed_radix_64_0256_243(b: &mut Bencher) { bench_mixed_radix_f64(b,  256, 243); }
+#[bench] fn mixed_radix_64_2048_3(b: &mut Bencher) { bench_mixed_radix_f64(b,  2048, 3); }
+#[bench] fn mixed_radix_64_2048_2187(b: &mut Bencher) { bench_mixed_radix_f64(b,  2048, 2187); }
+
 fn plan_butterfly_fft(len: usize) -> Arc<dyn Fft<f32>> {
     match len {
         2 => Arc::new(Butterfly2::new(FftDirection::Forward)),
@@ -344,6 +392,40 @@ fn bench_mixed_radix_small(b: &mut Bencher, width: usize, height: usize) {
 #[bench] fn mixed_radix_small_0004_5(b: &mut Bencher) { bench_mixed_radix_small(b,  4, 5); }
 #[bench] fn mixed_radix_small_0007_32(b: &mut Bencher) { bench_mixed_radix_small(b, 7, 32); }
 
+fn plan_butterfly_fft_f64(len: usize) -> Arc<dyn Fft<f64>> {
+    match len {
+        2 => Arc::new(Butterfly2::new(FftDirection::Forward)),
+        3 => Arc::new(Butterfly3::new(FftDirection::Forward)),
+        4 => Arc::new(Butterfly4::new(FftDirection::Forward)),
+        5 => Arc::new(Butterfly5::new(FftDirection::Forward)),
+        6 => Arc::new(Butterfly6::new(FftDirection::Forward)),
+        7 => Arc::new(Butterfly7::new(FftDirection::Forward)),
+        8 => Arc::new(Butterfly8::new(FftDirection::Forward)),
+        16 => Arc::new(Butterfly16::new(FftDirection::Forward)),
+        32 => Arc::new(Butterfly32::new(FftDirection::Forward)),
+        _ => panic!("Invalid butterfly size: {}", len),
+    }
+}
+
+/// Times just the FFT execution (not allocation and pre-calculation)
+/// for a given length, specific to the MixedRadixSmall algorithm
+fn bench_mixed_radix_small_f64(b: &mut Bencher, width: usize, height: usize) {
+
+    let width_fft = plan_butterfly_fft_f64(width);
+    let height_fft = plan_butterfly_fft_f64(height);
+
+    let fft : Arc<dyn Fft<_>> = Arc::new(MixedRadixSmall::new(width_fft, height_fft));
+
+    let mut signal = vec![Complex{re: 0_f64, im: 0_f64}; width * height];
+    let mut spectrum = signal.clone();
+    b.iter(|| {fft.process_with_scratch(&mut signal, &mut spectrum);} );
+}
+
+#[bench] fn mixed_radix_small_64_0002_3(b: &mut Bencher) { bench_mixed_radix_small_f64(b,  2, 3); }
+#[bench] fn mixed_radix_small_64_0003_4(b: &mut Bencher) { bench_mixed_radix_small_f64(b,  3, 4); }
+#[bench] fn mixed_radix_small_64_0004_5(b: &mut Bencher) { bench_mixed_radix_small_f64(b,  4, 5); }
+#[bench] fn mixed_radix_small_64_0007_32(b: &mut Bencher) { bench_mixed_radix_small_f64(b, 7, 32); }
+
 /// Times just the FFT execution (not allocation and pre-calculation)
 /// for a given length, specific to the Mixed-Radix Double Butterfly algorithm
 fn bench_good_thomas_small(b: &mut Bencher, width: usize, height: usize) {
@@ -363,6 +445,25 @@ fn bench_good_thomas_small(b: &mut Bencher, width: usize, height: usize) {
 #[bench] fn good_thomas_small_0004_5(b: &mut Bencher) { bench_good_thomas_small(b,  4, 5); }
 #[bench] fn good_thomas_small_0007_32(b: &mut Bencher) { bench_good_thomas_small(b, 7, 32); }
 
+/// Times just the FFT execution (not allocation and pre-calculation)
+/// for a given length, specific to the Mixed-Radix Double Butterfly algorithm
+fn bench_good_thomas_small_f64(b: &mut Bencher, width: usize, height: usize) {
+
+    let width_fft = plan_butterfly_fft_f64(width);
+    let height_fft = plan_butterfly_fft_f64(height);
+
+    let fft : Arc<dyn Fft<_>> = Arc::new(GoodThomasAlgorithmSmall::new(width_fft, height_fft));
+
+    let mut signal = vec![Complex{re: 0_f64, im: 0_f64}; width * height];
+    let mut spectrum = signal.clone();
+    b.iter(|| {fft.process_with_scratch(&mut signal, &mut spectrum);} );
+}
+
+#[bench] fn good_thomas_small_64_0002_3(b: &mut Bencher) { bench_good_thomas_small_f64(b,  2, 3); }
+#[bench] fn good_thomas_small_64_0003_4(b: &mut Bencher) { bench_good_thomas_small_f64(b,  3, 4); }
+#[bench] fn good_thomas_small_64_0004_5(b: &mut Bencher) { bench_good_thomas_small_f64(b,  4, 5); }
+#[bench] fn good_thomas_small_64_0007_32(b: &mut Bencher) { bench_good_thomas_small_f64(b, 7, 32); }
+
 
 /// Times just the FFT execution (not allocation and pre-calculation)
 /// for a given length, specific to Rader's algorithm

src/algorithm/mixed_radix.rs

@@ -324,13 +324,16 @@ impl<T: FftNum> MixedRadixSmall<T> {
         // STEP 2: perform FFTs of size `height`
         self.height_size_fft.process_with_scratch(scratch, buffer);
 
-        // STEP 3: Apply twiddle factors
-        for (element, twiddle) in scratch.iter_mut().zip(self.twiddles.iter()) {
-            *element = *element * twiddle;
-        }
-
-        // STEP 4: transpose again
-        unsafe { array_utils::transpose_small(self.height, self.width, scratch, buffer) };
+        // STEP 3 & 4: Apply twiddle factors and transpose
+        unsafe {
+            array_utils::transpose_small_twiddle(
+                self.height,
+                self.width,
+                scratch,
+                buffer,
+                &self.twiddles,
+            )
+        };
 
         // STEP 5: perform FFTs of size `width`
         self.width_size_fft
@@ -353,13 +356,16 @@ impl<T: FftNum> MixedRadixSmall<T> {
         // STEP 2: perform FFTs of size `height`
         self.height_size_fft.process_with_scratch(output, scratch);
 
-        // STEP 3: Apply twiddle factors
-        for (element, twiddle) in output.iter_mut().zip(self.twiddles.iter()) {
-            *element = *element * twiddle;
-        }
-
-        // STEP 4: transpose again
-        unsafe { array_utils::transpose_small(self.height, self.width, output, scratch) };
+        // STEP 3 & 4: Apply twiddle factors and transpose
+        unsafe {
+            array_utils::transpose_small_twiddle(
+                self.height,
+                self.width,
+                output,
+                scratch,
+                &self.twiddles,
+            )
+        };
 
         // STEP 5: perform FFTs of size `width`
         self.width_size_fft.process_with_scratch(scratch, output);
@@ -381,13 +387,16 @@ impl<T: FftNum> MixedRadixSmall<T> {
         // STEP 2: perform FFTs of size `height`
         self.height_size_fft.process_with_scratch(output, input);
 
-        // STEP 3: Apply twiddle factors
-        for (element, twiddle) in output.iter_mut().zip(self.twiddles.iter()) {
-            *element = *element * twiddle;
-        }
-
-        // STEP 4: transpose again
-        unsafe { array_utils::transpose_small(self.height, self.width, output, input) };
+        // STEP 3 & 4: Apply twiddle factors and transpose
+        unsafe {
+            array_utils::transpose_small_twiddle(
+                self.height,
+                self.width,
+                output,
+                input,
+                &self.twiddles,
+            )
+        };
 
         // STEP 5: perform FFTs of size `width`
         self.width_size_fft.process_with_scratch(input, output);

src/array_utils.rs

@@ -17,6 +17,36 @@ pub unsafe fn transpose_small<T: Copy>(width: usize, height: usize, input: &[T],
     }
 }
 
+pub unsafe fn transpose_small_twiddle<T: FftNum>(
+    width: usize,
+    height: usize,
+    input: &[Complex<T>],
+    output: &mut [Complex<T>],
+    twiddles: &[Complex<T>],
+) {
+    assert!(input.len() >= width * height);
+    assert!(output.len() >= width * height);
+    assert!(twiddles.len() >= width * height);
+
+    #[cfg(all(target_arch = "aarch64", feature = "neon"))]
+    {
+        if crate::neon::neon_utils::transpose_small_twiddle(width, height, input, output, twiddles)
+        {
+            return;
+        }
+    }
+
+    for y in 0..height {
+        for x in 0..width {
+            let in_idx = x + y * width;
+            let out_idx = y + x * height;
+            let val = *input.get_unchecked(in_idx);
+            let tw = *twiddles.get_unchecked(in_idx);
+            *output.get_unchecked_mut(out_idx) = val * tw;
+        }
+    }
+}
+
 #[allow(unused)]
 pub unsafe fn workaround_transmute<T, U>(slice: &[T]) -> &[U] {
     let ptr = slice.as_ptr() as *const U;
@@ -115,7 +145,7 @@ mod unit_tests {
     use num_traits::Zero;
 
     #[test]
-    fn test_transpose() {
+    fn test_transpose_small() {
         let sizes: Vec<usize> = (1..16).collect();
 
         for &width in &sizes {
@@ -141,6 +171,75 @@ mod unit_tests {
             }
         }
     }
+
+    #[test]
+    fn test_transpose_small_twiddle() {
+        let sizes: Vec<usize> = (1..16).collect();
+
+        for &width in &sizes {
+            for &height in &sizes {
+                let len = width * height;
+
+                // Test f32
+                let input_f32: Vec<Complex<f32>> = random_signal(len);
+                let twiddles_f32: Vec<Complex<f32>> = random_signal(len);
+                let mut output_f32 = vec![Zero::zero(); len];
+                unsafe {
+                    transpose_small_twiddle(
+                        width,
+                        height,
+                        &input_f32,
+                        &mut output_f32,
+                        &twiddles_f32,
+                    );
+                }
+                for x in 0..width {
+                    for y in 0..height {
+                        let expected = input_f32[x + y * width] * twiddles_f32[x + y * width];
+                        let actual = output_f32[y + x * height];
+                        assert!(
+                            (actual.re - expected.re).abs() < 1e-5
+                                && (actual.im - expected.im).abs() < 1e-5,
+                            "f32 mismatch at ({}, {}) for {}x{}",
+                            x,
+                            y,
+                            width,
+                            height
+                        );
+                    }
+                }
+
+                // Test f64
+                let input_f64: Vec<Complex<f64>> = random_signal(len);
+                let twiddles_f64: Vec<Complex<f64>> = random_signal(len);
+                let mut output_f64 = vec![Zero::zero(); len];
+                unsafe {
+                    transpose_small_twiddle(
+                        width,
+                        height,
+                        &input_f64,
+                        &mut output_f64,
+                        &twiddles_f64,
+                    );
+                }
+                for x in 0..width {
+                    for y in 0..height {
+                        let expected = input_f64[x + y * width] * twiddles_f64[x + y * width];
+                        let actual = output_f64[y + x * height];
+                        assert!(
+                            (actual.re - expected.re).abs() < 1e-10
+                                && (actual.im - expected.im).abs() < 1e-10,
+                            "f64 mismatch at ({}, {}) for {}x{}",
+                            x,
+                            y,
+                            width,
+                            height
+                        );
+                    }
+                }
+            }
+        }
+    }
 }
 
 // A utility that validates the following conditions, then calls chunk_fn() on each chunk of buffer. Passes the entire scratch buffer with each call.

src/lib.rs

@@ -440,7 +440,7 @@ pub use self::sse::sse_planner::FftPlannerSse;
 
 // Algorithms implemented to use Neon instructions. Only compiled on AArch64, and only compiled if the "neon" feature flag is set.
 #[cfg(all(target_arch = "aarch64", feature = "neon"))]
-mod neon;
+pub(crate) mod neon;
 
 // If we're not on AArch64, or if the "neon" feature was disabled, keep a stub implementation around that has the same API, but does nothing
 // That way, users can write code using the Neon planner and compile it on any platform

src/neon/mod.rs

@@ -8,7 +8,7 @@ pub mod neon_butterflies;
 pub mod neon_prime_butterflies;
 pub mod neon_radix4;
 
-mod neon_utils;
+pub(crate) mod neon_utils;
 
 pub mod neon_planner;