fft missing methods

[ptiede]

Reactant is missing a few use cases where an FFT is required, but the current extension is missing methods. Here are a few cases I have encountered.

Plan is defined based on cpu arguments and is essentially constant

using Reactant, FFTW

plan = plan_fft!(copy(x))

plan * x #works in place

xr = Reactant.to_rarray(x)

@compile p * xr
# Gives execution error because it is trying to use fftw methods
# Missing Base.:*(plan, x::RArray) method

Inplace ffts are not defined

using Reactant, FFTW

x = rand(ComplexF64, 64, 64)
@compile fft!(x)
# Missing fft!(x::RArray) 

ReactantFFT plans doesn't support dims

x = rand(ComplexF64, 64, 64, 4)
xr = Reactant.to_rarray(xr)
@compile plan_fft(xr, 1:2)
# ERROR: MethodError: no method matching plan_fft(::Reactant.TracedRArray{ComplexF64, 2}, ::UnitRange{Int64})

I noticed this while reviewing the new code from #1926. This should be easy enough to fix by changing the struct to capture the dims argument, or otherwise make it 1:ndims(x)

The second and third ones are easy to fix. For the first one, my naive guess would be to define the overload and use Reactant to compute the fft.

[wsmoses]

in place FFT's are also easy, we just should define it as copyto!(x, fft(x))

want to give these a go?

[ptiede]

Yep!