Support for multiple FFT backends via backend tokens/types

[nHackel]

We recently stumbled over the same issues as #32 with AbstractNFFTs. There we tried a backend token based solution which is quite similar to the way AbstractDifferentiation or KernelAbstractions handle support for multiple backends.

Furthermore, we've also added a concept of an active backend, which hides the token from the user and allows one to just load a backend package and call plan_(n)fft without thinking about backends. This is the approach Makie uses for backend selection. The active backend can be combined with setting ScopedValues, allowing each scope to have a "global" active backend.

A second variant without this active backend is shown in this PR and described towards the bottom of this comment.

In this PR I've ported the same solution to AbstractFFT which allows for both implicit and explicit backend selection:

julia> x = randn(ComplexF32, 10);

julia> # Active/current backend API, similar to Makie

julia> using FFTW

julia> typeof(plan_fft(x))
FFTW.cFFTWPlan{ComplexF32, -1, false, 1, Tuple{Int64}}

julia> using FFTA

julia> typeof(plan_fft(x))
FFTA.FFTAPlan_cx{ComplexF32, 1}

julia> FFTW.activate!();

julia> typeof(plan_fft(x))
FFTW.cFFTWPlan{ComplexF32, -1, false, 1, Tuple{Int64}}

julia> # Explicit backend, similar to AbstractDifferentiation or KernerlAbstractions

julia> typeof(plan_fft(FFTW.backend(), x))
FFTW.cFFTWPlan{ComplexF32, -1, false, 1, Tuple{Int64}}

julia> typeof(plan_fft(FFTA.backend(), x))
FFTA.FFTAPlan_cx{ComplexF32, 1}

julia> # Overwritting active backend with a dynamic scope

