怎么测试pycuda是否安装成功
时间: 2023-08-23 11:40:17 浏览: 48
可以使用以下代码测试 PyCUDA 是否安装成功:
```
import pycuda.autoinit
import pycuda.driver as drv
import numpy as np
# 创建一个随机的 4x4 浮点数矩阵
a = np.random.randn(4, 4).astype(np.float32)
# 将数据从主机内存复制到设备内存
a_gpu = drv.mem_alloc(a.nbytes)
drv.memcpy_htod(a_gpu, a)
# 创建一个 4x4 的浮点数矩阵,用于存储结果
b_gpu = drv.mem_alloc_like(a_gpu)
# 定义一个内核函数
kernel_code = """
__global__ void square(float *dest, float *a)
{
const int i = threadIdx.x;
const int j = threadIdx.y;
const int idx = i + j*4;
const float x = a[idx];
dest[idx] = x*x;
}
"""
# 编译内核函数
from pycuda.compiler import SourceModule
mod = SourceModule(kernel_code)
# 获取内核函数
square = mod.get_function("square")
# 调用内核函数
square(b_gpu, a_gpu, block=(4, 4, 1))
# 将结果从设备内存复制到主机内存
b = np.empty_like(a)
drv.memcpy_dtoh(b, b_gpu)
# 输出结果
print("a:\n", a)
print("b:\n", b)
```
这个程序会在 GPU 上执行一个简单的内核函数,计算输入矩阵中每个元素的平方,并将结果存储在输出矩阵中。如果程序能够正常运行并输出结果,那么 PyCUDA 应该已经安装成功了。