2D CUDA网格核函数中CuPy索引异常，与Numba实现结果不一致如何解决？

• 数据类型不匹配：CUDA内核中定义的三个数组参数类型均为float*，但实际传入的是int32类型的CuPy数组，类型不匹配导致内存访问和数值计算完全错误。
• 二维数组下标访问错误：C/C++原生语法不支持对一维指针直接使用[i,j]格式的二维下标访问，此处逗号会被识别为逗号运算符，最终实际仅取第二个值j作为下标，相当于访问一维数组的第j个元素，完全不符合预期的二维索引逻辑。CuPy传入RawKernel的数组都是扁平化的行优先连续内存指针，需要手动计算线性偏移量：i * 列数 + j。
• 网格尺寸计算存在隐患：计算bpg第二个维度时除以了tpb[0]，虽然本次测试中tpb两个维度都是16不会触发问题，但通用场景下应该除以tpb[1]，避免二维尺寸不同时出现线程覆盖不全的问题。

内核部分修正如下：

extern "C" __global__
void add_arrs(const int* x1, const int* x2, int* y, int N){
  int i = blockDim.x * blockIdx.x + threadIdx.x;
  int j = blockDim.y * blockIdx.y + threadIdx.y;

  if(i < N && j < N){
    int idx = i * N + j;
    y[idx] = x1[idx] + x2[idx];
  }
}

Python调用部分修正网格计算逻辑：

cp_add_arrs = cp.RawKernel(r'''
extern "C" __global__
void add_arrs(const int* x1, const int* x2, int* y, int N){
  int i = blockDim.x * blockIdx.x + threadIdx.x;
  int j = blockDim.y * blockIdx.y + threadIdx.y;

  if(i < N && j < N){
    int idx = i * N + j;
    y[idx] = x1[idx] + x2[idx];
  }
}
''', 'add_arrs')

x1 = cp.ones(25, dtype=cp.int32).reshape(5, 5)
x2 = cp.ones(25, dtype=cp.int32).reshape(5, 5)
y = cp.zeros((5, 5), dtype=cp.int32)
N = x1.shape[0]
tpb = (16, 16)
# 修正网格第二维度计算逻辑
bpg = (x1.shape[0] // tpb[0] + 1, x1.shape[1] // tpb[1] + 1)
cp_add_arrs(bpg, tpb, (x1, x2, y, cp.int32(N)))

运行后即可得到和Numba版本一致的全2数组结果。

内容的提问来源于stack exchange，提问作者andregalera