Add a SIMD RadixN for NEON, SSE and wasm_simd by HEnquist · Pull Request #179 · ejmahler/RustFFT

HEnquist · 2026-09-09T07:30:11Z

Adds a vectorised RadixN to the NEON, SSE and wasm_simd backends, written once against a shared
SimdVector trait. No planner changes: wiring it up lands in a separate PR so the plan changes
can be measured on their own.

src/simd/simd_radixn.rs: shared SimdRadixN<V, T>, generic over f32 and f64, base may be a
composite recipe that needs scratch. The three *_radixn.rs files are thin type aliases plus
tests
src/simd/simd_vector.rs: the SimdVector trait the algorithm is written against, implemented
once per vector type next to each backend's own vector trait impls
prime butterfly structs and constructors made pub for sibling modules, via the shared
generator template
split_cross_len and get_power_of factored into math_utils.rs, so plan.rs shares them
RadixN is only reachable from its own tests until the planner PR, hence the allow(dead_code)

The scalar planner spells out two pieces of arithmetic around its own base-choice rules, and the SIMD planners are about to need the same two. Both come out into building blocks that carry no policy: PrimeFactors::get_power_of(value) replaces the hand-written find_map over get_other_factors(). get_power_of_two and get_power_of_three already existed, this covers 5, 7 and anything later. RadixFactor::split_cross_len(len) is the 7/6/5/3-then-4s-last split of a cross-FFT length. It returns None where the scalar planner asserted, and plan.rs keeps the panic by calling .expect() on it. What is left in the planner is only its own decisions: which base, and whether Radix4 takes the length. Nothing about the algorithm set is shared, so a planner can use these without being tied to any other planner's choices. Recipes are unchanged: FftPlannerScalar dumped for every length from 1 to 20000, f32 and f64, all 40000 identical before and after.

The RadixN drivers about to be added reuse the size-7 butterfly for their radix-7 cross-FFT layer, so the generated structs and their constructors have to be visible outside their own module. Changes the shared template and regenerates all three backends, so the autogeneration check keeps passing. The neon, sse and wasm_simd modules are themselves private, so this exports nothing new from the crate.

Mirrors src/algorithm/radixn.rs: one flat transpose down to a base FFT, then in-place cross-FFT layers over a single packed twiddle array. The difference is that the layers use column butterflies, so a whole vector of columns goes through each butterfly call. The algorithm lives in src/simd_radixn.rs as SimdRadixN<V, T>, generic over a RadixNVector trait, so the other two SIMD backends can reuse it by supplying two trait impls. NEON is the first, at 224 lines of impl plus a type alias. - generic over f32 and f64. Twiddles are stored as vectors, and the trait pairs each vector type with its radix 3, 5, 6 and 7 butterflies. Radix 2 and 4 were already vector-generic. - the base may be a composite recipe with its own scratch, which the existing boilerplate macros hardcode to zero, so Fft is implemented on SimdRadixN directly instead. - an f32 vector holds two complex numbers, so every cross-FFT layer needs an even column count. An odd base folds in a spare factor of two, and odd lengths, which have none to spare, keep the mixed radix path. src/simd_planner.rs holds the planning arithmetic, again so the other two backends get it unchanged: design_radixn, design_butterfly_product, and complex_per_vector. Each planner owns a private Recipe enum, so these hand back plain numbers and the caller builds its own recipe. The planner's dispatch chain ends up in the same order src/plan.rs uses, with the butterfly pair search ahead of RadixN. That moves the pair search out of the final else, so it now sees lengths with trailing_zeros() >= 6 that it never used to. Two of them change plan, and both get faster. NEON, forward, 10*len buffer, M1, ns/iter, mean of two runs: len dtype old new speedup 320 f32 8730 7444 1.17x 320 f64 12235 10806 1.13x 576 f32 15690 13124 1.20x 576 f64 22265 19226 1.16x Recipes audited over every length from 1 to 20000 for f32 and f64: 26323 change, 26290 of them by gaining a RadixN, and all 33 of the rest are those same two plans propagating through nested designs. Nothing else moves. Measured on an M1, f64, at 1008, 1050, 1080, 1296, 10368 and 100800: 1.46x to 2.03x faster than the previous planner, and 1.13x to 1.53x faster than the best MixedRadix tree this backend could build before.

The shared layer in simd_radixn.rs and simd_planner.rs landed in a shape SSE can use unchanged, so this is additive apart from the two cfg gates in lib.rs, which grow an x86_64 arm. SSE is the most mechanical of the ports. SseVector already mirrors NeonVector method for method, the SseArray and SseArrayMut load and store traits match, and the butterfly structs for radix 3, 5, 6 and 7 have the same perform_fft_direct and perform_parallel_fft_direct shapes, so the two RadixNVector impls are the NEON ones with the names changed. sse_radixn.rs comes out at the same 224 lines as neon_radixn.rs. The planner gets the same treatment: a Recipe::RadixN variant, thin wrappers over simd_planner::design_radixn and design_butterfly_product, and the same dispatch chain order. Recipes audited over every length from 1 to 20000 for f32 and f64, and the change is exactly the one NEON saw: 26290 lengths gain a RadixN, and all 33 of the remaining differences are the butterfly-pair reorder at 320 and 576 propagating through nested designs. The resulting designs are identical to NEON's at all 40000, which is the check that the shared layer really is shared. Both reorder lengths measured faster on a Ryzen 7 250, forward, 10*len buffer, ns/iter, median of five runs: len dtype old new speedup 320 f32 9729 5642 1.72x 320 f64 11225 8123 1.38x 576 f32 18582 10532 1.76x 576 f64 31156 15097 2.06x That is a good deal more than the 1.13x to 1.20x the same plan change gave on an M1, so the size of the win is microarchitecture specific. Length 512, whose plan does not change, measures the same in both trees to within 0.5%, which rules out a build difference behind these numbers. A later run on the same machine put 576 f64 at 0.78x rather than 2.06x, so that one number is not settled and is being re-measured; the other three have been stable across runs.

The third and last SIMD backend. Same algorithm and the same planner branch as the other two, which the backends being structurally identical makes largely mechanical. Two things differ. The butterflies for radix 3, 5 and 6 are written against raw v128, while WasmVector32 and WasmVector64 are newtypes over it, so the column butterflies unwrap and rewrap around them; butterfly 7 already speaks the wrapper types. And wasm_simd_planner globs its own module, so mod.rs re-exports the new one. The planner tests copied from sse also get their names fixed to test_plan_wasm_simd_* and test_wasm_simd_*, matching what the other two call the same test bodies. All five RadixN tests are #[wasm_bindgen_test], not #[test]. The wasm-bindgen harness only collects the former, so a plain #[test] here is silently dropped on the one backend where these are slowest. wasm-pack test --node lists 73 passing and 1 ignored, the ignored one being the six-layer case. Recipes audited over every length from 1 to 20000 for f32 and f64, under node via wasm32-wasip1: the same 26290 lengths gain a RadixN and the same 33 are the butterfly-pair reorder at 320 and 576, with nothing else moving. The designs are identical to NEON's and SSE's at all 40000, so all three backends now plan these lengths the same way.

ejmahler · 2026-09-10T06:01:13Z

Hell yeah. I'm currently experiencing a burst of energy on working on FFT stuff. I'm writing a multidimensional FFT implementation at the moment, and I was going to get around to this afterwards. Thank you very much for tackling it! I'm busy this weekend but I will review it next week.

The main reason why I never made RadixN public initially is because I wasn't quite happy with the design. In particular, having to branch inside the loop seems to hurt performance, and as a result RadixN with all 4s is generally slower than Radix4. One thing I wanted to try was monomorphizing the various factor computation, both at the bit reversal step and the twiddle factor step. So for example, in the computation step instead of matching on the radix inside the loop, it would directly store a power2 value, power3 value, power4 value etc in the RadixN struct, and the FFT execution would be something like

for i in 0..self.power2 {
    // execute power2 steps
}
for i in 0..self.power3 {
    // execute power3 steps
}
for i in 0..self.power4 {
    // execute power4 steps
}
// more factors

And then just hardcoding the order they happen in based on benchmarking. I don't know if this would be faster, but if it was then i would want to change the constructor of RadixN to just take power2, power3, power4 values directly instead of letting callers pass an arbitrary array.

