文章目录

• • 一、简介
  • 二、numba

使用Python写CUDA程序有两种方式：

• Numba
• PyCUDA

numbapro现在已经不推荐使用了，功能被拆分并分别被集成到accelerate和Numba了。

一、简介

1. numba
   Numba通过及时编译机制（JIT）优化Python代码，Numba可以针对本机的硬件环境进行优化，同时支持CPU和GPU的优化，并且可以和Numpy集成，使Python代码可以在GPU上运行，只需在函数上方加上相关的指令标记，

2. PyCUDA
   PyCUDA的内核函数（kernel）其实就是使用C/C++编写的，通过动态编译为GPU微码，Python代码与GPU代码进行交互。

3. 对比
   numba使用一些指令标记某些函数进行加速（也可以使用Python编写内核函数），这一点类似于OpenACC，而PyCUDA需要自己写kernel，在运行时进行编译，底层是基于C/C++实现的。通过测试，这两种方式的加速比基本差不多。但是，numba更像是一个黑盒，不知道内部到底做了什么，而PyCUDA就显得很直观。因此，这两种方式具有不同的应用：

• 如果只是为了加速自己的算法而不关心CUDA编程，那么直接使用numba会更好。

• 如果为了学习、研究CUDA编程或者实验某一个算法在CUDA下的可行性，那么使用PyCUDA。

• 如果写的程序将来要移植到C/C++，那么就一定要使用PyCUDA了，因为使用PyCUDA写的kernel本身就是用CUDA C/C++写的。

二、numba

C/C++ 这类编译性语言最大的好处就是其编译过程是发生在运行之前的，源代码在调用前被编译器转换为可执行机器码，这样就节约了大量的时间。而 python 作为一种解释性语言，没法做到一次编译，后续可以直接运行，每次运行的时候都要重新将源代码通过解释器转化为机器码。这样一个好处就是非常容易 debug（ 这里要再次感叹一下 python 真不愧是新手友好型语言~）， 当然，这个问题自然也是有尝试解决的办法，一个很重要的技术就是 JIT （Just-in-time compilation）：JIT 即时编译技术是在运行时（runtime）将调用的函数或程序段编译成机器码载入内存，以加快程序的执行。说白了，就是在第一遍执行一段代码前，先执行编译动作，然后执行编译后的代码。

上面只是简单列出了两点，当然还有更多的原因，限于篇幅就不再具体介绍，numba 就是通过 JIT 加速了 python 代码。那么怎么使用 numba 加速我们的代码呢？我们可以看一些简单的例子：

numba 加速 python 的小例子
用 numba 加速 python 代码多简单方便呢，我们先来看看如何使用 numba 加速 python 代码：

如果让你用单纯的 python 计算一个矩阵所有元素的和，很容易可以写出下面的代码：

def cal_sum(a): 
    result = 0 
    for i in range(a.shape[0]): 
        for j in range(a.shape[1]): 
            result += a[i, j] 
    return result

当需要计算的矩阵很小的时候，貌似速度也不慢，可以接受，但是如果输入的矩阵大小为 (500, 500)，

a = np.random.random((500, 500)) 
%timeit cal_sum(a)

输出结果为：

47.8 ms ± 499 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

我们尝试加上 numba：

import numba  
 
@numba.jit(nopython=True) 
def cal_sum(a): 
    result = 0 
    for i in range(a.shape[0]): 
        for j in range(a.shape[1]): 
            result += a[i, j] 
    return result

输入同样大小的矩阵

a = np.random.random((500, 500)) 
%timeit cal_sum(a)

输出结果为：

236 µs ± 545 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)

注意在这里我们使用了 %itemit 测试运行时间（原因我们留到后面说），通过对比两个时间，我们可以发现通过 numba 获得了非常明显的加速效果！

我们来具体看一下如何用 numba 加速 python 代码：在实际使用过程中，numba 其实是以装饰器的形式加在 python 函数上的，用户可以不用关心到底 numba 是通过什么方法来优化代码，只需要调用就行。同时需要注意到 @jit 装饰器同时也有一个参数 nopython， 这个参数主要是来区分 numba 的运行模式，numba 其实有两种运行模式：一个是 nopython 模式，另一个就是 object模式。只有在nopython 模式下，才会获得最好的加速效果，如果 numba 发现你的代码里有它不能理解的东西，就会自动进入 object 模式，保证程序至少是能够运行的（当然这其实就失去了添加 numba 的意义）。如果我们将装饰器改为 @jit(nopython=True) 或者 @njit，numba 会假设你已经对所加速的函数非常了解，强制使用加速的方式，不会进入 object 模式，如编译不成功，则直接抛出异常。

当然说到这里，可能大家还是很困惑，numba 到底是怎么加速 python 代码的？

python 代码的编译过程包括四个阶段：词法分析 -> 语法分析 -> 生成字节码 -> 将字节码解释为机器码执行, 常见的 python 解释器的类型有 cpython、IPython、PyPy、Jython、IronPython，与其他解释器不同，numba 是使用 LLVM 编译技术来解释字节码的。

LLVM 是一个编译器，它采用字节码，并将其编译为机器码，编译过程涉及许多额外的传递，而
LLVM编译器可以优化字节码，例如某些频繁执行的模块，LLVM 可以将其作为 “hot code” 从而进行相应的优化，LLVM
工具链非常擅长优化字节码，它不仅可以编译 numba 的代码，还可以优化它。

