[AArch64 / NEON] Vectorize matrix transpose and fuse twiddle multiplication in MixedRadixSmall - #172
Conversation
|
I made the NEON implementation, so some quick feedback on that side. Could you add benchmarks to 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
The raw pointers are the bigger thing. The rest of the NEON code asserts what it assumes, The TypeId check itself is fine, that is how the crate specializes. The other half of the pattern is |
…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.
RustFFT NEON Transpose & Twiddle Fusion Benchmark ResultsBenchmark comparison measured directly on an Apple Silicon (aarch64 NEON) machine using
1. MixedRadixSmallFuses twiddle factor multiplication directly into the transposition pass, eliminating a separate memory load/store pass. Single Precision (
|
| Dimensions ( |
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 ( |
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 ( |
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 ( |
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 ( |
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 ( |
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 ( |
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 ( |
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
-
MixedRadixSmallTwiddle Fusion:- Fusing twiddle multiplication into the transpose step provides consistent +10% to +17% speedup on
f32and +10% to +17% speedup onf64.
- Fusing twiddle multiplication into the transpose step provides consistent +10% to +17% speedup on
-
GoodThomasAlgorithmSmallParity:- Preserving the direct scalar loop in
transpose_smallprevents tiling/dispatch overhead on tiny transforms, ensuring full performance parity with baseline across all dimensions.
- Preserving the direct scalar loop in
-
Broad Speedups at Scale (
MixedRadixandGoodThomasAlgorithm):-
MixedRadixachieves up to +31% speedup onf32and up to +17% speedup onf64. - Large transforms (
$2048 \times 2187 \approx 4.48\text{M points}$ ) run ~4.2 ms faster per FFT inf32(from 42.4 ms to 38.2 ms) and ~3.6 ms faster inf64(from 65.3 ms to 61.7 ms). - In
GoodThomasAlgorithm, the$2048 \times 2187$ size runs ~4.0 ms faster inf32(from 64.3 ms to 60.3 ms) and ~8.3 ms faster inf64(from 88.0 ms to 79.7 ms).
-
|
This is a lot better, thanks. Slices instead of raw pointers, asserts on the lengths, Two things I'd still like to see changed, and one observation about the size of it. The complex multiply helpers are duplicates. The same applies to the new traits. That matters for the size. The PR is +783 lines of src, and It matters more for what comes next. The SSE and wasm_simd planners both build the same generic 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 @ejmahler this next part is your call rather than mine. Building
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 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:
|
|
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.
B is clearly the better trade. It is this PR with the The code that becomes dead is not a judgement call, rustc lists it: Variant A is worth knowing about as the floor. 25 lines, no new @ejmahler on the transpose half, would it make more sense in the There may also be something to gain there without any SIMD. In 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. |
|
Thanks @HEnquist for your feedback and analysis. There's two key things I see worth pursuing here:
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 |
|
As for this PR, I'd like to see these things, which are basically just reiterating @HEnquist 's main feedback:
|
RustFFT NEON Transpose & Twiddle Fusion Benchmark ResultsBenchmark comparison measured directly on an Apple Silicon (aarch64 NEON) machine using
1. MixedRadixSmallFuses twiddle factor multiplication directly into the transposition pass using Single Precision (
|
| Dimensions ( |
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 ( |
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 ( |
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 ( |
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 ( |
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 ( |
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 ( |
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 ( |
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
-
MixedRadixSmallPerformance Preserved:- Retains a consistent +12% to +18% speedup in
f64and +3% to +8% speedup inf32by 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.
- Retains a consistent +12% to +18% speedup in
-
Clean Architectural Alignment with NEON Subsystem:
- Twiddle transpose now operates via
impl NeonArrayandimpl NeonArrayMut, callingNeonVector::mul_complex. - All duplicate complex multiply helpers,
TransposeKernel, andTransposeTwiddleKerneltraits have been eliminated.
- Twiddle transpose now operates via
-
Protection Against Large-Transform Cache Regressions:
- Reverting large matrix transpositions in
MixedRadixandGoodThomasAlgorithmback totranspose::transposeensures RustFFT avoids cache striding penalties on large dimensions (such as 512×512 or 1024×1024), maintaining planner stability and competitiveness againstRadix4.
- Reverting large matrix transpositions in
|
I reviewed the new changes, and it's good to go. It just needs a |
|
Thanks for the review! I ran cargo fmt. |
|
Merged. Thank you for submitting, and thanks again to @HEnquist for the feedback. |
Overview
This PR improves FFT performance on AArch64 (ARM NEON) platforms by optimizing memory-bound matrix operations in mixed-radix algorithms:
f32andf64inneon_utils, replacing the generic scalar transpose fallback for small matrices.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
ac8f1b1—neon: vectorize 2x2 block transpose in neon_utilsarray_utils::transpose.vtrn1q_f64/vtrn2q_f64andvzip1q_f32/vzip2q_f32.transpose::transposefor other architectures and non-float types.71f2a1f—neon: fuse twiddle multiplication into transpose pass in MixedRadixSmallMixedRadixSmall::perform_fft_inplaceandperform_fft_out_of_place.c82055b—neon: accelerate matrix transpose with tiled NEON vectorizationGoodThomasAlgorithmSmallandMixedRadixSmall.Benchmark Results (Apple Silicon / AArch64 NEON,
f64)Measured on Apple Silicon (
aarch64-apple-darwin), comparing upstreamv6.4.1(commit4758ab0) against this PR branch (commitc82055b) with 100,000 iterations per size:Correctness & Verification
cargo test --lib(122 passed; 0 failed).tests/accuracy.rs(all 4 forward/inversef32/f64roundtrip accuracy tests pass with identicalcargo test --doc(17 passed; 0 failed).#[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:
Benchmark Source Code (
examples/bench_mixed_radix.rs)