Polluted result using CUDA 3D rfft in Float32 data type

[doraemonho]

Hi, I am working on the 3D RFFT but find out the result of rfft using CUDA is a bit weird in Float32 .

I define a function F(k) in k-space as a gaussian function like e^{-(k-k0)^2}e^{-i\phi}, where k0 is some number and phi is a random number array between 0 and 2pi. Here is the implementation

function GetFk(grid;kf = 2,σ²= 1)
    kx,ky,kz = grid.kr,grid.l,grid.m;
    k⁻¹  = @. √(grid.invKrsq);
    k    = @. √(grid.Krsq);
    Fk   = @. √(exp(-(k.-kf)^2/σ²)/2/π)*k⁻¹;
    randN = typeof(Fk) <: Array ? Base.rand : CUDA.rand;
    eⁱᶿ¹ = exp.(im*randN(T,GPUgrid.nkr,GPUgrid.nl,GPUgrid.nm)*2*π);
    return Fk.*eⁱᶿ¹;
end

If I do a rfft on this function, the result of CUDA rfft will overlay a moiré like pattern like below:

Here is the implemention:

using PyPlot,CUDA, FourierFlows
using FFTW
using LinearAlgebra: mul!, ldiv!

T = Float32;
GPUgrid = ThreeDGrid(GPU(),128,2π;T=T);
GFkh = GetFk(GPUgrid;kf = 2,σ²= 1);

GFk  = CUDA.zeros(T,GPUgrid.nx,GPUgrid.ny,GPUgrid.nz);
ldiv!(GFk, GPUgrid.rfftplan, deepcopy(GFkh));

CPUgrid = ThreeDGrid(CPU(),128,2π;T=T);
CFkh = Array(GFkh);
CFk  = zeros(T,GPUgrid.nx,GPUgrid.ny,GPUgrid.nz);
ldiv!(CFk, CPUgrid.rfftplan, deepcopy(CFkh));

figure(figsize = (16,8))
subplot(121);title("GPU T =$T",size=16)
imshow(Array(GFk[:,:,1]));
subplot(122);title("CPU T =$T",size=16)
imshow(Array(CFk[:,:,1]))

However, if you change T from Float32 to Float64, the result would become normal again:

The problem itself is rare and would happen in this kind of function.

[navidcy]

Hi @doraemonho — Nice! I’ll try to have a look at it soon.

[navidcy]

Inside the GetFk function the GPUgrid is called instead the grid from the function’s argument. Is that the prob?

[doraemonho]

Thanks for taking a look!

Inside the GetFk function the GPUgrid is called instead the grid from the function’s argument. Is that the prob?

I rewrtie the function a bit to define the GPUgrid inside the function (showed below), but the result doesn't change.

function GetFk(N::Int,Lx::Number;T = Float32, dev = GPU(), kf = 2,σ²= 1)
    Fgrid    = ThreeDGrid(GPU(),N,Lx;T=T);
    kx,ky,kz = Fgrid.kr,Fgrid.l,Fgrid.m;
    k⁻¹  = @. √(Fgrid.invKrsq);
    k    = @. √(Fgrid.Krsq);
    Fk   = @. √(exp(-(k.-kf)^2/σ²)/2/π)*k⁻¹;
    randN = typeof(Fk) <: Array ? Base.rand : CUDA.rand;
    eⁱᶿ¹ = exp.(im*randN(T,Fgrid.nkr,Fgrid.nl,Fgrid.nm)*2*π);
    return Fk.*eⁱᶿ¹,Fgrid;
end

T = Float32;
N = 128
Lx= 2π
GFkh,GPUgrid = GetFk(N,Lx;T=T,kf = 8,σ²= 1);

GFk  = CUDA.zeros(T,GPUgrid.nx,GPUgrid.ny,GPUgrid.nz);
ldiv!(GFk, GPUgrid.rfftplan, deepcopy(GFkh));

CPUgrid = ThreeDGrid(CPU(),N,Lx;T=T);
CFkh = Array(GFkh);
CFk  = zeros(T,GPUgrid.nx,GPUgrid.ny,GPUgrid.nz);
ldiv!(CFk, CPUgrid.rfftplan, deepcopy(CFkh));

figure(figsize = (16,8))
subplot(121);title("GPU T =$T",size=16)
imshow(Array(GFk[:,:,1]));
subplot(122);title("CPU T =$T",size=16)
imshow(Array(CFk[:,:,1]))

[navidcy]

I’ll have a closer look in the morning from my laptop

[doraemonho]

It occurs to me that, the root of this problem is coming from the fft plan itself.
Now, I don't even use grid fft plan itself but define a rfft plan using FFTW, I can reproduce the polluted result

