Python GPU编程实战：利用Numba加速近邻表计算
本文深入解析如何利用Numba库在Python中实现GPU加速，重点介绍通过@cuda.jit修饰器优化近邻表计算的方法。通过对比CPU与GPU在分子动力学模拟场景下的运行时间，展示了GPU并行计算带来的显著性能提升，为Python高性能计算提供了实用方案。
技术背景
GPU加速已成为现代工业场景中广泛采用的技术手段，其核心优势在于极高的并行计算能力。在Python生态系统中，存在多种GPU并行优化方案，例如前文提及的cupy、pycuda以及numba.cuda，它们均是GPU加速领域的标志性Python库。本文重点推荐numba.cuda方案，原因在于：cupy虽然内置了大量已实现的函数，但在算法实现的灵活性上略显不足；而pycuda虽具备高灵活性和卓越性能，却要求在Python代码中嵌入C语言代码，这违背了Pythonic的编程理念。相比之下，numba.cuda方案仅需在Python函数前添加@numba.cuda.jit修饰器，即可使用纯正的Python语法实现GPU加速效果。
加速场景
首先需要明确的是，GPU并非适用于所有计算场景的加速。GPU的加速效果源于其高度并行化的特性，这意味着各个计算进程之间必须相互独立或无依赖关系。若某个进程的计算过程或结果依赖于其他进程的输出，则无法实现完全并行，只能采用串行处理。为了直观展示GPU加速的效果，本文引入分子动力学模拟领域中常见的问题：近邻表计算。
近邻表计算问题的定义如下：给定一组包含n个原子的系统，已知每个原子的三维坐标，并设定一个截断常数d0。当两个原子之间的距离di,j小于或等于d0时，判定这两个原子为近邻。最终目标是生成一个0-1矩阵Ai,j：若Ai,j=0，表示原子i与j不相邻；若Ai,j=1，则表示它们相邻。针对此场景，我们可以并行遍历n×n的空间，直接输出大小为An×n的近邻表。这是一个非常适合GPU加速的计算场景，下文首先展示不使用GPU加速时的常规实现方案。
# 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加速
针对上述近邻表计算场景，我们很容易想到利用GPU函数对neighbor_list进行改造。对于每一个距离计算di,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倍。
总结概要
对于Python开发者而言，性能瓶颈一直是痛点。若能以极其Pythonic的方式实现GPU加速，无疑是巨大的福音，Numba正是提供了这样一项基础功能。本文通过近邻表计算案例，展示了适用于GPU加速的计算场景特征：高并行化程度及函数被高频调用（如在分子动力学模拟的每一步中均需调用）。因此，这是最典型且最适合GPU加速的场景案例。
以上就是关于Python GPU编程实例近邻表计算的详细内容，更多Python GPU编程相关资料请关注本站其他文章。
免责声明：本文内容仅供参考，不代表本站观点。转载或引用请注明出处。
