Skip to content

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

Merged
ejmahler merged 8 commits into
ejmahler:masterfrom
Wang-Yue:master
Sep 24, 2026
Merged

ejmahler merged 8 commits into
ejmahler:masterfrom
Wang-Yue:master

Conversation

@Wang-Yue

@Wang-Yue Wang-Yue commented Sep 5, 2026

Copy link
Copy Markdown
Contributor

Overview

This PR improves FFT performance on AArch64 (ARM NEON) platforms by optimizing memory-bound matrix operations in mixed-radix algorithms:

  1. Vectorized Matrix Transpose: Implements NEON-accelerated tiled $2 \times 2$ complex transposition for f32 and f64 in neon_utils, replacing the generic scalar transpose fallback for small matrices.
  2. Fused Twiddle Multiplication + Transpose: In MixedRadixSmall, Steps 3 and 4 previously performed an elementwise twiddle multiply over the buffer followed by an out-of-place matrix transpose into scratch. We fuse these operations into a single pass (transpose_small_twiddle) that loads the source complex elements, multiplies them by the twiddle factors via NEON vector instructions, and stores them transposed in a single memory traversal.

Commits Breakdown

  1. ac8f1b1 — neon: vectorize 2x2 block transpose in neon_utils

    • Adds NEON acceleration hook in array_utils::transpose.
    • Implements SIMD vector transposition for $2 \times 2$ blocks of complex numbers using ARM64 vtrn1q_f64/vtrn2q_f64 and vzip1q_f32/vzip2q_f32.
    • Preserves portable fallback to transpose::transpose for other architectures and non-float types.
  2. 71f2a1f — neon: fuse twiddle multiplication into transpose pass in MixedRadixSmall

    • Fuses Step 3 (twiddle factor multiplication) and Step 4 (matrix transpose) in MixedRadixSmall::perform_fft_inplace and perform_fft_out_of_place.
    • Eliminates an entire intermediate memory read/write pass over the dataset, significantly improving L1 cache locality and reducing memory bandwidth.
  3. c82055b — neon: accelerate matrix transpose with tiled NEON vectorization

    • Generalizes the vector transpose into a blocked $2 \times 2$ kernel with scalar boundary/edge handlers for odd rows and columns.
    • Integrates the accelerated transpose with GoodThomasAlgorithmSmall and MixedRadixSmall.

Benchmark Results (Apple Silicon / AArch64 NEON, f64)

Measured on Apple Silicon (aarch64-apple-darwin), comparing upstream v6.4.1 (commit 4758ab0) against this PR branch (commit c82055b) with 100,000 iterations per size:

FFT Length Factors ($w \times h$) Upstream v6.4.1 This PR Speedup Latency Reduction
20 $5 \times 4$ 114.7 ns 94.7 ns 1.21x -17.4%
60 $12 \times 5$ 264.4 ns 206.8 ns 1.28x -21.8%
100 $10 \times 10$ 319.6 ns 245.9 ns 1.30x -23.1%
120 $15 \times 8$ 301.5 ns 228.8 ns 1.32x -24.1%
140 $20 \times 7$ 480.5 ns 443.5 ns 1.08x -7.7%
200 $20 \times 10$ 833.1 ns 699.2 ns 1.19x -16.1%
320 $16 \times 20$ 963.3 ns 913.6 ns 1.05x -5.2%

Correctness & Verification

  • Unit tests: cargo test --lib (122 passed; 0 failed).
  • Precision / Accuracy: tests/accuracy.rs (all 4 forward/inverse f32/f64 roundtrip accuracy tests pass with identical $\epsilon$ tolerances).
  • Documentation tests: cargo test --doc (17 passed; 0 failed).
  • Target isolation: All NEON intrinsics are gated behind #[cfg(all(target_arch = "aarch64", feature = "neon"))]. Non-ARM and non-NEON compilation targets remain untouched and continue using the portable crate routines.

Reproducing the Benchmark

The benchmarks can be reproduced using the included example script:

cargo run --release --example bench_mixed_radix
Benchmark Source Code (examples/bench_mixed_radix.rs)
use rustfft::num_complex::Complex;
use rustfft::num_traits::Zero;
use rustfft::FftPlanner;
use std::time::Instant;

