python用gpu运算的代码

在Python中使用GPU运算，常用的是通过CUDA（Compute Unified Device Architecture）进行计算。CUDA是NVIDIA公司推出的用于GPU加速计算的平台和编程模型。

下面是一个使用CUDA进行GPU加速计算的示例代码：

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

# 定义CUDA的代码
mod = SourceModule("""
    __global__ void add(int *a, int *b, int *c) {
        int tid = threadIdx.x + blockIdx.x * blockDim.x;
        c[tid] = a[tid] + b[tid];
    }
""")

# 定义输入数据
a = np.array([1, 2, 3, 4, 5, 6]).astype(np.int32)
b = np.array([6, 5, 4, 3, 2, 1]).astype(np.int32)
c = np.zeros_like(a)

# 获取CUDA函数
add_func = mod.get_function("add")

# 设置CUDA函数的参数并执行计算
block_size = 2
grid_size = (a.size + block_size - 1) // block_size
add_func(drv.In(a), drv.In(b), drv.Out(c), block=(block_size, 1, 1), grid=(grid_size, 1))

# 输出结果
print(c)

在上述代码中，首先定义了一个使用CUDA进行计算的函数`add`，然后定义了输入数据`a`和`b`，以及一个用于存储计算结果的数组`c`。接着，通过`SourceModule`函数将CUDA代码编译成可执行的CUDA函数。然后，通过`mod.get_function`函数获取CUDA函数，并将输入数据和输出数据作为参数传入该函数中。最后，通过`block`和`grid`参数指定并行计算的线程块和线程格的大小，执行CUDA函数并输出结果。

需要注意的是，使用CUDA进行GPU加速计算需要安装CUDA工具包，并且需要安装PyCUDA库进行Python的CUDA编程。