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的原子系統，每一個原子的三維坐標都是已知的，給定一個截斷常數d0，當兩個原子之間的距離di,j<=d0時，則認為這兩個原子是相鄰近的原子。那么最終我們需要給出一個0-1矩陣Ai,j，當Ai,j=0時，表示i,j兩個原子互不相鄰，反之則相鄰。那么對于這個問題場景，我們就可以并行化的遍歷n×n的空間，直接輸出An×n大小的近鄰表。這個計算場景是一個非常適合用GPU來加速的計算，以下我們先看一下不用GPU加速時的常規實現方案：

# cuda_neighbor_list.pyfrom numba import jitfrom numba import cudaimport numpy as np@jitdef 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] = 1neighbors[j][i] = 1 return neighborsif __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近鄰表。

對于上述的近鄰表計算的場景，我們很容易的想到這個neighbor_list函數可以用GPU的函數來進行改造。對于每一個di,j我們都可以啟動一個線程去執行計算，類似于CPU上的SIMD技術，GPU中的這項優化稱為SIMT。而在Python中改造成GPU函數的方法也非常簡單，只需要把函數前的修飾器改一下，去掉函數內部的for循環，就基本完成了，比如下面這個改造的近鄰表計算的案例：

# cuda_neighbor_list.pyfrom numba import jitfrom numba import cudaimport numpy as np@jitdef 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] = 1neighbors[j][i] = 1 return [email protected] 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 > 0if __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.6401469707489014sThe time cost of GPU with cuda.jit is: 0.19208502769470215sThe result error is: 0.0

需要說明的是，這里僅僅運行了一次的程序，而jit即時編譯的加速效果在第一次的運行中其實并不明顯，甚至還有一些速度偏慢，但是在后續過程的函數調用中，就能夠起到比較大的加速效果。所以這里的運行時間并沒有太大的代表性，比較有代表性的時間對比可以看如下的案例：

# cuda_neighbor_list.pyfrom numba import jitfrom numba import cudaimport numpy as np@jitdef 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] = 1neighbors[j][i] = 1 return [email protected] 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 > 0if __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 - time1print (’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.955506563186646sThe total time cost of GPU with cuda.jit is: 0.018685102462768555s

可以看到，在GPU加速后，相比于CPU的高性能運算，能夠提速達將近1000倍！

對于Pythoner而言，苦其性能已久。如果能夠用一種非常Pythonic的方法來實現GPU的加速效果，對于Pythoner而言無疑是巨大的好消息，Numba就為我們提供了這樣的一個基礎功能。本文通過一個近鄰表計算的案例，給出了適用于GPU加速的計算場景。這種計算場景可并行化的程度較高，而且函數會被多次用到（在分子動力學模擬的過程中，每一個step都會調用到這個函數），因此這是一種最典型的、最適用于GPU加速場景的案例。

以上就是關于Python的GPU編程實例近鄰表計算的講解的詳細內容，更多關于Python GPU編程實例的資料請關注好吧啦網其它相關文章！