Allow plans to have fewer dimensions than the data

[maximilian-gelbrecht]

I think it's a good idea to allow FFT plans to be applied to data that has more dimensions than the plan, as long as those are only the trailing dimensions. The expected behaviour in this case would be a batched FFT.
CUDA.jl PR JuliaGPU/CUDA.jl#1903 initially implemented this behaviour for CUFFT, but it was removed because it is not in accordance with AbstractFFTs.

As an example

using CUDA 
A = CUDA.rand(100); B = CUDA.rand(100,10); 
plan = CUDA.CUFFT.plan_fft(A) 
plan * A # works 
plan * B # doesn't work

In this case the expected behaviour would be (plan * B)[:,i] == (plan * B[:,i]).

You'd only need a single plan to apply the plan to data in different batch sizes as well. E.g. in the following situation:

A = rand(100,10); B = rand(100,20); C = rand(100) 

I'd like to have a plan that performs an FFT along the first dimension, but can apply to both A and B. It's a very common setup e.g. in ML applications, that the batch size is variable and not set when allocating the model.

[stevengj]

This is already supported: just pass the dims parameter to plan_fft. (FFTW works for any subset of the dimensions, not just trailing dimensions. Some other FFT implementations like MKL only support trailing dimensions.)

[stevengj]

Oh, I see, you want to plan a lower-dimensional FFT without the dims parameter and then apply it to a batch of FFTs.

This is do-able, in principle. The only wrinkle is that the data alignment may change for different batches, which could cause problems with FFTW where the plan is specialized for the array alignment (mod 16 bytes)

[maximilian-gelbrecht]

The dims parameter doesn't do what I want (as long as I am not absolutely misunderstanding it). Let me give another example, now with the dims specified.:

using AbstractFFTs, FFTW

A= rand(100); B=rand(100,10); C=rand(100,20);
planA = plan_fft(A, 1);
planB = plan_fft(B, 1);

planA * A # works 
planA * B # doesn't work 

planB * B # works
planB * C # doesn't work 

The second example with planB, is more for completeness. It is probably easier to limit this additional trailing dimensions, and not all trailing dimensions after the dims I guess.

To be honest, I am only really interested in this functionality for CUFFT. The CUFFT PR did already implement it, but it wasn't included into the final accepted one, due to the AbstractFFT design. It would be great if it could be added.