Add a SIMD RadixN for NEON, SSE and wasm_simd - #179
Conversation
|
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 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. |
|
I wanted to know how far behind Radix4 RadixN actually is, and was surprised at how close it is.
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 M1 only, and I did not test the reordering half. |
|
Also tried the reordering half. Swept every permutation of the factor multiset on NEON at 4032, 10368, 6300 and 100800, f32 and |
|
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:
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 |
|
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? |
|
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. |
| |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); |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
I was imagining making a version of cross_fft that runs out-of-place. Seems much easier to implement than an in-place transpose.
There was a problem hiding this comment.
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.
|
Done reviewing. This is a great change that I think could be a template for future changes we make. Examples:
|
|
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. |
Thanks! Fixes for the review comments are pushed, and I replied inline where there was something to Radix4 on the shared code looks like the natural next step, I can take that as a separate PR once |
A random sample over that range is doable. The cost model branch doesn't switch any planner |
|
I was reminding myself what didn't work about my attempted 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. |
|
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? |
HEnquist
changed the title
|
Makes sense with the planner work coming, so the planner integration is out: the three SIMD planners are back to master exactly and |
Checked it. With the prefixes normalised away, Dropping the Those two halves come apart though: collapsing the vector traits only needs the |
|
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. |
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.
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.
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.
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.
The in-place path transposed into scratch, ran the base FFT and the cross layers there, and ended with a full length copy back into the caller's chunk. Transpose into scratch and run the base FFT out of place from there back into the chunk instead. The cross layers then run in place on the chunk, and the result ends up where the caller wants it with no final copy. NEON on M1, old and new both built into one binary and timed interleaved, min of 15, two runs: the new path is 0.93x to 0.95x at 1008, 1080, 1296, 10368 and 100800, 0.89x to 0.91x at 64 and 128, and 0.81x to 0.84x with no factors at all, where the old path did two full length copies and the new one does one. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
f2c102a to
4d05564
Compare
|
Good call, that's better. The in-place path now transposes into scratch and runs the base out of place from there back into the caller's buffer, so every cross layer is in place again and the result still ends up where the caller wants it with no final copy. That drops the whole abstraction: no It's faster too. Old and new built into one binary and timed interleaved, min of 15, two runs, NEON on M1. Ratio is new over old, so below 1 is better:
That's against the state before the One cost: Rebased on master. The only conflict was the prime butterfly template, where you renamed And yes, it's ready. Marking it so. |
| for y in 0..height { | ||
| let row = x + y * width; | ||
| for (i, &r) in rev.iter().enumerate() { | ||
| unsafe { |
There was a problem hiding this comment.
Something we did for Radix4 and the scalar RadixN is add some analysis in the comments for why this is unsafe block is safe.
Since this is one of the "crunchiest" uses of unsafe we do where you can't just intuit that it's safe from the code, I think we should do the same here.
One thing that should help is, the comment block for the function says that D must divide the width, and I think we should add that to the assert on line 415. The chunks_exact on line 417 already does a divide by D, and since it's a compile time constant it should be easily reduced.
There was a problem hiding this comment.
For output.get_unchecked_mut(y + r * height): The value of r comes from the reversed_columns slice, and in the constructor of SimdRadixN all items in reversed_columns are asserted to be less than width, therefore the maximum value of r is width - 1. We can also see from the loop parameters that the maximum value of y is height - 1. Substituting these values into the expression y + r * height and simplifying, we get an overall maximum value of width * height - 1. We know that width * height == len, so the index will always be in range.
There was a problem hiding this comment.
For input.get_unchecked(row + i):
rowis defined asx + y * width, andxis defined asgroup * D, sorowisgroup * D + y * width.- The maximum value of
groupisreversed_columns.len() / D - 1. We know from the constructor thatreversed_columns.len() == width. Substituting that in, we get(width / D - 1) * D + y * width, and since we know D divides width, we can simplify this towidth + y * width - D - We can see from the loop condition that the maximum value of y is
height - 1. Substituting this in, we getwidth + (height - 1) * width - D, simplifying towidth * height - D - Finally, we can see from the loop parameters that the maximum value of
iisD - 1. Thus, the overall maximum value of this index iswidth * height - D+D - 1, simplifying towidth * height - 1 - We know that width * height == len, so the index will always be in range.
There was a problem hiding this comment.
IMO the way to do this would be to separate it into two unsafe blocks, and have those comments I wrote above, or equivalent, outside each one. With that done, the most-unsafe operation in this file will hopefully be clear to anyone who wants to audit it.
There was a problem hiding this comment.
Done. Two unsafe blocks now, with your analysis above each, and width % D == 0 added to the
assert.
The new assert can't fire: D is the first transpose factor and width is the product of all of
them, so it divides by construction. It's there for the reader and for the codegen.
|
I went over it one more time just to be as thorough as possible, and only found two things: The first is the comment I left above about the unsafe block, and the other is that now that FCMA is merged in, the FCMA path needs a radixn implementation. Thanks so much for pushing all of this through, once those two things are resolved I'll merge. |
Split the read and the write into separate unsafe blocks, each with the bound it relies on, and assert that D divides the width so the chunks_exact divide folds away as well.
The FCMA vector types are the plain NEON ones, and the NEON backend already implements SimdVector for those, so FCMA gets newtype wrappers to hang its own impls on. They only exist at the SimdVector boundary. Same recipes, scalar butterfly base, NEON on M1, fcma over neon: 1024 0.67 f64 0.80 f32 4096 0.66 f64 0.73 f32 1008 0.74 f64 0.79 f32 2520 0.78 f64 0.81 f32 12288 0.78 f64 0.79 f32
|
Both done. The FCMA one had a snag worth mentioning: FCMA's vector types are the NEON ones, and the NEON Worth it. Same recipes and the same scalar butterfly base for both, so the only difference is the
In line with what the rest of the fcma backend gets over neon. If the wrappers bother you, the alternative is parameterising the NEON |
|
CI is red because |
|
We already do the same kind of wrapper for wasm simd, I think that's totally fine. |
|
Done! Thanks so much for taking this on and iterating on all my feedback etc. I think the SimdVector concept will be one of the biggest leaps forward for this project in a long time. I want to start on SimdRadix4 right away. |
|
There a few a few other things we've talked about:
Are you planning on doing any of these, besides the estimating planner? I'm happy to do any or all of them, I just want to make sure we don't duplicate work. |
|
Thanks for merging, and for the review rounds, they made it a lot better. Glad the One assumption up front: nothing releases before all of this lands. If that's wrong the ordering Since you want to start on SimdRadix4 right away, that and the trait consolidation are yours. I'd The planner is the one I really want to land, but last. It reads op counts out of the algorithm That puts the prime half of #183 before it, since prime factors need no planner but do change
Then 7 unlocks the composite half of #183, since picking On the non-even bases in #183, that's more than a local augment: for f32 an odd base means an odd Happy to open issues for each with the dependencies noted. |
|
That breakdown works for me. Heads up, I've noticed a pretty big performance regression (10x+ computation time) on wasm simd for powers of 2, with the new simd radxin vs the old wasm simd radix4. I merged #185 which fixes at least one cause (bad code gen in the butterflies), and I found some cases in RadixN where things don't get inlined (for example the gather lambda in My WIP branch https://github.com/ejmahler/RustFFT/tree/simd-radix-4 has some benchmarks if you want to take a look, run I've reproduced it both on my windows 86_64 PC and on my M1 laptop. At this point I'm out of ideas for what could be causing it. My new simd radix4 implementation seems to have inherited the bad performance, because I wrote it by copying SimdRadixN and monomorphizing it for all-4 factors. |
2 participants
Adds a vectorised RadixN to the NEON, SSE and wasm_simd backends, written once against a shared
SimdVectortrait. No planner changes: wiring it up lands in a separate PR so the plan changescan be measured on their own.
src/simd/simd_radixn.rs: sharedSimdRadixN<V, T>, generic over f32 and f64, base may be acomposite recipe that needs scratch. The three
*_radixn.rsfiles are thin type aliases plustests
src/simd/simd_vector.rs: theSimdVectortrait the algorithm is written against, implementedonce per vector type next to each backend's own vector trait impls
pubfor sibling modules, via the sharedgenerator template
split_cross_lenandget_power_offactored intomath_utils.rs, soplan.rsshares themallow(dead_code)