python - 如何将二维数组传递到 pycuda 中的内核？

Question

我找到了答案 here ，但不清楚我是否应该 reshape 数组。在将它传递给 pycuda 内核之前，我是否需要将 2d 数组 reshape 为 1d？

Answer 1

无需 reshape 2D gpuarray 即可将其传递给 CUDA 内核。

正如我在您链接到的答案中所说，2D numpy 或 PyCUDA 数组只是倾斜线性内存的分配，默认情况下以行主要顺序存储。两者都有两个成员告诉您访问数组所需的一切 - shape 和 strides。例如:

In [8]: X=np.arange(0,15).reshape((5,3))

In [9]: print X.shape
(5, 3)

In [10]: print X.strides
(12, 4)

形状是不言自明的，步幅是以字节为单位的存储间距。内核代码的最佳实践是将 PyCUDA 提供的指针视为使用 cudaMallocPitch 分配的指针。并将 stride 的第一个元素视为内存中行的字节间距。一个简单的示例可能如下所示:

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

mod = SourceModule("""
__global__ void diag_kernel(float *dest, int stride, int N)
{
    const int tid = threadIdx.x + blockDim.x * blockIdx.x;

    if (tid < N) {
    float* p = (float*)((char*)dest + tid*stride) + tid;
        *p = 1.0f;
    }
}
""")

diag_kernel = mod.get_function("diag_kernel")

a = np.zeros((10,10), dtype=np.float32)
a_N = np.int32(a.shape[0])
a_stride = np.int32(a.strides[0])
a_bytes = a.size * a.dtype.itemsize
a_gpu = drv.mem_alloc(a_bytes)
drv.memcpy_htod(a_gpu, a)
diag_kernel(a_gpu, a_stride, a_N, block=(32,1,1))
drv.memcpy_dtoh(a, a_gpu)

print a

此处在设备上分配了一些内存，将归零的二维数组直接复制到该分配，并将内核的结果(用 1 填充对角线)复制回主机并打印。在此过程中的任何时候都没有必要展平或以其他方式修改 2D numpy 数据的形状或内存布局。结果是:

$ cuda-memcheck python ./gpuarray.py 
========= CUDA-MEMCHECK
[[ 1.  0.  0.  0.  0.  0.  0.  0.  0.  0.]
 [ 0.  1.  0.  0.  0.  0.  0.  0.  0.  0.]
 [ 0.  0.  1.  0.  0.  0.  0.  0.  0.  0.]
 [ 0.  0.  0.  1.  0.  0.  0.  0.  0.  0.]
 [ 0.  0.  0.  0.  1.  0.  0.  0.  0.  0.]
 [ 0.  0.  0.  0.  0.  1.  0.  0.  0.  0.]
 [ 0.  0.  0.  0.  0.  0.  1.  0.  0.  0.]
 [ 0.  0.  0.  0.  0.  0.  0.  1.  0.  0.]
 [ 0.  0.  0.  0.  0.  0.  0.  0.  1.  0.]
 [ 0.  0.  0.  0.  0.  0.  0.  0.  0.  1.]]
========= ERROR SUMMARY: 0 errors