迭代非方形二维数组时如何实现内存合并？

Question

我正在努力弄清楚 CUDA 中的内存合并。为了评估合并和未合并内存访问之间的性能差异，我实现了两个不同版本的内核，它们添加了两个 2D 矩阵：

from numba import cuda

@cuda.jit
def uncoalesced_matrix_add(a, b, out):
    x, y = cuda.grid(2)
    out[x][y] = a[x][y] + b[x][y]

@cuda.jit
def coalesced_matrix_add(a, b, out):
    x, y = cuda.grid(2)
    out[y][x] = a[y][x] + b[y][x]

当我用方矩阵测试上面的代码时，一切正常，我的意思是，两个内核产生相同的结果，合并版本明显更快：

import numpy as np

nrows, ncols = 2048, 2048
tpb = 32

threads_per_block = (tpb, tpb)
blocks = ((nrows + (tpb - 1))//tpb, (ncols + (tpb - 1))//tpb)

size = nrows*ncols

a = np.arange(size).reshape(nrows, ncols).astype(np.int32)
b = np.ones(shape=a.shape, dtype=np.int32)
out = np.empty_like(a).astype(np.int32)

d_a = cuda.to_device(a)
d_b = cuda.to_device(b)
d_out = cuda.to_device(out)

uncoalesced_matrix_add[blocks, threads_per_block](d_a, d_b, d_out)
slow = d_out.copy_to_host()

coalesced_matrix_add[blocks, threads_per_block](d_a, d_b, d_out)
fast = d_out.copy_to_host()

np.array_equal(slow, fast)
# True

但是，如果我更改 nrows = 1024 使矩阵不再是正方形，coalesced_matrix_add() 会抛出以下错误：

CudaAPIError: [700] Call to cuMemcpyDtoH results in UNKNOWN_CUDA_ERROR

我在这里错过了什么？

编辑

为了完整起见，我附上了一些分析。这些数据是通过使用@Robert Crovella 提出的 nrows = 1024 和 ncols = 2048:

提出的解决方法获得的

In [40]: %timeit uncoalesced_matrix_add[blocksu, threads_per_block](d_a, d_b, d_out)
289 µs ± 498 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [41]: %timeit coalesced_matrix_add[blocksc, threads_per_block](d_a, d_b, d_out)
164 µs ± 108 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Answer 1

当你制作数组 non-square 时，这个计算不再正确，因为你的两个内核中的 :

blocks = ((nrows + (tpb - 1))//tpb, (ncols + (tpb - 1))//tpb)

对于线程和块维度，首先是 x 索引维度，然后是 y 索引维度（指的是出现的 x 和 y在 in-kernel built-in 变量 x 和 y).

在你的第一个内核中的用法：

out[x][y] = a[x][y] + b[x][y]

我们希望 x 对行进行索引。这与您的网格定义一致。

在你的第二个内核中的用法：

out[y][x] = a[y][x] + b[y][x]

我们希望 y 对行进行索引。这与您的网格定义不一致。

结果是 out-of-bounds 在第二次内核调用时访问。矩形网格的方向与矩形数据的方向不匹配。

在正方形的情况下，这种反转并不重要，因为两个维度相同。

这是一个可能的“修复”：

$ cat t62.py
from numba import cuda
import numpy as np

@cuda.jit
def uncoalesced_matrix_add(a, b, out):
    x, y = cuda.grid(2)
    out[x][y] = a[x][y] + b[x][y]

@cuda.jit
def coalesced_matrix_add(a, b, out):
    x, y = cuda.grid(2)
    out[y][x] = a[y][x] + b[y][x]


nrows, ncols = 512, 1024
tpb = 32

threads_per_block = (tpb, tpb)
blocksu = ((nrows + (tpb - 1))//tpb, (ncols + (tpb - 1))//tpb)
blocksc = ((ncols + (tpb - 1))//tpb, (nrows + (tpb - 1))//tpb)

size = nrows*ncols

a = np.arange(size).reshape(nrows, ncols).astype(np.int32)
b = np.ones(shape=a.shape, dtype=np.int32)
out = np.empty_like(a).astype(np.int32)

d_a = cuda.to_device(a)
d_b = cuda.to_device(b)
d_out = cuda.to_device(out)

uncoalesced_matrix_add[blocksu, threads_per_block](d_a, d_b, d_out)
slow = d_out.copy_to_host()

coalesced_matrix_add[blocksc, threads_per_block](d_a, d_b, d_out)
fast = d_out.copy_to_host()

print(np.array_equal(slow, fast))
# True
$ python t62.py
True
$

另请注意，此网格大小调整策略仅适用于 whole-number 可被块大小整除的维度。

迭代非方形二维数组时如何实现内存合并？

How to achieve memory coalescing when iterating over non-square 2D arrays?

python

cuda

gpu

memory-access

numba

编辑