fn bench_len(len: usize, iters: usize) -> f64 {
    let mut planner = FftPlanner::<f64>::new();
    let fft = planner.plan_fft_forward(len);
    let mut buffer = vec![Complex::zero(); len];
    let mut scratch = vec![Complex::zero(); fft.get_inplace_scratch_len()];

    // Warm-up to ensure CPU frequency scaling and cache lines are hot
    for _ in 0..2000 {
        fft.process_with_scratch(&mut buffer, &mut scratch);
    }

    let start = Instant::now();
    for _ in 0..iters {
        fft.process_with_scratch(&mut buffer, &mut scratch);
        std::hint::black_box(&buffer);
    }
    let elapsed_ns = start.elapsed().as_nanos() as f64;
    elapsed_ns / (iters as f64)
}

fn main() {
    let lengths: &[(usize, &str)] = &[
        (20, "5x4"),
        (60, "12x5"),
        (100, "10x10"),
        (120, "15x8"),
        (140, "20x7"),
        (200, "20x10"),
        (320, "16x20"),
    ];

    let iters = 100_000;
    println!("Running mixed-radix microbenchmark ({} iterations per size, f64):", iters);
    println!("{:-<55}", "");
    println!("{:<8} | {:<15} | {:<15}", "Length", "Factors (w x h)", "Latency (ns)");
    println!("{:-<55}", "");

    for &(len, factors) in lengths {
        let ns = bench_len(len, iters);
        println!("{:<8} | {:<15} | {:>10.2} ns", len, factors, ns);
    }
    println!("{:-<55}", "");
}

@HEnquist

HEnquist commented Sep 6, 2026

Copy link
Copy Markdown
Contributor

I made the NEON implementation, so some quick feedback on that side.

Could you add benchmarks to benches/bench_rustfft.rs following the existing pattern, showing both the speedup and that nothing regresses? mixed_radix_*, mixed_radix_small_* and good_thomas_small_* already cover this code, though they are f32 only so f64 ones would need adding. I had a go at reproducing your numbers and got mixed results at the larger sizes, so it would be good to have this in the harness where anyone can run it.

Pasting the benchmark into the description as a code block is a strange way to do it. The description calls it the included example script and gives a cargo run --release --example bench_mixed_radix command, but examples/bench_mixed_radix.rs isn't in the diff, so that fails on a checkout of the branch until you copy the code out of the description by hand.

transpose_f64 and transpose_f32 are hard to follow. Four levels of nested while loops, index arithmetic written in scalars rather than complex values (in_r0 + x * 2, (cy + cx * height) * 2), and the two are near-identical 70 line copies. For unsafe code I want to be able to read it and convince myself the edge handling is right, and I can't here. A comment on what each loop covers, and something to remove the duplication between the two, would go a long way.

The raw pointers are the bigger thing. transpose_f64(p_in: *const f64, p_out: *mut f64, width, height) drops the slice lengths at the boundary, and the callers hand it input.as_ptr() as *const f64. So the requirement that both buffers hold width*height elements is never stated and never checked. Same for transpose_small_twiddle with three slices at once. array_utils::transpose is unsafe with nothing saying what it expects either, and it replaces transpose::transpose, which asserted the lengths.

The rest of the NEON code asserts what it assumes, assert_f32/assert_f64 on every entry point, and the helpers describe their lane layout in a comment above them. That is the difference, and I'd like the new code to do the same.

The TypeId check itself is fine, that is how the crate specializes. The other half of the pattern is array_utils::workaround_transmute, which reinterprets the slice and keeps the length, see neon_butterflies.rs. That is also the function this PR removes the #[allow(unused)] from, which brings back a dead code warning on cargo build --no-default-features. Going through it would keep &[Complex<f64>] in the signatures and let the kernels index with get_unchecked the way transpose_small did.

…mance

- 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.
@Wang-Yue

Wang-Yue commented Sep 7, 2026

Copy link
Copy Markdown
Contributor Author

RustFFT NEON Transpose & Twiddle Fusion Benchmark Results

Benchmark comparison measured directly on an Apple Silicon (aarch64 NEON) machine using cargo +nightly bench --bench bench_rustfft:

  • Baseline: v6.4.1 (commit 4758ab0) using the scalar transpose crate
  • Latest Commit: 0382a62 with tiled NEON transposition, twiddle-fused transposition in MixedRadixSmall, and direct scalar transpose_small

1. MixedRadixSmall

Fuses twiddle factor multiplication directly into the transposition pass, eliminating a separate memory load/store pass.

Single Precision (f32)

Dimensions ($W \times H$) Baseline v6.4.1 Latest Commit 0382a62 Speedup Improvement
2 × 3 20.10 ns 22.01 ns ~parity (sub-25ns measurement noise)
3 × 4 31.11 ns 25.73 ns 1.21× +17.3%
4 × 5 48.10 ns 40.11 ns 1.20× +16.6%
7 × 32 750.21 ns 682.44 ns 1.10× +9.0%

Double Precision (f64)

Dimensions ($W \times H$) Baseline v6.4.1 Latest Commit 0382a62 Speedup Improvement
2 × 3 20.71 ns 18.56 ns 1.12× +10.4%
3 × 4 32.27 ns 27.76 ns 1.16× +14.0%
4 × 5 51.89 ns 43.11 ns 1.20× +16.9%
7 × 32 761.34 ns 651.47 ns 1.17× +14.4%

2. MixedRadix

Standard 6-step mixed-radix FFT algorithm with 3 matrix transpositions per transform.

Single Precision (f32)

Dimensions ($W \times H$) Baseline v6.4.1 Latest Commit 0382a62 Speedup Improvement
2 × 3 22.02 ns 25.26 ns ~parity (sub-30ns noise)
3 × 4 31.79 ns 28.78 ns 1.10× +9.5%
4 × 5 46.36 ns 37.71 ns 1.23× +18.7%
7 × 32 408.39 ns 282.00 ns 1.45× +31.0%
32 × 27 2,872.80 ns 2,390.13 ns 1.20× +16.8%
256 × 243 352.16 µs 308.97 µs 1.14× +12.3%
2048 × 3 14.98 µs 13.02 µs 1.15× +13.0%
2048 × 2187 42.42 ms 38.21 ms 1.11× +9.9%

Double Precision (f64)

Dimensions ($W \times H$) Baseline v6.4.1 Latest Commit 0382a62 Speedup Improvement
2 × 3 22.15 ns 26.15 ns ~parity (sub-30ns noise)
3 × 4 34.93 ns 33.26 ns 1.05× +4.8%
4 × 5 53.21 ns 45.96 ns 1.16× +13.6%
7 × 32 568.77 ns 469.59 ns 1.21× +17.4%
32 × 27 3,714.75 ns 3,150.66 ns 1.18× +15.2%
256 × 243 482.31 µs 433.09 µs 1.11× +10.2%
2048 × 3 24.36 µs 23.03 µs 1.06× +5.4%
2048 × 2187 65.32 ms 61.72 ms 1.06× +5.5%

3. GoodThomasAlgorithm

Good-Thomas (prime-factor) algorithm utilizing matrix transposition without twiddles.

Single Precision (f32)

Dimensions ($W \times H$) Baseline v6.4.1 Latest Commit 0382a62 Speedup Improvement
2 × 3 18.95 ns 18.43 ns parity (1.03×)
3 × 4 33.29 ns 28.38 ns 1.17× +14.8%
4 × 5 43.80 ns 42.60 ns 1.03× +2.7%
7 × 32 375.69 ns 340.65 ns 1.10× +9.3%
32 × 27 2,897.71 ns 2,562.53 ns 1.13× +11.6%
256 × 243 397.05 µs 364.28 µs 1.09× +8.3%
2048 × 3 14.32 µs 14.66 µs parity (0.98×)
2048 × 2187 64.30 ms 60.27 ms 1.07× +6.3%

Double Precision (f64)

Dimensions ($W \times H$) Baseline v6.4.1 Latest Commit 0382a62 Speedup Improvement
2 × 3 20.61 ns 21.58 ns parity (0.96×)
3 × 4 32.27 ns 29.61 ns 1.09× +8.2%
4 × 5 48.61 ns 44.64 ns 1.09× +8.2%
7 × 32 491.68 ns 464.67 ns 1.06× +5.5%
32 × 27 3,596.20 ns 3,065.90 ns 1.17× +14.7%
256 × 243 501.41 µs 459.35 µs 1.09× +8.4%
2048 × 3 22.90 µs 23.16 µs parity (0.99×)
2048 × 2187 88.00 ms 79.73 ms 1.10× +9.4%

