Generalize construction of FFT plans on GPU

[navidcy]

Currently, we can construct FFT plans on the GPU via:

FourierFlows.jl/src/CuFourierFlows.jl

Lines 4 to 5 in 21cd276

	plan_flows_fft(a::CuArray, effort) = plan_fft(a)
	plan_flows_rfft(a::CuArray, effort) = plan_rfft(a)

This works great but it does not take into account optional dims arguments the user may give into plan_fft()/plan_rfft(); this is particularly relevant for MultilayerQG where FFTs are only done in the first two dimensions of the arrays. We need to generalize the gpuification of plan_fft()/plan_rfft() to account for dims.

[glwagner]

True. I'm not sure I completely understand the desired functionality. It does seem like a bit of a hack that the fft plans are included in the grid at all.

In MultilayerQG some additional plans are constructed, eg:

https://github.com/FourierFlows/GeophysicalFlows.jl/blob/0bb51f19ca47b12c10b6f885b7c14917c41bc6bb/src/multilayerqg.jl#L132

What I don't totally understand is why this only affects GPU functionality. Wouldn't the lack of this feature also affect CPU functionality? Why is this GPU specific?

[glwagner]

since

julia> p = plan_fft(CuArray(rand(3, 3, 3)), [1, 2], flags=FFTW.MEASURE)
FFTW forward plan for 3×3×3 array of Complex{Float64}
(dft-vrank>=1-x3/1
  (dft-rank>=2/1
    (dft-direct-3-x3 "n1fv_3_sse2")
    (dft-direct-3-x3 "n1fv_3_sse2")))

works, it seems we don't need plan_flows_fft anymore.

never mind, the above code is not correct (it returns an FFTW object for CPU FFTs) and we still need plan_flows_fft.

Note that the CPU definition of plan_flows_fft is here:

FourierFlows.jl/src/domains.jl

Line 1 in 21cd276

plan_flows_fft(a::Array, effort) = plan_fft(a; flags=effort)