24、利用GPU进行Python编程：PyCUDA与Numba的实践指南

利用GPU进行Python编程：PyCUDA与Numba的实践指南

无NVIDIA显卡使用CUDA的解决方案

直接使用CUDA及相关库（如PyCUDA）时，最大的阻碍之一是缺乏合适的硬件。不过，有一些解决办法，只是这些办法并非免费。
- 使用云服务 ：可以利用市场上现有的云服务，例如AWS，它允许使用配备NVIDIA硬件的服务器，且只需按实际使用时间付费。这是一种相对经济的入门方式，但需要花费时间学习如何设置和管理AWS基础设施。
- 购买NVIDIA显卡 ：为工作机器购买基于NVIDIA的显卡，这需要一笔固定的前期成本，尤其是最新一代、高规格的显卡价格可能较贵。

需要注意的是，虽然云服务短期内可能更经济，但有人曾在不知情的情况下让GPU实例长时间运行，最终产生的费用远超一张新显卡的总价。

PyCUDA简介

PyCUDA是一个Python库，它为Python开发者提供了访问NVIDIA CUDA并行计算API的能力。其官方文档可在 此处 找到，包含官方源代码和邮件列表。

PyCUDA的特性

• 对象清理 ：对象清理与对象的生命周期相关，因此编写的代码更不容易出现内存泄漏问题，长时间运行也不易崩溃。
• 抽象复杂性 ：抽象了诸如 pydve_compiler.SourceModule 和 pydva_gpuarray.GPUArray 等复杂操作。
• 自动错误检查 ：自动将遇到的错误转换为Python异常，无需编写繁琐的错误检查代码。
• 高性能 ：该库用C++编写，性能出色。

简单示例

以下是PyCUDA主页上的官方示例，展示了如何使用PyCUDA进行基本计算：

import pycuda.autoint
import pycuda.driver as drv
import numpy
from pycuda.compiler import SourceModule

mod = SourceModule("""
__global__ void multiply_them(float *dest, float *a, float *b) {
    const int i = threadIdx.x;
    dest[i] = a[i] * b[i];
}
""")

multiply_them = mod.get_function("multiply_them")

a = numpy.random.randn(100).astype(numpy.float32)
b = numpy.random.randn(100).astype(numpy.float32)
dest = numpy.zeros_like(a)

multiply_them(
    drv.Out(dest), drv.In(a), drv.In(b),
    block=(100, 1, 1), grid=(1, 1))

print(dest, a * b)

内核函数（Kernels）

内核函数是图形编程和各种语言中常见的概念。内核函数是一种特殊的GPU函数，由CPU代码调用。在Python程序中，内核函数通常如下所示：

mod = SourceModule("""
    __global__ void doubleify(float *a) {
        int idx = threadIdx.x + threadIdx.y * 3;
        a[idx] *= 2;
    }
""")

有一本优秀的在线书籍《The OpenCL Programming Book》详细介绍了OpenCL内核编程的基础知识，可在 此处 找到。

GPU数组

GPU数组是PyCUDA库的重要组成部分，它将使用GPU的复杂性进行了抽象，让我们可以像使用 numpy.ndarray 一样进行操作。其类定义如下：

class pycuda.gpuarray.GPUArray(shape, dtype, ..., allocator=None, order='C')

以下是一个快速示例，展示如何初始化一个GPU数组实例：

import numpy
import pycuda.autoint
import pycuda.gpuarray as gpuarray

a_gpu = gpuarray.to_gpu(numpy.random.randn(2, 2).astype(numpy.float32))
a_tripled = (3 * a_gpu).get()
print(a_tripled)

由于篇幅限制，无法像官方文档那样深入讲解，建议查看 官方文档 获取更多信息。

Numba简介

Numba是Continuum Analytics开发的Python编译器，它能让解释型语言实现高度并行化和强大的性能。官方pydata网站的 文档 提供了Numba的全面概述。

Numba的特性

• 即时代码生成 ：能够在运行时生成优化的机器代码。
• 原生代码生成 ：可针对CPU（默认）和GPU生成原生代码。
• 集成科学软件栈 ：与Python科学软件栈集成良好。
• 跨平台支持 ：支持Python 2和Python 3，可在三大主流操作系统上使用。