在第一次调用 numba 装饰的函数时，numba 将在调用期间推断参数类型，numba 会结合给定的参数类型将其编译为机器代码。这个过程是有一定的时间消耗的，但是一旦编译完成，numba 会为所呈现的特定类型的参数缓存函数的机器代码版本，如果再次使用相同的类型调用它，它可以重用缓存的机器代码而不必再次编译。

就上面举的简单的例子来说，使用 numpy 和 numba 加速基本效果差不多，但是在实际情况里面，不是所有的 for 循环代码都可以直接用 numpy 自带的函数实现。但是 numba 基本对所有的 for 循环代码都有非常好的加速效果，当然前提是 for 循环里面的代码必须是 numba 能够理解的。

而在从实际使用中，一般推荐将代码中密集的计算部分提取出来作为单独的函数实现，并使用 nopython 方式优化，这样可以保证我们能使用到 numba 的加速功能。其余部分还是使用 python 原生代码，这样一方面就可以做到在 numba 加速不明显或者无法加速的代码中调用各种函数实现自己的代码逻辑， 另一方面也能享受到 numba 的加速效果。

numba 加速 numpy 运算

上面说了 numba 一大亮点就是加速 for 循环，除此以外，numba 对 numpy 的运算也同样的有加速的效果。因为即使是numpy 也没有 numba 转换为机器码快，numba 尤其擅长加速 numpy 的基本运算 (如加法、相乘和平方等等)，其实准确来说如果 numpy 函数是对各个元素采用相同的操作的情况下，都会有比较好的效果。我们简单举一个 numba 加速 numpy运算的例子：

a = np.ones((1000, 1000), np.int64) * 5 
b = np.ones((1000, 1000), np.int64) * 10 
c = np.ones((1000, 1000), np.int64) * 15 
 
def add_arrays(a, b, c): 
    return np.square(a, b, c) 
 
@numba.njit 
def add_arrays_numba(a, b, c): 
    return np.square(a, b, c) 
 
# 第一次调用完成编译 
add_arrays_numba(a) 
 
# 函数被编译，机器代码被缓存 
start = time.time() 
add_arrays_numba(a) 
end = time.time() 
print("Elapsed (after compilation) = %s" % (end - start)) 
 
# 不使用 numba 加速 
start = time.time() 
add_arrays(a) 
end = time.time() 
print("Elapsed = %s" % (end - start))

Elapsed (after compilation) = 0.002088785171508789
Elapsed = 0.0031290054321289062

numba 使用 CUDA 加速

numba 更厉害的地方就在于，我们可以直接用 python 写 CUDA Kernel， 直接在 GPU 上编译和运行我们的 Python程序，numba 通过将 python 代码直接编译为遵循 CUDA 执行模型的 CUDA 内核和设备函数来支持 CUDA GPU 编程（但是实际上 numba 目前支持的 CUDA API 很少，希望开发团队能更肝一点~~~) ，对于不熟悉 CUDA 的同学，我们推荐大家看一下之前的一篇文章 PyTorch 源码解读之 cpp_extension：揭秘 C++/CUDA 算子实现和调用全流程 了解一下 CUDA 编程的基本概念。为了节省将 numpy 数组复制到指定设备，然后又将结果存储到 numpy数组中所浪费的时间，numba 提供了一些函数来声明并将数组送到指定设备来节省不必要的复制到 cpu 的时间。

常用内存分配函数：

• cuda.device_array()：在设备上分配一个空向量，类似于numpy.empty()；
• cuda.to_device()：将主机的数据拷贝到设备；
• cuda.copy_to_host()：将设备的数据拷贝回主机；

我们可以通过一个简单的矩阵相加的例子来看看通过 numba 使用 CUDA 加速的效果：

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

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


m = 5000
n = 4000

x = np.arange(m*n).reshape((m,n)).astype(np.int32)
y = np.arange(m*n).reshape((m,n)).astype(np.int32)

# 拷贝数据到设备端
x_device = cuda.to_device(x)
y_device = cuda.to_device(y)

# 在显卡设备上初始化一块用于存放GPU计算结果的空间
gpu_result1 = cuda.device_array((m,n))
gpu_result2 = cuda.device_array((m,n))
cpu_result = np.empty((m,n))

threads_per_block = 1024
blocks_per_grid = math.ceil(m*n / threads_per_block)
# 第一次调用包含编译时间
start = time()
matrix_add[blocks_per_grid, threads_per_block](x_device, y_device, gpu_result1, m, n)
cuda.synchronize()
print("gpu matrix add time (with compilation) " + str(time() - start))
start = time()
matrix_add[blocks_per_grid, threads_per_block](x_device, y_device, gpu_result2, m, n)
cuda.synchronize()
print("gpu matrix add time (after compilation)" + str(time() - start))
start = time()
cpu_result = np.add(x, y)
print("cpu matrix add time " + str(time() - start))

运行时间分别为：

gpu matrix add time (with compilation) 0.15977692604064941
gpu matrix add time (after compilation) 0.0005376338958740234
cpu matrix add time 0.023023128509521484

在通过 numba 进行 CUDA 加速的时候，主要是通过调用 @cuda.jit 装饰器实现，从结果可以看到 numba 通过调用 CUDA 明显加速了 python 程序。

---

参考：
http://www.gimoo.net/t/1811/5bef70fa0f4ae.html

numba：
https://zhuanlan.zhihu.com/p/454573654

pycuda：
https://zhuanlan.zhihu.com/p/91188887