我们如何在 CUDA 中访问 3D 数组的列?

How do we access the column of a 3D array in CUDA?

    mod=SourceModule("""

   __global__ void mat_ops(float *A,float *B)
  {   /*formula to get unique thread index*/
      int thrd= blockIdx.x*blockDim.x*blockDim.y+threadIdx.y*blockDim.x+threadIdx.x;
      B[]=A[];
   }    
   """)
        func = mod.get_function("mat_ops")
        func(A_k, B_k, grid=(3,1,1),block=(4,4,1))

我有两个 3D 数组 float *A 和 float *B,在这个 PyCUDA 内核中每个大小为 4 X 4 X 3。我在这里要做的是逐列遍历 3D 数组,而不是逐行遍历。我正在使用 2D 块的 1D 网格。我该怎么做呢 ?

为此,您需要向 CUDA 内核描述数组在内存中的布局,并且需要使用主机端提供的步幅在内核中进行正确的索引计算。一个简单的方法是在 CUDA 中定义一个小助手 class,它隐藏了大部分索引并提供了一个简单的索引语法。例如:

from pycuda import driver, gpuarray
from pycuda.compiler import SourceModule
import pycuda.autoinit
import numpy as np

mod=SourceModule("""

   struct stride3D
   {
       float* p;
       int s0, s1;

       __device__
       stride3D(float* _p, int _s0, int _s1) : p(_p), s0(_s0), s1(_s1) {};

       __device__
       float operator  () (int x, int y, int z) const { return p[x*s0 + y*s1 + z]; };

       __device__
       float& operator () (int x, int y, int z) { return p[x*s0 + y*s1 + z]; };
   };

   __global__ void mat_ops(float *A, int sA0, int sA1, float *B, int sB0, int sB1)
   {
       stride3D A3D(A, sA0, sA1);
       stride3D B3D(B, sB0, sB1);

       int xidx = blockIdx.x;
       int yidx = threadIdx.x;
       int zidx = threadIdx.y;

       B3D(xidx, yidx, zidx) = A3D(xidx, yidx, zidx);
   }    
   """)

A = 1 + np.arange(0, 4*4*3, dtype=np.float32).reshape(4,4,3)
B = np.zeros((5,5,5), dtype=np.float32)
A_k = gpuarray.to_gpu(A)
B_k = gpuarray.to_gpu(B)

astrides = np.array(A.strides, dtype=np.int32) // A.itemsize
bstrides = np.array(B.strides, dtype=np.int32) // B.itemsize

func = mod.get_function("mat_ops")
func(A_k, astrides[0], astrides[1], B_k, bstrides[0], bstrides[1], grid=(4,1,1),block=(4,3,1))
print(B_k[:4,:4,:3])

这里我选择了使源数组和目标数组的大小不同,只是为了表明代码是通用的,只要块大小足够,它就可以用于任何大小的数组。请注意,设备代码端没有数组边界检查,对于重要示例,您需要添加它。

另请注意,这对于 Fortran 和 C 有序的 numpy 数组应该都能正常工作,因为它直接使用 numpy 步幅值。但是,由于内存合并问题,性能将在 CUDA 方面受到影响。

注意:如果不扩展助手 class 以对所有维度采取步幅并更改内核以接受输入和输出数组的所有维度的步幅,这将不适用于 Fortran 和 C 排序。从性能的角度来看,最好为 Fortran 和 C 有序数组编写单独的助手 classes。