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中定义一个小的帮助程序类,该类可隐藏大量索引并提供简单的索引语法。例如:
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端受到影响。
注意:如果不扩展helper类以获取所有维度的步幅并且更改内核以接受输入和输出数组的所有维度的步幅,则这对于fortran和C排序都将无效。从性能的角度来看,最好为fortran和C有序数组编写单独的帮助器类。