Python的GPU程式設計例項——近鄰表計算
技術背景
GPU加速是現代工業各種場景中非常常用的一種技術，這得益於GPU計算的高度並行化。在Python中存在有多種GPU並行優化的解決方案，包括之前的部落格中提到的cupy、pycuda和numba.cuda，都是GPU加速的標誌性Python庫。這裡我們重點推numba.cuda這一解決方案，因為cupy的優勢在於實現好了的眾多的函式，在演算法實現的靈活性上還比較欠缺；而pycuda雖然提供了很好的靈活性和相當高的效能，但是這要求我們必須在Python的程式碼中插入C程式碼，這顯然是非常不Pythonic的解決方案。因此我們可以選擇numba.cuda這一解決方案，只要在Python函式前方加一個numba.cuda.jit的修飾器，就可以在Python中用最Python的程式設計語法，實現GPU的加速效果。
加速場景
我們需要先了解的是，GPU在什麼樣的計算場景下能夠實現加速的效果，很顯然的是，並不是所有的計算過程都能在GPU上表現出加速的效果。前面說道，GPU的加速作用，是源自於高度的並行化，所謂的並行，就要求程序之前互不干擾或者依賴。如果說一個程序的計算過程或者結果，依賴於另一個程序中的計算結果，那麼就無法實現完全的並行，只能使用序列的技術。這裡為了展示GPU加速的效果，我們就引入一個在分子動力學模擬領域中常見的問題：近鄰表的計算。
近鄰表計算的問題是這樣描述的：給定一堆數量為n的原子系統，每一個原子的三維座標都是已知的，給定一個截斷常數\(d_0\)，當兩個原子之間的距離\(d_{i,j}<=d_0\)
# cuda_neighbor_list.py from numba import jit from numba import cuda import numpy as np @jit def neighbor_list(crd, neighbors, data_length, cutoff): """CPU based neighbor list calculation. """ for i in range(data_length): for j in range(i+1, data_length): if np.linalg.norm(crd[i]-crd[j]) <= cutoff: neighbors[i][j] = 1 neighbors[j][i] = 1 return neighbors if __name__ == '__main__': np.random.seed(1) atoms = 2**2 cutoff = 0.5 crd = np.random.random((atoms,3)) adjacent = np.zeros((atoms, atoms)) adjacent = neighbor_list(crd, adjacent, atoms, cutoff) print (adjacent)
這是最常規的一種CPU上的實現方案，遍歷所有的原子，計算原子間距，然後填充近鄰表。這裡我們還使用到了numba.jit即時編譯的功能，這個功能是在執行到相關函式時再對其進行編譯的方法，在向量化的計算中有可能使用到晶片廠商所提供的SIMD的一些優化。當然，這裡都是CPU層面的執行和優化，執行結果如下：
$ python3 cuda_neighbor_list.py
[[0. 0. 0. 0.]
[0. 0. 1. 0.]
[0. 1. 0. 1.]
[0. 0. 1. 0.]]
這個輸出的結果就是一個0-1近鄰表。
基於Numba的GPU加速
對於上述的近鄰表計算的場景，我們很容易的想到這個neighbor_list函式可以用GPU的函式來進行改造。對於每一個\(d_{i,j}\)我們都可以啟動一個執行緒去執行計算，類似於CPU上的SIMD技術，GPU中的這項優化稱為SIMT。而在Python中改造成GPU函式的方法也非常簡單，只需要把函式前的修飾器改一下，去掉函式內部的for迴圈，就基本完成了，比如下面這個改造的近鄰表計算的案例：
# cuda_neighbor_list.py
from numba import jit
from numba import cuda
import numpy as np
@jit
def neighbor_list(crd, neighbors, data_length, cutoff):
"""CPU based neighbor list calculation.
"""
for i in range(data_length):
for j in range(i+1, data_length):
if np.linalg.norm(crd[i]-crd[j]) <= cutoff:
neighbors[i][j] = 1
neighbors[j][i] = 1
return neighbors
@cuda.jit
def cuda_neighbor_list(crd, neighbors, cutoff):
"""GPU based neighbor list calculation.
"""
i, j = cuda.grid(2)
dis = ((crd[i][0]-crd[j][0])**2+\
(crd[i][1]-crd[j][1])**2+\
(crd[i][2]-crd[j][2])**2)**0.5
neighbors[i][j] = dis <= cutoff[0] and dis > 0
if __name__ == '__main__':
import time
np.random.seed(1)
atoms = 2**5
cutoff = 0.5
cutoff_cuda = cuda.to_device(np.array([cutoff]).astype(np.float32))
crd = np.random.random((atoms,3)).astype(np.float32)
crd_cuda = cuda.to_device(crd)
adjacent = np.zeros((atoms, atoms)).astype(np.float32)
adjacent_cuda = cuda.to_device(adjacent)
time0 = time.time()
adjacent_c = neighbor_list(crd, adjacent, atoms, cutoff)
time1 = time.time()
cuda_neighbor_list[(atoms, atoms), (1, 1)](crd_cuda,
adjacent_cuda,
cutoff_cuda)
time2 = time.time()
adjacent_g = adjacent_cuda.copy_to_host()
print ('The time cost of CPU with numba.jit is: {}s'.format(\
time1-time0))
print ('The time cost of GPU with cuda.jit is: {}s'.format(\
time2-time1))
print ('The result error is: {}'.format(np.sum(adjacent_c-\
adjacent_g)))
需要說明的是，當前Numba並未支援所有的numpy的函式，因此有一些計算的功能需要我們自己去手動實現一下，比如計算一個Norm的值。這裡我們在輸出結果中不僅統計了結果的正確性，也給出了執行的時間：
$ python3 cuda_neighbor_list.py
The time cost of CPU with numba.jit is: 0.6401469707489014s
The time cost of GPU with cuda.jit is: 0.19208502769470215s
The result error is: 0.0
需要說明的是，這裡僅僅運行了一次的程式，而jit即時編譯的加速效果在第一次的執行中其實並不明顯，甚至還有一些速度偏慢，但是在後續過程的函式呼叫中，就能夠起到比較大的加速效果。所以這裡的執行時間並沒有太大的代表性，比較有代表性的時間對比可以看如下的案例：
# cuda_neighbor_list.py
from numba import jit
from numba import cuda
import numpy as np
@jit
def neighbor_list(crd, neighbors, data_length, cutoff):
"""CPU based neighbor list calculation.
"""
for i in range(data_length):
for j in range(i+1, data_length):
if np.linalg.norm(crd[i]-crd[j]) <= cutoff:
neighbors[i][j] = 1
neighbors[j][i] = 1
return neighbors
@cuda.jit
def cuda_neighbor_list(crd, neighbors, cutoff):
"""GPU based neighbor list calculation.
"""
i, j = cuda.grid(2)
dis = ((crd[i][0]-crd[j][0])**2+\
(crd[i][1]-crd[j][1])**2+\
(crd[i][2]-crd[j][2])**2)**0.5
neighbors[i][j] = dis <= cutoff[0] and dis > 0
if __name__ == '__main__':
import time
np.random.seed(1)
atoms = 2**10
cutoff = 0.5
cutoff_cuda = cuda.to_device(np.array([cutoff]).astype(np.float32))
crd = np.random.random((atoms,3)).astype(np.float32)
crd_cuda = cuda.to_device(crd)
adjacent = np.zeros((atoms, atoms)).astype(np.float32)
adjacent_cuda = cuda.to_device(adjacent)
time_c = 0.0
time_g = 0.0
for _ in range(100):
time0 = time.time()
adjacent_c = neighbor_list(crd, adjacent, atoms, cutoff)
time1 = time.time()
cuda_neighbor_list[(atoms, atoms), (1, 1)](crd_cuda,
adjacent_cuda,
cutoff_cuda)
time2 = time.time()
if _ != 0:
time_c += time1 - time0
time_g += time2 - time1
print ('The total time cost of CPU with numba.jit is: {}s'.format(\
time_c))
print ('The total time cost of GPU with cuda.jit is: {}s'.format(\
time_g))
這個案例中也沒有修改較多的地方，只是把一次計算的時間調整為多次計算的時間，並且忽略第一次計算過程中的即時編譯，最終輸出結果如下：
$ python3 cuda_neighbor_list.py
The total time cost of CPU with numba.jit is: 14.955506563186646s
The total time cost of GPU with cuda.jit is: 0.018685102462768555s
可以看到，在GPU加速後，相比於CPU的高效能運算，能夠提速達將近1000倍！
總結概要
對於Pythoner而言，苦其效能已久。如果能夠用一種非常Pythonic的方法來實現GPU的加速效果，對於Pythoner而言無疑是巨大的好訊息，Numba就為我們提供了這樣的一個基礎功能。本文通過一個近鄰表計算的案例，給出了適用於GPU加速的計算場景。這種計算場景可並行化的程度較高，而且函式會被多次用到（在分子動力學模擬的過程中，每一個step都會呼叫到這個函式），因此這是一種最典型的、最適用於GPU加速場景的案例。
版權宣告
本文首發連結為：https://www.cnblogs.com/dechinphy/p/cuda-neighbor.html
作者ID：DechinPhy
更多原著文章請參考：https://www.cnblogs.com/dechinphy/
打賞專用連結：https://www.cnblogs.com/dechinphy/gallery/image/379634.html
“留一手”加劇內卷，“講不清”浪費時間。