T = Float32;
GFkh_32,GPUgrid = GetFk(N,Lx;T=T,kf = 8,σ²= 1);
GFk32  = CUDA.zeros(T,GPUgrid.nx,GPUgrid.ny,GPUgrid.nz);
GPUrfftplan_F32 = plan_rfft(GFk32);
ldiv!(GFk32, GPUrfftplan_F32, deepcopy(GFkh_32));


figure(figsize = (16,8))
subplot(121);title("GPU T =$T",size=16)
imshow(Array(GFk32[:,:,1]));

T = Float64;
GFkh_64,GPUgrid = GetFk(N,Lx;T=T,kf = 8,σ²= 1);
GFk64  = CUDA.zeros(T,GPUgrid.nx,GPUgrid.ny,GPUgrid.nz);
GPUrfftplan_F64 = plan_rfft(GFk64);
ldiv!(GFk64, GPUrfftplan_F64, deepcopy(GFkh_64));

subplot(122);title("GPU T =$T",size=16)
imshow(Array(GFk64[:,:,1]))

I am not sure if this is a FFTW.jl 's or CUDA.jl 's problem.

[doraemonho]

Quick Update here
After a long trial and error, it seems that I am finally arriving to the root of problem.
It may has something to do with the fft plan of 2^n 3D real transform.
If you try switch the resolution of any 2^n 3D Map (128^3 in our case) to other even number (130^3).
The discrepancy between CPU and GPU rfft would vanish.
I suspect there may have a bug on the algorithm of implementing 2^n FP32 rfft but I am not sure if this is a FFTW problem or CUDA problem

[navidcy]

Can we make sure it's not a plotting issue?

I tried to do some FFTs back and forth and couldn't find any disagreement between different data types. Have a look:

Here's an example just with FFTW and CUDA

CUDA + FFTW

using CUDA
using FFTW
using LinearAlgebra: mul!, ldiv!
using Test