That doesn't need to happen in this PR, but it's probably what I will try next before release - if it makes it faster then i'll go with it. I wouldn't be surprised if it doesn't, RadixN is flat out doing more work than Radix4.

Move the factor dispatch outside the chunk loop, the way algorithm/radixn.rs already does it, so each layer runs one monomorphized loop over its chunks. Measured perf-neutral on NEON, this is for consistency.

HEnquist · 2026-09-10T20:49:17Z

I wanted to know how far behind Radix4 RadixN actually is, and was surprised at how close it is.
Radix4 vs an all-4s RadixN, same base and length, so the only difference is the generic layer loop:

len	NEON f64	NEON f32	scalar f64	scalar f32
1024	1.036x	1.031x	1.024x	1.023x
4096	1.028x	1.018x	1.016x	1.023x
16384	1.003x	0.995x	0.997x	1.010x
65536	1.008x	0.994x	0.999x	1.007x

Only 2-3% behind at the small sizes, and level from 16384 up.

I also tried moving the factor dispatch out of the chunk loop in the SIMD RadixN: 0.988x to 1.004x
over 8 planner recipes, so no win. Kept it anyway (9ce0980) to match the scalar file.

M1 only, and I did not test the reordering half.

HEnquist · 2026-09-10T21:16:58Z

Also tried the reordering half. split_cross_len already emits the factors grouped and in a fixed
order (7s, 6s, 5s, 3s, then a 2, then the 4s), so that part is already the status quo. Storing
power2/power3/power4 counts would be the same plan in a different shape.

Swept every permutation of the factor multiset on NEON at 4032, 10368, 6300 and 100800, f32 and
f64. Best to worst spread is 0.4% to 6.4%, but the current order is within 1.1% of the best in all
8 cases. So about 1% left, and only with a per-length search.

HEnquist · 2026-09-15T09:09:34Z

Heads up on something related. I've been working on a new estimating planner, and it's going a lot better than my old attempt in #38.

Different approach: it counts instructions by reading the source instead of measuring a per-machine table, plus a coarse memory term with a few fitted weights. Simpler and much more robust, with nothing big to re-measure per machine.

Worst-case regret against the best recipe I can enumerate, over 33 mixed-factor lengths:

	fixed planner	cost model
NEON f64	1.495	1.043
NEON f32	1.969	1.121
SSE f64	1.734	1.152
SSE f32	1.927	1.121

All four hold up on a held-out half. The lengths are deliberately adversarial, powers of two come out 1.000 for everyone.

Branch is at https://github.com/HEnquist/RustFFT/tree/counted_cost_spike, sitting on this PR's commits. It changes no planner behaviour yet. Full writeup in RESULTS.md.

If a smaller scoped version is more interesting, the clearest win is replacing MAX_RADER_PRIME_FACTOR with a cost comparison. No threshold on lpf(len-1) can be right: 1013 wants Bluestein's, 9661 wants Rader's, and both have lpf = 23.

ejmahler · 2026-09-16T05:52:54Z

That's great to hear! I'm curious to hear more about what you mean by reading the source.

Some regression is definitely ok - it's not like the current algorithm is optimal, it's just sitting in a local maximum of heuristics, and like you pointed out with the prime factors, there are some things it simply cannot do well. It would be interesting to see some statistical analysis of sizes that got better or worse. ie for the first million sizes, what's the 10th, 25th, 50th, 75th, 90th percentile change in run time? Although that sounds like it would take a very long time to measure, so maybe a random sampling in that range?

ejmahler · 2026-09-16T05:54:48Z

Ah, I see op-counts.md. That makes sense, and the fact that every operation is an explicit call must make it easier since you don't have to worry about hidden operator impls etc.

ejmahler · 2026-09-16T07:23:02Z

