爱站程序员基地技术背景
之前写过一篇讲述如何使用pycuda来在Python上写CUDA程序的博客。这个方案的特点在于完全遵循了CUDA程序的写法，只是支持了一些常用函数的接口，如果你需要自己写CUDA算子，那么就只能使用非常不Pythonic的写法。还有一种常见的方法是用cupy来替代numpy，相当于一个GPU版本的numpy。那么本文要讲述的是用numba自带的装饰器，来写一个非常Pythonic的CUDA程序。
CUDA的线程与块
GPU从计算逻辑来讲，可以认为是一个高并行度的计算阵列，我们可以想象成一个二维的像围棋棋盘一样的网格，每一个格子都可以执行一个单独的任务，并且所有的格子可以同时执行计算任务，这就是GPU加速的来源。那么刚才所提到的棋盘，每一列都认为是一个线程，并有自己的线程编号；每一行都是一个块，有自己的块编号。我们可以通过一些简单的程序来理解这其中的逻辑：
用GPU打印线程编号
# numba_cuda_test.pyfrom numba import cuda@cuda.jitdef gpu():print (\'threadIdx:\', cuda.threadIdx.x)if __name__ == \'__main__\':gpu[2,4]()
threadIdx: 0threadIdx: 1threadIdx: 2threadIdx: 3threadIdx: 0threadIdx: 1threadIdx: 2threadIdx: 3
用GPU打印块编号
# numba_cuda_test.pyfrom numba import cuda@cuda.jitdef gpu():print (\'blockIdx:\', cuda.blockIdx.x)if __name__ == \'__main__\':gpu[2,4]()
blockIdx: 0blockIdx: 0blockIdx: 0blockIdx: 0blockIdx: 1blockIdx: 1blockIdx: 1blockIdx: 1
用GPU打印块的维度
# numba_cuda_test.pyfrom numba import cuda@cuda.jitdef gpu():print (\'blockDim:\', cuda.blockDim.x)if __name__ == \'__main__\':gpu[2,4]()
blockDim: 4blockDim: 4blockDim: 4blockDim: 4blockDim: 4blockDim: 4blockDim: 4blockDim: 4
用GPU打印线程的维度
# numba_cuda_test.pyfrom numba import cuda@cuda.jitdef gpu():print (\'gridDim:\', cuda.gridDim.x)if __name__ == \'__main__\':gpu[2,4]()
gridDim: 2gridDim: 2gridDim: 2gridDim: 2gridDim: 2gridDim: 2gridDim: 2gridDim: 2
总结
我们可以用如下的一张图来总结刚才提到的GPU网格的概念，在上面的测试案例中，我们在GPU上划分一块2*4大小的阵列用于我们自己的计算，每一行都是一个块，每一列都是一个线程，所有的网格是同时执行计算的内容的（如果没有逻辑上的依赖的话）。
GPU所支持的最大并行度
我们可以用几个简单的程序来测试一下GPU的并行度，因为每一个GPU上的网格都可以独立的执行一个任务，因此我们认为可以分配多少个网格，就有多大的并行度。本机的最大并行应该是在$240$，因此假设我们给GPU分配$250$大小的网格，程序就会报错：
# numba_cuda_test.pyfrom numba import cuda@cuda.jitdef gpu():passif __name__ == \'__main__\':gpu[2**50,1]()print (\'Running Success!\')
运行结果如下：
Traceback (most recent call last):File "numba_cuda_test.py", line 10, in <module>gpu[2**50,1]()File "/home/dechin/.local/lib/python3.7/site-packages/numba/cuda/compiler.py", line 822, in __call__self.stream, self.sharedmem)File "/home/dechin/.local/lib/python3.7/site-packages/numba/cuda/compiler.py", line 966, in callkernel.launch(args, griddim, blockdim, stream, sharedmem)File "/home/dechin/.local/lib/python3.7/site-packages/numba/cuda/compiler.py", line 699, in launchcooperative=self.cooperative)File "/home/dechin/.local/lib/python3.7/site-packages/numba/cuda/cudadrv/driver.py", line 2100, in launch_kernelNone)File "/home/dechin/.local/lib/python3.7/site-packages/numba/cuda/cudadrv/driver.py", line 300, in safe_cuda_api_callself._check_error(fname, retcode)File "/home/dechin/.local/lib/python3.7/site-packages/numba/cuda/cudadrv/driver.py", line 335, in _check_errorraise CudaAPIError(retcode, msg)numba.cuda.cudadrv.driver.CudaAPIError: [1] Call to cuLaunchKernel results in CUDA_ERROR_INVALID_VALUE
而如果我们分配一个额定大小之内的网格，程序就可以正常的运行：
# numba_cuda_test.pyfrom numba import cuda@cuda.jitdef gpu():passif __name__ == \'__main__\':gpu[2**30,1]()print (\'Running Success!\')
这里加了一个打印输出：
Running Success!
需要注意的是，两个维度上的可分配大小是不一致的，比如本机的上限是分配230*210大小的空间用于计算：
# numba_cuda_test.pyfrom numba import cuda@cuda.jitdef gpu():passif __name__ == \'__main__\':gpu[2**30,2**10]()print (\'Running Success!\')
同样的，只要在允许的范围内都是可以执行成功的：
Running Success!
如果在本机上有多块GPU的话，还可以通过
select_device
的指令来选择执行指令的GPU编号：
# numba_cuda_test.pyfrom numba import cudacuda.select_device(1)import time@cuda.jitdef gpu():passif __name__ == \'__main__\':gpu[2**30,2**10]()print (\'Running Success!\')
如果两块GPU的可分配空间一致的话，就可以运行成功：
Running Success!
GPU的加速效果
前面我们经常提到一个词叫GPU加速，GPU之所以能够实现加速的效果，正源自于GPU本身的高度并行性。这里我们直接用一个数组求和的案例来说明GPU的加速效果，这个案例需要得到的结果是$b_j=a_j+b_j$，将求和后的值赋值在其中的一个输入数组之上，以节省一些内存空间。当然，如果这个数组还有其他的用途的话，是不能这样操作的。具体代码如下：
# gpu_add.pyfrom numba import cudacuda.select_device(1)import numpy as npimport time@cuda.jitdef gpu(a,b,DATA_LENGHTH):idx = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.xif idx < DATA_LENGHTH:b[idx] += a[idx]if __name__ == \'__main__\':np.random.seed(1)DATA_EXP_LENGTH = 20DATA_DIMENSION = 2**DATA_EXP_LENGTHnp_time = 0.0nb_time = 0.0for i in range(100):a = np.random.randn(DATA_DIMENSION).astype(np.float32)b = np.random.randn(DATA_DIMENSION).astype(np.float32)a_cuda = cuda.to_device(a)b_cuda = cuda.to_device(b)time0 = time.time()gpu[DATA_DIMENSION,4](a_cuda,b_cuda,DATA_DIMENSION)time1 = time.time()c = b_cuda.copy_to_host()time2 = time.time()d = np.add(a,b)time3 = time.time()if i == 0:print (\'The error between numba and numpy is: \', sum(c-d))continuenp_time += time3 - time2nb_time += time1 - time0print (\'The time cost of numba is: {}s\'.format(nb_time))print (\'The time cost of numpy is: {}s\'.format(np_time))
需要注意的是，基于Numba实现的Python的GPU加速程序，采用的jit即时编译的模式，也就是说，在运行调用到相关函数时，才会对其进行编译优化。换句话说，第一次执行这一条指令的时候，事实上达不到加速的效果，因为这个运行的时间包含了较长的一段编译时间。但是从第二次运行调用开始，就不需要重新编译，这时候GPU加速的效果就体现出来了，运行结果如下：
$ python3 gpu_add.py The error between numba and numpy is: 0.0The time cost of numba is: 0.018711328506469727sThe time cost of numpy is: 0.09502553939819336s
可以看到，即使是相比于Python中优化程度十分强大的的Numpy实现，我们自己写的GPU加速的程序也能够达到5倍的加速效果（在前面一篇博客中，针对于特殊计算场景，加速效果可达1000倍以上），而且可定制化程度非常之高。
总结概要
本文针对于Python中使用Numba的GPU加速程序的一些基本概念和实现的方法，比如GPU中的线程和模块的概念，以及给出了一个矢量加法的代码案例，进一步说明了GPU加速的效果。需要注意的是，由于Python中的Numba实现是一种即时编译的技术，因此第一次运算时的时间会明显较长，所以我们一般说GPU加速是指从第二步开始的运行时间。对于一些工业和学界常见的场景，比如分子动力学模拟中的系统演化，或者是深度学习与量子计算中的参数优化，都是相同维度参数多步运算的一个过程，非常适合使用即时编译的技术，配合以GPU高度并行化的加速效果，能够在实际工业和学术界的各种场景下发挥巨大的作用。
版权声明
本文首发链接为：https://www.cnblogs.com/dechinphy/p/nbc.html
作者ID：DechinPhy
更多原著文章请参考：https://www.cnblogs.com/dechinphy/
打赏专用链接：https://www.cnblogs.com/dechinphy/gallery/image/379634.html
腾讯云专栏同步：https://cloud.tencent.com/developer/column/91958
