关于Python的GPU编程实例近邻表计算的讲解

第关于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的原子系统，每一个原子的三维坐标都是已知的，给定一个截断常数d0，当两个原子之间的距离di,j=d0时，则认为这两个原子是相邻近的原子。那么最终我们需要给出一个0-1矩阵Ai,j，当Ai,j=0时，表示i,j两个原子互不相邻，反之则相邻。那么对于这个问题场景，我们就可以并行化的遍历n×n的空间，直接输出An×n大小的近邻表。这个计算场景是一个非常适合用GPU来加速的计算，以下我们先看一下不用GPU加速时的常规实现方案：

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

这是最常规的一种CPU上的实现方案，遍历所有的原子，计算原子间距，然后填充近邻表。这里我们还使用到了numba.jit即时编译的功能，这个功能是在执行到相关函数时再对其进行编译的方法，在矢量化的计算中有可能使用到芯片厂商所提供的SIMD的一些优化。当然，这里都是CPU层面的执行和优化，执行结果如下：

$python3cuda_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的函数来进行改造。对于每一个di,j我们都可以启动一个线程去执行计算，类似于CPU上的SIMD技术，GPU中的这项优化称为SIMT。而在Python中改造成GPU函数的方法也非常简单，只需要把函数前的修饰器改一下，去掉函数内部的for循环，就基本完成了，比如下面这个改造的近邻表计算的案例：

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

需要说明的是，当前Numba并未支持所有的numpy的函数，因此有一些计算的功能需要我们自己去手动实现一下，比如计算一个Norm的值。这里我们在输出结果中不仅统计了结果的正确性，也给出了运行的时间：

$python3cuda_neighbor_list.py

ThetimecostofCPUwithnumba.jitis:0.6401469707489014s

ThetimecostofGPUwithcuda.jitis:0.19208502769470215s

Theresulterroris:0.0

需要说明的是，这里仅仅运行了一次的程序，而jit即时编译的加速效果在第一次的运行中其实并不明显，甚至还有一些速度偏慢，但是在后续过程的函数调用中，就能够起到比较大的加速效果。所以这里的运行时间并没有太大的代表性，比较有代表性的时间对比可以看如下的案例：

#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