4. GoodThomasAlgorithmSmall

Performs small prime-factor FFTs without twiddle factors, utilizing direct scalar indexing for small matrix transpositions.

Single Precision (f32)

Dimensions ($W \times H$) Baseline v6.4.1 Latest Commit 0382a62 Speedup
2 × 3 14.75 ns 13.76 ns +6.7%
3 × 4 21.11 ns 21.59 ns ~parity
4 × 5 34.52 ns 34.89 ns ~parity
7 × 32 660.54 ns 666.41 ns ~parity

Double Precision (f64)

Dimensions ($W \times H$) Baseline v6.4.1 Latest Commit 0382a62 Speedup
2 × 3 12.68 ns 12.54 ns parity
3 × 4 22.08 ns 21.83 ns parity
4 × 5 37.57 ns 36.94 ns parity
7 × 32 651.54 ns 655.05 ns parity

Key Findings & Takeaways

  1. MixedRadixSmall Twiddle Fusion:
    • Fusing twiddle multiplication into the transpose step provides consistent +10% to +17% speedup on f32 and +10% to +17% speedup on f64.
  2. GoodThomasAlgorithmSmall Parity:
    • Preserving the direct scalar loop in transpose_small prevents tiling/dispatch overhead on tiny transforms, ensuring full performance parity with baseline across all dimensions.
  3. Broad Speedups at Scale (MixedRadix and GoodThomasAlgorithm):
    • MixedRadix achieves up to +31% speedup on f32 and up to +17% speedup on f64.
    • Large transforms ($2048 \times 2187 \approx 4.48\text{M points}$) run ~4.2 ms faster per FFT in f32 (from 42.4 ms to 38.2 ms) and ~3.6 ms faster in f64 (from 65.3 ms to 61.7 ms).
    • In GoodThomasAlgorithm, the $2048 \times 2187$ size runs ~4.0 ms faster in f32 (from 64.3 ms to 60.3 ms) and ~8.3 ms faster in f64 (from 88.0 ms to 79.7 ms).

@HEnquist

HEnquist commented Sep 7, 2026

Copy link
Copy Markdown
Contributor

This is a lot better, thanks. Slices instead of raw pointers, asserts on the lengths, workaround_transmute instead of a hand rolled cast, comments on the loops, and the f64 benches. All the things I asked about are addressed.

Two things I'd still like to see changed, and one observation about the size of it.

The complex multiply helpers are duplicates. neon_vector.rs already has NeonVector::mul_complex implemented for both float32x4_t and float64x2_t, and neon_complex_mul_f64 is the same three lines with left/right renamed to val/tw. neon_complex_mul_f32 is the same story. Only neon_complex_mul_single_f32 is genuinely new.

The same applies to the new traits. TransposeKernel and TransposeTwiddleKernel are an abstraction over the two NEON vector types, which is what NeonVector already is: it has COMPLEX_PER_VECTOR, ScalarType, load_complex, store_complex and mul_complex. Writing the kernels against NeonVector should let both new traits and both duplicated multiplies go away.

That matters for the size. The PR is +783 lines of src, and neon_utils.rs goes from 301 to 946 lines. Some of that is unavoidable, but a good part is scaffolding that duplicates what is already in neon_vector.rs.

It matters more for what comes next. The SSE and wasm_simd planners both build the same generic MixedRadix, MixedRadixSmall and GoodThomasAlgorithmSmall from src/algorithm, so they go through the same array_utils::transpose and transpose_small_twiddle this PR changes. An SSE version is the obvious follow up, and the pieces are already there: sse_utils has transpose_complex_2x2_f32, mul_complex_f32 and mul_complex_f64, and SseNum has COMPLEX_PER_VECTOR. Nothing needs an instruction SSE lacks. But if this lands in its current shape, the SSE port means a second copy of the two kernel traits and a third and fourth copy of the complex multiply. Better to settle the shape now. AVX is not affected, it has its own mixed radix implementations and doesn't use these.

On the numbers, I ran the planner benches on this head, f32 and f64, master (4758ab0) against 0382a62. 28 of 92 improve by more than 10%, best around 1.25x, nothing regresses, and the 21 power of two benches are flat at 0.993x mean since those go through Radix4. That is a good result. It is also a fair amount of new aarch64 only code for gains that land on non power of two sizes. Worth saying that the fused twiddle multiply is the portable half of it, the scalar fallback in array_utils already gives SSE, wasm and scalar builds most of that win for free.

