python - 二维 Cuda 网格内核中的 Cupy 索引?
问题描述
我正在尝试开始使用 Cupy 进行一些 Cuda 编程。我需要编写自己的内核。但是,我正在为 2D 内核苦苦挣扎。似乎 Cupy 没有按我预期的方式工作。这是 Numba Cuda 中的 2D 内核的一个非常简单的示例:
import cupy as cp
from numba import cuda
@cuda.jit
def nb_add_arrs(x1, x2, y):
i, j = cuda.grid(2)
if i < y.shape[0] and j < y.shape[1]:
y[i, j] = x1[i, j] + x2[i, j]
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)
# Grid and block sizes
tpb = (16, 16)
bpg = (x1.shape[0] // tpb[0] + 1, x1.shape[1] // tpb[0] + 1)
# Call kernel
nb_add_arrs[bpg, tpb](x1, x2, y)
结果正如预期的那样:
y
[[2 2 2 2 2]
[2 2 2 2 2]
[2 2 2 2 2]
[2 2 2 2 2]
[2 2 2 2 2]]
然而,当我尝试在 Cupy 中做这个简单的内核时,我没有得到相同的结果。
cp_add_arrs = cp.RawKernel(r'''
extern "C" __global__
void add_arrs(const float* x1, const float* x2, float* y, int N){
int i = blockDim.x * blockIdx.x + threadIdx.x;
int j = blockDim.y * blockIdx.y + threadIdx.y;
if(i < N && j < N){
y[i, j] = x1[i, j] + x2[i, j];
}
}
''', '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]
# Grid and block sizes
tpb = (16, 16)
bpg = (x1.shape[0] // tpb[0] + 1, x1.shape[1] // tpb[0] + 1)
# Call kernel
cp_add_arrs(bpg, tpb, (x1, x2, y, cp.int32(N)))
y
[[0 0 0 0 0]
[0 0 0 0 0]
[0 0 0 0 0]
[0 0 0 0 0]
[0 0 0 0 0]]
有人可以帮我弄清楚为什么吗?
解决方案
C 中的内存以行优先顺序存储。所以,我们需要按照这个顺序进行索引。此外,由于我传递的是 int 数组,因此我更改了内核的参数类型。这是代码:
cp_add_arrs = cp.RawKernel(r'''
extern "C" __global__
void add_arrs(int* x1, 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){
y[j + i*N] = x1[j + i*N] + x2[j + i*N];
}
}
''', '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]
# Grid and block sizes
tpb = (16, 16)
bpg = (x1.shape[0] // tpb[0] + 1, x1.shape[1] // tpb[0] + 1)
# Call kernel
cp_add_arrs(bpg, tpb, (x1, x2, y, cp.int32(N)))
y
array([[2, 2, 2, 2, 2],
[2, 2, 2, 2, 2],
[2, 2, 2, 2, 2],
[2, 2, 2, 2, 2],
[2, 2, 2, 2, 2]], dtype=int32)
推荐阅读
- java - 无法使用 Apache Camel 3.1.0 将 XML 转换为 JSON
- ruby-on-rails - SocketError:执行 docker-compose 文件时无法打开与 docker 服务的 TCP 连接
- python - Pytest - 获取使用的固定装置和创建的数据的可见性
- javascript - 将 JSON 文件传递到 php 服务器,然后将其中的数据保存到 Mysql DB
- sql - 如何摆脱动态 SQL (T-SQL)
- c - 确定数组[1...n] 是否为最大堆
- css - 为什么我的 CSS 网格行会相互重叠?
- lisp - 以其他格式返回日期和时间 AutoCAD
- python - ModuleNotFoundError:没有名为“Foundation”的模块
- variables - 在 Ansible 中使用虚线 YAML 变量文件