for T in [Float64, Float32]
    nx, ny, nz = 6, 8, 10
    nk, nl, nm = Int(nx//2)+1, ny, nz

    a = randn(T, (nx, ny, nz))
    a_cu = CuArray(a)

    ah = randn(Complex{T}, (nk, nl, nm))
    ah_cu = CuArray(ah)

    plan = plan_rfft(a)
    plan_cu = plan_rfft(a_cu)

    bh = zeros(Complex{T}, (nk, nl, nm))
    bh_cu = CuArray(bh)

    c = zeros(T, (nx, ny, nz))
    c_cu = CuArray(c)

    mul!(bh, plan, a)
    ldiv!(c, plan, deepcopy(bh))

    mul!(bh_cu, plan_cu, a_cu)
    ldiv!(c_cu, plan_cu, deepcopy(bh_cu))

    @testset "CPU test for $T" begin
        @test bh == rfft(a)
        @test c == irfft(bh, nx)
    end

    @testset "GPU test for $T" begin
        @test bh_cu == rfft(a_cu)
        @test c_cu == irfft(bh_cu, nx)
    end
end

and here's an example with CUDA + FourierFlows

CUDA + FourierFlows

using CUDA
using LinearAlgebra: mul!, ldiv!
using FourierFlows
using Test

for T in [Float64, Float32]
    nx, ny, nz = 6, 8, 10

    grid = ThreeDGrid(nx, nx, ny, ny, nz, nz; T)
    grid_cu = ThreeDGrid(GPU(), nx, nx, ny, ny, nz, nz; T)

    a = randn(T, (grid.nx, grid.ny, grid.nz))
    a_cu = CuArray(a)

    ah = randn(Complex{eltype(grid)}, (grid.nkr, grid.nl, grid.nm))
    ah_cu = CuArray(ah)

    plan = grid.rfftplan
    plan_cu = grid_cu.rfftplan

    bh = zeros(Complex{eltype(grid)}, (grid.nkr, grid.nl, grid.nm))
    bh_cu = CuArray(bh)

    c = zeros(eltype(grid), (grid.nx, grid.ny, grid.nz))
    c_cu = CuArray(c)

    mul!(bh, plan, a)
    ldiv!(c, plan, deepcopy(bh))

    mul!(bh_cu, plan_cu, a_cu)
    ldiv!(c_cu, plan_cu, deepcopy(bh_cu))

    @testset "CPU test for $T" begin
        @test bh == rfft(a)
        @test c == irfft(bh, nx)
    end

    @testset "GPU test for $T" begin
        @test bh_cu == rfft(a_cu)
        @test c_cu == irfft(bh_cu, nx)
    end
end

All test pass for me. Same if I choose nx = ny = nz = 128.

[doraemonho]

Can we make sure it's not a plotting issue?

I tried to do some FFTs back and forth and couldn't find any disagreement between different data types. Have a look:

Here's an example just with FFTW and CUDA

CUDA + FFTW
and here's an example with CUDA + FourierFlows

CUDA + FourierFlows
All test pass for me. Same if I choose nx = ny = nz = 128.

Thanks for checking!
Same result for me but I think this problem would only occurs if we satisfy the following condition :

1. Doing rfft from k-space to real space
2. The function that we are transferring is special ( like the spectral filtered e^{i\phi} k^{-2} )

I quite certain the problem is not coming the plot as the radial spectrum of FP32-CUDA found a spike at large k number.
I cannot check the result using your method as the FFT result is a little different (~1e-18) between GPU and CPU result.
However, I come up a alternative by using \approx instead of == inside the test. The Float64 is fine but fail in Float32.

using CUDA
using FFTW
using LinearAlgebra: mul!, ldiv!
using FourierFlows
using Test

nx, ny, nz = 128, 128, 128;       # define size of the array
θ = rand(θ,div(nx,2)+1, ny, nz)*2π; # define \theta \in [0,2\pi]

for T in [Float64,Float32]
    grid = ThreeDGrid(nx, nx, ny, ny, nz, nz; T)             
    grid_cu = ThreeDGrid(GPU(), nx, nx, ny, ny, nz, nz; T)

    θT  = convert(Array{T,3},θ)
    eⁱᶿk⁻² = @. exp.(im*θT)*grid.invKrsq;     #define the spectral filtered function in CPU
    eⁱᶿk⁻²_cu = CuArray(eⁱᶿk⁻²);               #define the spectral filtered function in GPU

    plan = grid.rfftplan
    plan_cu = grid_cu.rfftplan

    c = zeros(eltype(grid), (grid.nx, grid.ny, grid.nz))
    c_cu = CuArray(c);
    
    ldiv!(c, plan, deepcopy(eⁱᶿk⁻²))    
    ldiv!(c_cu, plan_cu, deepcopy(eⁱᶿk⁻²_cu))
    c_cu = Array(c_cu); #move the result from GPU-> CPU
    
    test_title = "eⁱᶿk⁻² CPU & GPUtest for " * string(T)

    @testset "$test_title" begin
        #Compare the result using ≈
        @test c_cu  ≈  c
    end 
end

[navidcy]

Can you print the types of eⁱᶿk⁻² and eⁱᶿk⁻²_cu to make sure they are correct type?

[navidcy]

Before and after the conversion to Array for eⁱᶿk⁻²_cu

[doraemonho]

Before and after the conversion to Array for eⁱᶿk⁻²_cu

do you mean before/after the rift or before the testing process?
This is my type checking code and the result, as you can see, FP64 is fine but FP32 just fail

using CUDA
using FFTW
using LinearAlgebra: mul!, ldiv!
using FourierFlows
using Test

nx, ny, nz = 128, 128, 128;       # define size of the array
θ = rand(div(nx,2)+1, ny, nz)*2π; # define \theta \in [0,2\pi]

for T in [Float64,Float32]
    grid = ThreeDGrid(nx, nx, ny, ny, nz, nz; T)             
    grid_cu = ThreeDGrid(GPU(), nx, nx, ny, ny, nz, nz; T)

    θT  = convert(Array{T,3},θ)
    eⁱᶿk⁻² = @. exp.(im*θT)*grid.invKrsq;     #define the spectral filtered function in CPU
    eⁱᶿk⁻²_cu = CuArray(eⁱᶿk⁻²);              #define the spectral filtered function in GPU
    
    println("before the rfft")
    @show typeof(eⁱᶿk⁻²)
    @show typeof(eⁱᶿk⁻²_cu)
    
    plan = grid.rfftplan
    plan_cu = grid_cu.rfftplan

    c = zeros(eltype(grid), (grid.nx, grid.ny, grid.nz))
    c_cu = CuArray(c);
    
    
    ldiv!(c, plan, deepcopy(eⁱᶿk⁻²))    
    ldiv!(c_cu, plan_cu, deepcopy(eⁱᶿk⁻²_cu))
    c_cu = Array(c_cu); #move the result from GPU-> CPU
    
    
    println("before the test")
    @show typeof(c)
    @show typeof(c_cu)
    
    test_title = "eⁱᶿk⁻² CPU & GPU test for " * string(T)

    @testset "$test_title" begin
        #Compare the result using ≈
        @test c_cu  ≈  c
    end 
end

[navidcy]

@doraemonho sorry for slacking... I went out on leave and now slowly coming back.

Did you have any progress on this or is it still outstanding?

[doraemonho]

@doraemonho sorry for slacking... I went out on leave and now slowly coming back.

Did you have any progress on this or is it still outstanding?

@navidcy no worries, I still have this issue.
The way of circumventing it is to change the resolution from 2^n to 10*n or switch to FP64 calculation. The problem only occurs in this kind of function

[doraemonho]

Just an update from what we were : -)
This may be a CuFFT bug that would occur after the version CuFFT v.10.3.
I reported it to NVIDIA and just heard back from them. They confirmed the issue and their CuFFT team is taking a look into it right now.
Not sure when will be solved but I will keep updating when I get back from them