@ejmahler this next part is your call rather than mine.

Building MixedRadix directly at large f64 sizes regresses: 512x512 is 0.60x, 16x1024 is 0.54x, 1024x1024 is 0.78x, 64x16384 is 0.82x. The planner doesn't produce those shapes, so I did not find anything user visible. What concerns me more is what it does to MixedRadix relative to Radix4 at the same length. Ratio of MixedRadix to Radix4, so lower is better for MixedRadix:

length master this PR
f32 256 0.97x 0.94x
f32 1024 1.05x 1.02x
f32 65536 1.13x 1.08x
f32 1048576 1.24x 1.62x
f64 65536 1.36x 1.43x
f64 1048576 1.70x 2.09x

At the smaller sizes MixedRadix gains ground, and at f32 256 it is actually faster than Radix4. That part is good. At the large sizes it loses ground hard: f32 at 1M goes from 24% behind Radix4 to 62% behind, and f64 from 70% to 109%. The two used to be reasonably close, and MixedRadix could win in places. Widening that gap costs flexibility in the planner later even if nothing changes today, and I would rather not trade it away for gains elsewhere.

To me it points at TILE_SIZE of 32 and the y outer / x inner loop order being wrong at large sizes, especially for Complex<f64> where a 128-bit register holds exactly one value and a tile is 16 kB. The write striding is the expensive part. Whatever gets picked there would port straight to SSE, since __m128d has the same geometry.

One caveat on all of these numbers. They come from an M1, and everything above is on the performance cores, which have 128 kB of L1 data cache and a lot of memory bandwidth. That is about the friendliest environment there is for a 32x32 tile. I did spot check the worst cases on the efficiency cores and it looked similar rather than worse, but those still sit behind the same fast memory system, so take that for what it is and not more. A Cortex-A53 or A55 with 32 kB of L1 and a much narrower path to memory is a different situation, and RustFFT gets run on Raspberry Pis and phones. Where a 16 kB tile starts to hurt depends on L1 size and on how expensive the strided writes are, so this needs numbers from weaker hardware before it can really be evaluated, the gains as much as the regressions.

Smaller things:

  • The length asserts run at three levels, array_utils, neon_utils::transpose*, and again in the _f64/_f32 functions. That is 6 checks per transpose and 9 per twiddle transpose, in release builds. sse_vector.rs uses a single debug_assert! for the same job. The smallest benches regress a little (mixed_radix_64_0002_3 32.1 to 37.9 ns) and making the inner two levels debug_assert! recovers part of that.
  • transpose_f64 calls assert_f64::<f64>(), which is always true since the function is already monomorphic on f64.
  • The doc comment on transpose_small lost its last two lines and is now a dangling fragment.

@HEnquist

HEnquist commented Sep 7, 2026

Copy link
Copy Markdown
Contributor

I got curious about how much of the win actually depends on the transpose change, so I built two cut down versions and ran all 92 planner benches on each, same machine and same commits as before.

variant src diff benches improving >10% mean worst
A: scalar fused twiddle only, no NEON at all +25/-21 7 of 92 1.026x 0.99x
B: A plus the NEON twiddle kernel, transpose::transpose kept everywhere +470/-50 23 of 92 1.057x 0.99x
C: this PR +783/-62 28 of 92 1.064x 0.93x

B is clearly the better trade. It is this PR with the array_utils::transpose calls in mixed_radix.rs and good_thomas_algorithm.rs reverted to transpose::transpose, so the transpose itself is untouched and only the fused twiddle multiply gets a NEON kernel. That keeps 89% of the mean gain and 23 of the 28 double digit wins for 60% of the code, and it has no regressions at all, worst case 0.99x against 0.93x for the full version. It also leaves MixedRadix's standing against Radix4 exactly as it is today, which was my main worry.

The code that becomes dead is not a judgement call, rustc lists it: array_utils::transpose, neon_utils::transpose, TransposeKernel with both its impls, transpose_tiled, transpose_f64 and transpose_f32. Removing those and the test that covers them leaves a clean build with the 122 tests passing. B still carries TransposeTwiddleKernel and the three duplicated multiply helpers, so building it on NeonVector as suggested above should shrink it a good deal further.

