From ac8f1b1e3d3fb47ea03196474a1d990951bf8d0c Mon Sep 17 00:00:00 2001 From: Wang-Yue Date: Sun, 6 Sep 2026 01:21:45 +0900 Subject: [PATCH 1/7] neon: vectorize 2x2 block transpose in neon_utils --- src/array_utils.rs | 14 +++++++--- src/lib.rs | 2 +- src/neon/mod.rs | 2 +- src/neon/neon_utils.rs | 61 ++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 74 insertions(+), 5 deletions(-) diff --git a/src/array_utils.rs b/src/array_utils.rs index bfae7df0..2848c26d 100644 --- a/src/array_utils.rs +++ b/src/array_utils.rs @@ -6,9 +6,16 @@ use std::ops::{Deref, DerefMut}; /// Given an array of size width * height, representing a flattened 2D array, /// transpose the rows and columns of that 2D array into the output /// benchmarking shows that loop tiling isn't effective for small arrays (in the range of 50x50 or smaller) -pub unsafe fn transpose_small(width: usize, height: usize, input: &[T], output: &mut [T]) { - for x in 0..width { - for y in 0..height { +pub unsafe fn transpose_small(width: usize, height: usize, input: &[T], output: &mut [T]) { + #[cfg(all(target_arch = "aarch64", feature = "neon"))] + { + if crate::neon::neon_utils::transpose_small(width, height, input, output) { + return; + } + } + + for y in 0..height { + for x in 0..width { let input_index = x + y * width; let output_index = y + x * height; @@ -17,6 +24,7 @@ pub unsafe fn transpose_small(width: usize, height: usize, input: &[T], } } + #[allow(unused)] pub unsafe fn workaround_transmute(slice: &[T]) -> &[U] { let ptr = slice.as_ptr() as *const U; diff --git a/src/lib.rs b/src/lib.rs index 2bcb30dd..89ee732a 100644 --- a/src/lib.rs +++ b/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 diff --git a/src/neon/mod.rs b/src/neon/mod.rs index 20b936cb..e7c7d2e6 100644 --- a/src/neon/mod.rs +++ b/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; diff --git a/src/neon/neon_utils.rs b/src/neon/neon_utils.rs index 0bc8db2e..a89d4c14 100644 --- a/src/neon/neon_utils.rs +++ b/src/neon/neon_utils.rs @@ -1,4 +1,5 @@ use core::arch::aarch64::*; +use num_complex::Complex; // __ __ _ _ _________ _ _ _ // | \/ | __ _| |_| |__ |___ /___ \| |__ (_) |_ @@ -246,6 +247,66 @@ impl Rotate90F64 { } } +pub unsafe fn transpose_small( + width: usize, + height: usize, + input: &[T], + output: &mut [T], +) -> bool { + use std::any::TypeId; + if TypeId::of::() == TypeId::of::>() { + let p_in = input.as_ptr() as *const f64; + let p_out = output.as_mut_ptr() as *mut f64; + + let mut y = 0; + while y + 2 <= height { + let in_row0 = (y * width) * 2; + let in_row1 = ((y + 1) * width) * 2; + let mut x = 0; + while x + 2 <= width { + let in_idx0 = in_row0 + x * 2; + let in_idx1 = in_row1 + x * 2; + let a0 = vld1q_f64(p_in.add(in_idx0)); + let a1 = vld1q_f64(p_in.add(in_idx0 + 2)); + let b0 = vld1q_f64(p_in.add(in_idx1)); + let b1 = vld1q_f64(p_in.add(in_idx1 + 2)); + + let out_idx0 = (y + x * height) * 2; + let out_idx1 = (y + (x + 1) * height) * 2; + vst1q_f64(p_out.add(out_idx0), a0); + vst1q_f64(p_out.add(out_idx0 + 2), b0); + vst1q_f64(p_out.add(out_idx1), a1); + vst1q_f64(p_out.add(out_idx1 + 2), b1); + + x += 2; + } + while x < width { + let in_idx0 = in_row0 + x * 2; + let in_idx1 = in_row1 + x * 2; + let a0 = vld1q_f64(p_in.add(in_idx0)); + let b0 = vld1q_f64(p_in.add(in_idx1)); + let out_idx0 = (y + x * height) * 2; + vst1q_f64(p_out.add(out_idx0), a0); + vst1q_f64(p_out.add(out_idx0 + 2), b0); + x += 1; + } + y += 2; + } + while y < height { + let in_row = (y * width) * 2; + for x in 0..width { + let in_idx = in_row + x * 2; + let out_idx = (y + x * height) * 2; + let a = vld1q_f64(p_in.add(in_idx)); + vst1q_f64(p_out.add(out_idx), a); + } + y += 1; + } + return true; + } + false +} + #[cfg(test)] mod unit_tests { use super::*; From 71f2a1fa389fceb0879484808ddba6e897add998 Mon Sep 17 00:00:00 2001 From: Wang-Yue Date: Sun, 6 Sep 2026 01:22:53 +0900 Subject: [PATCH 2/7] neon: fuse twiddle multiplication into transpose pass in MixedRadixSmall --- src/algorithm/mixed_radix.rs | 51 +++++++++------- src/array_utils.rs | 25 +++++++- src/neon/neon_utils.rs | 115 +++++++++++++++++++++++++++++++++++ 3 files changed, 169 insertions(+), 22 deletions(-) diff --git a/src/algorithm/mixed_radix.rs b/src/algorithm/mixed_radix.rs index 317d4cb5..a4630b11 100644 --- a/src/algorithm/mixed_radix.rs +++ b/src/algorithm/mixed_radix.rs @@ -324,13 +324,16 @@ impl MixedRadixSmall { // 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 MixedRadixSmall { // 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 MixedRadixSmall { // 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); diff --git a/src/array_utils.rs b/src/array_utils.rs index 2848c26d..50fc48e7 100644 --- a/src/array_utils.rs +++ b/src/array_utils.rs @@ -24,8 +24,31 @@ pub unsafe fn transpose_small(width: usize, height: usize, in } } +pub unsafe fn transpose_small_twiddle( + width: usize, + height: usize, + input: &[Complex], + output: &mut [Complex], + twiddles: &[Complex], +) { + #[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(slice: &[T]) -> &[U] { let ptr = slice.as_ptr() as *const U; let len = slice.len(); diff --git a/src/neon/neon_utils.rs b/src/neon/neon_utils.rs index a89d4c14..636812ff 100644 --- a/src/neon/neon_utils.rs +++ b/src/neon/neon_utils.rs @@ -1,5 +1,6 @@ use core::arch::aarch64::*; use num_complex::Complex; +use crate::FftNum; // __ __ _ _ _________ _ _ _ // | \/ | __ _| |_| |__ |___ /___ \| |__ (_) |_ @@ -307,6 +308,120 @@ pub unsafe fn transpose_small( false } +#[inline(always)] +unsafe fn neon_complex_mul_f64( + val: float64x2_t, + tw: float64x2_t, +) -> float64x2_t { + let temp = vcombine_f64(vneg_f64(vget_high_f64(val)), vget_low_f64(val)); + let sum = vmulq_laneq_f64::<0>(val, tw); + vfmaq_laneq_f64::<1>(sum, temp, tw) +} + +pub unsafe fn transpose_small_twiddle( + width: usize, + height: usize, + input: &[Complex], + output: &mut [Complex], + twiddles: &[Complex], +) -> bool { + use std::any::TypeId; + if TypeId::of::() == TypeId::of::() { + let p_in = input.as_ptr() as *const f64; + let p_out = output.as_mut_ptr() as *mut f64; + let p_tw = twiddles.as_ptr() as *const f64; + + let mut y = 0; + while y + 2 <= height { + let in_row0 = (y * width) * 2; + let in_row1 = ((y + 1) * width) * 2; + let mut x = 0; + while x + 2 <= width { + let in_idx0 = in_row0 + x * 2; + let in_idx1 = in_row1 + x * 2; + + let a0 = vld1q_f64(p_in.add(in_idx0)); + let tw_a0 = vld1q_f64(p_tw.add(in_idx0)); + let a1 = vld1q_f64(p_in.add(in_idx0 + 2)); + let tw_a1 = vld1q_f64(p_tw.add(in_idx0 + 2)); + + let b0 = vld1q_f64(p_in.add(in_idx1)); + let tw_b0 = vld1q_f64(p_tw.add(in_idx1)); + let b1 = vld1q_f64(p_in.add(in_idx1 + 2)); + let tw_b1 = vld1q_f64(p_tw.add(in_idx1 + 2)); + + let res_a0 = neon_complex_mul_f64(a0, tw_a0); + let res_a1 = neon_complex_mul_f64(a1, tw_a1); + let res_b0 = neon_complex_mul_f64(b0, tw_b0); + let res_b1 = neon_complex_mul_f64(b1, tw_b1); + + let out_idx0 = (y + x * height) * 2; + let out_idx1 = (y + (x + 1) * height) * 2; + + vst1q_f64(p_out.add(out_idx0), res_a0); + vst1q_f64(p_out.add(out_idx0 + 2), res_b0); + vst1q_f64(p_out.add(out_idx1), res_a1); + vst1q_f64(p_out.add(out_idx1 + 2), res_b1); + + x += 2; + } + while x < width { + let in_idx0 = in_row0 + x * 2; + let in_idx1 = in_row1 + x * 2; + + let a0 = vld1q_f64(p_in.add(in_idx0)); + let tw_a0 = vld1q_f64(p_tw.add(in_idx0)); + let b0 = vld1q_f64(p_in.add(in_idx1)); + let tw_b0 = vld1q_f64(p_tw.add(in_idx1)); + + let res_a0 = neon_complex_mul_f64(a0, tw_a0); + let res_b0 = neon_complex_mul_f64(b0, tw_b0); + + let out_idx0 = (y + x * height) * 2; + vst1q_f64(p_out.add(out_idx0), res_a0); + vst1q_f64(p_out.add(out_idx0 + 2), res_b0); + x += 1; + } + y += 2; + } + while y < height { + let in_row = (y * width) * 2; + for x in 0..width { + let in_idx = in_row + x * 2; + let out_idx = (y + x * height) * 2; + let a = vld1q_f64(p_in.add(in_idx)); + let tw = vld1q_f64(p_tw.add(in_idx)); + let res = neon_complex_mul_f64(a, tw); + vst1q_f64(p_out.add(out_idx), res); + } + y += 1; + } + return true; + } else if TypeId::of::() == TypeId::of::() { + let p_in = input.as_ptr() as *const f32; + let p_out = output.as_mut_ptr() as *mut f32; + let p_tw = twiddles.as_ptr() as *const f32; + for y in 0..height { + for x in 0..width { + let in_idx = (x + y * width) * 2; + let out_idx = (y + x * height) * 2; + let left = vld1_f32(p_in.add(in_idx)); + let right = vld1_f32(p_tw.add(in_idx)); + let left_q = vcombine_f32(left, left); + let right_q = vcombine_f32(right, right); + let temp1 = vtrn1q_f32(right_q, right_q); + let temp2 = vtrn2q_f32(right_q, vnegq_f32(right_q)); + let temp3 = vmulq_f32(temp2, left_q); + let temp4 = vrev64q_f32(temp3); + let res = vfmaq_f32(temp4, temp1, left_q); + vst1_f32(p_out.add(out_idx), vget_low_f32(res)); + } + } + return true; + } + false +} + #[cfg(test)] mod unit_tests { use super::*; From c82055bbbe40d19cbb0d57f0a07edefa6baf0fae Mon Sep 17 00:00:00 2001 From: Wang-Yue Date: Sun, 6 Sep 2026 01:43:00 +0900 Subject: [PATCH 3/7] neon: accelerate matrix transpose with tiled NEON vectorization --- src/algorithm/good_thomas_algorithm.rs | 7 +- src/algorithm/mixed_radix.rs | 19 ++- src/array_utils.rs | 69 ++++++--- src/neon/neon_utils.rs | 187 +++++++++++++++++++------ 4 files changed, 202 insertions(+), 80 deletions(-) diff --git a/src/algorithm/good_thomas_algorithm.rs b/src/algorithm/good_thomas_algorithm.rs index 2122310a..5fb63b34 100644 --- a/src/algorithm/good_thomas_algorithm.rs +++ b/src/algorithm/good_thomas_algorithm.rs @@ -4,7 +4,6 @@ use std::sync::Arc; use num_complex::Complex; use num_integer::Integer; use strength_reduce::StrengthReducedUsize; -use transpose; use crate::array_utils; use crate::{common::FftNum, FftDirection}; @@ -237,7 +236,7 @@ impl GoodThomasAlgorithm { .process_with_scratch(scratch, width_scratch); // transpose - transpose::transpose(scratch, buffer, self.width, self.height); + unsafe { array_utils::transpose(scratch, buffer, self.width, self.height) }; // run FFTs of size 'height' self.height_size_fft @@ -262,7 +261,7 @@ impl GoodThomasAlgorithm { let (scratch, inner_scratch) = scratch.split_at_mut(self.len()); // transpose - transpose::transpose(output, scratch, self.width, self.height); + unsafe { array_utils::transpose(output, scratch, self.width, self.height) }; // run FFTs of size 'height' self.height_size_fft @@ -291,7 +290,7 @@ impl GoodThomasAlgorithm { .process_with_scratch(output, width_scratch); // transpose - transpose::transpose(output, input, self.width, self.height); + unsafe { array_utils::transpose(output, input, self.width, self.height) }; // run FFTs of size 'height' let height_scratch = if scratch.len() > output.len() { diff --git a/src/algorithm/mixed_radix.rs b/src/algorithm/mixed_radix.rs index a4630b11..cb6bcc9b 100644 --- a/src/algorithm/mixed_radix.rs +++ b/src/algorithm/mixed_radix.rs @@ -3,7 +3,6 @@ use std::sync::Arc; use num_complex::Complex; use num_traits::Zero; -use transpose; use crate::array_utils; use crate::{common::FftNum, twiddles, FftDirection}; @@ -130,7 +129,7 @@ impl MixedRadix { let (scratch, inner_scratch) = scratch.split_at_mut(self.len()); // STEP 1: transpose - transpose::transpose(buffer, scratch, self.width, self.height); + unsafe { array_utils::transpose(buffer, scratch, self.width, self.height) }; // STEP 2: perform FFTs of size `height` let height_scratch = if inner_scratch.len() > buffer.len() { @@ -147,14 +146,14 @@ impl MixedRadix { } // STEP 4: transpose again - transpose::transpose(scratch, buffer, self.height, self.width); + unsafe { array_utils::transpose(scratch, buffer, self.height, self.width) }; // STEP 5: perform FFTs of size `width` self.width_size_fft .process_outofplace_with_scratch(buffer, scratch, inner_scratch); // STEP 6: transpose again - transpose::transpose(scratch, buffer, self.width, self.height); + unsafe { array_utils::transpose(scratch, buffer, self.width, self.height) }; } fn perform_fft_immut( @@ -164,7 +163,7 @@ impl MixedRadix { scratch_raw: &mut [Complex], ) { // STEP 1: transpose - transpose::transpose(input, output, self.width, self.height); + unsafe { array_utils::transpose(input, output, self.width, self.height) }; // STEP 2: perform FFTs of size `height` self.height_size_fft @@ -178,14 +177,14 @@ impl MixedRadix { let (scratch, inner_scratch) = scratch_raw.split_at_mut(self.len()); // STEP 4: transpose again - transpose::transpose(output, scratch, self.height, self.width); + unsafe { array_utils::transpose(output, scratch, self.height, self.width) }; // STEP 5: perform FFTs of size `width` self.width_size_fft .process_with_scratch(scratch, inner_scratch); // STEP 6: transpose again - transpose::transpose(scratch, output, self.width, self.height); + unsafe { array_utils::transpose(scratch, output, self.width, self.height) }; } fn perform_fft_out_of_place( @@ -197,7 +196,7 @@ impl MixedRadix { // SIX STEP FFT: // STEP 1: transpose - transpose::transpose(input, output, self.width, self.height); + unsafe { array_utils::transpose(input, output, self.width, self.height) }; // STEP 2: perform FFTs of size `height` let height_scratch = if scratch.len() > input.len() { @@ -214,7 +213,7 @@ impl MixedRadix { } // STEP 4: transpose again - transpose::transpose(output, input, self.height, self.width); + unsafe { array_utils::transpose(output, input, self.height, self.width) }; // STEP 5: perform FFTs of size `width` let width_scratch = if scratch.len() > output.len() { @@ -226,7 +225,7 @@ impl MixedRadix { .process_with_scratch(input, width_scratch); // STEP 6: transpose again - transpose::transpose(input, output, self.width, self.height); + unsafe { array_utils::transpose(input, output, self.width, self.height) }; } } boilerplate_fft!( diff --git a/src/array_utils.rs b/src/array_utils.rs index 50fc48e7..e0c6615d 100644 --- a/src/array_utils.rs +++ b/src/array_utils.rs @@ -3,25 +3,35 @@ use crate::Complex; use crate::FftNum; use std::ops::{Deref, DerefMut}; +/// Transpose the input array into the output array. +/// /// Given an array of size width * height, representing a flattened 2D array, -/// transpose the rows and columns of that 2D array into the output -/// benchmarking shows that loop tiling isn't effective for small arrays (in the range of 50x50 or smaller) -pub unsafe fn transpose_small(width: usize, height: usize, input: &[T], output: &mut [T]) { +/// transpose the rows and columns of that 2D array into the output. +pub unsafe fn transpose( + input: &[T], + output: &mut [T], + width: usize, + height: usize, +) { #[cfg(all(target_arch = "aarch64", feature = "neon"))] { - if crate::neon::neon_utils::transpose_small(width, height, input, output) { + if crate::neon::neon_utils::transpose(width, height, input, output) { return; } } - for y in 0..height { - for x in 0..width { - let input_index = x + y * width; - let output_index = y + x * height; + transpose::transpose(input, output, width, height); +} - *output.get_unchecked_mut(output_index) = *input.get_unchecked(input_index); - } - } +/// Given an array of size width * height, representing a flattened 2D array, +/// transpose the rows and columns of that 2D array into the output +pub unsafe fn transpose_small( + width: usize, + height: usize, + input: &[T], + output: &mut [T], +) { + transpose(input, output, width, height); } pub unsafe fn transpose_small_twiddle( @@ -147,25 +157,44 @@ mod unit_tests { #[test] fn test_transpose() { - let sizes: Vec = (1..16).collect(); + let sizes: Vec = (1..35).collect(); for &width in &sizes { for &height in &sizes { let len = width * height; - let input: Vec> = random_signal(len); - let mut output = vec![Zero::zero(); len]; - - unsafe { transpose_small(width, height, &input, &mut output) }; + // Test f32 + let input_f32: Vec> = random_signal(len); + let mut output_f32 = vec![Zero::zero(); len]; + unsafe { transpose(&input_f32, &mut output_f32, width, height) }; + for x in 0..width { + for y in 0..height { + assert_eq!( + input_f32[x + y * width], + output_f32[y + x * height], + "f32: x = {}, y = {}, width = {}, height = {}", + x, + y, + width, + height + ); + } + } + // Test f64 + let input_f64: Vec> = random_signal(len); + let mut output_f64 = vec![Zero::zero(); len]; + unsafe { transpose(&input_f64, &mut output_f64, width, height) }; for x in 0..width { for y in 0..height { assert_eq!( - input[x + y * width], - output[y + x * height], - "x = {}, y = {}", + input_f64[x + y * width], + output_f64[y + x * height], + "f64: x = {}, y = {}, width = {}, height = {}", x, - y + y, + width, + height ); } } diff --git a/src/neon/neon_utils.rs b/src/neon/neon_utils.rs index 636812ff..fea2f9cf 100644 --- a/src/neon/neon_utils.rs +++ b/src/neon/neon_utils.rs @@ -248,7 +248,141 @@ impl Rotate90F64 { } } -pub unsafe fn transpose_small( +#[inline(always)] +unsafe fn transpose_f64(p_in: *const f64, p_out: *mut f64, width: usize, height: usize) { + const TILE_SIZE: usize = 32; + + let mut by = 0; + while by < height { + let block_h = (height - by).min(TILE_SIZE); + let mut bx = 0; + while bx < width { + let block_w = (width - bx).min(TILE_SIZE); + + let mut y = 0; + while y + 2 <= block_h { + let cy = by + y; + let in_r0 = (cy * width + bx) * 2; + let in_r1 = ((cy + 1) * width + bx) * 2; + + let mut x = 0; + while x + 2 <= block_w { + let cx = bx + x; + let a0 = vld1q_f64(p_in.add(in_r0 + x * 2)); + let a1 = vld1q_f64(p_in.add(in_r0 + x * 2 + 2)); + + let b0 = vld1q_f64(p_in.add(in_r1 + x * 2)); + let b1 = vld1q_f64(p_in.add(in_r1 + x * 2 + 2)); + + let out_c0 = (cy + cx * height) * 2; + let out_c1 = (cy + (cx + 1) * height) * 2; + + vst1q_f64(p_out.add(out_c0), a0); + vst1q_f64(p_out.add(out_c0 + 2), b0); + + vst1q_f64(p_out.add(out_c1), a1); + vst1q_f64(p_out.add(out_c1 + 2), b1); + + x += 2; + } + while x < block_w { + let cx = bx + x; + let a0 = vld1q_f64(p_in.add(in_r0 + x * 2)); + let b0 = vld1q_f64(p_in.add(in_r1 + x * 2)); + + let out_c0 = (cy + cx * height) * 2; + + vst1q_f64(p_out.add(out_c0), a0); + vst1q_f64(p_out.add(out_c0 + 2), b0); + + x += 1; + } + y += 2; + } + while y < block_h { + let cy = by + y; + let in_r = (cy * width + bx) * 2; + for x in 0..block_w { + let cx = bx + x; + let a0 = vld1q_f64(p_in.add(in_r + x * 2)); + let out_c = (cy + cx * height) * 2; + vst1q_f64(p_out.add(out_c), a0); + } + y += 1; + } + + bx += TILE_SIZE; + } + by += TILE_SIZE; + } +} + +#[inline(always)] +unsafe fn transpose_f32(p_in: *const f32, p_out: *mut f32, width: usize, height: usize) { + const TILE_SIZE: usize = 32; + + let mut by = 0; + while by < height { + let block_h = (height - by).min(TILE_SIZE); + let mut bx = 0; + while bx < width { + let block_w = (width - bx).min(TILE_SIZE); + + let mut y = 0; + while y + 2 <= block_h { + let cy = by + y; + let in_r0 = (cy * width + bx) * 2; + let in_r1 = ((cy + 1) * width + bx) * 2; + + let mut x = 0; + while x + 2 <= block_w { + let cx = bx + x; + let row0 = vld1q_f32(p_in.add(in_r0 + x * 2)); + let row1 = vld1q_f32(p_in.add(in_r1 + x * 2)); + + let transposed = transpose_complex_2x2_f32(row0, row1); + + let out_c0 = (cy + cx * height) * 2; + let out_c1 = (cy + (cx + 1) * height) * 2; + + vst1q_f32(p_out.add(out_c0), transposed[0]); + vst1q_f32(p_out.add(out_c1), transposed[1]); + + x += 2; + } + while x < block_w { + let cx = bx + x; + let a0 = vld1_f32(p_in.add(in_r0 + x * 2)); + let b0 = vld1_f32(p_in.add(in_r1 + x * 2)); + + let out_c0 = (cy + cx * height) * 2; + + vst1_f32(p_out.add(out_c0), a0); + vst1_f32(p_out.add(out_c0 + 2), b0); + + x += 1; + } + y += 2; + } + while y < block_h { + let cy = by + y; + let in_r = (cy * width + bx) * 2; + for x in 0..block_w { + let cx = bx + x; + let a0 = vld1_f32(p_in.add(in_r + x * 2)); + let out_c = (cy + cx * height) * 2; + vst1_f32(p_out.add(out_c), a0); + } + y += 1; + } + + bx += TILE_SIZE; + } + by += TILE_SIZE; + } +} + +pub unsafe fn transpose( width: usize, height: usize, input: &[T], @@ -258,51 +392,12 @@ pub unsafe fn transpose_small( if TypeId::of::() == TypeId::of::>() { let p_in = input.as_ptr() as *const f64; let p_out = output.as_mut_ptr() as *mut f64; - - let mut y = 0; - while y + 2 <= height { - let in_row0 = (y * width) * 2; - let in_row1 = ((y + 1) * width) * 2; - let mut x = 0; - while x + 2 <= width { - let in_idx0 = in_row0 + x * 2; - let in_idx1 = in_row1 + x * 2; - let a0 = vld1q_f64(p_in.add(in_idx0)); - let a1 = vld1q_f64(p_in.add(in_idx0 + 2)); - let b0 = vld1q_f64(p_in.add(in_idx1)); - let b1 = vld1q_f64(p_in.add(in_idx1 + 2)); - - let out_idx0 = (y + x * height) * 2; - let out_idx1 = (y + (x + 1) * height) * 2; - vst1q_f64(p_out.add(out_idx0), a0); - vst1q_f64(p_out.add(out_idx0 + 2), b0); - vst1q_f64(p_out.add(out_idx1), a1); - vst1q_f64(p_out.add(out_idx1 + 2), b1); - - x += 2; - } - while x < width { - let in_idx0 = in_row0 + x * 2; - let in_idx1 = in_row1 + x * 2; - let a0 = vld1q_f64(p_in.add(in_idx0)); - let b0 = vld1q_f64(p_in.add(in_idx1)); - let out_idx0 = (y + x * height) * 2; - vst1q_f64(p_out.add(out_idx0), a0); - vst1q_f64(p_out.add(out_idx0 + 2), b0); - x += 1; - } - y += 2; - } - while y < height { - let in_row = (y * width) * 2; - for x in 0..width { - let in_idx = in_row + x * 2; - let out_idx = (y + x * height) * 2; - let a = vld1q_f64(p_in.add(in_idx)); - vst1q_f64(p_out.add(out_idx), a); - } - y += 1; - } + transpose_f64(p_in, p_out, width, height); + return true; + } else if TypeId::of::() == TypeId::of::>() { + let p_in = input.as_ptr() as *const f32; + let p_out = output.as_mut_ptr() as *mut f32; + transpose_f32(p_in, p_out, width, height); return true; } false From 709bbb4283c68eda1a8738d354780e6458cd4c1d Mon Sep 17 00:00:00 2001 From: Wang-Yue Date: Mon, 7 Sep 2026 16:01:40 +0900 Subject: [PATCH 4/7] Refactor array transposition functions for improved safety and performance - Removed unsafe keyword from transpose functions and added input/output length assertions. - Updated unit tests to reflect changes in the transpose function signatures. - Enhanced NEON utility functions for complex number multiplication and transposition. - Introduced templated kernel structures for handling different data types (f32, f64) in transposition. - Implemented tiled matrix transpose to optimize cache usage and SIMD operations. - Added support for twiddle factor multiplication during small matrix transpositions. --- benches/bench_rustfft.rs | 101 ++++ src/algorithm/good_thomas_algorithm.rs | 6 +- src/algorithm/mixed_radix.rs | 18 +- src/array_utils.rs | 16 +- src/neon/neon_utils.rs | 768 ++++++++++++++++++------- 5 files changed, 696 insertions(+), 213 deletions(-) diff --git a/benches/bench_rustfft.rs b/benches/bench_rustfft.rs index af9c6422..c00562a5 100644 --- a/benches/bench_rustfft.rs +++ b/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> = 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> = 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> { 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> { + 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> = 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> = 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 diff --git a/src/algorithm/good_thomas_algorithm.rs b/src/algorithm/good_thomas_algorithm.rs index 5fb63b34..4801170c 100644 --- a/src/algorithm/good_thomas_algorithm.rs +++ b/src/algorithm/good_thomas_algorithm.rs @@ -236,7 +236,7 @@ impl GoodThomasAlgorithm { .process_with_scratch(scratch, width_scratch); // transpose - unsafe { array_utils::transpose(scratch, buffer, self.width, self.height) }; + array_utils::transpose(scratch, buffer, self.width, self.height); // run FFTs of size 'height' self.height_size_fft @@ -261,7 +261,7 @@ impl GoodThomasAlgorithm { let (scratch, inner_scratch) = scratch.split_at_mut(self.len()); // transpose - unsafe { array_utils::transpose(output, scratch, self.width, self.height) }; + array_utils::transpose(output, scratch, self.width, self.height); // run FFTs of size 'height' self.height_size_fft @@ -290,7 +290,7 @@ impl GoodThomasAlgorithm { .process_with_scratch(output, width_scratch); // transpose - unsafe { array_utils::transpose(output, input, self.width, self.height) }; + array_utils::transpose(output, input, self.width, self.height); // run FFTs of size 'height' let height_scratch = if scratch.len() > output.len() { diff --git a/src/algorithm/mixed_radix.rs b/src/algorithm/mixed_radix.rs index cb6bcc9b..a8d30d86 100644 --- a/src/algorithm/mixed_radix.rs +++ b/src/algorithm/mixed_radix.rs @@ -129,7 +129,7 @@ impl MixedRadix { let (scratch, inner_scratch) = scratch.split_at_mut(self.len()); // STEP 1: transpose - unsafe { array_utils::transpose(buffer, scratch, self.width, self.height) }; + array_utils::transpose(buffer, scratch, self.width, self.height); // STEP 2: perform FFTs of size `height` let height_scratch = if inner_scratch.len() > buffer.len() { @@ -146,14 +146,14 @@ impl MixedRadix { } // STEP 4: transpose again - unsafe { array_utils::transpose(scratch, buffer, self.height, self.width) }; + array_utils::transpose(scratch, buffer, self.height, self.width); // STEP 5: perform FFTs of size `width` self.width_size_fft .process_outofplace_with_scratch(buffer, scratch, inner_scratch); // STEP 6: transpose again - unsafe { array_utils::transpose(scratch, buffer, self.width, self.height) }; + array_utils::transpose(scratch, buffer, self.width, self.height); } fn perform_fft_immut( @@ -163,7 +163,7 @@ impl MixedRadix { scratch_raw: &mut [Complex], ) { // STEP 1: transpose - unsafe { array_utils::transpose(input, output, self.width, self.height) }; + array_utils::transpose(input, output, self.width, self.height); // STEP 2: perform FFTs of size `height` self.height_size_fft @@ -177,14 +177,14 @@ impl MixedRadix { let (scratch, inner_scratch) = scratch_raw.split_at_mut(self.len()); // STEP 4: transpose again - unsafe { array_utils::transpose(output, scratch, self.height, self.width) }; + array_utils::transpose(output, scratch, self.height, self.width); // STEP 5: perform FFTs of size `width` self.width_size_fft .process_with_scratch(scratch, inner_scratch); // STEP 6: transpose again - unsafe { array_utils::transpose(scratch, output, self.width, self.height) }; + array_utils::transpose(scratch, output, self.width, self.height); } fn perform_fft_out_of_place( @@ -196,7 +196,7 @@ impl MixedRadix { // SIX STEP FFT: // STEP 1: transpose - unsafe { array_utils::transpose(input, output, self.width, self.height) }; + array_utils::transpose(input, output, self.width, self.height); // STEP 2: perform FFTs of size `height` let height_scratch = if scratch.len() > input.len() { @@ -213,7 +213,7 @@ impl MixedRadix { } // STEP 4: transpose again - unsafe { array_utils::transpose(output, input, self.height, self.width) }; + array_utils::transpose(output, input, self.height, self.width); // STEP 5: perform FFTs of size `width` let width_scratch = if scratch.len() > output.len() { @@ -225,7 +225,7 @@ impl MixedRadix { .process_with_scratch(input, width_scratch); // STEP 6: transpose again - unsafe { array_utils::transpose(input, output, self.width, self.height) }; + array_utils::transpose(input, output, self.width, self.height); } } boilerplate_fft!( diff --git a/src/array_utils.rs b/src/array_utils.rs index e0c6615d..09d97486 100644 --- a/src/array_utils.rs +++ b/src/array_utils.rs @@ -7,15 +7,18 @@ use std::ops::{Deref, DerefMut}; /// /// Given an array of size width * height, representing a flattened 2D array, /// transpose the rows and columns of that 2D array into the output. -pub unsafe fn transpose( +pub fn transpose( input: &[T], output: &mut [T], width: usize, height: usize, ) { + assert!(input.len() >= width * height); + assert!(output.len() >= width * height); + #[cfg(all(target_arch = "aarch64", feature = "neon"))] { - if crate::neon::neon_utils::transpose(width, height, input, output) { + if unsafe { crate::neon::neon_utils::transpose(width, height, input, output) } { return; } } @@ -41,6 +44,10 @@ pub unsafe fn transpose_small_twiddle( output: &mut [Complex], twiddles: &[Complex], ) { + 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) { @@ -59,6 +66,7 @@ pub unsafe fn transpose_small_twiddle( } } +#[allow(unused)] pub unsafe fn workaround_transmute(slice: &[T]) -> &[U] { let ptr = slice.as_ptr() as *const U; let len = slice.len(); @@ -166,7 +174,7 @@ mod unit_tests { // Test f32 let input_f32: Vec> = random_signal(len); let mut output_f32 = vec![Zero::zero(); len]; - unsafe { transpose(&input_f32, &mut output_f32, width, height) }; + transpose(&input_f32, &mut output_f32, width, height); for x in 0..width { for y in 0..height { assert_eq!( @@ -184,7 +192,7 @@ mod unit_tests { // Test f64 let input_f64: Vec> = random_signal(len); let mut output_f64 = vec![Zero::zero(); len]; - unsafe { transpose(&input_f64, &mut output_f64, width, height) }; + transpose(&input_f64, &mut output_f64, width, height); for x in 0..width { for y in 0..height { assert_eq!( diff --git a/src/neon/neon_utils.rs b/src/neon/neon_utils.rs index fea2f9cf..4763bf41 100644 --- a/src/neon/neon_utils.rs +++ b/src/neon/neon_utils.rs @@ -185,8 +185,12 @@ pub unsafe fn duplicate_hi_f32(values: float32x4_t) -> float32x4_t { )) } -// transpose a 2x2 complex matrix given as [x0, x1], [x2, x3] -// result is [x0, x2], [x1, x3] +// Transpose a 2x2 complex f32 matrix in NEON registers: +// left: [a0.re, a0.im, a1.re, a1.im] (row 0: a0 = col 0, a1 = col 1) +// right: [b0.re, b0.im, b1.re, b1.im] (row 1: b0 = col 0, b1 = col 1) +// Returns: +// [0]: [a0.re, a0.im, b0.re, b0.im] (col 0: a0 = row 0, b0 = row 1) +// [1]: [a1.re, a1.im, b1.re, b1.im] (col 1: a1 = row 0, b1 = row 1) #[inline(always)] pub unsafe fn transpose_complex_2x2_f32(left: float32x4_t, right: float32x4_t) -> [float32x4_t; 2] { let temp02 = extract_lo_lo_f32(left, right); @@ -248,132 +252,397 @@ impl Rotate90F64 { } } +// Multiply two complex f64 numbers (1 complex number per 128-bit register): +// val: [val.re, val.im] +// tw: [tw.re, tw.im] +// Returns: [val.re * tw.re - val.im * tw.im, val.re * tw.im + val.im * tw.re] #[inline(always)] -unsafe fn transpose_f64(p_in: *const f64, p_out: *mut f64, width: usize, height: usize) { - const TILE_SIZE: usize = 32; +pub unsafe fn neon_complex_mul_f64(val: float64x2_t, tw: float64x2_t) -> float64x2_t { + // temp: [-val.im, val.re] + let temp = vcombine_f64(vneg_f64(vget_high_f64(val)), vget_low_f64(val)); + // sum: [val.re * tw.re, val.im * tw.re] + let sum = vmulq_laneq_f64::<0>(val, tw); + // sum + temp * tw.im: [val.re * tw.re - val.im * tw.im, val.im * tw.re + val.re * tw.im] + vfmaq_laneq_f64::<1>(sum, temp, tw) +} - let mut by = 0; - while by < height { - let block_h = (height - by).min(TILE_SIZE); - let mut bx = 0; - while bx < width { - let block_w = (width - bx).min(TILE_SIZE); +// Pairwise multiply two pairs of complex f32 numbers (2 complex numbers per 128-bit register): +// val: [a0.re, a0.im, a1.re, a1.im] +// tw: [t0.re, t0.im, t1.re, t1.im] +// Returns: [ +// a0.re * t0.re - a0.im * t0.im, a0.re * t0.im + a0.im * t0.re, +// a1.re * t1.re - a1.im * t1.im, a1.re * t1.im + a1.im * t1.re, +// ] +#[inline(always)] +pub unsafe fn neon_complex_mul_f32(val: float32x4_t, tw: float32x4_t) -> float32x4_t { + let temp1 = vtrn1q_f32(tw, tw); + let temp2 = vtrn2q_f32(tw, vnegq_f32(tw)); + let temp3 = vmulq_f32(temp2, val); + let temp4 = vrev64q_f32(temp3); + vfmaq_f32(temp4, temp1, val) +} - let mut y = 0; - while y + 2 <= block_h { - let cy = by + y; - let in_r0 = (cy * width + bx) * 2; - let in_r1 = ((cy + 1) * width + bx) * 2; +// Multiply a single complex f32 number (64-bit vector): +// val: [a.re, a.im] +// tw: [t.re, t.im] +// Returns: [a.re * t.re - a.im * t.im, a.re * t.im + a.im * t.re] +#[inline(always)] +pub unsafe fn neon_complex_mul_single_f32(val: float32x2_t, tw: float32x2_t) -> float32x2_t { + let val_q = vcombine_f32(val, val); + let tw_q = vcombine_f32(tw, tw); + vget_low_f32(neon_complex_mul_f32(val_q, tw_q)) +} - let mut x = 0; - while x + 2 <= block_w { - let cx = bx + x; - let a0 = vld1q_f64(p_in.add(in_r0 + x * 2)); - let a1 = vld1q_f64(p_in.add(in_r0 + x * 2 + 2)); +trait TransposeKernel { + type Elem: Copy; + unsafe fn transpose_2x2( + input: &[Self::Elem], + output: &mut [Self::Elem], + in_r0: usize, + in_r1: usize, + out_c0: usize, + out_c1: usize, + ); + unsafe fn transpose_1x2( + input: &[Self::Elem], + output: &mut [Self::Elem], + in_r0: usize, + in_r1: usize, + out_c0: usize, + ); + unsafe fn transpose_1x1( + input: &[Self::Elem], + output: &mut [Self::Elem], + in_r: usize, + out_c: usize, + ); +} - let b0 = vld1q_f64(p_in.add(in_r1 + x * 2)); - let b1 = vld1q_f64(p_in.add(in_r1 + x * 2 + 2)); +trait TransposeTwiddleKernel { + type Elem: Copy; + unsafe fn step_2x2( + input: &[Self::Elem], + output: &mut [Self::Elem], + twiddles: &[Self::Elem], + in_r0: usize, + in_r1: usize, + out_c0: usize, + out_c1: usize, + ); + unsafe fn step_1x2( + input: &[Self::Elem], + output: &mut [Self::Elem], + twiddles: &[Self::Elem], + in_r0: usize, + in_r1: usize, + out_c0: usize, + ); + unsafe fn step_1x1( + input: &[Self::Elem], + output: &mut [Self::Elem], + twiddles: &[Self::Elem], + in_r: usize, + out_c: usize, + ); +} - let out_c0 = (cy + cx * height) * 2; - let out_c1 = (cy + (cx + 1) * height) * 2; +struct KernelF64; +struct KernelF32; - vst1q_f64(p_out.add(out_c0), a0); - vst1q_f64(p_out.add(out_c0 + 2), b0); +impl TransposeKernel for KernelF64 { + type Elem = Complex; - vst1q_f64(p_out.add(out_c1), a1); - vst1q_f64(p_out.add(out_c1 + 2), b1); + #[inline(always)] + unsafe fn transpose_2x2( + input: &[Complex], + output: &mut [Complex], + in_r0: usize, + in_r1: usize, + out_c0: usize, + out_c1: usize, + ) { + let a0 = vld1q_f64(input.get_unchecked(in_r0) as *const _ as *const f64); + let a1 = vld1q_f64(input.get_unchecked(in_r0 + 1) as *const _ as *const f64); + let b0 = vld1q_f64(input.get_unchecked(in_r1) as *const _ as *const f64); + let b1 = vld1q_f64(input.get_unchecked(in_r1 + 1) as *const _ as *const f64); + + vst1q_f64(output.get_unchecked_mut(out_c0) as *mut _ as *mut f64, a0); + vst1q_f64(output.get_unchecked_mut(out_c0 + 1) as *mut _ as *mut f64, b0); + vst1q_f64(output.get_unchecked_mut(out_c1) as *mut _ as *mut f64, a1); + vst1q_f64(output.get_unchecked_mut(out_c1 + 1) as *mut _ as *mut f64, b1); + } - x += 2; - } - while x < block_w { - let cx = bx + x; - let a0 = vld1q_f64(p_in.add(in_r0 + x * 2)); - let b0 = vld1q_f64(p_in.add(in_r1 + x * 2)); + #[inline(always)] + unsafe fn transpose_1x2( + input: &[Complex], + output: &mut [Complex], + in_r0: usize, + in_r1: usize, + out_c0: usize, + ) { + let a0 = vld1q_f64(input.get_unchecked(in_r0) as *const _ as *const f64); + let b0 = vld1q_f64(input.get_unchecked(in_r1) as *const _ as *const f64); + + vst1q_f64(output.get_unchecked_mut(out_c0) as *mut _ as *mut f64, a0); + vst1q_f64(output.get_unchecked_mut(out_c0 + 1) as *mut _ as *mut f64, b0); + } - let out_c0 = (cy + cx * height) * 2; + #[inline(always)] + unsafe fn transpose_1x1( + input: &[Complex], + output: &mut [Complex], + in_r: usize, + out_c: usize, + ) { + *output.get_unchecked_mut(out_c) = *input.get_unchecked(in_r); + } +} - vst1q_f64(p_out.add(out_c0), a0); - vst1q_f64(p_out.add(out_c0 + 2), b0); +impl TransposeTwiddleKernel for KernelF64 { + type Elem = Complex; - x += 1; - } - y += 2; - } - while y < block_h { - let cy = by + y; - let in_r = (cy * width + bx) * 2; - for x in 0..block_w { - let cx = bx + x; - let a0 = vld1q_f64(p_in.add(in_r + x * 2)); - let out_c = (cy + cx * height) * 2; - vst1q_f64(p_out.add(out_c), a0); - } - y += 1; - } + #[inline(always)] + unsafe fn step_2x2( + input: &[Complex], + output: &mut [Complex], + twiddles: &[Complex], + in_r0: usize, + in_r1: usize, + out_c0: usize, + out_c1: usize, + ) { + let a0 = vld1q_f64(input.get_unchecked(in_r0) as *const _ as *const f64); + let tw_a0 = vld1q_f64(twiddles.get_unchecked(in_r0) as *const _ as *const f64); + let a1 = vld1q_f64(input.get_unchecked(in_r0 + 1) as *const _ as *const f64); + let tw_a1 = vld1q_f64(twiddles.get_unchecked(in_r0 + 1) as *const _ as *const f64); + + let b0 = vld1q_f64(input.get_unchecked(in_r1) as *const _ as *const f64); + let tw_b0 = vld1q_f64(twiddles.get_unchecked(in_r1) as *const _ as *const f64); + let b1 = vld1q_f64(input.get_unchecked(in_r1 + 1) as *const _ as *const f64); + let tw_b1 = vld1q_f64(twiddles.get_unchecked(in_r1 + 1) as *const _ as *const f64); + + let res_a0 = neon_complex_mul_f64(a0, tw_a0); + let res_a1 = neon_complex_mul_f64(a1, tw_a1); + let res_b0 = neon_complex_mul_f64(b0, tw_b0); + let res_b1 = neon_complex_mul_f64(b1, tw_b1); + + vst1q_f64(output.get_unchecked_mut(out_c0) as *mut _ as *mut f64, res_a0); + vst1q_f64(output.get_unchecked_mut(out_c0 + 1) as *mut _ as *mut f64, res_b0); + vst1q_f64(output.get_unchecked_mut(out_c1) as *mut _ as *mut f64, res_a1); + vst1q_f64(output.get_unchecked_mut(out_c1 + 1) as *mut _ as *mut f64, res_b1); + } - bx += TILE_SIZE; - } - by += TILE_SIZE; + #[inline(always)] + unsafe fn step_1x2( + input: &[Complex], + output: &mut [Complex], + twiddles: &[Complex], + in_r0: usize, + in_r1: usize, + out_c0: usize, + ) { + let a0 = vld1q_f64(input.get_unchecked(in_r0) as *const _ as *const f64); + let tw_a0 = vld1q_f64(twiddles.get_unchecked(in_r0) as *const _ as *const f64); + let b0 = vld1q_f64(input.get_unchecked(in_r1) as *const _ as *const f64); + let tw_b0 = vld1q_f64(twiddles.get_unchecked(in_r1) as *const _ as *const f64); + + let res_a0 = neon_complex_mul_f64(a0, tw_a0); + let res_b0 = neon_complex_mul_f64(b0, tw_b0); + + vst1q_f64(output.get_unchecked_mut(out_c0) as *mut _ as *mut f64, res_a0); + vst1q_f64(output.get_unchecked_mut(out_c0 + 1) as *mut _ as *mut f64, res_b0); + } + + #[inline(always)] + unsafe fn step_1x1( + input: &[Complex], + output: &mut [Complex], + twiddles: &[Complex], + in_r: usize, + out_c: usize, + ) { + let a = vld1q_f64(input.get_unchecked(in_r) as *const _ as *const f64); + let tw = vld1q_f64(twiddles.get_unchecked(in_r) as *const _ as *const f64); + let res = neon_complex_mul_f64(a, tw); + vst1q_f64(output.get_unchecked_mut(out_c) as *mut _ as *mut f64, res); } } +impl TransposeKernel for KernelF32 { + type Elem = Complex; + + #[inline(always)] + unsafe fn transpose_2x2( + input: &[Complex], + output: &mut [Complex], + in_r0: usize, + in_r1: usize, + out_c0: usize, + out_c1: usize, + ) { + // Load row 0: 2 Complex into 1 float32x4_t: [a0, a1] + let row0 = vld1q_f32(input.get_unchecked(in_r0) as *const _ as *const f32); + // Load row 1: 2 Complex into 1 float32x4_t: [b0, b1] + let row1 = vld1q_f32(input.get_unchecked(in_r1) as *const _ as *const f32); + + // Transpose to get col 0 [a0, b0] and col 1 [a1, b1] + let transposed = transpose_complex_2x2_f32(row0, row1); + + vst1q_f32(output.get_unchecked_mut(out_c0) as *mut _ as *mut f32, transposed[0]); + vst1q_f32(output.get_unchecked_mut(out_c1) as *mut _ as *mut f32, transposed[1]); + } + + #[inline(always)] + unsafe fn transpose_1x2( + input: &[Complex], + output: &mut [Complex], + in_r0: usize, + in_r1: usize, + out_c0: usize, + ) { + let a0 = vld1_f32(input.get_unchecked(in_r0) as *const _ as *const f32); + let b0 = vld1_f32(input.get_unchecked(in_r1) as *const _ as *const f32); + + let col0 = vcombine_f32(a0, b0); + vst1q_f32(output.get_unchecked_mut(out_c0) as *mut _ as *mut f32, col0); + } + + #[inline(always)] + unsafe fn transpose_1x1( + input: &[Complex], + output: &mut [Complex], + in_r: usize, + out_c: usize, + ) { + *output.get_unchecked_mut(out_c) = *input.get_unchecked(in_r); + } +} + +impl TransposeTwiddleKernel for KernelF32 { + type Elem = Complex; + + #[inline(always)] + unsafe fn step_2x2( + input: &[Complex], + output: &mut [Complex], + twiddles: &[Complex], + in_r0: usize, + in_r1: usize, + out_c0: usize, + out_c1: usize, + ) { + let row0 = vld1q_f32(input.get_unchecked(in_r0) as *const _ as *const f32); + let tw_row0 = vld1q_f32(twiddles.get_unchecked(in_r0) as *const _ as *const f32); + + let row1 = vld1q_f32(input.get_unchecked(in_r1) as *const _ as *const f32); + let tw_row1 = vld1q_f32(twiddles.get_unchecked(in_r1) as *const _ as *const f32); + + let res0 = neon_complex_mul_f32(row0, tw_row0); + let res1 = neon_complex_mul_f32(row1, tw_row1); + + let transposed = transpose_complex_2x2_f32(res0, res1); + + vst1q_f32(output.get_unchecked_mut(out_c0) as *mut _ as *mut f32, transposed[0]); + vst1q_f32(output.get_unchecked_mut(out_c1) as *mut _ as *mut f32, transposed[1]); + } + + #[inline(always)] + unsafe fn step_1x2( + input: &[Complex], + output: &mut [Complex], + twiddles: &[Complex], + in_r0: usize, + in_r1: usize, + out_c0: usize, + ) { + let a0 = vld1_f32(input.get_unchecked(in_r0) as *const _ as *const f32); + let tw_a0 = vld1_f32(twiddles.get_unchecked(in_r0) as *const _ as *const f32); + let b0 = vld1_f32(input.get_unchecked(in_r1) as *const _ as *const f32); + let tw_b0 = vld1_f32(twiddles.get_unchecked(in_r1) as *const _ as *const f32); + + let res_a0 = neon_complex_mul_single_f32(a0, tw_a0); + let res_b0 = neon_complex_mul_single_f32(b0, tw_b0); + + let col0 = vcombine_f32(res_a0, res_b0); + vst1q_f32(output.get_unchecked_mut(out_c0) as *mut _ as *mut f32, col0); + } + + #[inline(always)] + unsafe fn step_1x1( + input: &[Complex], + output: &mut [Complex], + twiddles: &[Complex], + in_r: usize, + out_c: usize, + ) { + let a = vld1_f32(input.get_unchecked(in_r) as *const _ as *const f32); + let tw = vld1_f32(twiddles.get_unchecked(in_r) as *const _ as *const f32); + let res = neon_complex_mul_single_f32(a, tw); + vst1_f32(output.get_unchecked_mut(out_c) as *mut _ as *mut f32, res); + } +} + +/// Tiled matrix transpose for 2D array of Complex elements. +/// +/// Divides the matrix into TILE_SIZE x TILE_SIZE blocks to maximize cache locality. +/// Within each tile, processes 2x2 complex sub-blocks using SIMD vector instructions, +/// followed by remainder column and remainder row handling. #[inline(always)] -unsafe fn transpose_f32(p_in: *const f32, p_out: *mut f32, width: usize, height: usize) { +unsafe fn transpose_tiled( + input: &[K::Elem], + output: &mut [K::Elem], + width: usize, + height: usize, +) { const TILE_SIZE: usize = 32; + // Loop 1: iterate over row tiles of height TILE_SIZE let mut by = 0; while by < height { let block_h = (height - by).min(TILE_SIZE); + + // Loop 2: iterate over column tiles of width TILE_SIZE let mut bx = 0; while bx < width { let block_w = (width - bx).min(TILE_SIZE); + // Loop 3: process 2 rows at a time within the tile let mut y = 0; while y + 2 <= block_h { let cy = by + y; - let in_r0 = (cy * width + bx) * 2; - let in_r1 = ((cy + 1) * width + bx) * 2; + let in_r0 = cy * width + bx; + let in_r1 = (cy + 1) * width + bx; + // Loop 4: process 2 columns at a time (2x2 complex sub-block) let mut x = 0; while x + 2 <= block_w { let cx = bx + x; - let row0 = vld1q_f32(p_in.add(in_r0 + x * 2)); - let row1 = vld1q_f32(p_in.add(in_r1 + x * 2)); - - let transposed = transpose_complex_2x2_f32(row0, row1); - - let out_c0 = (cy + cx * height) * 2; - let out_c1 = (cy + (cx + 1) * height) * 2; - - vst1q_f32(p_out.add(out_c0), transposed[0]); - vst1q_f32(p_out.add(out_c1), transposed[1]); + let out_c0 = cy + cx * height; + let out_c1 = cy + (cx + 1) * height; + K::transpose_2x2(input, output, in_r0 + x, in_r1 + x, out_c0, out_c1); x += 2; } - while x < block_w { - let cx = bx + x; - let a0 = vld1_f32(p_in.add(in_r0 + x * 2)); - let b0 = vld1_f32(p_in.add(in_r1 + x * 2)); - - let out_c0 = (cy + cx * height) * 2; - vst1_f32(p_out.add(out_c0), a0); - vst1_f32(p_out.add(out_c0 + 2), b0); - - x += 1; + // Handle remainder column when block_w is odd (2 rows x 1 column) + if x < block_w { + let cx = bx + x; + let out_c0 = cy + cx * height; + K::transpose_1x2(input, output, in_r0 + x, in_r1 + x, out_c0); } + y += 2; } - while y < block_h { + + // Handle remainder row when block_h is odd (1 row x block_w columns) + if y < block_h { let cy = by + y; - let in_r = (cy * width + bx) * 2; + let in_r = cy * width + bx; for x in 0..block_w { let cx = bx + x; - let a0 = vld1_f32(p_in.add(in_r + x * 2)); - let out_c = (cy + cx * height) * 2; - vst1_f32(p_out.add(out_c), a0); + let out_c = cy + cx * height; + K::transpose_1x1(input, output, in_r + x, out_c); } - y += 1; } bx += TILE_SIZE; @@ -382,35 +651,132 @@ unsafe fn transpose_f32(p_in: *const f32, p_out: *mut f32, width: usize, height: } } +pub unsafe fn transpose_f64( + input: &[Complex], + output: &mut [Complex], + width: usize, + height: usize, +) { + super::neon_common::assert_f64::(); + assert!(input.len() >= width * height); + assert!(output.len() >= width * height); + + transpose_tiled::(input, output, width, height); +} + +pub unsafe fn transpose_f32( + input: &[Complex], + output: &mut [Complex], + width: usize, + height: usize, +) { + super::neon_common::assert_f32::(); + assert!(input.len() >= width * height); + assert!(output.len() >= width * height); + + transpose_tiled::(input, output, width, height); +} + pub unsafe fn transpose( width: usize, height: usize, input: &[T], output: &mut [T], ) -> bool { + assert!(input.len() >= width * height); + assert!(output.len() >= width * height); + use std::any::TypeId; if TypeId::of::() == TypeId::of::>() { - let p_in = input.as_ptr() as *const f64; - let p_out = output.as_mut_ptr() as *mut f64; - transpose_f64(p_in, p_out, width, height); + let input: &[Complex] = crate::array_utils::workaround_transmute(input); + let output: &mut [Complex] = crate::array_utils::workaround_transmute_mut(output); + transpose_f64(input, output, width, height); return true; } else if TypeId::of::() == TypeId::of::>() { - let p_in = input.as_ptr() as *const f32; - let p_out = output.as_mut_ptr() as *mut f32; - transpose_f32(p_in, p_out, width, height); + let input: &[Complex] = crate::array_utils::workaround_transmute(input); + let output: &mut [Complex] = crate::array_utils::workaround_transmute_mut(output); + transpose_f32(input, output, width, height); return true; } false } +/// Fused complex twiddle multiplication and matrix transpose for small sizes. +/// +/// Computes `output[y + x * height] = input[x + y * width] * twiddles[x + y * width]`. +/// Processes 2x2 complex blocks with SIMD multiplication and transposition, +/// followed by remainder column and row handling. #[inline(always)] -unsafe fn neon_complex_mul_f64( - val: float64x2_t, - tw: float64x2_t, -) -> float64x2_t { - let temp = vcombine_f64(vneg_f64(vget_high_f64(val)), vget_low_f64(val)); - let sum = vmulq_laneq_f64::<0>(val, tw); - vfmaq_laneq_f64::<1>(sum, temp, tw) +unsafe fn transpose_small_twiddle_impl( + input: &[K::Elem], + output: &mut [K::Elem], + twiddles: &[K::Elem], + width: usize, + height: usize, +) { + // Loop 1: process 2 rows at a time + let mut y = 0; + while y + 2 <= height { + let in_r0 = y * width; + let in_r1 = (y + 1) * width; + + // Loop 2: process 2 columns at a time (2x2 complex sub-block) + let mut x = 0; + while x + 2 <= width { + let out_c0 = y + x * height; + let out_c1 = y + (x + 1) * height; + + K::step_2x2(input, output, twiddles, in_r0 + x, in_r1 + x, out_c0, out_c1); + x += 2; + } + + // Remainder column when width is odd (2 rows x 1 column) + if x < width { + let out_c0 = y + x * height; + K::step_1x2(input, output, twiddles, in_r0 + x, in_r1 + x, out_c0); + } + + y += 2; + } + + // Remainder row when height is odd (1 row x width columns) + if y < height { + let in_r = y * width; + for x in 0..width { + let out_c = y + x * height; + K::step_1x1(input, output, twiddles, in_r + x, out_c); + } + } +} + +pub unsafe fn transpose_small_twiddle_f64( + input: &[Complex], + output: &mut [Complex], + twiddles: &[Complex], + width: usize, + height: usize, +) { + super::neon_common::assert_f64::(); + assert!(input.len() >= width * height); + assert!(output.len() >= width * height); + assert!(twiddles.len() >= width * height); + + transpose_small_twiddle_impl::(input, output, twiddles, width, height); +} + +pub unsafe fn transpose_small_twiddle_f32( + input: &[Complex], + output: &mut [Complex], + twiddles: &[Complex], + width: usize, + height: usize, +) { + super::neon_common::assert_f32::(); + assert!(input.len() >= width * height); + assert!(output.len() >= width * height); + assert!(twiddles.len() >= width * height); + + transpose_small_twiddle_impl::(input, output, twiddles, width, height); } pub unsafe fn transpose_small_twiddle( @@ -420,98 +786,22 @@ pub unsafe fn transpose_small_twiddle( output: &mut [Complex], twiddles: &[Complex], ) -> bool { + assert!(input.len() >= width * height); + assert!(output.len() >= width * height); + assert!(twiddles.len() >= width * height); + use std::any::TypeId; if TypeId::of::() == TypeId::of::() { - let p_in = input.as_ptr() as *const f64; - let p_out = output.as_mut_ptr() as *mut f64; - let p_tw = twiddles.as_ptr() as *const f64; - - let mut y = 0; - while y + 2 <= height { - let in_row0 = (y * width) * 2; - let in_row1 = ((y + 1) * width) * 2; - let mut x = 0; - while x + 2 <= width { - let in_idx0 = in_row0 + x * 2; - let in_idx1 = in_row1 + x * 2; - - let a0 = vld1q_f64(p_in.add(in_idx0)); - let tw_a0 = vld1q_f64(p_tw.add(in_idx0)); - let a1 = vld1q_f64(p_in.add(in_idx0 + 2)); - let tw_a1 = vld1q_f64(p_tw.add(in_idx0 + 2)); - - let b0 = vld1q_f64(p_in.add(in_idx1)); - let tw_b0 = vld1q_f64(p_tw.add(in_idx1)); - let b1 = vld1q_f64(p_in.add(in_idx1 + 2)); - let tw_b1 = vld1q_f64(p_tw.add(in_idx1 + 2)); - - let res_a0 = neon_complex_mul_f64(a0, tw_a0); - let res_a1 = neon_complex_mul_f64(a1, tw_a1); - let res_b0 = neon_complex_mul_f64(b0, tw_b0); - let res_b1 = neon_complex_mul_f64(b1, tw_b1); - - let out_idx0 = (y + x * height) * 2; - let out_idx1 = (y + (x + 1) * height) * 2; - - vst1q_f64(p_out.add(out_idx0), res_a0); - vst1q_f64(p_out.add(out_idx0 + 2), res_b0); - vst1q_f64(p_out.add(out_idx1), res_a1); - vst1q_f64(p_out.add(out_idx1 + 2), res_b1); - - x += 2; - } - while x < width { - let in_idx0 = in_row0 + x * 2; - let in_idx1 = in_row1 + x * 2; - - let a0 = vld1q_f64(p_in.add(in_idx0)); - let tw_a0 = vld1q_f64(p_tw.add(in_idx0)); - let b0 = vld1q_f64(p_in.add(in_idx1)); - let tw_b0 = vld1q_f64(p_tw.add(in_idx1)); - - let res_a0 = neon_complex_mul_f64(a0, tw_a0); - let res_b0 = neon_complex_mul_f64(b0, tw_b0); - - let out_idx0 = (y + x * height) * 2; - vst1q_f64(p_out.add(out_idx0), res_a0); - vst1q_f64(p_out.add(out_idx0 + 2), res_b0); - x += 1; - } - y += 2; - } - while y < height { - let in_row = (y * width) * 2; - for x in 0..width { - let in_idx = in_row + x * 2; - let out_idx = (y + x * height) * 2; - let a = vld1q_f64(p_in.add(in_idx)); - let tw = vld1q_f64(p_tw.add(in_idx)); - let res = neon_complex_mul_f64(a, tw); - vst1q_f64(p_out.add(out_idx), res); - } - y += 1; - } + let input: &[Complex] = crate::array_utils::workaround_transmute(input); + let output: &mut [Complex] = crate::array_utils::workaround_transmute_mut(output); + let twiddles: &[Complex] = crate::array_utils::workaround_transmute(twiddles); + transpose_small_twiddle_f64(input, output, twiddles, width, height); return true; } else if TypeId::of::() == TypeId::of::() { - let p_in = input.as_ptr() as *const f32; - let p_out = output.as_mut_ptr() as *mut f32; - let p_tw = twiddles.as_ptr() as *const f32; - for y in 0..height { - for x in 0..width { - let in_idx = (x + y * width) * 2; - let out_idx = (y + x * height) * 2; - let left = vld1_f32(p_in.add(in_idx)); - let right = vld1_f32(p_tw.add(in_idx)); - let left_q = vcombine_f32(left, left); - let right_q = vcombine_f32(right, right); - let temp1 = vtrn1q_f32(right_q, right_q); - let temp2 = vtrn2q_f32(right_q, vnegq_f32(right_q)); - let temp3 = vmulq_f32(temp2, left_q); - let temp4 = vrev64q_f32(temp3); - let res = vfmaq_f32(temp4, temp1, left_q); - vst1_f32(p_out.add(out_idx), vget_low_f32(res)); - } - } + let input: &[Complex] = crate::array_utils::workaround_transmute(input); + let output: &mut [Complex] = crate::array_utils::workaround_transmute_mut(output); + let twiddles: &[Complex] = crate::array_utils::workaround_transmute(twiddles); + transpose_small_twiddle_f32(input, output, twiddles, width, height); return true; } false @@ -555,18 +845,102 @@ mod unit_tests { } #[test] - fn test_pack() { - unsafe { - let nbr2 = vld1q_f32([5.0, 6.0, 7.0, 8.0].as_ptr()); - let nbr1 = vld1q_f32([1.0, 2.0, 3.0, 4.0].as_ptr()); - let first = extract_lo_lo_f32(nbr1, nbr2); - let second = extract_hi_hi_f32(nbr1, nbr2); - let first = std::mem::transmute::; 2]>(first); - let second = std::mem::transmute::; 2]>(second); - let first_expected = [Complex::new(1.0, 2.0), Complex::new(5.0, 6.0)]; - let second_expected = [Complex::new(3.0, 4.0), Complex::new(7.0, 8.0)]; - assert_eq!(first, first_expected); - assert_eq!(second, second_expected); + fn test_transpose_neon_f32_f64() { + use num_traits::Zero; + for width in [1, 2, 3, 4, 5, 7, 8, 16, 31, 32, 33, 40] { + for height in [1, 2, 3, 4, 5, 7, 8, 16, 31, 32, 33, 40] { + let len = width * height; + + // f32 + let input_f32: Vec> = (0..len) + .map(|i| Complex::new(i as f32, (i * 2) as f32)) + .collect(); + let mut out_f32 = vec![Complex::zero(); len]; + unsafe { transpose(width, height, &input_f32, &mut out_f32) }; + for y in 0..height { + for x in 0..width { + assert_eq!( + input_f32[x + y * width], + out_f32[y + x * height], + "f32 mismatch at ({}, {}) for {}x{}", + x, y, width, height + ); + } + } + + // f64 + let input_f64: Vec> = (0..len) + .map(|i| Complex::new(i as f64, (i * 2) as f64)) + .collect(); + let mut out_f64 = vec![Complex::zero(); len]; + unsafe { transpose(width, height, &input_f64, &mut out_f64) }; + for y in 0..height { + for x in 0..width { + assert_eq!( + input_f64[x + y * width], + out_f64[y + x * height], + "f64 mismatch at ({}, {}) for {}x{}", + x, y, width, height + ); + } + } + } + } + } + + #[test] + fn test_transpose_small_twiddle_neon() { + use num_traits::Zero; + for width in [1, 2, 3, 4, 5, 7, 8, 16, 31, 32, 33] { + for height in [1, 2, 3, 4, 5, 7, 8, 16, 31, 32, 33] { + let len = width * height; + + // f32 + let input_f32: Vec> = (0..len) + .map(|i| Complex::new((i % 7) as f32, (i % 5) as f32)) + .collect(); + let twiddles_f32: Vec> = (0..len) + .map(|i| Complex::new((i % 3) as f32, (i % 4) as f32)) + .collect(); + let mut out_f32 = vec![Complex::zero(); len]; + unsafe { + transpose_small_twiddle(width, height, &input_f32, &mut out_f32, &twiddles_f32); + } + for y in 0..height { + for x in 0..width { + let expected = input_f32[x + y * width] * twiddles_f32[x + y * width]; + let actual = out_f32[y + x * height]; + assert!( + (actual.re - expected.re).abs() < 1e-5 && (actual.im - expected.im).abs() < 1e-5, + "f32 twiddle mismatch at ({}, {}) for {}x{}: expected {:?}, got {:?}", + x, y, width, height, expected, actual + ); + } + } + + // f64 + let input_f64: Vec> = (0..len) + .map(|i| Complex::new((i % 7) as f64, (i % 5) as f64)) + .collect(); + let twiddles_f64: Vec> = (0..len) + .map(|i| Complex::new((i % 3) as f64, (i % 4) as f64)) + .collect(); + let mut out_f64 = vec![Complex::zero(); len]; + unsafe { + transpose_small_twiddle(width, height, &input_f64, &mut out_f64, &twiddles_f64); + } + for y in 0..height { + for x in 0..width { + let expected = input_f64[x + y * width] * twiddles_f64[x + y * width]; + let actual = out_f64[y + x * height]; + assert!( + (actual.re - expected.re).abs() < 1e-10 && (actual.im - expected.im).abs() < 1e-10, + "f64 twiddle mismatch at ({}, {}) for {}x{}: expected {:?}, got {:?}", + x, y, width, height, expected, actual + ); + } + } + } } } } From 0382a62037a72ae852db8edce4610399c48f1e9d Mon Sep 17 00:00:00 2001 From: Wang-Yue Date: Mon, 7 Sep 2026 16:20:23 +0900 Subject: [PATCH 5/7] Optimize transpose_small function for improved performance and safety --- src/array_utils.rs | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/src/array_utils.rs b/src/array_utils.rs index 09d97486..746f6bea 100644 --- a/src/array_utils.rs +++ b/src/array_utils.rs @@ -27,14 +27,15 @@ pub fn transpose( } /// Given an array of size width * height, representing a flattened 2D array, -/// transpose the rows and columns of that 2D array into the output -pub unsafe fn transpose_small( - width: usize, - height: usize, - input: &[T], - output: &mut [T], -) { - transpose(input, output, width, height); +pub unsafe fn transpose_small(width: usize, height: usize, input: &[T], output: &mut [T]) { + for x in 0..width { + for y in 0..height { + let input_index = x + y * width; + let output_index = y + x * height; + + *output.get_unchecked_mut(output_index) = *input.get_unchecked(input_index); + } + } } pub unsafe fn transpose_small_twiddle( From 77c7b2fb546261fe0885d63553e2faa6deafb3bd Mon Sep 17 00:00:00 2001 From: Wang-Yue Date: Tue, 8 Sep 2026 15:34:11 +0900 Subject: [PATCH 6/7] neon: integrate fused twiddle transpose with NeonArray and revert big transposes --- src/algorithm/good_thomas_algorithm.rs | 7 +- src/algorithm/mixed_radix.rs | 19 +- src/array_utils.rs | 95 ++-- src/neon/neon_utils.rs | 630 +++++-------------------- 4 files changed, 182 insertions(+), 569 deletions(-) diff --git a/src/algorithm/good_thomas_algorithm.rs b/src/algorithm/good_thomas_algorithm.rs index 4801170c..2122310a 100644 --- a/src/algorithm/good_thomas_algorithm.rs +++ b/src/algorithm/good_thomas_algorithm.rs @@ -4,6 +4,7 @@ use std::sync::Arc; use num_complex::Complex; use num_integer::Integer; use strength_reduce::StrengthReducedUsize; +use transpose; use crate::array_utils; use crate::{common::FftNum, FftDirection}; @@ -236,7 +237,7 @@ impl GoodThomasAlgorithm { .process_with_scratch(scratch, width_scratch); // transpose - array_utils::transpose(scratch, buffer, self.width, self.height); + transpose::transpose(scratch, buffer, self.width, self.height); // run FFTs of size 'height' self.height_size_fft @@ -261,7 +262,7 @@ impl GoodThomasAlgorithm { let (scratch, inner_scratch) = scratch.split_at_mut(self.len()); // transpose - array_utils::transpose(output, scratch, self.width, self.height); + transpose::transpose(output, scratch, self.width, self.height); // run FFTs of size 'height' self.height_size_fft @@ -290,7 +291,7 @@ impl GoodThomasAlgorithm { .process_with_scratch(output, width_scratch); // transpose - array_utils::transpose(output, input, self.width, self.height); + transpose::transpose(output, input, self.width, self.height); // run FFTs of size 'height' let height_scratch = if scratch.len() > output.len() { diff --git a/src/algorithm/mixed_radix.rs b/src/algorithm/mixed_radix.rs index a8d30d86..a4630b11 100644 --- a/src/algorithm/mixed_radix.rs +++ b/src/algorithm/mixed_radix.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use num_complex::Complex; use num_traits::Zero; +use transpose; use crate::array_utils; use crate::{common::FftNum, twiddles, FftDirection}; @@ -129,7 +130,7 @@ impl MixedRadix { let (scratch, inner_scratch) = scratch.split_at_mut(self.len()); // STEP 1: transpose - array_utils::transpose(buffer, scratch, self.width, self.height); + transpose::transpose(buffer, scratch, self.width, self.height); // STEP 2: perform FFTs of size `height` let height_scratch = if inner_scratch.len() > buffer.len() { @@ -146,14 +147,14 @@ impl MixedRadix { } // STEP 4: transpose again - array_utils::transpose(scratch, buffer, self.height, self.width); + transpose::transpose(scratch, buffer, self.height, self.width); // STEP 5: perform FFTs of size `width` self.width_size_fft .process_outofplace_with_scratch(buffer, scratch, inner_scratch); // STEP 6: transpose again - array_utils::transpose(scratch, buffer, self.width, self.height); + transpose::transpose(scratch, buffer, self.width, self.height); } fn perform_fft_immut( @@ -163,7 +164,7 @@ impl MixedRadix { scratch_raw: &mut [Complex], ) { // STEP 1: transpose - array_utils::transpose(input, output, self.width, self.height); + transpose::transpose(input, output, self.width, self.height); // STEP 2: perform FFTs of size `height` self.height_size_fft @@ -177,14 +178,14 @@ impl MixedRadix { let (scratch, inner_scratch) = scratch_raw.split_at_mut(self.len()); // STEP 4: transpose again - array_utils::transpose(output, scratch, self.height, self.width); + transpose::transpose(output, scratch, self.height, self.width); // STEP 5: perform FFTs of size `width` self.width_size_fft .process_with_scratch(scratch, inner_scratch); // STEP 6: transpose again - array_utils::transpose(scratch, output, self.width, self.height); + transpose::transpose(scratch, output, self.width, self.height); } fn perform_fft_out_of_place( @@ -196,7 +197,7 @@ impl MixedRadix { // SIX STEP FFT: // STEP 1: transpose - array_utils::transpose(input, output, self.width, self.height); + transpose::transpose(input, output, self.width, self.height); // STEP 2: perform FFTs of size `height` let height_scratch = if scratch.len() > input.len() { @@ -213,7 +214,7 @@ impl MixedRadix { } // STEP 4: transpose again - array_utils::transpose(output, input, self.height, self.width); + transpose::transpose(output, input, self.height, self.width); // STEP 5: perform FFTs of size `width` let width_scratch = if scratch.len() > output.len() { @@ -225,7 +226,7 @@ impl MixedRadix { .process_with_scratch(input, width_scratch); // STEP 6: transpose again - array_utils::transpose(input, output, self.width, self.height); + transpose::transpose(input, output, self.width, self.height); } } boilerplate_fft!( diff --git a/src/array_utils.rs b/src/array_utils.rs index 746f6bea..eb6502c0 100644 --- a/src/array_utils.rs +++ b/src/array_utils.rs @@ -3,30 +3,9 @@ use crate::Complex; use crate::FftNum; use std::ops::{Deref, DerefMut}; -/// Transpose the input array into the output array. -/// -/// Given an array of size width * height, representing a flattened 2D array, -/// transpose the rows and columns of that 2D array into the output. -pub fn transpose( - input: &[T], - output: &mut [T], - width: usize, - height: usize, -) { - assert!(input.len() >= width * height); - assert!(output.len() >= width * height); - - #[cfg(all(target_arch = "aarch64", feature = "neon"))] - { - if unsafe { crate::neon::neon_utils::transpose(width, height, input, output) } { - return; - } - } - - transpose::transpose(input, output, width, height); -} - /// Given an array of size width * height, representing a flattened 2D array, +/// transpose the rows and columns of that 2D array into the output +/// benchmarking shows that loop tiling isn't effective for small arrays (in the range of 50x50 or smaller) pub unsafe fn transpose_small(width: usize, height: usize, input: &[T], output: &mut [T]) { for x in 0..width { for y in 0..height { @@ -165,8 +144,36 @@ mod unit_tests { use num_traits::Zero; #[test] - fn test_transpose() { - let sizes: Vec = (1..35).collect(); + fn test_transpose_small() { + let sizes: Vec = (1..16).collect(); + + for &width in &sizes { + for &height in &sizes { + let len = width * height; + + let input: Vec> = random_signal(len); + let mut output = vec![Zero::zero(); len]; + + unsafe { transpose_small(width, height, &input, &mut output) }; + + for x in 0..width { + for y in 0..height { + assert_eq!( + input[x + y * width], + output[y + x * height], + "x = {}, y = {}", + x, + y + ); + } + } + } + } + } + + #[test] + fn test_transpose_small_twiddle() { + let sizes: Vec = (1..16).collect(); for &width in &sizes { for &height in &sizes { @@ -174,36 +181,38 @@ mod unit_tests { // Test f32 let input_f32: Vec> = random_signal(len); + let twiddles_f32: Vec> = random_signal(len); let mut output_f32 = vec![Zero::zero(); len]; - transpose(&input_f32, &mut output_f32, width, height); + unsafe { + transpose_small_twiddle(width, height, &input_f32, &mut output_f32, &twiddles_f32); + } for x in 0..width { for y in 0..height { - assert_eq!( - input_f32[x + y * width], - output_f32[y + x * height], - "f32: x = {}, y = {}, width = {}, height = {}", - x, - y, - width, - 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> = random_signal(len); + let twiddles_f64: Vec> = random_signal(len); let mut output_f64 = vec![Zero::zero(); len]; - transpose(&input_f64, &mut output_f64, width, height); + unsafe { + transpose_small_twiddle(width, height, &input_f64, &mut output_f64, &twiddles_f64); + } for x in 0..width { for y in 0..height { - assert_eq!( - input_f64[x + y * width], - output_f64[y + x * height], - "f64: x = {}, y = {}, width = {}, height = {}", - x, - y, - width, - 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 ); } } diff --git a/src/neon/neon_utils.rs b/src/neon/neon_utils.rs index 4763bf41..3bf16306 100644 --- a/src/neon/neon_utils.rs +++ b/src/neon/neon_utils.rs @@ -1,6 +1,7 @@ use core::arch::aarch64::*; use num_complex::Complex; use crate::FftNum; +use crate::neon::neon_vector::{NeonArray, NeonArrayMut, NeonVector}; // __ __ _ _ _________ _ _ _ // | \/ | __ _| |_| |__ |___ /___ \| |__ (_) |_ @@ -252,465 +253,97 @@ impl Rotate90F64 { } } -// Multiply two complex f64 numbers (1 complex number per 128-bit register): -// val: [val.re, val.im] -// tw: [tw.re, tw.im] -// Returns: [val.re * tw.re - val.im * tw.im, val.re * tw.im + val.im * tw.re] -#[inline(always)] -pub unsafe fn neon_complex_mul_f64(val: float64x2_t, tw: float64x2_t) -> float64x2_t { - // temp: [-val.im, val.re] - let temp = vcombine_f64(vneg_f64(vget_high_f64(val)), vget_low_f64(val)); - // sum: [val.re * tw.re, val.im * tw.re] - let sum = vmulq_laneq_f64::<0>(val, tw); - // sum + temp * tw.im: [val.re * tw.re - val.im * tw.im, val.im * tw.re + val.re * tw.im] - vfmaq_laneq_f64::<1>(sum, temp, tw) -} - -// Pairwise multiply two pairs of complex f32 numbers (2 complex numbers per 128-bit register): -// val: [a0.re, a0.im, a1.re, a1.im] -// tw: [t0.re, t0.im, t1.re, t1.im] -// Returns: [ -// a0.re * t0.re - a0.im * t0.im, a0.re * t0.im + a0.im * t0.re, -// a1.re * t1.re - a1.im * t1.im, a1.re * t1.im + a1.im * t1.re, -// ] -#[inline(always)] -pub unsafe fn neon_complex_mul_f32(val: float32x4_t, tw: float32x4_t) -> float32x4_t { - let temp1 = vtrn1q_f32(tw, tw); - let temp2 = vtrn2q_f32(tw, vnegq_f32(tw)); - let temp3 = vmulq_f32(temp2, val); - let temp4 = vrev64q_f32(temp3); - vfmaq_f32(temp4, temp1, val) -} - -// Multiply a single complex f32 number (64-bit vector): -// val: [a.re, a.im] -// tw: [t.re, t.im] -// Returns: [a.re * t.re - a.im * t.im, a.re * t.im + a.im * t.re] -#[inline(always)] -pub unsafe fn neon_complex_mul_single_f32(val: float32x2_t, tw: float32x2_t) -> float32x2_t { - let val_q = vcombine_f32(val, val); - let tw_q = vcombine_f32(tw, tw); - vget_low_f32(neon_complex_mul_f32(val_q, tw_q)) -} - -trait TransposeKernel { - type Elem: Copy; - unsafe fn transpose_2x2( - input: &[Self::Elem], - output: &mut [Self::Elem], - in_r0: usize, - in_r1: usize, - out_c0: usize, - out_c1: usize, - ); - unsafe fn transpose_1x2( - input: &[Self::Elem], - output: &mut [Self::Elem], - in_r0: usize, - in_r1: usize, - out_c0: usize, - ); - unsafe fn transpose_1x1( - input: &[Self::Elem], - output: &mut [Self::Elem], - in_r: usize, - out_c: usize, - ); -} - -trait TransposeTwiddleKernel { - type Elem: Copy; - unsafe fn step_2x2( - input: &[Self::Elem], - output: &mut [Self::Elem], - twiddles: &[Self::Elem], - in_r0: usize, - in_r1: usize, - out_c0: usize, - out_c1: usize, - ); - unsafe fn step_1x2( - input: &[Self::Elem], - output: &mut [Self::Elem], - twiddles: &[Self::Elem], - in_r0: usize, - in_r1: usize, - out_c0: usize, - ); - unsafe fn step_1x1( - input: &[Self::Elem], - output: &mut [Self::Elem], - twiddles: &[Self::Elem], - in_r: usize, - out_c: usize, - ); -} - -struct KernelF64; -struct KernelF32; - -impl TransposeKernel for KernelF64 { - type Elem = Complex; - - #[inline(always)] - unsafe fn transpose_2x2( - input: &[Complex], - output: &mut [Complex], - in_r0: usize, - in_r1: usize, - out_c0: usize, - out_c1: usize, - ) { - let a0 = vld1q_f64(input.get_unchecked(in_r0) as *const _ as *const f64); - let a1 = vld1q_f64(input.get_unchecked(in_r0 + 1) as *const _ as *const f64); - let b0 = vld1q_f64(input.get_unchecked(in_r1) as *const _ as *const f64); - let b1 = vld1q_f64(input.get_unchecked(in_r1 + 1) as *const _ as *const f64); - - vst1q_f64(output.get_unchecked_mut(out_c0) as *mut _ as *mut f64, a0); - vst1q_f64(output.get_unchecked_mut(out_c0 + 1) as *mut _ as *mut f64, b0); - vst1q_f64(output.get_unchecked_mut(out_c1) as *mut _ as *mut f64, a1); - vst1q_f64(output.get_unchecked_mut(out_c1 + 1) as *mut _ as *mut f64, b1); - } - - #[inline(always)] - unsafe fn transpose_1x2( - input: &[Complex], - output: &mut [Complex], - in_r0: usize, - in_r1: usize, - out_c0: usize, - ) { - let a0 = vld1q_f64(input.get_unchecked(in_r0) as *const _ as *const f64); - let b0 = vld1q_f64(input.get_unchecked(in_r1) as *const _ as *const f64); - - vst1q_f64(output.get_unchecked_mut(out_c0) as *mut _ as *mut f64, a0); - vst1q_f64(output.get_unchecked_mut(out_c0 + 1) as *mut _ as *mut f64, b0); - } - - #[inline(always)] - unsafe fn transpose_1x1( - input: &[Complex], - output: &mut [Complex], - in_r: usize, - out_c: usize, - ) { - *output.get_unchecked_mut(out_c) = *input.get_unchecked(in_r); - } -} - -impl TransposeTwiddleKernel for KernelF64 { - type Elem = Complex; - - #[inline(always)] - unsafe fn step_2x2( - input: &[Complex], - output: &mut [Complex], - twiddles: &[Complex], - in_r0: usize, - in_r1: usize, - out_c0: usize, - out_c1: usize, - ) { - let a0 = vld1q_f64(input.get_unchecked(in_r0) as *const _ as *const f64); - let tw_a0 = vld1q_f64(twiddles.get_unchecked(in_r0) as *const _ as *const f64); - let a1 = vld1q_f64(input.get_unchecked(in_r0 + 1) as *const _ as *const f64); - let tw_a1 = vld1q_f64(twiddles.get_unchecked(in_r0 + 1) as *const _ as *const f64); - - let b0 = vld1q_f64(input.get_unchecked(in_r1) as *const _ as *const f64); - let tw_b0 = vld1q_f64(twiddles.get_unchecked(in_r1) as *const _ as *const f64); - let b1 = vld1q_f64(input.get_unchecked(in_r1 + 1) as *const _ as *const f64); - let tw_b1 = vld1q_f64(twiddles.get_unchecked(in_r1 + 1) as *const _ as *const f64); - - let res_a0 = neon_complex_mul_f64(a0, tw_a0); - let res_a1 = neon_complex_mul_f64(a1, tw_a1); - let res_b0 = neon_complex_mul_f64(b0, tw_b0); - let res_b1 = neon_complex_mul_f64(b1, tw_b1); - - vst1q_f64(output.get_unchecked_mut(out_c0) as *mut _ as *mut f64, res_a0); - vst1q_f64(output.get_unchecked_mut(out_c0 + 1) as *mut _ as *mut f64, res_b0); - vst1q_f64(output.get_unchecked_mut(out_c1) as *mut _ as *mut f64, res_a1); - vst1q_f64(output.get_unchecked_mut(out_c1 + 1) as *mut _ as *mut f64, res_b1); - } - - #[inline(always)] - unsafe fn step_1x2( - input: &[Complex], - output: &mut [Complex], - twiddles: &[Complex], - in_r0: usize, - in_r1: usize, - out_c0: usize, - ) { - let a0 = vld1q_f64(input.get_unchecked(in_r0) as *const _ as *const f64); - let tw_a0 = vld1q_f64(twiddles.get_unchecked(in_r0) as *const _ as *const f64); - let b0 = vld1q_f64(input.get_unchecked(in_r1) as *const _ as *const f64); - let tw_b0 = vld1q_f64(twiddles.get_unchecked(in_r1) as *const _ as *const f64); - - let res_a0 = neon_complex_mul_f64(a0, tw_a0); - let res_b0 = neon_complex_mul_f64(b0, tw_b0); - - vst1q_f64(output.get_unchecked_mut(out_c0) as *mut _ as *mut f64, res_a0); - vst1q_f64(output.get_unchecked_mut(out_c0 + 1) as *mut _ as *mut f64, res_b0); - } - - #[inline(always)] - unsafe fn step_1x1( - input: &[Complex], - output: &mut [Complex], - twiddles: &[Complex], - in_r: usize, - out_c: usize, - ) { - let a = vld1q_f64(input.get_unchecked(in_r) as *const _ as *const f64); - let tw = vld1q_f64(twiddles.get_unchecked(in_r) as *const _ as *const f64); - let res = neon_complex_mul_f64(a, tw); - vst1q_f64(output.get_unchecked_mut(out_c) as *mut _ as *mut f64, res); - } -} - -impl TransposeKernel for KernelF32 { - type Elem = Complex; - - #[inline(always)] - unsafe fn transpose_2x2( - input: &[Complex], - output: &mut [Complex], - in_r0: usize, - in_r1: usize, - out_c0: usize, - out_c1: usize, - ) { - // Load row 0: 2 Complex into 1 float32x4_t: [a0, a1] - let row0 = vld1q_f32(input.get_unchecked(in_r0) as *const _ as *const f32); - // Load row 1: 2 Complex into 1 float32x4_t: [b0, b1] - let row1 = vld1q_f32(input.get_unchecked(in_r1) as *const _ as *const f32); - - // Transpose to get col 0 [a0, b0] and col 1 [a1, b1] - let transposed = transpose_complex_2x2_f32(row0, row1); - - vst1q_f32(output.get_unchecked_mut(out_c0) as *mut _ as *mut f32, transposed[0]); - vst1q_f32(output.get_unchecked_mut(out_c1) as *mut _ as *mut f32, transposed[1]); - } - - #[inline(always)] - unsafe fn transpose_1x2( - input: &[Complex], - output: &mut [Complex], - in_r0: usize, - in_r1: usize, - out_c0: usize, - ) { - let a0 = vld1_f32(input.get_unchecked(in_r0) as *const _ as *const f32); - let b0 = vld1_f32(input.get_unchecked(in_r1) as *const _ as *const f32); - - let col0 = vcombine_f32(a0, b0); - vst1q_f32(output.get_unchecked_mut(out_c0) as *mut _ as *mut f32, col0); - } - - #[inline(always)] - unsafe fn transpose_1x1( - input: &[Complex], - output: &mut [Complex], - in_r: usize, - out_c: usize, - ) { - *output.get_unchecked_mut(out_c) = *input.get_unchecked(in_r); - } -} - -impl TransposeTwiddleKernel for KernelF32 { - type Elem = Complex; - - #[inline(always)] - unsafe fn step_2x2( - input: &[Complex], - output: &mut [Complex], - twiddles: &[Complex], - in_r0: usize, - in_r1: usize, - out_c0: usize, - out_c1: usize, - ) { - let row0 = vld1q_f32(input.get_unchecked(in_r0) as *const _ as *const f32); - let tw_row0 = vld1q_f32(twiddles.get_unchecked(in_r0) as *const _ as *const f32); - - let row1 = vld1q_f32(input.get_unchecked(in_r1) as *const _ as *const f32); - let tw_row1 = vld1q_f32(twiddles.get_unchecked(in_r1) as *const _ as *const f32); - - let res0 = neon_complex_mul_f32(row0, tw_row0); - let res1 = neon_complex_mul_f32(row1, tw_row1); - - let transposed = transpose_complex_2x2_f32(res0, res1); - - vst1q_f32(output.get_unchecked_mut(out_c0) as *mut _ as *mut f32, transposed[0]); - vst1q_f32(output.get_unchecked_mut(out_c1) as *mut _ as *mut f32, transposed[1]); - } - - #[inline(always)] - unsafe fn step_1x2( - input: &[Complex], - output: &mut [Complex], - twiddles: &[Complex], - in_r0: usize, - in_r1: usize, - out_c0: usize, - ) { - let a0 = vld1_f32(input.get_unchecked(in_r0) as *const _ as *const f32); - let tw_a0 = vld1_f32(twiddles.get_unchecked(in_r0) as *const _ as *const f32); - let b0 = vld1_f32(input.get_unchecked(in_r1) as *const _ as *const f32); - let tw_b0 = vld1_f32(twiddles.get_unchecked(in_r1) as *const _ as *const f32); - - let res_a0 = neon_complex_mul_single_f32(a0, tw_a0); - let res_b0 = neon_complex_mul_single_f32(b0, tw_b0); - - let col0 = vcombine_f32(res_a0, res_b0); - vst1q_f32(output.get_unchecked_mut(out_c0) as *mut _ as *mut f32, col0); - } - - #[inline(always)] - unsafe fn step_1x1( - input: &[Complex], - output: &mut [Complex], - twiddles: &[Complex], - in_r: usize, - out_c: usize, - ) { - let a = vld1_f32(input.get_unchecked(in_r) as *const _ as *const f32); - let tw = vld1_f32(twiddles.get_unchecked(in_r) as *const _ as *const f32); - let res = neon_complex_mul_single_f32(a, tw); - vst1_f32(output.get_unchecked_mut(out_c) as *mut _ as *mut f32, res); - } -} - -/// Tiled matrix transpose for 2D array of Complex elements. +/// Fused complex twiddle multiplication and matrix transpose for small f64 sizes. /// -/// Divides the matrix into TILE_SIZE x TILE_SIZE blocks to maximize cache locality. -/// Within each tile, processes 2x2 complex sub-blocks using SIMD vector instructions, -/// followed by remainder column and remainder row handling. +/// Computes `output[y + x * height] = input[x + y * width] * twiddles[x + y * width]`. +/// Processes 2x2 complex blocks with NEON SIMD vector instructions, +/// followed by remainder column and row handling. #[inline(always)] -unsafe fn transpose_tiled( - input: &[K::Elem], - output: &mut [K::Elem], +pub unsafe fn transpose_small_twiddle_f64( + input: impl NeonArray, + mut output: impl NeonArrayMut, + twiddles: impl NeonArray, width: usize, height: usize, ) { - const TILE_SIZE: usize = 32; - - // Loop 1: iterate over row tiles of height TILE_SIZE - let mut by = 0; - while by < height { - let block_h = (height - by).min(TILE_SIZE); - - // Loop 2: iterate over column tiles of width TILE_SIZE - let mut bx = 0; - while bx < width { - let block_w = (width - bx).min(TILE_SIZE); - - // Loop 3: process 2 rows at a time within the tile - let mut y = 0; - while y + 2 <= block_h { - let cy = by + y; - let in_r0 = cy * width + bx; - let in_r1 = (cy + 1) * width + bx; - - // Loop 4: process 2 columns at a time (2x2 complex sub-block) - let mut x = 0; - while x + 2 <= block_w { - let cx = bx + x; - let out_c0 = cy + cx * height; - let out_c1 = cy + (cx + 1) * height; - - K::transpose_2x2(input, output, in_r0 + x, in_r1 + x, out_c0, out_c1); - x += 2; - } + // Loop 1: process 2 rows at a time + let mut y = 0; + while y + 2 <= height { + let in_r0 = y * width; + let in_r1 = (y + 1) * width; - // Handle remainder column when block_w is odd (2 rows x 1 column) - if x < block_w { - let cx = bx + x; - let out_c0 = cy + cx * height; - K::transpose_1x2(input, output, in_r0 + x, in_r1 + x, out_c0); - } + // Loop 2: process 2 columns at a time (2x2 complex block) + let mut x = 0; + while x + 2 <= width { + let out_c0 = y + x * height; + let out_c1 = y + (x + 1) * height; - y += 2; - } + let a0 = input.load_complex(in_r0 + x); + let tw_a0 = twiddles.load_complex(in_r0 + x); + let a1 = input.load_complex(in_r0 + x + 1); + let tw_a1 = twiddles.load_complex(in_r0 + x + 1); - // Handle remainder row when block_h is odd (1 row x block_w columns) - if y < block_h { - let cy = by + y; - let in_r = cy * width + bx; - for x in 0..block_w { - let cx = bx + x; - let out_c = cy + cx * height; - K::transpose_1x1(input, output, in_r + x, out_c); - } - } + let b0 = input.load_complex(in_r1 + x); + let tw_b0 = twiddles.load_complex(in_r1 + x); + let b1 = input.load_complex(in_r1 + x + 1); + let tw_b1 = twiddles.load_complex(in_r1 + x + 1); - bx += TILE_SIZE; + let res_a0 = NeonVector::mul_complex(a0, tw_a0); + let res_a1 = NeonVector::mul_complex(a1, tw_a1); + let res_b0 = NeonVector::mul_complex(b0, tw_b0); + let res_b1 = NeonVector::mul_complex(b1, tw_b1); + + output.store_complex(res_a0, out_c0); + output.store_complex(res_b0, out_c0 + 1); + output.store_complex(res_a1, out_c1); + output.store_complex(res_b1, out_c1 + 1); + + x += 2; } - by += TILE_SIZE; - } -} -pub unsafe fn transpose_f64( - input: &[Complex], - output: &mut [Complex], - width: usize, - height: usize, -) { - super::neon_common::assert_f64::(); - assert!(input.len() >= width * height); - assert!(output.len() >= width * height); + // Remainder column when width is odd (2 rows x 1 column) + if x < width { + let out_c0 = y + x * height; - transpose_tiled::(input, output, width, height); -} + let a0 = input.load_complex(in_r0 + x); + let tw_a0 = twiddles.load_complex(in_r0 + x); + let b0 = input.load_complex(in_r1 + x); + let tw_b0 = twiddles.load_complex(in_r1 + x); -pub unsafe fn transpose_f32( - input: &[Complex], - output: &mut [Complex], - width: usize, - height: usize, -) { - super::neon_common::assert_f32::(); - assert!(input.len() >= width * height); - assert!(output.len() >= width * height); + let res_a0 = NeonVector::mul_complex(a0, tw_a0); + let res_b0 = NeonVector::mul_complex(b0, tw_b0); - transpose_tiled::(input, output, width, height); -} + output.store_complex(res_a0, out_c0); + output.store_complex(res_b0, out_c0 + 1); + } -pub unsafe fn transpose( - width: usize, - height: usize, - input: &[T], - output: &mut [T], -) -> bool { - assert!(input.len() >= width * height); - assert!(output.len() >= width * height); + y += 2; + } - use std::any::TypeId; - if TypeId::of::() == TypeId::of::>() { - let input: &[Complex] = crate::array_utils::workaround_transmute(input); - let output: &mut [Complex] = crate::array_utils::workaround_transmute_mut(output); - transpose_f64(input, output, width, height); - return true; - } else if TypeId::of::() == TypeId::of::>() { - let input: &[Complex] = crate::array_utils::workaround_transmute(input); - let output: &mut [Complex] = crate::array_utils::workaround_transmute_mut(output); - transpose_f32(input, output, width, height); - return true; + // Remainder row when height is odd (1 row x width columns) + if y < height { + let in_r = y * width; + for x in 0..width { + let in_idx = in_r + x; + let out_idx = y + x * height; + let a = input.load_complex(in_idx); + let tw = twiddles.load_complex(in_idx); + let res = NeonVector::mul_complex(a, tw); + output.store_complex(res, out_idx); + } } - false } -/// Fused complex twiddle multiplication and matrix transpose for small sizes. +/// Fused complex twiddle multiplication and matrix transpose for small f32 sizes. /// /// Computes `output[y + x * height] = input[x + y * width] * twiddles[x + y * width]`. -/// Processes 2x2 complex blocks with SIMD multiplication and transposition, +/// Processes 2x2 complex blocks with NEON SIMD vector instructions, /// followed by remainder column and row handling. #[inline(always)] -unsafe fn transpose_small_twiddle_impl( - input: &[K::Elem], - output: &mut [K::Elem], - twiddles: &[K::Elem], +pub unsafe fn transpose_small_twiddle_f32( + input: impl NeonArray, + mut output: impl NeonArrayMut, + twiddles: impl NeonArray, width: usize, height: usize, ) { @@ -720,20 +353,43 @@ unsafe fn transpose_small_twiddle_impl( let in_r0 = y * width; let in_r1 = (y + 1) * width; - // Loop 2: process 2 columns at a time (2x2 complex sub-block) + // Loop 2: process 2 columns at a time (2x2 complex block) let mut x = 0; while x + 2 <= width { let out_c0 = y + x * height; let out_c1 = y + (x + 1) * height; - K::step_2x2(input, output, twiddles, in_r0 + x, in_r1 + x, out_c0, out_c1); + let row0 = input.load_complex(in_r0 + x); + let tw_row0 = twiddles.load_complex(in_r0 + x); + let row1 = input.load_complex(in_r1 + x); + let tw_row1 = twiddles.load_complex(in_r1 + x); + + let res0 = NeonVector::mul_complex(row0, tw_row0); + let res1 = NeonVector::mul_complex(row1, tw_row1); + + let [col0, col1] = transpose_complex_2x2_f32(res0, res1); + + output.store_complex(col0, out_c0); + output.store_complex(col1, out_c1); + x += 2; } // Remainder column when width is odd (2 rows x 1 column) if x < width { let out_c0 = y + x * height; - K::step_1x2(input, output, twiddles, in_r0 + x, in_r1 + x, out_c0); + + let a0 = vget_low_f32(input.load_partial_lo_complex(in_r0 + x)); + let b0 = vget_low_f32(input.load_partial_lo_complex(in_r1 + x)); + let val = vcombine_f32(a0, b0); + + let tw_a0 = vget_low_f32(twiddles.load_partial_lo_complex(in_r0 + x)); + let tw_b0 = vget_low_f32(twiddles.load_partial_lo_complex(in_r1 + x)); + let tw = vcombine_f32(tw_a0, tw_b0); + + let res = NeonVector::mul_complex(val, tw); + + output.store_complex(res, out_c0); } y += 2; @@ -743,42 +399,16 @@ unsafe fn transpose_small_twiddle_impl( if y < height { let in_r = y * width; for x in 0..width { - let out_c = y + x * height; - K::step_1x1(input, output, twiddles, in_r + x, out_c); + let in_idx = in_r + x; + let out_idx = y + x * height; + let a = input.load_partial_lo_complex(in_idx); + let tw = twiddles.load_partial_lo_complex(in_idx); + let res = NeonVector::mul_complex(a, tw); + output.store_partial_lo_complex(res, out_idx); } } } -pub unsafe fn transpose_small_twiddle_f64( - input: &[Complex], - output: &mut [Complex], - twiddles: &[Complex], - width: usize, - height: usize, -) { - super::neon_common::assert_f64::(); - assert!(input.len() >= width * height); - assert!(output.len() >= width * height); - assert!(twiddles.len() >= width * height); - - transpose_small_twiddle_impl::(input, output, twiddles, width, height); -} - -pub unsafe fn transpose_small_twiddle_f32( - input: &[Complex], - output: &mut [Complex], - twiddles: &[Complex], - width: usize, - height: usize, -) { - super::neon_common::assert_f32::(); - assert!(input.len() >= width * height); - assert!(output.len() >= width * height); - assert!(twiddles.len() >= width * height); - - transpose_small_twiddle_impl::(input, output, twiddles, width, height); -} - pub unsafe fn transpose_small_twiddle( width: usize, height: usize, @@ -786,9 +416,9 @@ pub unsafe fn transpose_small_twiddle( output: &mut [Complex], twiddles: &[Complex], ) -> bool { - assert!(input.len() >= width * height); - assert!(output.len() >= width * height); - assert!(twiddles.len() >= width * height); + debug_assert!(input.len() >= width * height); + debug_assert!(output.len() >= width * height); + debug_assert!(twiddles.len() >= width * height); use std::any::TypeId; if TypeId::of::() == TypeId::of::() { @@ -845,46 +475,18 @@ mod unit_tests { } #[test] - fn test_transpose_neon_f32_f64() { - use num_traits::Zero; - for width in [1, 2, 3, 4, 5, 7, 8, 16, 31, 32, 33, 40] { - for height in [1, 2, 3, 4, 5, 7, 8, 16, 31, 32, 33, 40] { - let len = width * height; - - // f32 - let input_f32: Vec> = (0..len) - .map(|i| Complex::new(i as f32, (i * 2) as f32)) - .collect(); - let mut out_f32 = vec![Complex::zero(); len]; - unsafe { transpose(width, height, &input_f32, &mut out_f32) }; - for y in 0..height { - for x in 0..width { - assert_eq!( - input_f32[x + y * width], - out_f32[y + x * height], - "f32 mismatch at ({}, {}) for {}x{}", - x, y, width, height - ); - } - } - - // f64 - let input_f64: Vec> = (0..len) - .map(|i| Complex::new(i as f64, (i * 2) as f64)) - .collect(); - let mut out_f64 = vec![Complex::zero(); len]; - unsafe { transpose(width, height, &input_f64, &mut out_f64) }; - for y in 0..height { - for x in 0..width { - assert_eq!( - input_f64[x + y * width], - out_f64[y + x * height], - "f64 mismatch at ({}, {}) for {}x{}", - x, y, width, height - ); - } - } - } + fn test_pack() { + unsafe { + let nbr2 = vld1q_f32([5.0, 6.0, 7.0, 8.0].as_ptr()); + let nbr1 = vld1q_f32([1.0, 2.0, 3.0, 4.0].as_ptr()); + let first = extract_lo_lo_f32(nbr1, nbr2); + let second = extract_hi_hi_f32(nbr1, nbr2); + let first = std::mem::transmute::; 2]>(first); + let second = std::mem::transmute::; 2]>(second); + let first_expected = [Complex::new(1.0, 2.0), Complex::new(5.0, 6.0)]; + let second_expected = [Complex::new(3.0, 4.0), Complex::new(7.0, 8.0)]; + assert_eq!(first, first_expected); + assert_eq!(second, second_expected); } } From 6fd6af526fd3fc456b9097a20f824f06867caed0 Mon Sep 17 00:00:00 2001 From: Wang-Yue Date: Thu, 17 Sep 2026 14:40:47 +0900 Subject: [PATCH 7/7] Run cargo fmt --- src/array_utils.rs | 35 ++++++++++++++++++++++++++++------- src/neon/neon_utils.rs | 24 ++++++++++++++++++------ 2 files changed, 46 insertions(+), 13 deletions(-) diff --git a/src/array_utils.rs b/src/array_utils.rs index eb6502c0..54ecf4b7 100644 --- a/src/array_utils.rs +++ b/src/array_utils.rs @@ -30,7 +30,8 @@ pub unsafe fn transpose_small_twiddle( #[cfg(all(target_arch = "aarch64", feature = "neon"))] { - if crate::neon::neon_utils::transpose_small_twiddle(width, height, input, output, twiddles) { + if crate::neon::neon_utils::transpose_small_twiddle(width, height, input, output, twiddles) + { return; } } @@ -184,16 +185,26 @@ mod unit_tests { let twiddles_f32: Vec> = random_signal(len); let mut output_f32 = vec![Zero::zero(); len]; unsafe { - transpose_small_twiddle(width, height, &input_f32, &mut output_f32, &twiddles_f32); + 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, + (actual.re - expected.re).abs() < 1e-5 + && (actual.im - expected.im).abs() < 1e-5, "f32 mismatch at ({}, {}) for {}x{}", - x, y, width, height + x, + y, + width, + height ); } } @@ -203,16 +214,26 @@ mod unit_tests { let twiddles_f64: Vec> = random_signal(len); let mut output_f64 = vec![Zero::zero(); len]; unsafe { - transpose_small_twiddle(width, height, &input_f64, &mut output_f64, &twiddles_f64); + 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, + (actual.re - expected.re).abs() < 1e-10 + && (actual.im - expected.im).abs() < 1e-10, "f64 mismatch at ({}, {}) for {}x{}", - x, y, width, height + x, + y, + width, + height ); } } diff --git a/src/neon/neon_utils.rs b/src/neon/neon_utils.rs index 3bf16306..a309355b 100644 --- a/src/neon/neon_utils.rs +++ b/src/neon/neon_utils.rs @@ -1,7 +1,7 @@ +use crate::neon::neon_vector::{NeonArray, NeonArrayMut, NeonVector}; +use crate::FftNum; use core::arch::aarch64::*; use num_complex::Complex; -use crate::FftNum; -use crate::neon::neon_vector::{NeonArray, NeonArrayMut, NeonVector}; // __ __ _ _ _________ _ _ _ // | \/ | __ _| |_| |__ |___ /___ \| |__ (_) |_ @@ -513,9 +513,15 @@ mod unit_tests { let expected = input_f32[x + y * width] * twiddles_f32[x + y * width]; let actual = out_f32[y + x * height]; assert!( - (actual.re - expected.re).abs() < 1e-5 && (actual.im - expected.im).abs() < 1e-5, + (actual.re - expected.re).abs() < 1e-5 + && (actual.im - expected.im).abs() < 1e-5, "f32 twiddle mismatch at ({}, {}) for {}x{}: expected {:?}, got {:?}", - x, y, width, height, expected, actual + x, + y, + width, + height, + expected, + actual ); } } @@ -536,9 +542,15 @@ mod unit_tests { let expected = input_f64[x + y * width] * twiddles_f64[x + y * width]; let actual = out_f64[y + x * height]; assert!( - (actual.re - expected.re).abs() < 1e-10 && (actual.im - expected.im).abs() < 1e-10, + (actual.re - expected.re).abs() < 1e-10 + && (actual.im - expected.im).abs() < 1e-10, "f64 twiddle mismatch at ({}, {}) for {}x{}: expected {:?}, got {:?}", - x, y, width, height, expected, actual + x, + y, + width, + height, + expected, + actual ); } }