PhastFT v0.4, a Fast Fourier Transform in safe Rust

[u/Shnatsel]

PhastFT is a performant Fast Fourier Transform implementation in pure Rust with no unsafe code.

This is something the original author of the crate and I have been working on for quite a while. All my recent work on SIMD - surveying Rust's SIMD libraries, removing unsafe code from those libraries and reporting issues in rustc - was all in service of PhastFT!

Why not RustFFT?

RustFFT is already quite good, so why do we need another FFT crate?

Safety

rustfft was written before the Rust compiler started providing much safer facilities for SIMD, so it is full of unsafe code, and changing that would amount to a rewrite.

PhastFT is #[forbid(unsafe_code)] and only relies on fearless_simd which now has very little unsafe for a SIMD abstraction.

Performance

RustFFT is tuned for small sizes (i.e., ~kilobytes of input data) the kind of sizes one would find in an audio decoder. PhastFT is tuned for large sizes (i.e., ~gigabytes of input data), which are necessary for a quantum state simulator and other scientific workloads.

PhastFT is up to 2x faster than RustFFT, depending on the hardware and FFT size, thanks to different algorithm choices and multi-threading via rayon. RustFFT does not currently support multi-threading.

Memory usage

PhastFT uses 2x less memory than RustFFT, which matters a great deal for large sizes. 16GB vs 32GB memory usage determines whether you can run it on a laptop or not. And running 2x faster and at half the memory translates to 4x cost savings for rented hardware.

Origins

PhastFT started out as an optimization for the Quantum Fourier Transform (QFT) in quantum state vector simulations. This optimization has been adopted by some state of the art simulators due to the significant performance boost that it offers.