julia> with(fft_backend=>FFTA.backend() do 
    println(typeof(plan_fft(x)) # FFTA inside, while FFTW outside
end
FFTA.FFTAPlan_cx{ComplexF32, 1}

Backends need to implement the following:

module MyFFTPackage

using AbstractFFTs
struct MyFFTBackend <: AbstractFFTBackend end
backend() = MyFFTBackend()
activate!() = AbstractFFTs.set_active_backend!(MyFFTPackage)

function __init__()
    activate!()
end

# required specializations
plan_fft(::MyFFTBackend , x; kws...) = ...
plan_bfft(::MyFFTBackend , x; kws...) = ...
plan_rfft(::MyFFTBackend , x; kws...) = ...
plan_brfft(::MyFFTBackend , x, dims; kws...) = ...
end

I've also implemented the necessary API changes for FFTW, FFTA and RustFFT.

Some limitations I observed during the implementation were:

• Method ambiguities for the plan_x functions which promote to complexfloat if a backend does not restrict T in ::AbstractArray{T}. The high-level call avoids this by promoting element types before delegating to the active backend. The low-level form can still be ambiguous in theory, but the backends tested (FFTW, FFTA, RustFFT) are unaffected.
• Return type inference for plan_fft(x) changes based on the number of loaded backends, which is just inherent to supporting multiple-backends. However, plan_fft(backend, x) should still be type-stable and inferable.

We'd be happy to hear feedback on this proposal

---

The second variant does not have an active backend and instead removes all plan_x(x::AbstractArray) functions in favour of the plan_x(::AbstractFFTBackend, x::AbsttractArray) versions. It's therefore a breaking change affecting all backends and all packages using a backend.

Furthermore, this variant allows backends to define not-exported plan_x functions which provide the "old" plan_x(::AbstractArray) functionality. For this to work the backend cant directly load and reexport plan_x from AbstractFFTs which gives the following interface:

# Package
module MyPackage
using AbstractFFTs

f(...; backend::FFTBackend) = ... fft(backend, ...) ...
end

# User with explicit backend have to load AbstractFFTs as well
using AbstractFFTs, FFTA
fft(FFTA.backend(), ...)

# User without backend can just rely on a backend but must be fully qualified
using FFTA
FFTA.fft(...)

[giordano]

A major caveat is that this is thread-unsafe. Maybe a task-local storage or a scope value would be better options? They may come at a performance penalty though, not sure. Alternatively, not having a global state would be better.

[nHackel]

Good point! I also tried playing around a bit with ScopedValues. I just didn't see a way where I could preserve the "default" behaviour of plan_fft with it, since the user would always need to do:

with(fft_backend=>FFTW.backend()) do
    p = plan_fft(x)
    # ...
end

I did see ScopedSettings.jl, which would allow for the desired behaviour but I'm not sure if such an additional non-base dependency is justified.

Another option would be to have both a "global" active backend and a dynamic backend with the later being a ScopedValue. Then the "high-level" API would need to first check if the dynamic scoped value is set, if not then check the global one and if neither is set throw the no-backend-error.

I'm not sure if packages that want to actively switch between backends, should target the high-level interface. The current behaviour if multiple FFT backends are loaded is also not well-defined at the moment

[codecov]

Codecov Report

❌ Patch coverage is 70.96774% with 9 lines in your changes missing coverage. Please review.
✅ Project coverage is 93.27%. Comparing base (04a14f5) to head (a25ba9e).

Files with missing lines	Patch %	Lines
src/definitions.jl	70.96%	9 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #146      +/-   ##
==========================================
- Coverage   94.83%   93.27%   -1.56%     
==========================================
  Files           4        4              
  Lines         445      461      +16     
==========================================
+ Hits          422      430       +8     
- Misses         23       31       +8     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[devmotion]

I ran into this issue of AbstractFFTs multiple times, and IMO some way to switch between backends similar to DifferentiationInterface is very much needed at this point, as nowadays there exist alternatives to FFTW.jl and FFTW.jl is problematic for proprietary software on Apple silicon since MKL is not available.

I actually wonder: Is a global setting needed?

[giordano]

I actually wonder: Is a global setting needed?

Yeah, I feel like this could be a breaking change where the backend token is a mandatory argument and that's it, it's all local, no global state.

[nHackel]

Initially we added the global state to avoid producing breaking changes to the plan_x interface. I wasn't sure how large the community buy-in was on this pain-point of the interface, since such a breaking change would also affect the backends and packages using the backends.

I somewhat like the ability of a user to do a "drop-in" replacement of the used FFT backend if an intermediate package just depends on AbstractFFTs.

Of course that ability could also be implemented by the intermediate package with an backend keyword argument, a scoped-value or a global state in the intermediate package. And a user might chose an FFT backend which doesn't implement the transformations required by the intermediate package. (On a related note we could do traits on the backend tokens)

[nHackel]

I've tried out a variant of the active backend based on a ScopedValue, which would additionally offer the following interface:

julia> FFTW.activate!()
julia> with(fft_backend=>FFTA.backend()) do
           println(typeof(plan_fft(x)))
       end
FFTA.FFTAPlan_cx{ComplexF32, 1}

For thread-safety within a given dynamic scope, I can add a lock to the underlying BackendReference.

I'm a bit unsure about an adoption of the breaking change without such a "default" backend state, since the change would affect a lot of packages. That would force many packages to for example thread a new keyword through every call chain that touches FFTs. Packages that (want to) depend only on AbstractFFTs also couldn’t sensibly pick a default

[tknopp]

Yeah, I feel like this could be a breaking change where the backend token is a mandatory argument and that's it, it's all local, no global state.

I don't see how this would work though. A mandatory token would mean that all code that internally uses an FFT function needs to provide a keyword argument to make this configurable. I don't think this is realistic because the FFT is used in so many packages (100+?) and often only internally.

Regarding thread safety I wonder how realistic it is, that one wants to use different FFT implementations on different threads. In my experience it is more that one wants to statically change the FFT for the runtime of a program.
But if we can get thread safety using scopedValues then it seems we have everything in place we need.

[devmotion]

I've tried out a variant of the active backend based on a ScopedValue, which would additionally offer the following interface:

I'm not convinced yet that this should be done. It seems much simpler to just write plan_fft(FFTABackend(), x).

The AD ecosystem is switching to DifferentiationInterface and passing around AD backend types to all kind of functions (e.g. in Turing or SciML), so I don't immediately see why something similar couldn't be feasible for FFT backend types as well.

[tknopp]

How do you want solve this e.g. in this situation:
https://github.com/JuliaDSP/DSP.jl/blob/7fc3842c13fe6d3977bfea0426db920b3f1a7dad/src/estimation.jl#L95

FFT is here only an implementation detail. Should all functions now get an FFTbackend as karg?

[jusack]

Maybe I misunderstand the problem, but how would the simplest use-case for an end-user know which backend to use without any global state? The basic functions are implemented in AbstractFFTs, create a plan and apply it

using FFTW
fft(randn(64))

[devmotion]

FFT is here only an implementation detail. Should all functions now get an FFTbackend as karg?

If you want to make it configurable, you could make it a keyword argument. Another alternative is to hardcode a backend, e.g., if some functionality is needed that is only available in e.g. FFTW. A third option would be to make the backend configurable on the package level. However, probably I'd prefer the first two options in most cases.

Maybe I misunderstand the problem, but how would the simplest use-case for an end-user know which backend to use without any global state?

Users have to load a FFT package anyway as AbstractFFTs is not sufficient to compute FFTs (and actually all FFT packages commit type piracy, and depending on the order in which you load FFT packages one or the other backend might overwrite the other). So they (or the package authors) already have to decide which backend to use anyways, already nowadays. If there's no global state, you'd just pass the backend to the fft call as well, e.g.,

using FFTW
fft(FFTWBackend(), randn(64))

[jusack]

I would strongly argue that from a user point of view it should be possible to just do fft(x) after using FFTW, without passing an explicit backend, because by loading FFTW I already chose a package that supplies me with an FFT (and probably do not even know the details of AbstractFFTs).

I totally agree that loading another package that decided to depend on e.g. FFTA should not change the implementation of the fft in my code, but that is the current problem that this PR would solve, right?

[devmotion]

Maybe instead of a global setting we could copy another design choice from the AD ecosystem, in addition to the unified interface based on backend instances: Every FFT package defines its own custom public but unexported API.

Then

• if a user or package author wants to easily switch between different FFT backends, he or she would use the AbstractFFTs API + FFT backend types
• if a user or package author wants to use a specific backend (IIRC FFTA and RustFFT currently don't support all features of FFTW), he or she would use the backend-specific API

It would be up to each package to decide on its internal API but I suppose similar to the AD ecosystem where you have ForwardDiff.gradient, ReverseDiff.gradient, etc. likely it would be something like FFTW.fft, FFTA.fft etc.

So in the first case you might as a package author write something like

using AbstractFFTs

# could define a default backend
f(...; backend::FFTBackend) = ... fft(backend, ...) ...

and as a user

using FFTW
using FFTA

fft(FFTWBackend(), ...)
fft(FFTABackend(), ...)

Whereas in the second case you'd just write

using FFTW
FFTW.fft(...)

[nHackel]

Maybe instead of a global setting we could copy another design choice from the AD ecosystem, in addition to the unified interface based on backend instances: Every FFT package defines its own custom public but unexported API.

I'll give that version quick try and see if I stumble over anything in the backend packages

[nHackel]

I actually don't know how to structure using/import and export in that setting. For that interface we want to have the following setup:

1. AbstractFFTs defines and exports functions for fft, plan_fft, ifft, plan_ifft, and so on
2. Backends extend and reexport those functions
3. Furthremore backends define their own functions with the same name as the exported functions from AbstractFFTs. This function is not exported and only accessible via Backend.fft(...)

Maybe I am missing something, but I don't know how to satisfy all three conditions at once, since 2.) makes the binding visible to a backend which means it can't create its own function with the same name.

As far as I can tell the approach of AD works because there none of the functions are exported. If we drop the reexports from 2 we would get:

# Package
module MyPackage
using AbstractFFTs

f(...; backend::FFTBackend) = ... fft(backend, ...) ...
end

# User with explicit backend have to load AbstractFFTs as well
using AbstractFFTs, FFTA
fft(FFTA.backend(), ...)

# User without backend can just rely on a backend but be fully qualified
using FFTA
FFTA.fft(...)

I've implemented the changes here for AbstractFFTs, FFTA and FFTW. To avoid spamming PRs or splitting the discussion, I just did the PR against my forks for now. In this setup backends have to be careful not to import plan_x or lose the ability to define their own function with the same name

[dannys4]

for what my opinion is worth, I think that (if this PR goes through) adopting, e.g., FFTW.fft(...) would be nice because then it would be easy to fix legacy packages that are borderline unmaintained; FFTW is a very "old" package.

[tknopp]

• if a user or package author wants to easily switch between different FFT backends, he or she would use the AbstractFFTs API + FFT backend types
• if a user or package author wants to use a specific backend (IIRC FFTA and RustFFT currently don't support all features of FFTW), he or she would use the backend-specific API

That is basically what this PR proposes. Each package of course can have its internal API and it does not need to align with the interface of AbstractFFTs. This PR targets the shared API and allows that the concrete implementation do not "claim" the AbstractFFTs interface but that they allow to make it switchable.

For me the discussion on an alternative interface with explicite tokens or the use of a fully qualified name is somewhat orthogonal to the problem, this PR aims to resolve. Here it is mainly about what AbstractFFTs.fft(x::AbstractArray) (...) should do. And we here have basically three options:

1. leave as is
2. providing a global token based API (this PR)
3. removing this function
   Or do I miss any alternative?

[Taaitaaiger]

(IIRC FFTA and RustFFT currently don't support all features of FFTW),

This is true. RustFFT is very limited in what it supports, and I don't expect that to change. Even real FFTs are not supported out of the box, but require an additional crate, and I'm still on the fence whether RustFFT.jl should also provide that functionality, or to offer it as a separate package.

How does that affect switchability? It seems to me that globally setting an implementation will lead to issues unless it's possible to specify what features the backend actually supports, and the user what features they need.

[nHackel]

How does that affect switchability? It seems to me that globally setting an implementation will lead to issues unless it's possible to specify what features the backend actually supports, and the user what features they need.

I'd say that while that could result in errors, those errors are already there with the current interface. Loading two packages with CPU FFTs leads to while not "undetermined" certainly unexpected behaviour such as "unexpected" plans being loaded or ambiguity.

Whatever version this PR ends up with at least allows packages to specify or even pin specific versions of an FFT backend and ignore the other backends.

Additionally, since we then have a dedicated type from each backend, one could define traits describing the capabilities of a backend