如何在cuda python中获取当前工作网格索引

Question

您好我正在尝试了解cuda内核的每一步。获得数据占用的所有网格索引都会很好。我的代码是添加2个向量，用python numba编写。

n = 10 
x = np.arange(n).astype(np.float32)
y = x + 1   

设置网格中的线程和块数

threads_per_block = 8
blocks_per_grid = 2

内核

def kernel_manual_add(x, y, out):
    threads_number = cuda.blockDim.x
    block_number = cuda.gridDim.x

    thread_index = cuda.threadIdx.x
    block_index = cuda.blockIdx.x

    grid_index = thread_index + block_index * threads_number  
    threads_range = threads_number * block_number

    for i in range(grid_index, x.shape[0], threads_range):
        out[i] = x[i] + y[i]

初始化内核：

kernel_manual_add[blocks_per_grid, threads_per_block](x, y, out)

当我尝试打印出grid_index时，我得到所有输入索引2 * 8.

如何获取用于计算数据的网格索引（其中10个）？

Answer 1

编写内核的规范方法就是这样

@cuda.jit
def kernel_manual_add(x, y, out):

    i = cuda.grid(1)
    if i < x.shape[0]:
        out[i] = x[i] + y[i]

您必须至少运行与输入数组中的元素一样多的线程。这里没有魔力，你需要在调用内核之前手动计算网格和块尺寸。有关建议，请参阅here和here。