在Python中使用PyCuda.compilerSourceModule()函数进行GPU加速
发布时间:2024-01-08 21:45:35
在Python中使用PyCuda的compiler.SourceModule()函数可以将CUDA C代码编译为CUDA内核,并生成可以在Python中调用的PyCuda模块。
以下是一个使用PyCuda进行GPU加速的简单示例代码:
首先,确保已经安装了PyCuda库,并导入所需的模块和库:
import pycuda.autoinit import pycuda.driver as drv from pycuda.compiler import SourceModule import numpy as np
接下来,我们定义一个简单的向量加法函数,该函数将在GPU上执行:
# 定义CUDA C代码
cuda_code = """
__global__ void vector_add(float *a, float *b, float *c, int n)
{
int tid = blockIdx.x * blockDim.x + threadIdx.x;
if (tid < n)
c[tid] = a[tid] + b[tid];
}
"""
# 使用SourceModule编译CUDA C代码
mod = SourceModule(cuda_code)
# 获取CUDA内核函数
vector_add = mod.get_function("vector_add")
然后,我们创建输入向量和输出向量,并将它们分配到GPU内存中:
# 定义输入向量和输出向量的大小 n = 1000000 # 在主机上生成输入向量a和b a = np.random.randn(n).astype(np.float32) b = np.random.randn(n).astype(np.float32) # 在主机上分配输出向量c c = np.zeros_like(a)
接下来,我们将输入向量和输出向量从主机内存复制到GPU内存:
# 在设备上分配内存 dev_a = drv.mem_alloc(a.nbytes) dev_b = drv.mem_alloc(b.nbytes) dev_c = drv.mem_alloc(c.nbytes) # 将输入向量a和b从主机内存复制到设备内存 drv.memcpy_htod(dev_a, a) drv.memcpy_htod(dev_b, b)
然后,我们定义要在GPU上执行的内核的块大小和网格大小,并调用内核函数进行计算:
# 定义块大小和网格大小 block_size = (256, 1, 1) grid_size = (int(n/block_size[0]) + 1, 1) # 调用内核函数进行计算 vector_add(dev_a, dev_b, dev_c, np.int32(n), block=block_size, grid=grid_size)
最后,我们将输出向量从GPU内存复制回主机内存,并在主机上打印结果:
# 将输出向量c从设备内存复制到主机内存 drv.memcpy_dtoh(c, dev_c) # 打印结果 print(c)
这是一个简单的示例,展示了如何使用PyCuda进行GPU加速。要实现更复杂的计算,可以编写更复杂的CUDA C代码,并将其编译为PyCuda模块以在Python中调用。
