Python Numba实现GPU加速

 

Python与GPU
Python作为解释型语言，.py文件一般是没法直接用GPU加速的，关于Python与GPU的结合点，以及GPU、CPU、CUDA、多核、并行、机器码…等底层实现，参考：

《Python程序如何用GPU加速：Tesla、CUDA、Numba》

《计算机底层运转机制：多核、缓存、CPU、CU、ALU、Cache》

《编译型语言与解释型语言如何在计算机底层运行》

Numba加速Python：在CPU下
Numba加速Python，可以基于CPU也可以基于GPU，关于CPU下的加速，参见《Python代码在CPU下加速：Numba入门》

Numba加速Python：在GPU下
Numba还可以使用GPU进行加速，目前支持英伟达的CUDA和AMD的ROC。

CUDA是英伟达提供给开发者的一个GPU编程框架。

CPU+主存，称为主机（Host）；GPU+显存，称为设备（Device）。

CPU无法直接读取显存数据，GPU无法直接读取主存数据；主机与设备必须通过总线（Bus）相互通讯。

安装CUDA及Numba
在进行GPU编程之前，先确认是否安装了CUDA工具箱。

# linux环境
echo $CUDA_HOME # 检查CUDA环境变量
# 查看显卡情况
nvidia-smi
1
2
3
4
可以用Anaconda里的conda命令安装CUDA：

$ conda install cudatoolkit
1
安装Numba库：

$ conda install numba
1
检查CUDA和Numba是否安装成功：

from numba import cuda
print(cuda.gpus)
1
2
一切顺利，得到<Managed Device 0>...

CUDA程序执行时会独占一张卡，如果机器上有多张GPU卡，CUDA默认会选用0号卡。

CUDA_VISIBLE_DEVICES='5'

此环境变量用来选择某张卡。

目前手上没有GPU，没法测试上述代码。

Numba模拟器
如果手头暂时没有GPU设备，Numba提供了一个模拟器，供用户学习和调试，只需要在命令行里添加一个环境变量。

Windows：

SET NUMBA_ENABLE_CUDASIM = 1
1
GPU程序与CPU程序的区别
CPU顺序执行

CPU程序，顺序执行流程：

初始化；
CPU计算；
得到计算结果；
GPU执行
在CUDA编程中，CPU和主存被称为主机（Host），GPU被称为设备（Device）。

GPU计算流程：

初始化，并将必要的数据拷贝到GPU设备的显存上；
CPU调用GPU函数，启动GPU多个核心同时进行计算；
CPU与GPU异步计算；
将GPU计算结果拷贝回主机端，得到计算结果；
GPU程序gpu_print.py解析
from numba import cuda

def cpu_print():
    print("print by cpu.")

@cuda.jit
def gpu_print():
    # GPU核函数
    print("print by gpu.")

def main():
    gpu_print[1, 2]()
    cuda.synchronize()
    cpu_print()

if __name__ == "__main__":
    main()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
使用CUDA_VISIBLE_DEVICES='0' python gpu_print.py执行上述代码，得到的结果是：

print by gpu.
print by gpu.
print by cpu.
1
2
3
解析上述代码，与传统的CPU运行Python不同：

使用from numba import cuda引入cuda库
在GPU函数上添加@cuda.jit装饰符，表示该函数是一个在GPU设备上运行的函数，GPU函数又被称为==核函数==。
主函数调用GPU核函数时，需要添加如[1, 2]这样的执行配置，这个配置是在告知GPU以多大的并行粒度同时进行计算。gpu_print[1, 2]()表示同时开启2个线程并行地执行gpu_print函数，函数将被并行地执行2次。下文会深入探讨如何设置执行配置。
GPU核函数的启动方式是异步的：启动GPU函数后，CPU不会等待GPU函数执行完毕才执行下一行代码。必要时，需要调用cuda.synchronize()，告知CPU等待GPU执行完核函数后，再进行CPU端后续计算。这个过程被称为同步，也就是GPU执行流程图中的红线部分。如果不调用cuda.synchronize()函数，执行结果也将改变，print by cpu将先被打印。虽然GPU函数在前，但是程序并没有等待GPU函数执行完，而是继续执行后面的cpu_print函数，由于CPU调用GPU有一定的延迟，反而后面的cpu_print先被执行，因此cpu_print的结果先被打印了出来。
Thread层次结构
CUDA将核函数（GPU函数）所定义的运算称为线程（Thread），多个线程组成一个块（Block），多个块组成网格（Grid）。

