按索引访问 GPU 缓冲区

Question

注意：我的问题是关于 Apple 的 Metal API，但我认为这个概念足够通用，也可以转化为其他 GPU 框架。

我的目标：在M x N 矩阵A 的每一行中添加一个1 x N 行向量b。

我的内核，简化为我遇到问题的部分：

kernel void vmadd(const device float* A [[ buffer(0) ]],
                  const device float* b [[ buffer(1) ]],
                  device float* C [[ buffer(2) ]],
                  constant ushort& aWidth [[ buffer(3) ]],
                  ushort2 gid [[ thread_position_in_grid ]]) {

    int idx = gid.y * aWidth + gid.x; // Compute absolute index in C
    C[idx] = A[idx] + b[gid.x];

}

假设：我的理解是gid是C中单个元素的位置：gid.x是列，gid.y是行。如果这不是真的，请有人纠正我。

现在，如果我用 8 x 8 个零填充 A：

A = 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 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 0 0 0 0 0 0
    0 0 0 0 0 0 0 0

和b 像这样：

b = 1 2 3 4 5 6 7 8

然后在执行之后，C 应该是一个 8 x 8 矩阵，其中每一行是 1 2 3 4 5 6 7 8。

相反，我得到了这个：

C = 1 2 3 4 5 6 7 8
    1 2 3 4 5 6 7 8
    1 2 3 4 5 6 7 8
    0 0 0 0 0 0 0 0
    1 2 3 4 5 6 7 8
    1 2 3 4 5 6 7 8
    1 2 3 4 5 6 7 8
    0 0 0 0 0 0 0 0

我正在根据 Apple 的建议 here 调整线程组的大小：

let w = computePipeline.threadExecutionWidth
let h = computePipeline.maxTotalThreadsPerThreadgroup / w
let threadsPerThreadgroup = MTLSizeMake(w, h, 1)
let threadgroupsPerGrid = MTLSize(width: (cWidth + w - 1) / w,
                                  height: (cHeight + h - 1) / h,
                                  depth: 1)

在我的机器上，每个组产生(64, 16, 1) 线程，每个网格产生(1, 1, 1) 线程组。

但是，如果我手动将threadsPerThreadgroup 设置为(1, 1, 1) 并将threadgroupsPerGrid 设置为(8, 8, 1)，我会在C 中得到正确的结果。

问题：

我相当肯定我的问题与线程组大小和管理缓冲区的方式有关，但我对 GPU 编程还很陌生，所以我并不完全理解它。

为什么减小线程组的大小会产生正确的结果？更一般地说，为什么线程组大小对这个计算有任何影响？

在我看来，如果gid 总是对应于C 中的索引，并且我要求b[gid.x]，那么该值应该可用。我知道数据通常会被分解以适合线程组缓存 - 如果是这种情况，它遵循什么规则，我该如何解释？

Answer 1

我的理解是gid是C中单个元素的位置：gid.x是列，gid.y是行。如果这不是真的，请有人纠正我。

这并不完全正确。 gid 是 网格 中的位置。

因为网格恰好是 64x16，所以计算函数将在 8x8 矩阵（A 和 C）和 8 元素向量（b）之外的位置调用。发生这种情况时，读取A 可能会访问错误的行，甚至超出A 的末尾。同样，b 的读取将超出其末尾。

例如，考虑gid 何时为 (8, 0)。 idx 将是 8。您将阅读 A[8]，它实际上位于 (0, 1)。你会读到b[8]，它已经结束了。这在技术上是未定义的，但实际上很可能是 0 对于相对较短长度的缓冲区。您将写信给同样位于 (0, 1) 的 C[8]。这大致与函数调用同时发生，该函数调用应该写在 (0, 1) 并且有一场关于哪个占上风的竞赛。

您的函数应该在接近开始时测试gid 是否超出范围，如果超出范围，请提前返回：

if (any(gid > aWidth))
    return;

（假设A 和C 始终为正方形，因此可以根据单个值检查宽度和高度。）

您可以尝试调整 threadsPerThreadgroup 和 threadgroupsPerGrid 的计算，以使网格的大小与您的矩阵的大小完全相同，但在所有情况下正确执行此操作可能会很乏味。也就是说，您当然可以避免 threadsPerThreadgroup 太大：

let w = min(computePipeline.threadExecutionWidth, cWidth)
let h = min(computePipeline.maxTotalThreadsPerThreadgroup / w, cHeight)

但是您仍然需要在计算函数中进行检查，因为总网格仍然可能太大。例如，假设 computePipeline.threadExecutionWidth 至少为 8，computePipeline.maxTotalThreadsPerThreadgroup 至少为 60。那么，w 将是 8，但 h 将是 7。那么，threadgroupsPerGrid 将是 (1, 2, 1) 并且总网格大小为 8x14x1，再次大于您的矩阵。