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

Question

    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))

I have two 3D arrays float *A and float *B, each of size 4 X 4 X 3 in this PyCUDA kernel. What I am trying to do here is, traverse the 3D array column by column, instead of row by row. I am making use of a 1D Grid of 2D blocks. How do I do this ?

Answer 1

To do this, you need to describe to layout of the array in memory to the CUDA kernel, and you need the correct indexing calculations in the kernel using the host side provided strides. A simple way to do this is to define a small helper class in CUDA which hides the bulk of the indexing and provides a simple indexing syntax. For example:

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])

Here I have chosen to make the source and destination arrays different sizes, just to show that the code is general and will work for any size arrays as long as the block size is sufficient. Note that there is no array bounds checking here on the device code side, you will need to add that for non-trivial examples.

Note also that this should work correctly both for fortran and C ordered numpy arrays, because it uses the numpy stride values directly. Performance will be effected on the CUDA side because of memory coalescing issues, however.

Note: this won't work for both fortran and C ordering without extending the helper class to take strides for all dimensions and changing the kernel to accept strides for all dimensions of the input and output arrays. From a performance perspective it would be better to write separate helper classes for fortran and C ordered arrays.