这样的方式，一个grid可以定义成千上万个线程，也就解决了并行执行上万次操作的问题。

例如，把前面的程序改为并行执行8次：可以用2个block，每个block中有4个thread。

原来的代码可以改为gpu_print[2,4](),方括号内第一个数字表示整个grid有多少个blcok，第二个数字表示，每个block有多少个thread。

线程（Thread）是一个编程上的软件概念。从硬件角度，thread运行在一个CUDA核心上，多个thread组成的blcok运行在Streaming Multiprocessor上，多个block组成的grid运行在一个GPU显卡上。

CUDA提供了一系列内置变量，以记录thread核block的大小以及索引下标。

以[2,4]这样的配置为例：blockDim.x变量表示block的大小是4，即每个block有4个thread，threadIdx.x变量是一个从0到blockDim.x-1（4-1=3）的索引下标，记录这是第几个thread；gridDim.x变量表示grid的大小是2，即表示grid的大小是2，即每个grid有2个block，blockIdx.x变量是一个从0到gridDim.x-1(2-1=1)的索引下标，记录这是第几个block。

某个thread在整个grid中的编号位置为：threadIdx.x + blockIdx.x * blockDim.x

更新上述代码：

from numba import cuda

def cpu_print(N):
    for i in range(0, N):
        print(i)

@cuda.jit
def gpu_print(N):
    idx = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x 
    if (idx < N):
        print(idx)

def main():
    print("gpu print:")
    gpu_print[2, 4](8)
    cuda.synchronize()
    print("cpu print:")
    cpu_print(8)

if __name__ == "__main__":
    main()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
执行结果为：

gpu print:
0
1
2
3
4
5
6
7
cpu print:
0
1
2
3
4
5
6
7
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
这里GPU函数在每个CUDA thread中打印了当前thread的编号，起到了CPU函数for循环同样的作用。

因为for循环中的计算内容互相不依赖，也就是说，某次循环只是专心做自己的事情，循环第i次不影响循环第j次的计算，所以这样互相不依赖的for循环非常适合放到CUDA thread里做并行计算。

在实际使用中，我们一般将CPU代码中互相不依赖的的for循环适当替换成CUDA代码。

这份代码打印了8个数字，核函数有一个参数N，N = 8，假如我们只想打印5个数字呢？当前的执行配置共2 * 4 = 8个线程，线程数8与要执行的次数5不匹配，不过我们已经在代码里写好了if (idx < N)的判断语句，判断会帮我们过滤不需要的计算。我们只需要把N = 5传递给gpu_print函数中就好，CUDA仍然会启动8个thread，但是大于等于N的thread不进行计算。注意，当线程数与计算次数不一致时，一定要使用这样的判断语句，以保证某个线程的计算不会影响其他线程的数据。

Block大小设置
不同的执行配置会影响GPU程序的速度，一般需要多次调试才能找到较好的执行配置。实际编程中，执行配置[gridDim, blcokDim]应参考下面的方法：

block运行在SM上，不同硬件架构（Turing、Volta、Pascal…）的CUDA核心数不同，一般需要根据当前硬件来配置block的大小blockDim（执行配置中的第二个参数）。一个block中的thread数最好是32、128、256的倍数。限于硬件设置，当前block大小不能超过1024.
grid的大小gridDim（执行配置中第一个参数），即一个grid中block的个数可以由总次数N除以blockDim，向上取整。
举例说明，我们想并行启动1000个thread，可以将blockDim设置为128，1000/128=7.8，向上取整为8。得到执行配置为gpuWork[8,128](),CUDA实际共启用了8*128=1024个thread，实际只使用了前1000个thread进行计算，多余的24个thread不进行计算。

blockDim：block中thread的数量；一个block中threadIdx最大不超过blockDim；
gridDim：grid中block的数量；一个grid中blockIdx最大不超过gridDim；
上述讨论中，block和grid都是一维的，实际可以是二维三维。

内存分配
GPU计算时直接从显存中读取数据，因此每当计算时要将数据从主存拷贝到显存上，用CUDA的术语来说就是要把数据从主机端拷贝到设备端。

CUDA强大之处在于它能自动将数据从主机和设备间相互拷贝，不需要程序员在代码中写明。

这种方法对编程者来说非常方便，不必对原有的CPU代码做大量改动。

举例说明：

from numba import cuda
import numpy as np
import math
from time import time
# 核函数
@cuda.jit
def gpu_add(a, b, result, n):
    idx = cuda.threadIdx.x + cuda.blockDim.x * cuda.blockIdx.x
    if idx < n:
        result[idx] = a[idx] + b[idx]