LLVM简介

LLVM是一个由模块化和可重用的编译器及工具链技术组成的项目。它最初是一个研究项目，现在已得到广泛认可。LLVM专注于生成最优的低级代码（中间代码或二进制代码），主要用C++编写，是众多语言和项目的基础，如Ada、Fortran、Python和Ruby等。对于想构建自己编译器的人，推荐阅读 这篇文章 。

跨硬件兼容性

Numba非常灵活，它支持将Python代码编译为可在CPU或GPU硬件上运行的代码，且本节中的示例可在多种不同类型的GPU上运行，不一定局限于NVIDIA显卡。

Python编译模式

在深入了解Numba库之前，有必要了解标准CPython程序和Numba Python程序的编译和执行方式的关键区别。主要有两种编译模式：
- 即时编译（JiT） ：可以消除Python代码的解释器开销，从而显著提高程序的运行速度。虽然Numba采用了JiT编译，但它只能让代码性能更接近编译型语言（如C或C++），而不能完全达到其性能水平。
- 提前编译（AoT） ：将函数编译为磁盘上的二进制对象，可独立分发和执行。这种方式无需解释代码，机器只需运行预构建的二进制文件，而无需同时进行编译和执行。

使用Numba的步骤

Stanley Seibert在一次演讲中介绍了有效使用Numba的五个步骤：
1. 创建基准测试用例 ：创建一个现实的基准测试用例，以获取系统在负载下的实际性能指标，而不仅仅是使用标准的单元测试库。
2. 使用性能分析工具 ：使用性能分析工具（如cProfile）对基准测试进行分析。
3. 识别热点代码 ：找出代码中执行时间较长的热点部分。
4. 使用装饰器 ：根据需要为关键函数使用 @numba.jit 和 @numba.vectorize 装饰器。
5. 重新运行基准测试 ：重新运行基准测试并分析结果，以确定是否提高了程序的性能。你可以在 YouTube 上观看他的原始演讲。

Anaconda的安装

在使用Numba之前，需要从Continuum Analytics网站安装Anaconda包，链接为 此处 。Anaconda是Python数据科学生态系统中强大且受尊重的一部分，它是开源的，提供了Python和R编程语言的高性能发行版。R是数据科学家和量化分析师处理大型数据集时的首选语言。Anaconda还自带包和依赖管理工具Conda，目前包含超过1000个数据科学特定的包。

编写基本的Numba Python程序

以下是一个简单的Python程序，包含一个返回两个值之和的函数：

def f(x, y):
    return x + y

f(2, 3)

可以通过导入 jit 装饰器并对函数进行装饰，为其添加懒编译功能：

from numba import jit

@jit
def f(x, y):
    return x + y

f(2, 3)

添加装饰器后，函数将被编译。

编译选项

@jit 装饰器接受多个关键字参数，可用于明确告诉编译器如何编译函数。
- nopython模式 ：Numba有两种编译模式，nopython模式和对象模式。nopython模式生成的代码不直接访问Python C API，因此执行速度非常快。

@jit(nopython=True)
def func(x, y):
    return x + y

• nogil选项 ：如果Numba能将函数编译为本地代码（如nopython模式），可以指定 nogil 标志。进入这些函数时，可释放全局解释器锁（GIL），使函数能与其他Python线程并发运行，从而提高性能。

@jit(nogil=True)
def func(x, y):
    return x + y

• cache选项 ：设置 cache 选项后，Numba会将特定函数的编译代码存储在磁盘缓存文件中，这样函数编译后，每次程序执行时无需重新编译。

@jit(cache=True)
def func(x, y):
    return x + y

• parallel选项 ： parallel 选项是一个实验性功能，旨在自动并行化指定函数中具有并行语义的操作，但必须与nopython选项一起使用。

@jit(nopython=True, parallel=True)
def func(x, y):
    return x + y

Numba的局限性

使用Numba构建软件系统时，需要注意其局限性。在某些情况下，类型推断可能无法实现，例如以下示例：

@jit(nopython=True)
def f(x, y):
    return x + y

f(1, (2,))

当传入不同类型的值时，Numba无法推断其类型，会抛出错误。

Numba在CUDA GPU和AMD APU上的应用

了解了Numba的基础知识后，接下来可以探索如何在GPU上使用它。从Numba 0.21版本开始，支持在异构系统架构（HSA）上编程。HSA旨在结合CPU和GPU的性能，并为它们提供共享内存空间。

综上所述，PyCUDA和Numba为Python开发者提供了强大的工具，使他们能够利用GPU的并行计算能力，提高程序的性能。通过合理使用这些工具和技术，可以在处理计算密集型任务时取得更好的效果。

利用GPU进行Python编程：PyCUDA与Numba的实践指南

对比PyCUDA与Numba

为了更清晰地了解PyCUDA和Numba的特点，我们可以通过以下表格进行对比：
| 特性 | PyCUDA | Numba |
| — | — | — |
| 硬件依赖 | 主要依赖NVIDIA GPU | 支持CPU和多种GPU，不限于NVIDIA |
| 代码复杂度 | 需要编写CUDA内核代码，相对复杂 | 通过装饰器简化代码，易于上手 |
| 性能优化 | 手动管理GPU资源，可精细优化 | 自动进行代码优化，有一定局限性 |
| 适用场景 | 适合对GPU编程有深入了解，需要精细控制的场景 | 适合快速实现并行计算，对性能有一定要求的场景 |

实际应用案例

下面通过一个实际的计算密集型任务，展示如何使用PyCUDA和Numba提高程序性能。假设我们要计算两个大型矩阵的乘法。

使用纯Python实现

import numpy as np

def matrix_multiplication_python(A, B):
    rows_A = len(A)
    cols_A = len(A[0])
    rows_B = len(B)
    cols_B = len(B[0])

    if cols_A != rows_B:
        raise ValueError("Number of columns in A must be equal to number of rows in B.")

    C = [[0 for row in range(cols_B)] for col in range(rows_A)]
    for i in range(rows_A):
        for j in range(cols_B):
            for k in range(cols_A):
                C[i][j] += A[i][k] * B[k][j]
    return C

# 生成两个大型矩阵
A = np.random.rand(100, 100)
B = np.random.rand(100, 100)

# 计算矩阵乘法
result_python = matrix_multiplication_python(A, B)

使用PyCUDA实现

import pycuda.driver as cuda
import pycuda.autoinit
from pycuda.compiler import SourceModule
import numpy as np

mod = SourceModule("""
    __global__ void matrix_multiply(float *A, float *B, float *C, int rows_A, int cols_A, int cols_B) {
        int idx = threadIdx.x + blockIdx.x * blockDim.x;
        int idy = threadIdx.y + blockIdx.y * blockDim.y;

        if (idx < cols_B && idy < rows_A) {
            float sum = 0;
            for (int k = 0; k < cols_A; k++) {
                sum += A[idy * cols_A + k] * B[k * cols_B + idx];
            }
            C[idy * cols_B + idx] = sum;
        }
    }
""")

matrix_multiply = mod.get_function("matrix_multiply")

# 生成两个大型矩阵
A = np.random.rand(100, 100).astype(np.float32)
B = np.random.rand(100, 100).astype(np.float32)
C = np.zeros((100, 100)).astype(np.float32)

# 执行矩阵乘法
matrix_multiply(
    cuda.In(A), cuda.In(B), cuda.Out(C),
    np.int32(A.shape[0]), np.int32(A.shape[1]), np.int32(B.shape[1]),
    block=(16, 16, 1), grid=(int((100 + 15) / 16), int((100 + 15) / 16))
)

print(C)

使用Numba实现

import numpy as np
from numba import jit