+                |chunk, scratch| {
+                    let (self_scratch, inner_scratch) = scratch.split_at_mut(self.len());
+                    self.perform_fft_out_of_place(chunk, self_scratch, inner_scratch);
+                    chunk.copy_from_slice(self_scratch);


Because the actual implementations are so simple and because we don't have to duplicate a macro to fix it, imo this would be a good time to eliminate this copy, and make a dedicated perform_fft_inplace function.

Oh, hmm, now that I think about it it's more complicated than I thought, we'd need two cross_fft functions, or for cross_fft to take impl LoadStore<T> both of which would be involved tasks. From what I benchamrked in the past, the copy isn't expensive, What do you think?

In RadixN the cross FFTs already run in place, so the copy doesn't come from cross_fft. It comes
from the transpose, which has to read and write separate buffers. Scalar RadixN does the same
copy through boilerplate_fft_oop!.

I measured what it costs on M1, in-place vs out-of-place on the same plan, lengths 1008 to 100800:

f32 f64

in-place / out-of-place 1.018x to 1.054x 1.046x to 1.080x

So not free, but not big. Getting rid of it would need an in-place transpose, and I'd rather try
that as a separate experiment.

I was imagining making a version of cross_fft that runs out-of-place. Seems much easier to implement than an in-place transpose.

That works nicely, thanks. Only the last layer needs it: a layer reads all of a column's rows before it writes any of them, and writes them at the indices it read them from, so it doesn't care that source and destination are different buffers. cross_layer is now generic over a LayerBuffer, InPlace or OutOfPlace, so there's still one copy of the loop and one of the radix dispatch.

NEON on M1, min of two runs, 128 to 100800, f32 and f64: in place 1.03x to 1.09x faster, out of place 0.99x to 1.02x, so the refactor costs nothing on the paths that were already out of place. In-place now edges past out-of-place at the small sizes, since it touches one buffer instead of two.

ejmahler · 2026-09-16T07:57:17Z

Done reviewing. This is a great change that I think could be a template for future changes we make.

Examples:

At minimum, I think converting the various simd radix4 implementations to this shared code style would be a slam dunk, although that should be a separate PR.
I also was holding off on implementing bluestein's algorithm for SIMD because I didn't want to deal with several copies of it floating around, but this architecture would allow for a single implementation, which would be much more palatable.
We could honestly go as far as to port the scalar versions of these algorithms to use these shared-code versions, by creating a scalar adapter to SimdVector, where SimdVector for f64 is just [Complex<f64>; 1] and SimdVector for f32 is just [Complex<f32>;2]. I'm still turning this one over in my head, maybe it's not good idea, but it's worth evaluating at least.

ejmahler · 2026-09-16T09:25:04Z

After thinking more, the last bullet point has a stumbling block because the scalar algorithms support things that aren't f32 and f64. I think we can still make it work though, it just won't be as clean of a mapping as i was hoping.

factor_transpose recomputes every column's reversed index on each call, with an out-of-line reverse_remainders call per column and two hardware divides, and chunks_exact_mut adds one more divide per cross layer. None of that scales with the length, so it is a large share of a short FFT, and more so on x86 where a 64-bit divide takes tens of cycles. Compute the reversed columns once in new() and walk the layer chunks with split_at_mut. The per-element work is unchanged. factor_transpose itself is left alone, since the scalar RadixN still uses it.

…imdVector The trait gets its own simd_vector.rs so other algorithms can be written against it. Each backend's impls move next to its own vector trait impls in *_vector.rs, together with the fft_helper forwarding macro, leaving *_radixn.rs with just the type alias and tests.

It is factoring arithmetic, so it belongs with PrimeFactors rather than on RadixFactor in common.rs.

… up front The perform methods were one line forwards from the Fft closures. The column loop now computes how many pairs and whether a column is left over before it starts, instead of testing vcol + 2 <= num_vector_columns.

The comment said from_fn was newer than the MSRV, but it has been stable since 1.63 and the MSRV is 1.77. Perf-neutral on NEON: 0.995x to 1.002x over 10 lengths from 120 to 100800, f32 and f64.

HEnquist · 2026-09-16T20:23:33Z

Done reviewing. This is a great change that I think could be a template for future changes we make.

Thanks! Fixes for the review comments are pushed, and I replied inline where there was something to
discuss.

Radix4 on the shared code looks like the natural next step, I can take that as a separate PR once
this one lands. A single SIMD Bluestein's after that would be nice too.

HEnquist · 2026-09-16T20:23:34Z

for the first million sizes, what's the 10th, 25th, 50th, 75th, 90th percentile change in run
time? [...] maybe a random sampling in that range?

A random sample over that range is doable. The cost model branch doesn't switch any planner
decisions yet, so the percentiles will come with the planner PR once it's ready. I'll try to get
it into shape for a draft PR asap.

ejmahler · 2026-09-20T22:48:40Z

I was reminding myself what didn't work about my attempted SimdVector trait in the past. It was specifically the SimdNum trait and its VectorType associated type. You can't constrain it to be SimdVector in the parent trait and then specialize it to be NeonVector in the child type.

But thinking through that more, if we had SimdVector and the various other platform vectors inheriting from that, what wouldn't we put in the parent trait? It sure seems like the answer is nothing. So now what I'm thinking is - what if we had a single SimdVector trait and eliminated the platform vector traits altogether? We could consolidate the various SseArray etc types as well.

In the past, I wanted to keep them separate because different platforms might have different needs for what was fastest etc but the longer this goes on the less likely it seems that we would ever want platform-specific differences like that.

It might not like the fact that we're creating multiple implementations of a hypothetical SimdNum trait where nothing differs but the VectorType, but I'm wondering if we can get away with just eliminating the SseNum etc traits and not even having a SimdNum trait.

ejmahler · 2026-09-20T22:53:49Z

As for sequencing, I'd like to get this merged without integration into the planner yet, and radix4 merging + integration into the planner can be separate PRs. Merging the platform traits can come after those two are ready. Does that sound good to you?

The planner side lands in its own PR, so the plan changes can be measured on their own. The three SIMD planners go back to master exactly, and simd_planner.rs goes with them. RadixN is then only reachable from its own tests, hence the allow(dead_code) on src/simd. All three backends build warning free and the tests are unchanged.

HEnquist · 2026-09-21T13:53:04Z

Makes sense with the planner work coming, so the planner integration is out: the three SIMD planners are back to master exactly and simd_planner.rs goes with them, leaving just the algorithm and the shared SimdVector code.

The in-place path ended with a full length copy from the scratch buffer back into the caller's chunk. A cross layer reads all of a column's rows before it writes any of them, and writes them at the indices it read them from, so the last layer can write straight into the chunk instead. cross_layer is now generic over a LayerBuffer, either InPlace or OutOfPlace, so there is still one copy of the loop and one copy of the per-radix dispatch. NEON, M1, min of two runs, at 128, 1008, 1080, 1296, 10368 and 100800, f32 and f64: in place 1.033x to 1.090x faster, out of place 0.992x to 1.021x, so the refactor costs nothing on the paths that were already out of place.

HEnquist · 2026-09-21T14:22:26Z

what wouldn't we put in the parent trait? It sure seems like the answer is nothing.

Checked it. With the prefixes normalised away, NeonVector and SseVector are byte-identical, 18
methods each. WasmVector differs by one line, fn unwrap(self) -> v128;. Plus three identical
copies of Rotation90<V>. So nothing to keep apart.

Dropping the *Num traits looks doable too. A single SimdNum can't exist, since f32 would need
VectorType to be float32x4_t, __m128 and WasmVector32 at once. But VectorType has only two
consumers: the Radix4 structs, and the NeonArray/SseArray traits. Both go away in the direction
you already want, Radix4 as SimdRadix4<V, T> and the array traits parameterised by the vector
instead of the scalar. The butterflies never touch it.

Those two halves come apart though: collapsing the vector traits only needs the *Num bounds
changed, so it isn't blocked on Radix4, while dropping the *Num and *Array traits is. Either way
it wants its own PR, the rename alone hits ~3600 generated lines per backend.

ejmahler · 2026-09-22T01:41:00Z

I have one more idea for the inplace vs out of place thing: I think it would result in less new code if, in the in-place path, we computed the base FFT out of place instead of in place. It would accomplish the same goal, while allowing us to avoid the new abstraction of the LayerWalk etc. I have no idea if it would be faster or not.

Other than that, this is still marked as a draft - do you think it's ready? I don't see any reason not to merge it once the conflicts are resolved.

HEnquist added 5 commits September 9, 2026 08:03

Match the scalar RadixN dispatch shape

9ce0980

Move the factor dispatch outside the chunk loop, the way algorithm/radixn.rs already does it, so each layer runs one monomorphized loop over its chunks. Measured perf-neutral on NEON, this is for consistency.