def main():
    # 数据初始化
    n = 20000000
    x = np.arange(n).astype(np.int32)
    y = 2 * x

    gpu_result = np.zeros(n)
    cpu_result = np.zeros(n)

    threads_per_block = 1024
    blocks_per_grid = math.ceil(n / threads_per_block)
    start = time()
    # 执行函数
    gpu_add[blocks_per_grid, threads_per_block](x, y, gpu_result, n)
    cuda.synchronize()
    print("gpu vector add time " + str(time() - start))
    start = time()
    cpu_result = np.add(x, y)
    print("cpu vector add time " + str(time() - start))

    if (np.array_equal(cpu_result, gpu_result)):
        print("result correct")

if __name__ == "__main__":
    main()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
上述代码运行后，GPU代码要比CPU代码慢10多倍。主要原因有：

向量加法的这个计算比较简单，CPU的numpy已经优化到了极致，无法突出GPU的优势，我们要解决实际问题往往比这个复杂得多，当解决复杂问题时，优化后的GPU代码将远快于CPU代码。
这份代码使用CUDA默认的统一内存管理机制，没有对数据的拷贝做优化。CUDA的统一内存系统是当GPU运行到某块数据发现不在设备端时，再去主机端中将数据拷贝过来，当执行完核函数后，又将所有的内存拷贝回主存。在上面的代码中，输入的两个向量是只读的，没必要再拷贝回主存。
这份代码没有做流水线优化。CUDA并非同时计算2千万个数据，一般分批流水线工作：一边对2000万中的某批数据进行计算，一边将下一批数据从主存拷贝过来。计算占用的是CUDA核心，数据拷贝占用的是总线，所需资源不同，互相不存在竞争关系。这种机制被称为流水线。
原因2中本该程序员动脑思考的问题交给了CUDA解决，增加了时间开销，所以CUDA非常方便的统一内存模型缺点是计算速度慢。针对原因2，我们可以继续优化这个程序，告知GPU哪些数据需要拷贝到设备，哪些需要拷贝回主机。

from numba import cuda # 从numba调用cuda
import numpy as np
import math
from time import time

@cuda.jit
def gpu_add(a, b, result, n):
    idx = cuda.threadIdx.x + cuda.blockDim.x * cuda.blockIdx.x
    if idx < n :
        result[idx] = a[idx] + b[idx]

def main():
    n = 20000000
    x = np.arange(n).astype(np.int32)
    y = 2 * x

    # 拷贝数据到设备端
    x_device = cuda.to_device(x)
    y_device = cuda.to_device(y)
    # 在显卡设备上初始化一块用于存放GPU计算结果的空间
    gpu_result = cuda.device_array(n)
    cpu_result = np.empty(n)

    threads_per_block = 1024
    blocks_per_grid = math.ceil(n / threads_per_block)
    start = time()
    gpu_add[blocks_per_grid, threads_per_block](x_device, y_device, gpu_result, n)
    cuda.synchronize()
    print("gpu vector add time " + str(time() - start))
    start = time()
    cpu_result = np.add(x, y)
    print("cpu vector add time " + str(time() - start))

    if (np.array_equal(cpu_result, gpu_result.copy_to_host())):
        print("result correct!")

if __name__ == "__main__":
    main()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
此时，GPU速度就比CPU快了。

常用内存分配函数：

cuda.device_array()：在设备上分配一个空向量，类似于numpy.empty()；
cuda.to_device()：将主机的数据拷贝到设备；
cuda.copy_to_host()：将设备的数据拷贝回主机；
总结
Python的Numba库可以调用CUDA进行GPU编程；

CPU端被称为主机，GPU端被称为设备；

运行在GPU上的函数被称为核函数；调用核函数需要有执行配置，用来告诉CUDA以多大的并行力度来计算；

使用GPU编程需要合理的将数据在主机和设备间互相拷贝。

上图展示了CUDA编程的基本流程：

初始化，并将必要的数据从CPU主存拷贝到GPU显存上；
使用某个执行配置，以一定的并行粒度调用CUDA核函数；
CPU和GPU异步计算；
将GPU计算结果拷贝回主机。
References
GPU加速02:超详细Python Cuda零基础入门教程，没有显卡也能学！ 原文讲的非常好，本文只是抄写学习
————————————————
版权声明：本文为CSDN博主「quantLearner」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/The_Time_Runner/article/details/103364257