用CUDA寫出比Numpy更快的規約求和函數
技術背景
在前面的幾篇博客中我們介紹了在Python中使用Numba來寫CUDA程序的一些基本操作和方法，並且展示了GPU加速的實際效果。在可並行化的算法中，比如計算兩個矢量的加和，或者是在分子動力學模擬領域中的查找近鄰表等等，都是可以直接並行的算法，而且實現起來難度不大。而有一種情況是，如果我們要計算的內容的線程之間互相存在依賴，比方說最常見的，計算一個矩陣所有元素的和。
CUDA的atomic運算
正如前面所提到的問題，如何去計算一個矩陣所有元素之和呢？具體問題可以表述為：
\]
對於此類的問題，如果我們像普通的CUDA並行操作一樣，直接創建一個S變量，然後直接在線程和分塊上直接把每一個矩陣元素加到這個S變量中，那麼會出現一種情況：在線程同步時，存在衝突的線程是無法同時加和成功的，也就是說，這種情況下雖然程序不會報錯，但是得到的結果是完全錯誤的。對於此類情況，CUDA官方給出了atomic運算這樣的方案，可以保障線程之間不被干擾：
import numpy as np
from numba import cuda
from numba import vectorize
cuda.select_device(1)
@cuda.jit
def ReducedSum(arr, result):
i, j = cuda.grid(2)
cuda.atomic.add(result, 0, arr[i][j])
if __name__ == '__main__':
import time
np.random.seed(2)
data_length = 2**10
arr = np.random.random((data_length,data_length)).astype(np.float32)
print (arr)
arr_cuda = cuda.to_device(arr)
np_time = 0.0
nb_time = 0.0
for i in range(100):
res = np.array([0],dtype=np.float32)
res_cuda = cuda.to_device(res)
time0 = time.time()
ReducedSum[(data_length,data_length),(1,1)](arr_cuda,res_cuda)
time1 = time.time()
res = res_cuda.copy_to_host()[0]
time2 = time.time()
np_res = np.sum(arr)
time3 = time.time()
if i == 0:
print ('The error rate is: ', abs(np_res-res)/res)
continue
np_time += time3 - time2
nb_time += time1 - time0
print ('The time cost of numpy is: {}s'.format(np_time))
print ('The time cost of numba is: {}s'.format(nb_time))
這裡需要重點關注的就是用CUDA實現的簡單函數ReducedSum，這個函數中調用了CUDA的atomic.add方法，用這個方法直接替代系統內置的加法，就完成了所有的操作。我們將這個函數的運行時間去跟np.sum函數做一個對比，結果如下：
$ python3 cuda_reduced_sum.py
[[0.4359949 0.02592623 0.5496625 ... 0.3810055 0.6834749 0.5225032 ]
[0.62763107 0.3184925 0.5822277 ... 0.89322233 0.7845663 0.4595605 ]
[0.9666947 0.16615923 0.6931703 ... 0.29497907 0.63724256 0.06265242]
...
[0.96224505 0.36741972 0.6673239 ... 0.3115176 0.7561843 0.9396167 ]
[0.781736 0.28829736 0.38047555 ... 0.15837361 0.00392629 0.6236886 ]
[0.03247315 0.3664344 0.00369871 ... 0.0205253 0.15924706 0.8655231 ]]
The error rate is: 4.177044e-06
The time cost of numpy is: 0.027491092681884766s
The time cost of numba is: 0.01042938232421875s
在GPU的計算中，會有一定的精度損失，比如這裡的誤差率就在1e-06級別，但是運行的速度要比numpy的實現快上2倍！
總結概要
我們知道GPU加速在可並行化程度比較高的算法中，能夠發揮出比較大的作用，展示出明顯的加速效果，而對於一些線程之間存在依賴這樣的場景就不一定能夠起到很大的加速作用。CUDA官方針對此類問題，提供了atomic的內置函數解決方案，包含有求和、求最大值等常用函數。而這些函數的特點就在於，線程與線程之間需要有一個時序的依賴關係。就比如說求最大值的函數，它會涉及到不同線程之間的輪詢。經過測試，CUDA的這種atomic的方案，實現起來非常方便，性能也很樂觀，相比於自己動手實現一個不斷切割、遞歸的規約函數，還是要容易快捷的多。
版權聲明
本文首發鏈接為：//www.cnblogs.com/dechinphy/p/gpu-sum.html
作者ID：DechinPhy
更多原著文章請參考：//www.cnblogs.com/dechinphy/
打賞專用鏈接：//www.cnblogs.com/dechinphy/gallery/image/379634.html
騰訊雲專欄同步：//cloud.tencent.com/developer/column/91958
