IDCT for dim=3 in 3D Poisson solver does not work on the GPU. about oceananigans.jl HOT 4 CLOSED

@ali-ramadhan could the NVidia cufftw interface help ( https://docs.nvidia.com/cuda/cufft/index.html#fftw-supported-interface ) as a step. My reading is that this does things on GPU, so could avoid copy, but may not be as super optimal as native cuFFT. Maybe you already tried that - the advertising on the box seems to imply you just switch a few headers, link and existing FFTW works, but maybe there is small print?

from oceananigans.jl.

@christophernhill Yes I believe I am using the FFTW interface to do the FFT/IFFTs on the GPUs.

Where I got stuck is that the existing solver uses FFTW.r2r!(f.data, FFTW.REDFT10, 3) for the DCT and FFTW.r2r!(ϕ.data, FFTW.REDFT01, 3) for the IDCT, but under the "FFTW Interface to cuFFT" table, Read-to-Real Transforms and Read-to-Real Transform Kinds are "Not supported".

So I implemented a DCT using the N FFT (following this post)
https://github.com/ali-ramadhan/Oceananigans.jl/blob/93aa0038b3126470f263475d648bceb9562bbe91/src/spectral_solvers.jl#L359-L370

and an IDCT using the IFFT (following these course notes)
https://github.com/ali-ramadhan/Oceananigans.jl/blob/93aa0038b3126470f263475d648bceb9562bbe91/src/spectral_solvers.jl#L372-L387

They work in isolation (they produce the same output as FFTW.r2r! but do not work when I substitute them into the Poisson solver. I'm playing around with it right now, hoping it's just a bug on my part.

from oceananigans.jl.

Hmmm, now I look through the notes again, maybe the IDCT I implemented only works for real inputs. That's exactly what we need though. I'll keep playing around.

from oceananigans.jl.

Upon closer inspection of the REDFT01 (DCT-III) transform in FFTW, I should have been dividing the first element X₀ by 2 in our DCT-III implementation using the IFFT. The correct IDCT is thus

function idct_dim3_gpu!(f)
    Nx, Ny, Nz = size(f)
    
    bfactors = exp.(collect(1im*π*(0:Nz-1) / (2*Nz)))
    bfactors[1] *= 0.5

    f .*= cu(repeat(reshape(bfactors, 1, 1, Nz), Nx, Ny, 1))
    ifft!(f, 3)
    
    f .= cu(reshape(permutedims(cat(f[:, :, 1:Int(Nz/2)], f[:, :, end:-1:Int(Nz/2)+1]; dims=4), (1, 2, 4, 3)), Nx, Ny, Nz))
    # @. f = real(f)  # Don't do it here. We'll do it when assigning real(ϕ) to pNHS to save some measly FLOPS.
    
    nothing
end

This can all be optimized, but at least a working Poisson solver is a good step.

from oceananigans.jl.