Variant A is worth knowing about as the floor. 25 lines, no new unsafe, no arch gating, and it works on x86, wasm and scalar builds too since it is just the fused loop in array_utils. 40% of the gain for 3% of the code.

@ejmahler on the transpose half, would it make more sense in the transpose crate? It is generic matrix transposition with nothing FFT specific about it, so done there it would help every architecture and every user of the crate, and RustFFT would pick it up with a version bump instead of carrying 500 lines of aarch64 code. The twiddle fusion can't move there, it is inherently an FFT thing, and going by the table that is where nearly all the value is anyway.

There may also be something to gain there without any SIMD. In out_of_place.rs, BLOCK_SIZE, SMALL_LEN and MEDIUM_LEN are all counted in elements rather than bytes, so a 16x16 block is 4 kB of Complex<f32> but 8 kB of Complex<f64>, and the small/medium/large thresholds land in different places in bytes depending on the element type. Making those aware of size_of::<T>() looks like a small change that would help everyone, and it is aimed at the same effect behind the f64 numbers I posted earlier. B plus something like that would probably get the last bit that B gives up.

Put together, B plus a transpose crate change looks strictly better than what is here now: fewer lines in RustFFT, no regressions anywhere, MixedRadix keeps its footing against Radix4, and the transpose gains land on every architecture instead of aarch64 only. I would rather see it split that way than merge the current version and tune it afterwards.

@ejmahler

ejmahler commented Sep 8, 2026 •

Copy link
Copy Markdown
Owner

Thanks @HEnquist for your feedback and analysis.

There's two key things I see worth pursuing here:

  1. The transposes would benefit from SIMD.
  2. The main bottleneck for mixed radix, especially the big ones, is memory traffic, so merging the twiddle pass with the middle transpose has a lot of potential.

There's a tradeoff of code size vs speed, especially if we combine #1 and #2. Fusing the twiddles with the transposes would probably help even more for big FFTs than for small (as long as it's tuned properly for cache etc), but it'd require a custom implementation for every architecture, which again idk if i have the stomach for.

I agree that B is a good compromise: Keep using transpose::transpose for the big FFTs, and optimize mixed radix small, since the implementation of that is significantly simpler. We can implement one fused twiddle-transpose function per arch, then write a wrapper function that just dynamically checks platform like this PR does. And separately, investigate using SIMD in transpose::transpose in an encapsulated way.

Since we're looking at this, I also realize that I've never tried switching from 6 step FFT to 4 step FFT for mixed radix small: Load the columns strided into a contiguous scratch buffer, do the column FFT, apply twiddles, and write back strided. This would skip the first and middle transpose altogether. That's probably better for scalar and simd f64, although for simd f32 we'd need some kind of specialized version that keeps the data unrolled 2x. But that would probably require a specialized implementation per architecture again, because we would want the data layout to be compatible with our perform_parallel_fft_contiguous functions, not the main FFT trait layout.

@ejmahler

ejmahler commented Sep 8, 2026

Copy link
Copy Markdown
Owner

