CUDA.jl cuFFT underperforming against CuPy cuFFT

[dreycenfoiles]

Describe the bug

cuFFT's with Julia are underperforming when compared with CuPy and I consistently see a ~2x performance gap. Below is an example run.

BenchmarkTools.Trial: 7144 samples with 1 evaluation.
 Range (min … max):  475.800 μs … 34.983 ms  ┊ GC (min … max): 0.00% … 8.24%
 Time  (median):     620.500 μs              ┊ GC (median):    0.00%
 Time  (mean ± σ):   688.121 μs ±  1.282 ms  ┊ GC (mean ± σ):  0.61% ± 0.33%

fft_func            :    CPU:  166.498 us   +/-64.221 (min:  112.600 / max: 1298.200) us     GPU-0:  336.961 us   +/-62.120 (min:  238.592 / max: 1489.920) u

To reproduce

The Minimal Working Example (MWE) for this bug:

Julia code I ran.

using CUDA 
using CUDA.CUFFT 
using BenchmarkTools

A = CUDA.rand(Float32, 500,500)

function fft_func(A)
    return fft(A) 
end

@benchmark @CUDA.sync fft_func(A)

Python code I ran.

import cupyx.scipy.fft as cufft
import cupy as cp 
from cupyx.profiler import benchmark

A = cp.random.random((500,500)).astype(cp.float32)

def fft_func(A):

    return cufft.fftn(A)

print(benchmark(fft_func, (A,), n_repeat=10000))

Expected behavior

Given that both CuPy and CUDA.jl should call out to the same cuFFT routines, I would expect their runtimes to be almost identical.

Version info

Details on Julia

Julia Version 1.8.1
Commit afb6c60d69 (2022-09-06 15:09 UTC)
Platform Info:
OS: Windows (x86_64-w64-mingw32)
CPU: 6 × Intel(R) Core(TM) i5-9400 CPU @ 2.90GHz
WORD_SIZE: 64
LIBM: libopenlibm
LLVM: libLLVM-13.0.1 (ORCJIT, skylake)
Threads: 1 on 6 virtual cores
Environment:
JULIA_EDITOR = code
JULIA_NUM_THREADS =

Details on CUDA:

CUDA toolkit 11.7, artifact installation
NVIDIA driver 516.94.0, for CUDA 11.7
CUDA driver 11.7

Libraries:

• CUBLAS: 11.10.1
• CURAND: 10.2.10
• CUFFT: 10.7.1
• CUSOLVER: 11.3.5
• CUSPARSE: 11.7.3
• CUPTI: 17.0.0
• NVML: 11.0.0+516.94
• CUDNN: 8.30.2 (for CUDA 11.5.0)
• CUTENSOR: 1.4.0 (for CUDA 11.5.0)

Toolchain:

• Julia: 1.8.1
• LLVM: 13.0.1
• PTX ISA support: 3.2, 4.0, 4.1, 4.2, 4.3, 5.0, 6.0, 6.1, 6.3, 6.4, 6.5, 7.0, 7.1, 7.2
• Device capability support: sm_35, sm_37, sm_50, sm_52, sm_53, sm_60, sm_61, sm_62, sm_70, sm_72, sm_75, sm_80, sm_86

1 device:
0: NVIDIA GeForce GTX 1050 Ti (sm_61, 3.366 GiB / 4.000 GiB available)

[vchuravy]

If you want to investigate. One avenue might be to look at a it with NSight systems and check that indeed we end up calling the same kernels.

[maleadt]

x-ref https://discourse.julialang.org/t/cupy-cufft-2x-faster-than-cuda-jl-cufft

[doraemonho]

I am not sure if this is the hint but I hope this may help: Link of Cupy.fft detail

As they said from the doc, starting from cupy v8, they had enabled the built-in plan cache for assisting the compution. So, they will recycle the cache every time you used the fft function. We can alway check by using the cp.fft.config.get_plan_cache() function. Like the code I showed below (sorry for not using benchmark, my version of cupy doesn't support the function)

import time
import cupy as cp 
from cupyx.scipy.fft import get_fft_plan
import cupyx.scipy.fft as cufft

def fft_func(A):
    return cufft.fftn(A)

start_time = time.time();
with cp.cuda.Device(1):
    cache = cp.fft.config.get_plan_cache()
    print(cache.show_info())
    A = cp.random.random((1024,1024)).astype(cp.float32)
    for i in range(100000):
        fft_func(A);    
    print(cache.show_info())

totol_time = (time.time() - start_time);
print("--- %s seconds ---" % totol_time)

print("end\n")

Output :

------------------- cuFFT plan cache (device 1) -------------------
cache enabled? True
current / max size   : 0 / 16 (counts)
current / max memsize: 0 / (unlimited) (bytes)
hits / misses: 0 / 0 (counts)

cached plans (most recently used first):

None
------------------- cuFFT plan cache (device 1) -------------------
cache enabled? True
current / max size   : 1 / 16 (counts)
current / max memsize: 512 / (unlimited) (bytes)
hits / misses: 99999 / 1 (counts)

cached plans (most recently used first):
key: ((1024, 1024), (1024, 1024), 1, 1, (1024, 1024), 1, 1, 41, 1, 'C', 1, None), plan type: PlanNd, memory usage: 512

None
--- 8.673956394195557 seconds ---
end

So, you can see the fft function use the cache every time.
However for julia, from my understanding, everytime if you run fft, you would reallocate an CuArray and cause a performance drop. In order to avoid this, we may need to first declare a fftplan and work our the forward transform.

using CUDA
using FFTW
using BenchmarkTools
using LinearAlgebra: mul!, ldiv!
device!(1)

function loopfft!(plan,A,B;N=100)
    for i = 1:N::Int
         mul!(B,fftplan,A)
    end
    return nothing;
end

nx,ny = 1024,1024;
T = ComplexF32
ArrayType = CuArray
fftplan = plan_fft(ArrayType{T, 2}(undef, nx, ny))

A = ArrayType{T, 2}(undef, nx, ny)
A = CUDA.randn!(A)
B = copy(A);
@CUDA.time loopfft!(fftplan,A,B;N=100000)

Output

 6.160003 seconds (32.73 k CPU allocations: 2.346 MiB)

In such setup, julia is faster than the python's version but I dont think this is a fair comparison.(as I run a long for loop in python....) If my guess is true, the key massage is that the difference of time is coming from the way of implementation between python and julia.

[maleadt]

Please test #1734