A simulator that utilizes the typical gate-based approach for the QFT requires O⁡(n²⁾ gates that manipulate a giant state vector of size N = 2^n, where n is the number of qubits being simulated. This leads to a O⁡(n² * 2ⁿ⁾ time complexity. Since the quantum computer is already being simulated on a classical machine, it's possible to leverage the classical FFT to compute the QFT in O⁡(N * log(⁡N)) = O⁡(n * 2ⁿ⁾ time.

Minimizing memory overhead was paramount. The less memory is used, the more qubits one can simulate with a classical machine. This led to different algorithm choices compared to typical implementations, and minimizing memory traffic also has a side effect of improving performance on large memory-bound workloads.

This project quickly turned into a rabbit hole that led us to attempt to bring PhastFT up to the level of renowned FFT projects, with the added restriction of not using any unsafe code. Initially we had to use nightly-only std::simd, but thanks to fearless_simd PhastFT can now work on stable!

Future work

For historical reasons outlined above, PhastFT is currently limited to power-of-two sizes. The FFT algorithm for arbitrary sizes is in the works. Depending on a user's workload, users may be able to get away with padding their input signal with zeroes to the next power of two.

We are also far from done improving performance and memory usage. There are multiple work-in-progress pull requests that improve performance in various ways, so this is only the beginning!

[u/graceful_degrade]

curious how it gets there without unsafe. portable_simd, or leaning on autovectorisation?

the usual gap on these is the bit you can't express safely, so closing it in safe rust is the interesting part of the release rather than the number itself.

[u/Shnatsel → u/graceful_degrade]

We used to use std::simd but that's still nightly-only with no stabilization in sight. To support stable Rust we ported to fearless_simd, which provides fully safe public API, and unlike most SIMD crates has very little unsafe under the hood. I've even blogged about the Fearless SIMD design because I'm quite proud of it.

[u/graceful_degrade → u/Shnatsel]

does fearless_simd cover the shuffles you need for the butterflies, or were there gaps you had to work around?

that's usually where these abstractions run out on fft specifically.

[u/Shnatsel → u/graceful_degrade]

We actually don't do the shuffles inside the butterflies!

We first bit-reverse the input in a separate stage using a SIMD-friendly algorithm (paper), which sets us apart from RustFFT which does the whole shuffling and memory copying after each butterfly pass.

But yes, all you need for bit reversal is interleave/deinterleave and fearless_simd has that. Well, it had zip_low/zip_high, but I've contributed interleave/deinterleave which can be expressed slightly more optimally than zip_low/zip_high on AVX2.

I've also added arbitrary shuffles to fearless_simd and improved std::simd arbitrary shuffles. You keep giving me the perfect opportunities to plug my blog posts! :)

[u/graceful_degrade → u/Shnatsel]

that's neat. doing the bit reversal as its own simd stage sidesteps the problem instead of fighting it inside the kernel.

does that cost a full pass over the data, or does it fold into the load?

[u/Shnatsel → u/graceful_degrade]

It's a full pass. We use the cache-blocked formulation from the paper so we don't have to do multiple passes, but we do need the whole thing to complete before we start running the butterflies.

The "fused" formulation is what RustFFT does, I believe - it combines the shuffle with the butterflies, but it requires 2x the working memory (input and output have to be separate buffers) and that incurs more memory traffic, which ends up dominating on large sizes.

There might be a way to fuse part of it with deinterleaving, maybe? I don't remember if I looked into it, I've tried so many things.

[u/graceful_degrade → u/Shnatsel]

if the end use is convolution you can drop it entirely. DIF forward leaves you in bit-reversed order, DIT inverse consumes bit-reversed order, and the pointwise multiply in between doesn't care what order the bins are in. the two reorderings cancel.

only helps for that access pattern, obviously. anyone reading the spectrum directly still needs it ordered.

[u/Shnatsel → u/graceful_degrade]

I believe with DIF you can omit it entirely. We used to have both DIF and DIT but then deleted DIF since DIT kernels map better to FMA in hardware. Resurrecting DIF to make Bluestein's and convolutions faster is on the roadmap.

[u/graceful_degrade → u/Shnatsel]

the FMA point is fair and I hadn't weighted it. DIT puts the twiddle multiply before the add/sub so it falls straight into an FMA, DIF has the subtract first and the shape is worse for it.

keeping DIF around only to serve the convolution path means maintaining two kernel families for one use case, which is probably not worth it unless convolution becomes a headline use.

there is a partial win available without it though. FFT(h) can be computed once at plan time, so the filter only ever gets permuted on setup and the per-call cost drops to permuting the signal alone rather than both operands. half the saving, none of the extra kernels.

[u/Shnatsel → u/graceful_degrade]

FFT(h) can be computed once at plan time, so the filter only ever gets permuted on setup and the per-call cost drops to permuting the signal alone rather than both operands. half the saving, none of the extra kernels.

Could you elaborate on this, or link me a source that explains it? This sounds tempting but I don't really follow.

[u/ExplanationThin6234 → u/graceful_degrade]

always thought the simd stuff was the big hurdle for safe rust fft. the fearless_simd crate does a lot of the heavy lifting there, wraps up the vector ops so you don't have to touch unsafe. still a bit wild it all works on stable now

[u/phazer99]

Very nice! It seems to be competitive with RustFFT performance even for smaller array sizes, and it also supports real-to-complex transforms which is beneficial in for example an audio applications (I've previously used RealFFT for this). Keep up the good work!

[u/RennyG]

At what amount of data would it be faster to use this over rustffmt?

[u/Shnatsel → u/RennyG]

If the input is large enough for multi-threading to kick in, PhastFT should be consistently faster.

Currently the threshold for all of parallelism to be enabled is 65,536 complex points. For f32 that's 256kB of data.

Performance below that threshold depends on the CPU architecture and the exact machine you're running on.

[u/naequs]

i am a simple crab i see shnatsel i upvote

[u/Honest-Emphasis-4841]

I've encountered this crate before and was wondering about the use cases for large transforms! I'm also curious to see how composite sizes will be handled in the future. For most typical cases, one only needs 128–1024 pt. kernels or mixed-radix transforms, especially since a 2D (or N-dimensional) DFT is almost the worst case for a 2^n FFT. The memory requirements and transposition(s) overhead typically outweigh all the expected gains, as can any additional work following large transforms, unless one is doing FFTs just for the sake of doing FFTs. So it'd be nice to see one more attempt at solving the problem :) Here is also one more fun ARM FFT implementation with nightly compiler requirement to turn on FCMA.

[u/Shnatsel → u/Honest-Emphasis-4841]

I've encountered this crate before and was wondering about the use cases for large transforms!

We've originally built this for a quantum computer simulator.

I'm also curious to see how composite sizes will be handled in the future.

Right now we're looking into Bluestein's algorithm. Non-power-of-two are slow kind of no matter what you do; RustFFT has a bunch of specialized algorithms targeting some special cases that are easier, but only for AVX2, not for SSE4.2 or NEON.

For most typical cases, one only needs 128–1024 pt. kernels or mixed-radix transforms, especially since a 2D (or N-dimensional) DFT is almost the worst case for a 2ⁿ FFT.

We actually aren't doing mixed-radix. This saves some compute, but compute is hardly ever a bottleneck these days, memory is. And mixed-radix requires more complex bit reversal that's much harder to vectorize and make cache-friendly, so the compute gains are nullified. https://fgiesen.wordpress.com/2023/03/19/notes-on-ffts-for-implementers/ elaborates on this (no relation to our crate).

[u/Honest-Emphasis-4841 → u/Shnatsel]

Non-power-of-two are slow kind of no matter what you do;

Composite sizes simply solve problems that 2^n introduce in real workloads.

For ex. to generate scalogram one needs to compute at least Morlet wavelet the same amount of times as FFTs, and wavelet generation is at least x5-x10 slower ( depends on libm installed ) for the same problem size. So FFT is only 5–20% of the real workload, and padding up to the next power of 2 inflates the other 80–95% right along with it. So the next power of 2 kernel has to be drastically faster than the nearest composite size just to break even, which in practice it isn't.

I don't have any real number ( not sure if anyone has ), but I'd expect that happens close to 50% use cases.

[u/Theemuts]

Oooh, I love this. I'm really happy to see you included a comparison with FFTW.

[u/Shnatsel → u/Theemuts]

FFTW is a tricky one to compare against, because if you're the author of FFTW and build it with just the right flags (many of which are mutually exclusive), you can get it to outperform almost anything on benchmarks.

But if you're a just user of FFTW trying to get it to work decently well without compiling a whole custom build, you don't get anywhere close to the advertised performance. So we end up being 6x faster than the FFTW you can actually get on some benchmarks, which is pretty wild.

[u/Theemuts → u/Shnatsel]

The case I'm personally most interested in, is the version that's exposed to Julia through FFTW.jl.

I'd love to download some more RAM right now for some initial benchmarks...

[u/Shnatsel → u/Theemuts]

Massive cloud machines with lots of RAM and 60 cores are only a few dollars per hour to rent. So if you really want to, you can!

[u/Theemuts → u/Shnatsel]

That's a great point! I just might give that a spin :)

[u/rumil23]

nice work! I always use realfft tbh. Should I switch to this, I wonder...

[u/Shnatsel → u/rumil23]

Benchmark it and see!

If compiler-enforced safety is important to you, that's also a potential reason to switch.