As for this PR, I'd like to see these things, which are basically just reiterating @HEnquist 's main feedback:

  1. The neon SIMD code was written "outside of the ecosystem" of the existing code so to speak, which will hurt the maintainability of the code if it's kept this way. The multiply and transpose should use the existing utilities, especially the multiply. I'm open to changing multiply implementations if the new one is faster, otherwise everything should use the existing one. In order to get access to the multiply function, the functions (like transpose_small_twiddle_f64') will have to take impl NeonArrayandimpl NeonArrayMut` instances as parameters instead of slices, which as a bonus will be another step in

  2. Switch the big good thomas and mixed radix impls back to use transpose::transpose. Since we're dropping the full simd transposes and only keeping the small ones, it looks like it will reduce the LOC count to just have one function do f32 and a separate one do f64, and we don't need the transpose kernel trait at all, because we won't need to abstract away the data size.

@Wang-Yue

Wang-Yue commented Sep 8, 2026

Copy link
Copy Markdown
Contributor Author

RustFFT NEON Transpose & Twiddle Fusion Benchmark Results

Benchmark comparison measured directly on an Apple Silicon (aarch64 NEON) machine using cargo +nightly bench --bench bench_rustfft:

  • Baseline (v6.4.1): Upstream master (4758ab0) using the scalar transpose crate
  • Previous Commit (0382a62): Full tiled NEON transpose across all algorithms and twiddle-fused transposition in MixedRadixSmall
  • Latest Commit (77c7b2f): Maintainer-requested "Variant B"
    • Reverted large MixedRadix and GoodThomasAlgorithm matrix transpositions back to transpose::transpose
    • Removed dead/duplicate transpose kernels and traits
    • Retained fused twiddle transpose in MixedRadixSmall, integrated with NeonArray / NeonArrayMut and NeonVector::mul_complex
    • Replaced inner length asserts with debug_assert!

1. MixedRadixSmall

Fuses twiddle factor multiplication directly into the transposition pass using transpose_small_twiddle.

Single Precision (f32)

Dimensions ($W \times H$) Baseline v6.4.1 Commit 0382a62 Latest 77c7b2f Latest vs Baseline Latest vs 0382a62
2 × 3 20.10 ns 22.01 ns 19.06 ns +5.2% +13.4% (overhead reduced)
3 × 4 31.11 ns 25.73 ns 29.52 ns +5.1% -14.7%
4 × 5 48.10 ns 40.11 ns 44.06 ns +8.4% -9.8%
7 × 32 750.21 ns 682.44 ns 730.40 ns +2.6% -7.0%

Double Precision (f64)

Dimensions ($W \times H$) Baseline v6.4.1 Commit 0382a62 Latest 77c7b2f Latest vs Baseline Latest vs 0382a62
2 × 3 20.71 ns 18.56 ns 18.18 ns +12.2% +2.0% (overhead reduced)
3 × 4 32.27 ns 27.76 ns 28.35 ns +12.1% -2.1%
4 × 5 51.89 ns 43.11 ns 42.75 ns +17.6% +0.8%
7 × 32 761.34 ns 651.47 ns 647.29 ns +15.0% +0.6%

2. MixedRadix

Standard 6-step mixed-radix FFT algorithm with 3 matrix transpositions per transform. Reverted to transpose::transpose in 77c7b2f to prevent large-size cache striding regressions.

Single Precision (f32)

Dimensions ($W \times H$) Baseline v6.4.1 Commit 0382a62 Latest 77c7b2f Latest vs Baseline Latest vs 0382a62
2 × 3 22.02 ns 25.26 ns 21.72 ns ~parity +14.0%
3 × 4 31.79 ns 28.78 ns 32.37 ns ~parity -12.5%
4 × 5 46.36 ns 37.71 ns 45.97 ns ~parity -21.9%
7 × 32 408.39 ns 282.00 ns 407.59 ns ~parity -44.5%
32 × 27 2,872.80 ns 2,390.13 ns 2,686.57 ns +6.5% -12.4%
256 × 243 352.16 µs 308.97 µs 321.02 µs +8.8% -3.9%
2048 × 3 14.98 µs 13.02 µs 15.42 µs ~parity -18.4%
2048 × 2187 42.42 ms 38.21 ms 40.34 ms +4.9% -5.6%

Double Precision (f64)

Dimensions ($W \times H$) Baseline v6.4.1 Commit 0382a62 Latest 77c7b2f Latest vs Baseline Latest vs 0382a62
2 × 3 22.15 ns 26.15 ns 23.87 ns ~parity +8.7%
3 × 4 34.93 ns 33.26 ns 36.49 ns ~parity -9.7%
4 × 5 53.21 ns 45.96 ns 54.79 ns ~parity -19.2%
7 × 32 568.77 ns 469.59 ns 586.23 ns ~parity -24.8%
32 × 27 3,714.75 ns 3,150.66 ns 3,323.18 ns +10.5% -5.5%
256 × 243 482.31 µs 433.09 µs 435.76 µs +9.7% -0.6%
2048 × 3 24.36 µs 23.03 µs 24.42 µs ~parity -6.0%
2048 × 2187 65.32 ms 61.72 ms 62.79 ms +3.9% -1.7%

3. GoodThomasAlgorithm

Good-Thomas (prime-factor) algorithm utilizing matrix transposition without twiddles. Reverted to transpose::transpose in 77c7b2f.

Single Precision (f32)

Dimensions ($W \times H$) Baseline v6.4.1 Commit 0382a62 Latest 77c7b2f Latest vs Baseline Latest vs 0382a62
2 × 3 18.95 ns 18.43 ns 18.45 ns +2.6% ~parity
3 × 4 33.29 ns 28.38 ns 30.64 ns +8.0% -8.0%
4 × 5 43.80 ns 42.60 ns 43.95 ns ~parity -3.2%
7 × 32 375.69 ns 340.65 ns 372.15 ns ~parity -9.2%
32 × 27 2,897.71 ns 2,562.53 ns 2,715.83 ns +6.3% -6.0%
256 × 243 397.05 µs 364.28 µs 368.06 µs +7.3% -1.0%
2048 × 3 14.32 µs 14.66 µs 14.31 µs ~parity +2.4%
2048 × 2187 64.30 ms 60.27 ms 59.89 ms +6.9% +0.6%

Double Precision (f64)

Dimensions ($W \times H$) Baseline v6.4.1 Commit 0382a62 Latest 77c7b2f Latest vs Baseline Latest vs 0382a62
2 × 3 20.61 ns 21.58 ns 21.03 ns ~parity +2.5%
3 × 4 32.27 ns 29.61 ns 32.07 ns ~parity -8.3%
4 × 5 48.61 ns 44.64 ns 49.00 ns ~parity -9.8%
7 × 32 491.68 ns 464.67 ns 487.82 ns ~parity -5.0%
32 × 27 3,596.20 ns 3,065.90 ns 3,041.37 ns +15.4% +0.8%
256 × 243 501.41 µs 459.35 µs 437.47 µs +12.8% +4.8%
2048 × 3 22.90 µs 23.16 µs 22.98 µs ~parity +0.8%
2048 × 2187 88.00 ms 79.73 ms 82.42 ms +6.3% -3.4%

4. GoodThomasAlgorithmSmall

Small prime-factor FFTs without twiddles, utilizing direct scalar indexing in transpose_small.

Single Precision (f32)

Dimensions ($W \times H$) Baseline v6.4.1 Commit 0382a62 Latest 77c7b2f Status
2 × 3 14.75 ns 13.76 ns 13.82 ns Parity / +6.3%
3 × 4 21.11 ns 21.59 ns 21.67 ns Parity
4 × 5 34.52 ns 34.89 ns 35.05 ns Parity
7 × 32 660.54 ns 666.41 ns 662.49 ns Parity

Double Precision (f64)

Dimensions ($W \times H$) Baseline v6.4.1 Commit 0382a62 Latest 77c7b2f Status
2 × 3 12.68 ns 12.54 ns 13.10 ns Parity
3 × 4 22.08 ns 21.83 ns 22.59 ns Parity
4 × 5 37.57 ns 36.94 ns 36.27 ns Parity
7 × 32 651.54 ns 655.05 ns 645.05 ns Parity

Key Takeaways

  1. MixedRadixSmall Performance Preserved:

    • Retains a consistent +12% to +18% speedup in f64 and +3% to +8% speedup in f32 by fusing twiddle factor multiplication directly with matrix transposition.
    • Reducing redundant assertion checks from release builds recovered small-size overhead, making tiny transforms (2×3) +13% faster than 0382a62.
  2. Clean Architectural Alignment with NEON Subsystem:

    • Twiddle transpose now operates via impl NeonArray and impl NeonArrayMut, calling NeonVector::mul_complex.
    • All duplicate complex multiply helpers, TransposeKernel, and TransposeTwiddleKernel traits have been eliminated.
  3. Protection Against Large-Transform Cache Regressions:

    • Reverting large matrix transpositions in MixedRadix and GoodThomasAlgorithm back to transpose::transpose ensures RustFFT avoids cache striding penalties on large dimensions (such as 512×512 or 1024×1024), maintaining planner stability and competitiveness against Radix4.

@ejmahler

Copy link
Copy Markdown
Owner

I reviewed the new changes, and it's good to go. It just needs a cargo fmt to pass the remaining check.

@Wang-Yue

Copy link
Copy Markdown
Contributor Author

Thanks for the review! I ran cargo fmt.

@ejmahler
ejmahler merged commit 3f88e31 into ejmahler:master Sep 24, 2026
19 checks passed
@ejmahler

Copy link
Copy Markdown
Owner

Merged. Thank you for submitting, and thanks again to @HEnquist for the feedback.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants