CUDA Python 编程实例

以下是一些使用Python进行CUDA编程的实例：

1. 使用Numba库的实例

Numba是一个开源的JIT（Just In Time）编译器，它允许用户将Python和NumPy代码转换为可在GPU上执行的CUDA内核。下面是一个简单的向量加法示例：

import numpy as np
from numba import cuda

# 定义CUDA内核函数
@cuda.jit
def vector_add(a, b, out):
    i = cuda.grid(1)
    if i < a.shape:
        out[i] = a[i] + b[i]

# 准备数据
n = 1024 * 1024
a = np.random.random(n).astype(np.float32)
b = np.random.random(n).astype(np.float32)
out = np.zeros_like(a)

# 获取GPU设备并分配内存
device_a = cuda.to_device(a)
device_b = cuda.to_device(b)
device_out = cuda.device_array_like(out)

# 设置线程块和网格尺寸
threads_per_block = 256
blocks_per_grid = (n + (threads_per_block - 1)) // threads_per_block

# 在GPU上执行内核函数
vector_add[blocks_per_grid, threads_per_block](device_a, device_b, device_out)

# 将结果从GPU复制回主机内存
device_out.copy_to_host(out)

# 现在out包含a和b对应元素相加的结果

2. 使用PyCUDA库的实例

PyCUDA是另一个用于编写和运行CUDA内核的Python接口，提供了一种更底层的方式来操作CUDA API。以下是一个使用PyCUDA的简单向量加法示例：

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

# CUDA内核源代码
kernel_code = """
__global__ void add(float *a, float *b, float *c)
{
    int idx = threadIdx.x + blockIdx.x * blockDim.x;
    if (idx < gridDim.x * blockDim.x)
        c[idx] = a[idx] + b[idx];
}
"""

# 初始化环境
cuda.init()
device = cuda.Device(0) # 使用第一个GPU
ctx = device.make_context()

# 创建输入和输出数组，并将它们传输到GPU
a = np.random.randn(10000).astype(np.float32)
b = np.random.randn(10000).astype(np.float32)
c = np.empty_like(a)
a_gpu = cuda.mem_alloc(a.nbytes)
b_gpu = cuda.mem_alloc(b.nbytes)
c_gpu = cuda.mem_alloc(c.nbytes)
cuda.memcpy_htod(a_gpu, a)
cuda.memcpy_htod(b_gpu, b)

# 编译内核
mod = SourceModule(kernel_code)
add_func = mod.get_function("add")

# 设置线程块和网格尺寸
block_dim = 256
grid_dim = (len(a) + block_dim - 1) // block_dim

# 执行内核函数
add_func(a_gpu, b_gpu, c_gpu, block=block_dim, grid=grid_dim)

# 将结果从GPU复制回主机内存
cuda.memcpy_dtoh(c, c_gpu)

# 清理资源
ctx.pop() # 离开当前上下文

# 现在c包含a和b对应元素相加的结果

注意事项

• 这些代码示例假设您已经正确安装了Numba或PyCUDA库及其依赖项，并且系统上有一块支持CUDA的NVIDIA显卡。

• 实际使用时，请根据实际情况调整数组大小和线程配置。

• CUDA编程涉及并行计算和资源管理，因此在实际应用中可能需要进行性能优化和错误处理。

以上实例展示了如何在Python中使用Numba和PyCUDA库进行CUDA编程，以实现向量加法操作。这些库简化了CUDA编程的复杂性，使得开发人员能够更轻松地利用GPU的并行计算能力。