Python: GPU 加速 FFT DFT 計算

前言

之前讀碩班常常要用 FFT 來處理光學的問題，通常都是會用 MATLAB 來撰寫程式，如果矩陣比較大也可以直接用 MATLAB 的 gpuArray 指令，來用 GPU 加速運算，使用起來非常方便，但是到了 Python 就沒有內建這些加速功能，好在 Python 的優點就是有非常完整的套件可以幫助我完成一些原本在 MATLAB 的工作。

這邊會介紹兩種工具，分別是 CuPy 以及 PyVkFFT。

CuPy

CuPy 是兼容 NumPy、SciPy 的 GPU 加速函式庫。簡單講就是 NumPy 的數學運算都可以直接轉換成 CuPy 的 API，所以要將原有的程式改寫成 GPU 加速會非常方便。
順帶一提，我以前就是用 NumPy 去自己手刻 CNN 出來，然後直接把相關計算改成 CuPy 就完成 GPU 加速了，而且也有基於這個工具的神經網路框架 Chainer。

我目前看一本書叫做 Deep Learning 3：用Python進行深度學習框架的開發實作 (連結)，也是使用 CuPy 來加速的樣子。

那 NumPy 到底是怎麼直接轉 CuPy 呢?

原本用 NumPy 去使用 FFT 會類似下面這樣的寫法

import numpy as np
A = np.ones((10,10))
B = np.fft.fft2(A)

只要把 import numpy as np 改成 import cupy as cp，然後把原本 np 的部分改成 cp 即可代換。幾乎都不用改動太多地方，連參數都是仿 NumPy，所以是很容易轉換的！！！

import cupy as cp
A = cp.ones((10,10))
B = cp.fft.fft2(A)

甚至，還可以將 NumPy Array 跟 CuPy Array 互轉 (我的感覺就像是 MATLAB 的 gpuArray 一樣)

import numpy as np
import cupy as cp
A_cpu = np.ones((10,10))
A_gpu = cp.asarray(A_cpu)   #類似 gpuArray (MATLAB)
B_gpu = cp.fft.fft2(A_gpu)
B_cpu = B_gpu.get()         #類似 gather (MATLAB))

也可以混用 NumPy、CuPy API

import numpy as np
import cupy as cp
A = cp.ones((10,10))
B_gpu = np.fft.fft2(A_gpu)  #會自動從 GPU 轉回 CPU

安裝

安裝指令也很簡單，只要根據你安裝的 cuda 版本去輸入對應指令即可。

詳細安裝版本對應可以參考官網的說明 (連結)

要先安裝 cuda toolkit，再去裝 CuPy。

Command (depend on cuda version)

$ pip install cupy-cuda92
$ pip install cupy-cuda100
$ pip install cupy-cuda101
$ pip install cupy-cuda102
$ pip install cupy-cuda110
$ pip install cupy-cuda111
$ pip install cupy-cuda112
$ pip install cupy-cuda113
$ pip install cupy-cuda114

上面的指令應該會隨著時間而慢慢變多，目前這些版本應該支援大多數電腦了，寫這篇文章時 Cuda 最新版本為 v11.5，CuPy 支援到 v11.4，更新也算蠻快的了。

PyVkFFT

這是基於 VkFFT GPU 加速計算的 Python API，專門用來計算 FFT 的相關運算。
VkFFT 是使用 C 語言撰寫，適用 Vulkan/CUDA/HIP/OpenCL 等技術進行加速。
有個很關鍵的事情，它是利用 CuPy 來進行 CUDA 加速，所以我的相關計算都可以無縫接軌，全部靠 CuPy 完成。

這個我其實我還沒很熟悉，我只是要用它來幫我計算 DFT 而已 (因為 CuPy 沒有 DFT)，而且光是安裝成功就搞很久，就沒深入研究。

安裝

這東西安裝起來很複雜，可能是因為我在 Windows 下安裝，所以問題一堆…

我是 Windows + Anaconda3 去建立 Python 環境

1. 安裝 Visual Studio 2019
2. 安裝 CUDA Toolkit 11.0 (其他版本沒試過)
3. conda activate your_env (進入 conda 虛擬環境)
4. pip install pybind11、pip install mako
5. pip install pyopencl-2021.2.9+cl12-cp37-cp37m-win_amd64.whl (這要去下載檔案)
6. 加入 ‘cl.exe’ 的路徑到 ‘環境變數’，目前路徑是 C:\Program Files (x86)\Microsoft Visual Studio\2019\Community\VC\Tools\MSVC\14.29.30037\bin\Hostx64\x64

‘cl.exe’ 路徑可能會隨著 VS 更新而改變。

7. pip install cupy-cuda110 (我是 cuda 11.0 來驅動 pyvkfft)
8. pip install pyvkfft

這 8 個步驟我搞了快一個禮拜才搞出來，而且不確定是否每台電腦都可以成功，我也懶得驗證了XD

pyvkfft fft 範例

小改官方的，原本它是拿 pycuda 來做，我改成 cupy 的

import numpy as np
import cupy as cp
from pyvkfft.fft import fftn, ifftn
d0 = cp.asarray(np.random.uniform(0,1,(200,200)).astype(np.complex64))
# This will compute the fft to a new GPU array
d1 = fftn(d0)
# An in-place transform can also be done by specifying the destination
d0 = fftn(d0, d0)
# Or an out-of-place transform to an existing array (the destination array is always returned)
d1 = fftn(d0, d1)
# Here will compute ifft to a new GPU array
d2 = ifftn(d1)

pyvkfft dft 範例

import numpy as np
import cupy as cp
from pyvkfft.cuda import VkFFTApp
d0 = cp.asarray(np.random.uniform(0,1,(200,200)).astype(np.float64))
# Set VkFFTApp for DFT
dft_app = VkFFTApp(d0.shape, d0.dtype, ndim=2, dct=2)
# Here will compute the dft
d0 = dft_app.fft(d0)
# Here will compute the idft
d0 = dft_app.ifft(d0)

結語

目前就是我知道可以在 Python 加速 FFT 的兩種方式，給像我一樣只是程式小白的人參考參考。至於加速多快，我自己目前感覺是一定比 CPU 快，快多少應該就要看硬體規格了XD。