用CUDA寫出比Numpy更快的規約求和函式
技術背景
在前面的幾篇部落格中我們介紹了在Python中使用Numba來寫CUDA程式的一些基本操作和方法,並且展示了GPU加速的實際效果。在可並行化的演算法中,比如計算兩個向量的加和,或者是在分子動力學模擬領域中的查詢近鄰表等等,都是可以直接並行的演算法,而且實現起來難度不大。而有一種情況是,如果我們要計算的內容的執行緒之間互相存在依賴,比方說最常見的,計算一個矩陣所有元素的和。
CUDA的atomic運算
正如前面所提到的問題,如何去計算一個矩陣所有元素之和呢?具體問題可以表述為:
\[S=\sum_{i,j}A_{i,j} \]對於此類的問題,如果我們像普通的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的方案,實現起來非常方便,效能也很樂觀,相比於自己動手實現一個不斷切割、遞迴的規約函式,還是要容易快捷的多。
版權宣告
本文首發連結為:https://www.cnblogs.com/dechinphy/p/gpu-sum.html
作者ID:DechinPhy
更多原著文章請參考:https://www.cnblogs.com/dechinphy/
打賞專用連結:https://www.cnblogs.com/dechinphy/gallery/image/379634.html
“留一手”加劇內卷,“講不清”浪費時間。