python cuda多流传输样例

以下是一个简单的Python CUDA多流传输样例：

import numpy as np
import pycuda.driver as cuda
import pycuda.gpuarray as gpuarray
import pycuda.autoinit

# 定义主机数据
h_a = np.random.randn(10000).astype(np.float32)
h_b = np.random.randn(10000).astype(np.float32)
h_c_cpu = np.zeros_like(h_a)

# 定义设备数据
d_a = gpuarray.to_gpu(h_a)
d_b = gpuarray.to_gpu(h_b)
d_c = gpuarray.zeros_like(d_a)

# 定义流
streams = []
for i in range(4):
    streams.append(cuda.Stream())

# 多流计算
for i in range(4):
    start = i * 2500
    end = start + 2500
    d_a_part = d_a[start:end]
    d_b_part = d_b[start:end]
    d_c_part = d_c[start:end]
    d_c_part += d_a_part * d_b_part
    d_c_part.get(h_c_cpu[start:end], stream=streams[i])

# 等待所有流计算完成
for stream in streams:
    stream.synchronize()

# 验证结果
h_c_gpu = d_c.get()
assert np.allclose(h_c_cpu, h_c_gpu)
print("Results verified!")

该样例将主机端的两个数组（h_a和h_b）传输到设备端的两个数组（d_a和d_b），并使用四个流计算它们的点积（d_c = d_a * d_b）。

要实现多流传输，我们首先需要创建几个流，并使用流的方法传输不同部分的数据。在这个例子中，我们将数据分成四个等大小的部分，并使用四个流分别计算它们的点积。在每个流上，我们使用GPUArray的切片来获取对应的部分数组，并在该流上计算点积。然后，我们使用get方法将计算结果传输回主机端，同时指定使用该流。

最后，我们需要等待所有流完成计算，并检查计算结果是否正确。

这个样例只是一个简单的示例，实际上，使用多流传输可以更好地利用GPU的并行性，提高计算效率。