N-dimensional batch for cuFFT

Accelerated Computing GPU-Accelerated Libraries
1

Let’s say I’ve got a 4D FORTRAN array and I’m interested in performing multiple 1D FFTs just along the third dimension. Say the dimensions of the array are NX,NY,NZ and NN; where NX is the inner most index and NZ is the size of the transform. As far as I understand, such task can only be done (efficiently) as follows:

real*8 :: IN(NX,NY,NZ,NN),OUT(NX,NY,NZ,NN)
integer :: n(1),inembed(1),onembed(1),idist,odist,istride,ostride,batch
integer :: ierr,planh

n(1) = NZ
inembed(1) = NZ
onembed(1) = NZ
idist = 1
odist = 1
istride = NX*NY
ostride = NX*NY
batch = NX*NY

! Create the plan
ierr = cufftPlanMany(planh,1,n,inembed,istride,idist,onembed,ostride,odist,CUFFT_D2Z,batch)

! Call cuFFT NN times (this would be streamed cuFFT calls)
do i=1,NN
  ierr = ierr + cufftExecD2Z(planh,IN(:,:,:,i),OUT(:,:,:,i))
end do

For simplicity, I have omitted the multiple plan creation necessary to async the multiple FFT calls. On the other hand, all these transforms could be collapsed into a single FFT call if the idist, odist and batch parameters were arrays. Thus, the above example would be:

real*8 :: IN(NX,NY,NZ,NN),OUT(NX,NY,NZ,NN)
integer :: n(1),inembed(1),onembed(1),idist,odist,istride,ostride,batch
integer :: ierr,planh

n(1) = NZ
inembed(1) = NZ
onembed(1) = NZ
idist(:) = (/1,NX*NY*NZ/)
odist(:) = (/1,NX*NY*NZ/)
istride = NX*NY
ostride = NX*NY
batch(:) = (/NX*NY,NN/)

! Create the plan
ierr = cufftPlanMany(planh,1,n,inembed,istride,idist,onembed,ostride,odist,CUFFT_D2Z,batch)

! Call cuFFT
err = ierr + cufftExecD2Z(planh,IN,OUT)

Is such thing possible, or are the multiple streamed cuFFT calls the only solution to this?

Thanks!