關(guān)于Python的GPU編程實(shí)例近鄰表計算的講解

第關(guān)于Python的GPU編程實(shí)例近鄰表計算的講解GPU加速是現(xiàn)代工業(yè)各種場景中非常常用的一種技術(shù)，這得益于GPU計算的高度并行化。在Python中存在有多種GPU并行優(yōu)化的解決方案，包括之前的博客中提到的cupy、pycuda和numba.cuda，都是GPU加速的標(biāo)志性Python庫。這里我們重點(diǎn)推numba.cuda這一解決方案，因?yàn)閏upy的優(yōu)勢在于實(shí)現(xiàn)好了的眾多的函數(shù)，在算法實(shí)現(xiàn)的靈活性上還比較欠缺；而pycuda雖然提供了很好的靈活性和相當(dāng)高的性能，但是這要求我們必須在Python的代碼中插入C代碼，這顯然是非常不Pythonic的解決方案。因此我們可以選擇numba.cuda這一解決方案，只要在Python函數(shù)前方加一個numba.cuda.jit的修飾器，就可以在Python中用最Python的編程語法，實(shí)現(xiàn)GPU的加速效果。

加速場景

我們需要先了解的是，GPU在什么樣的計算場景下能夠?qū)崿F(xiàn)加速的效果，很顯然的是，并不是所有的計算過程都能在GPU上表現(xiàn)出加速的效果。前面說道，GPU的加速作用，是源自于高度的并行化，所謂的并行，就要求進(jìn)程之前互不干擾或者依賴。如果說一個進(jìn)程的計算過程或者結(jié)果，依賴于另一個進(jìn)程中的計算結(jié)果，那么就無法實(shí)現(xiàn)完全的并行，只能使用串行的技術(shù)。這里為了展示GPU加速的效果，我們就引入一個在分子動力學(xué)模擬領(lǐng)域中常見的問題：近鄰表的計算。

近鄰表計算的問題是這樣描述的：給定一堆數(shù)量為n的原子系統(tǒng)，每一個原子的三維坐標(biāo)都是已知的，給定一個截斷常數(shù)d0，當(dāng)兩個原子之間的距離di,j=d0時，則認(rèn)為這兩個原子是相鄰近的原子。那么最終我們需要給出一個0-1矩陣Ai,j，當(dāng)Ai,j=0時，表示i,j兩個原子互不相鄰，反之則相鄰。那么對于這個問題場景，我們就可以并行化的遍歷n×n的空間，直接輸出An×n大小的近鄰表。這個計算場景是一個非常適合用GPU來加速的計算，以下我們先看一下不用GPU加速時的常規(guī)實(shí)現(xiàn)方案：

#cuda_neighbor_list.py

fromnumbaimportjit

fromnumbaimportcuda

importnumpyasnp

defneighbor_list(crd,neighbors,data_length,cutoff):

"""CPUbasedneighborlistcalculation.

foriinrange(data_length):

forjinrange(i+1,data_length):

ifnp.linalg.norm(crd[i]-crd[j])=cutoff:

neighbors[i][j]=1

neighbors[j][i]=1

returnneighbors

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)

這是最常規(guī)的一種CPU上的實(shí)現(xiàn)方案，遍歷所有的原子，計算原子間距，然后填充近鄰表。這里我們還使用到了numba.jit即時編譯的功能，這個功能是在執(zhí)行到相關(guān)函數(shù)時再對其進(jìn)行編譯的方法，在矢量化的計算中有可能使用到芯片廠商所提供的SIMD的一些優(yōu)化。當(dāng)然，這里都是CPU層面的執(zhí)行和優(yōu)化，執(zhí)行結(jié)果如下：

$python3cuda_neighbor_list.py

[[0.0.0.0.]

[0.0.1.0.]

[0.1.0.1.]

[0.0.1.0.]]

這個輸出的結(jié)果就是一個0-1近鄰表。

基于Numba的GPU加速

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

#cuda_neighbor_list.py

fromnumbaimportjit

fromnumbaimportcuda

importnumpyasnp

defneighbor_list(crd,neighbors,data_length,cutoff):

"""CPUbasedneighborlistcalculation.

foriinrange(data_length):

forjinrange(i+1,data_length):

ifnp.linalg.norm(crd[i]-crd[j])=cutoff:

neighbors[i][j]=1

neighbors[j][i]=1

returnneighbors

@cuda.jit

defcuda_neighbor_list(crd,neighbors,cutoff):

"""GPUbasedneighborlistcalculation.

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]anddis0

if__name__=='__main__':

importtime

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('ThetimecostofCPUwithnumba.jitis:{}s'.format(\

time1-time0))

print('ThetimecostofGPUwithcuda.jitis:{}s'.format(\

time2-time1))

print('Theresulterroris:{}'.format(np.sum(adjacent_c-\

adjacent_g)))

需要說明的是，當(dāng)前Numba并未支持所有的numpy的函數(shù)，因此有一些計算的功能需要我們自己去手動實(shí)現(xiàn)一下，比如計算一個Norm的值。這里我們在輸出結(jié)果中不僅統(tǒng)計了結(jié)果的正確性，也給出了運(yùn)行的時間：

$python3cuda_neighbor_list.py

ThetimecostofCPUwithnumba.jitis:0.6401469707489014s

ThetimecostofGPUwithcuda.jitis:0.19208502769470215s

Theresulterroris:0.0

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

#cuda_neighbor_list.py

fromnumbaimportjit

fromnumbaimportcuda

importnumpyasnp

defneighbor_list(crd,neighbors,data_length,cutoff):

"""CPUbasedneighborlistcalculation.

foriinrange(data_length):

forjinrange(i+1,data_length):

ifnp.linalg.norm(crd[i]-crd[j])=cutoff:

neighbors[i][j]=1

neighbors[j][i]=1

returnneighbors

@cuda.jit

defcuda_neighbor_list(crd,neighbors,cutoff):

"""GPUbasedneighborlistcalculation.

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]anddis0

if__name__=='__main__':

importtime

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_inrange(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('ThetotaltimecostofCPUwithnumba.jitis:{}s'.format(\

time_c))

print('ThetotaltimecostofGPUwithcuda.jitis:{}s'.format(\

ti