@jit(nopython=True, parallel=True)
def matrix_multiplication_numba(A, B):
    rows_A = A.shape[0]
    cols_A = A.shape[1]
    rows_B = B.shape[0]
    cols_B = B.shape[1]

    if cols_A != rows_B:
        raise ValueError("Number of columns in A must be equal to number of rows in B.")

    C = np.zeros((rows_A, cols_B))
    for i in range(rows_A):
        for j in range(cols_B):
            for k in range(cols_A):
                C[i, j] += A[i, k] * B[k, j]
    return C

# 生成两个大型矩阵
A = np.random.rand(100, 100).astype(np.float32)
B = np.random.rand(100, 100).astype(np.float32)

# 计算矩阵乘法
result_numba = matrix_multiplication_numba(A, B)
print(result_numba)

性能测试与分析

为了比较纯Python、PyCUDA和Numba实现矩阵乘法的性能，我们可以使用 timeit 模块进行测试。以下是测试代码：

import timeit

# 纯Python实现的性能测试
python_time = timeit.timeit(lambda: matrix_multiplication_python(A, B), number=10)
print(f"Pure Python time: {python_time} seconds")

# PyCUDA实现的性能测试
pycuda_time = timeit.timeit(lambda: matrix_multiply(
    cuda.In(A), cuda.In(B), cuda.Out(C),
    np.int32(A.shape[0]), np.int32(A.shape[1]), np.int32(B.shape[1]),
    block=(16, 16, 1), grid=(int((100 + 15) / 16), int((100 + 15) / 16))
), number=10)
print(f"PyCUDA time: {pycuda_time} seconds")

# Numba实现的性能测试
numba_time = timeit.timeit(lambda: matrix_multiplication_numba(A, B), number=10)
print(f"Numba time: {numba_time} seconds")

通过性能测试，我们可以得到以下结果（具体时间可能因硬件环境而异）：
| 实现方式 | 执行时间（秒） |
| — | — |
| 纯Python | 较长 |
| PyCUDA | 较短 |
| Numba | 较短 |

从结果可以看出，纯Python实现的矩阵乘法性能最差，而PyCUDA和Numba的实现性能明显优于纯Python。PyCUDA需要手动管理GPU资源，编写CUDA内核代码，适合对GPU编程有深入了解的开发者；Numba通过装饰器简化了代码，易于上手，适合快速实现并行计算。

总结与建议

在利用GPU进行Python编程时，PyCUDA和Numba是两个非常有用的工具。以下是一些总结和建议：
- 选择合适的工具 ：如果对GPU编程有深入了解，需要精细控制GPU资源，建议使用PyCUDA；如果希望快速实现并行计算，对性能有一定要求，建议使用Numba。
- 注意硬件依赖 ：PyCUDA主要依赖NVIDIA GPU，而Numba支持CPU和多种GPU，不限于NVIDIA。在选择工具时，需要考虑硬件环境。
- 了解工具的局限性 ：PyCUDA需要编写复杂的CUDA内核代码，学习成本较高；Numba在某些情况下可能无法进行类型推断，导致性能下降。在使用时，需要了解这些局限性，并采取相应的措施。
- 进行性能测试 ：在实际应用中，需要对不同的实现方式进行性能测试，选择性能最优的方案。

通过合理使用PyCUDA和Numba，Python开发者可以充分利用GPU的并行计算能力，提高程序的性能，处理更复杂的计算密集型任务。

操作流程总结

为了方便大家使用PyCUDA和Numba进行GPU编程，以下是一个操作流程总结：

graph LR
    classDef startend fill:#F5EBFF,stroke:#BE8FED,stroke-width:2px
    classDef process fill:#E5F6FF,stroke:#73A6FF,stroke-width:2px

    A([开始]):::startend --> B{选择工具}:::process
    B -->|需要精细控制| C(PyCUDA):::process
    B -->|快速实现并行计算| D(Numba):::process
    C --> E(安装PyCUDA和相关依赖):::process
    D --> F(安装Numba和Anaconda):::process
    E --> G(编写CUDA内核代码):::process
    F --> H(编写Python代码并使用装饰器):::process
    G --> I(执行代码并优化性能):::process
    H --> I
    I --> J([结束]):::startend

按照以上流程，你可以逐步掌握使用PyCUDA和Numba进行GPU编程的方法，提高程序的性能。