代码收藏家技术教程 2024-07-29

Python — — GPU编程

要想将Python程序运行在GPU上，我们可以使用numba库或者使用cupy库来实现GPU编程。

壹、numba

Numba 是一个开源的 JIT (Just-In-Time) 编译器，它可以将 Python 代码转换成机器代码以提高性能。Numba 特别适用于需要高性能计算的科学计算和数值计算任务。也就是说可以将python程序编译为机器码，使其可以像c/c++、Java一样快速的运行。同样Numba不仅可以加速 CPU 上的 Python 代码，还可以利用 GPU 进行加速。

安装Numba：

pip install numba

一、机器码编程

1. 函数编写：

Numba 的核心功能是 @jit 装饰器，它可以将 Python 函数编译成优化的机器代码。

from numba import jit

@jit(nopython=True)
def my_function(x):
    return x * x

x = 112.0
print(my_function(x))

指定传递参数类型以及返回值类型，nopython表示不使用python编译而直接编译为机器码：

from numba import jit

@jit('float64(float64)', nopython=True)  # 指定输入和输出类型，括号内的是参数类型，括号外的是返回值类型
def my_function(x):
    return x * x

x = 112.0
print(my_function(x))

从 Numba 0.15.1 版本开始，你可以使用 Python 类型注解来指定函数的参数类型：

from numba import jit

@jit
def my_function(x: float) -> float:
    return x * x

2. 使用Numba函数：

使用 Numba 函数，我们可以像使用普通函数一样使用jit修饰过的函数：

result = my_function(10.5)
print(result)  # 输出 110.25

Numba 特别适合于在 NumPy 数组上进行操作。你可以使用 NumPy 数组作为 Numba 函数的参数：

from numba import njit
import numpy as np


@njit
def parallel_function(arr):
    return arr * 2


arr = np.arange(10)
result = parallel_function(arr)
print(result)

3. 使用 Numba 的并行功能：

Numba 提供了并行执行的功能，可以使用 @njit 装饰器来替代 @jit，它会自动并行化循环：

from numba import njit
import numpy as np


@njit
def parallel_function(arr):
    return arr * 2


arr = np.arange(10)
result = parallel_function(arr)
print(result)

二、CUDA编程

1. 引入CUDA 模块：

from numba import cuda

2. 定义 GPU 核函数：

使用 @cuda.jit 装饰器定义 GPU 核函数，这与 CPU 加速中使用的 @jit 类似，但 @cuda.jit 指定了函数将在 GPU 上执行：

@cuda.jit
def gpu_kernel(x, y):
    # 核函数体，使用 CUDA 线程索引进行计算
    # 例如: position = cuda.grid(1)
    # if position < len(x):
    #     y[position] = x[position] * x[position]

position = cuda.grid(1)：其中cuda.grid(1)用于确定当前线程在执行的整个网格（grid）中的位置，这里的参数1表示一维的GPU网格索引，如果是cuda.grid(2)则表示二维的GPU网格索引。

CUDA 的执行模型的概念：

线程（Thread）：执行计算的最小单元。

块（Block）：一组线程，它们可以共享数据并通过共享内存进行通信。

网格（Grid）：由多个块组成，用于实现更大范围的并行性。

上面的代码表示的是对每一个元素分配一个GPU线程，通过cuda.grid(1)来获取每一个线程，本质上也是获取每一个元素，然后再进行运算操作，通常情况下希望数组长度至少与线程数相等。因为如果线程总数大于数组长度，就会有多余的线程没有执行任何操作。例如，如果数组 x 只有 5 个元素，但配置了 32 个线程，那么只有前 5 个线程会计算和存储结果，其余 27 个线程将不会执行任何操作。

3. 设置执行配置：

GPU 核函数需要执行配置来确定并行执行的线程数和块数。这通过在函数调用时使用方括号指定：

threads_per_block = 256
blocks_per_grid = (n + (threads_per_block - 1)) // threads_per_block
gpu_kernel[blocks_per_grid, threads_per_block](x, y)	# 给cuda.jit修饰的函数分配资源，并传入参数x 和 y

threads_per_block = 256：定义了每个块内的线程的个数，这里是256，如果是二维数组，那么需要使用元组的方式来进行定义，如：threads_per_block = (16, 16)。

blocks_per_grid = (n + (threads_per_block - 1)) // threads_per_block：定义了整个网格（grid）中的块数量。它也是一个元组，n表示数组的长度，(n + (threads_per_block - 1)) // threads_per_block这种运算相当于一个向上取整的操作，保证了数组中的每一个元素都能分配一个GPU线程，因为一个原则是：线程数量要大于等于数组的个数。如果是二维数组需要这样定义网格中块的数量：blocks_per_grid = (m // threads_per_block[0], n // threads_per_block[1])，其中m表示行数，n表示列数。

4. 数据传输：

在 GPU 上执行计算之前，需要将数据从 CPU 内存传输到 GPU 内存，这通常使用 cuda.to_device() 方法完成：

x_device = cuda.to_device(x)
y_device = cuda.to_device(y)

5. 在 GPU 上分配内存：

如果 GPU 上的核函数需要额外的存储空间，可以使用 cuda.device_array() 在 GPU 上分配内存：

result_device = cuda.device_array_like(x_device)

6. 同步执行:

GPU 核函数的执行是异步的，可能需要调用 cuda.synchronize() 来确保 CPU 等待 GPU 计算完成：

cuda.synchronize()

7. 将结果从 GPU 传回 CPU：

计算完成后，使用 copy_to_host() 方法将 GPU 上的结果复制回 CPU 内存：

result = result_device.copy_to_host()

8. 实例一：

二维数组的GPU运算

import numpy as np
from numba import cuda


@cuda.jit
def matrix_add(A, B, C, m, n):
    row, col = cuda.grid(2)
    if row < m and col < n:
        C[row, col] = A[row, col] + B[row, col]


m, n = 1024, 1024
A = np.random.rand(m, n).astype(np.float32)
B = np.random.rand(m, n).astype(np.float32)

C = np.zeros_like(A)   # 创建与A形状相同的0数组
threads_per_block = (16, 16)
blocks_per_grid = (m // threads_per_block[0], n // threads_per_block[1])
matrix_add[blocks_per_grid, threads_per_block](A, B, C, m, n)
print(C)

9. 实例二：

GPU显存与主机内存之间的通信

from numba import cuda
import numpy as np

# 在主机上创建一个NumPy数组
host_array = np.array([1, 2, 3, 4], dtype=np.int32)

# 使用cuda.to_device将主机数组复制到GPU
device_array = cuda.to_device(host_array)

# 确保数据传输完成
cuda.synchronize()

# 使用.copy_to_host()方法将GPU数组复制回主机数组
host_result = device_array.copy_to_host()
print(host_result)

# 释放GPU内存
del device_array

10. 实例三：

一维数组的GPU运算

from numba import cuda
import numpy as np

# 定义一个简单的cuda内核
@cuda.jit()
def add_kernel(x, y, z, n):
    i = cuda.grid(1) 
    if i < n:  		# 确保不会超出数组边界
        z[i] = x[i] + y[i]

# 主函数
def main():
    n = 256
    x = cuda.device_array(n, dtype=np.int32)   # 直接在GPU上创建数据，占用GPU显存
    y = cuda.device_array(n, dtype=np.int32)
    z = cuda.device_array(n, dtype=np.int32)
    # 初始化数据
    for i in range(n):
        x[i] = i
        y[i] = 2 * i

    # 计算线程块大小和网格大小, 线程块是一组可以同时执行的线程集合
    threadsperblock = 32  		# 这意味着每个线程块将包含256个线程。
    blockspergrid = (n + (threadsperblock - 1)) // threadsperblock  # 定义每个网格内的块的个数
    # 启动内核
    add_kernel[blockspergrid, threadsperblock](x, y, z, n)
    # 将结果从GPU复制回主机
    result = z.copy_to_host()
    print(result)


if __name__ == '__main__':
    main()

贰、cupy

CuPy 是一个与 NumPy 兼容的库，提供了 NumPy 相同的多维数组 API，但是所有的数值计算都在 GPU 上执行。CuPy 底层使用 CUDA，但是 API 更简洁，使用起来比直接使用 CUDA 更加方便。
使用cupy时，我们首先需要将CUDA的环境给配置好，包括CUDA Toolkit

一、安装cupy：

pip install cupy

二、使用cupy：

1. 导入 CuPy：

import cupy as cp

2. 创建 CuPy 数组

可以使用与 NumPy 类似的函数来创建 CuPy 数组。CuPy 数组是在 GPU 上的多维数组。

# 创建一个全零数组
x = cp.zeros((3, 3))

# 创建一个全一数组
y = cp.ones((2, 2))

# 从 Python 列表创建数组
z = cp.array([[1, 2], [3, 4]])

3. 基本运算

# 矩阵乘法
result = cp.dot(x, z)

# 元素乘法
elementwise_product = x * y

# 元素加法
sum_result = x + z

# 计算平方根
sqrt_result = cp.sqrt(x)

4. 利用 GPU 加速

# 计算数组的总和
total = x.sum()

# 计算数组的均值
mean_value = x.mean()

5. 与 NumPy 的互操作性

# 将 NumPy 数组转换为 CuPy 数组
numpy_array = np.random.rand(10)
cupy_array = cp.array(numpy_array)

# 将 CuPy 数组转换回 NumPy 数组
numpy_array_again = cupy_array.get()

6. 使用随机数生成

# 生成随机数数组
random_array = cp.random.rand(3, 3)

# 生成符合正态分布的随机数数组
normal_array = cp.random.normal(0, 1, (3, 3))

7. 广播

# 广播示例
a = cp.array([1, 2, 3])
b = cp.array([[1], [2], [3]])
result = a + b  # 结果是一个 3x3 的数组

8. 索引和切片

# 获取第二行
second_row = z[1]

# 切片操作
upper_triangle = z[cp.triu(cp.ones((3, 3), dtype=cp.bool_))]

9. 内存管理

CuPy 使用 GPU 内存，当不再需要 CuPy 数组时，应该释放它们以避免内存泄漏。

del x, y, z
cp.get_default_memory_pool().free_all_blocks()

10. 实例一：

import cupy as cp

# 创建一个Cupy数组（自动在GPU上）
x = cp.array([1., 2., 3., 4., 5.])
y = cp.sqrt(x)
print(y)

# 将Cupy数组转换回Numpy数组（如果有需要的话）
z = cp.asnumpy(y)

print(z)

11. 实例二：

import cupy as cp

# 创建两个随机的浮点型 CuPy 数组，相当于 NumPy 中的矩阵
A = cp.random.rand(3, 3).astype('float32')  # 3x3 矩阵
B = cp.random.rand(3, 3).astype('float32')  # 另一个 3x3 矩阵

# 执行矩阵乘法
C = cp.dot(A, B)  # 或者使用 @ 操作符 C = A @ B

# 打印结果
print("矩阵 A:\n", A)
print("矩阵 B:\n", B)
print("矩阵 A 和 B 的乘积:\n", C)

# 将 CuPy 数组转换回 NumPy 数组（如果需要）
import numpy as np
numpy_C = C.get()  # 将 CuPy 数组转换为 NumPy 数组

# 执行一些基本的 NumPy 操作，比如求和
sum_C = cp.sum(C)  # 在 GPU 上计算 C 的总和

# 打印 C 的总和
print("矩阵 C 的总和:", sum_C)

# 释放不再使用的 CuPy 数组以节省 GPU 内存
del A, B, C
cp.get_default_memory_pool().free_all_blocks()

作者：_Soy_Milk