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

更新時(shí)間：2021年08月31日 17:20:27 作者：DechinPhy

本文主要講解了Python的GPU編程實(shí)例近鄰表計(jì)算，通過(guò)一個(gè)近鄰表計(jì)算的案例，給出了適用于GPU加速的計(jì)算場(chǎng)景。需要了解的朋友可以參考一下這篇文章

技術(shù)背景

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

加速場(chǎng)景

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

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

# 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)

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

$ python3 cuda_neighbor_list.py
[[0. 0. 0. 0.]
[0. 0. 1. 0.]
[0. 1. 0. 1.]
[0. 0. 1. 0.]]

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

基于Numba的GPU加速

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

# 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)))

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

$ 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

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

# 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))

這個(gè)案例中也沒(méi)有修改較多的地方，只是把一次計(jì)算的時(shí)間調(diào)整為多次計(jì)算的時(shí)間，并且忽略第一次計(jì)算過(guò)程中的即時(shí)編譯，最終輸出結(jié)果如下：

$ 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的高性能運(yùn)算，能夠提速達(dá)將近1000倍！

總結(jié)概要

對(duì)于Pythoner而言，苦其性能已久。如果能夠用一種非常Pythonic的方法來(lái)實(shí)現(xiàn)GPU的加速效果，對(duì)于Pythoner而言無(wú)疑是巨大的好消息，Numba就為我們提供了這樣的一個(gè)基礎(chǔ)功能。本文通過(guò)一個(gè)近鄰表計(jì)算的案例，給出了適用于GPU加速的計(jì)算場(chǎng)景。這種計(jì)算場(chǎng)景可并行化的程度較高，而且函數(shù)會(huì)被多次用到（在分子動(dòng)力學(xué)模擬的過(guò)程中，每一個(gè)step都會(huì)調(diào)用到這個(gè)函數(shù)），因此這是一種最典型的、最適用于GPU加速場(chǎng)景的案例。

以上就是關(guān)于Python的GPU編程實(shí)例近鄰表計(jì)算的講解的詳細(xì)內(nèi)容，更多關(guān)于Python GPU編程實(shí)例的資料請(qǐng)關(guān)注腳本之家其它相關(guān)文章！

您可能感興趣的文章: