CuPy의 FFT를 가속화하는 숨겨진 기능

추가

CuPy v7에서 plan을 context manager로 취급하는 기능이 추가되었으므로이 기사의 방법보다 그곳을 사용하는 것이 좋습니다.

소개

CuPy 에 v4에서 FFT가 추가되었습니다.
이를 통해 NumPy와 동일한 인터페이스에서 cuFFT를 사용할 수 있습니다.
그러나 NumPy와 인터페이스를 정렬하기 위해 cuFFT의 성능을 다 사용하지 않을 수 있습니다.
이 기사에서는 CuPy의 내부 구현을 이용하여 FFT를 더욱 가속화하는 방법을 소개합니다.
참고 이 방법은 사양이 되지 않은 내부 구현을 이용하므로 minor update에서도 사용할 수 없게 될 가능성이 있습니다.

환경

CPU: Intel Core i7-8700K
GPU: GeForce GTX 1080 Ti
CUDA: 9.1
NumPy: 1.15.2 (+MKL)
CuPy: 4.4.1

길이 1024의 단정밀도의 배열을 100개씩 100회 FFT했을 때의 속도를 계측합니다.

측정용 소스 코드는 여기

NumPy

100×1024의 ndarray를 만들어 100회 FFT합니다.

x_cpu = numpy.random.rand(100, 1024) + 1j * numpy.random.rand(100, 1024)
x_cpu = x_cpu.astype(numpy.complex64)

with timer('numpy.fft.fft'):
    for i in range(100):
        numpy.fft.fft(x_cpu)

numpy.fft.fft: 93.943119 ms

CuPy

CuPy는 NumPy와 동일한 인터페이스를 가지므로 기본적으로 numpy를 cupy로 바꾸는 것만으로 GPU를 사용하는 코드입니다.

x_gpu = cupy.asarray(x_cpu)

with timer('cupy.fft.fft'):
    for i in range(100):
        cupy.fft.fft(x_gpu)

cupy.fft.fft: 24.476290 ms

벌써 이 시점에서 4배 약 빠르고 있습니다만, 실은 이 코드에서는 cuFFT의 plan을 만드는 처리가 병목이 되고 있습니다.
이번과 같이 같은 사이즈의 FFT를 여러 번 실시하는 경우는 plan을 사용 돌리는 것이 빨라집니다.

cupy.cuda.cufft.Plan1d

NumPy와 동일한 인터페이스에서 plan을 사용할 수 없으므로 CuPy의 내부 구현에서 plan을 관리하는 클래스를 사용합니다.

with timer('cupy.cuda.cufft.Plan1d'):
    plan = cufft.Plan1d(1024, cufft.CUFFT_C2C, 100)
    for i in range(100):
        y_gpu = plan.get_output_array(x_gpu)
        plan.fft(x_gpu, y_gpu, cufft.CUFFT_FORWARD)

cupy.cuda.cufft.Plan1d: 1.344442 ms

한자리 더 빨라졌습니다.
한편, 인터페이스는 여러 인수가 늘어나 이해하기 어려워지고 있습니다.
이 예는 아직 몹시 편이며, 직접 Plan1d 를 호출하면 구현이 상당히 복잡해지는 일이 있습니다.
그러한 경우에도 cupy.fft 는 형이나 사이즈나 축등을 잘 조정해 주기 때문에 편리합니다.

배열의 길이를 변경한 경우

배열의 길이가 짧을 때 고속화되고 있습니다.
데이터가 커지면 FFT를 실행하는 것이 시간이 걸리게 되어 plan을 만드는 오버헤드가 보이지 않게 됩니다.

요약

내부 실장을 이용해 Plan1d 를 직접 사용하는 편이 빨라진다.
코드는 다소 어려워지고 사양이 되어 있는 기능도 아니기 때문에, 속도를 그다지 필요로 하지 않는다면 cupy.fft 를 사용하는 편이 좋다.