使用 ComponentArrays 和 CUDA 在 Julia 中对大小不等的矩阵进行批量 GPU 矩阵-矩阵乘法

Question

我有两个 ComponentArray，每个都存储 L 个大小兼容的矩阵，我想将它们相乘并存储在第三个 ComponentArray 中。通过使用 for 循环迭代键，这在 CPU 上很容易做到，但我想使用 CUDA.jl 在 GPU 上并行执行所有内容。

这是一个在 CPU 上产生所需结果的示例。

using ComponentArrays

a = ComponentArray(L0 = rand(Float32, 2, 2), L1 = rand(Float32, 3, 3))
b = ComponentArray(L0 = rand(Float32, 2, 2), L1 = rand(Float32, 3, 3))
c = ComponentArray(L0 = zeros(Float32, 2, 2), L1 = zeros(Float32, 3, 3))
names = valkeys(a)

for k in names
    @view(c[k]) .= @view(a[k]) * @view(b[k])
end

现在，如果我们将所有内容移至 GPU：

using CUDA

ag = cu(a)
bg = cu(b)
cg = cu(c)

我们不能再仅仅使用常规的 for 循环进行迭代，因为这非常慢。如下图所示：

for k in names
    @view(cg[k]) .= @view(ag[k]) * @view(bg[k])
end

它返回一个警告，表明使用了标量索引。

问题在于 CUDA 数组的矩阵乘法调用 CUDA 内核，而在常规 CPU for 循环中这样做是一个很大的禁忌。正确的方法是调用一个 CUDA 内核，该内核为每个组件一次执行所有乘法和加法。

现在，库 NNlib.jl 有一个函数 batched_mul！本质上是这样做的，但它不适用于 2D 矩阵的 ComponentArray。相反，它适用于 3D 张量。因此，我可以通过将 ag 和 bg 转换为 3D 张量来使用它，这样在第三维上进行索引相当于对 ag 和 bg 的分量进行索引。然而，这仅在所有组件尺寸相同的情况下才有效。所以我必须将一些零连接到每个数组的 L0 分量。对于我的真实示例，这会浪费大量内存。

那么，有没有像batched_mul这样的东西！对于 GPU 组件数组？我对此非常怀疑，并希望必须编写一个 CUDA 内核。但是，cuBLAS 矩阵乘法经过了令人难以置信的优化和复杂。我是否必须从头开始编写整个 CUDA 内核，包括所有基本的矩阵乘法细节？或者，我是否可以编写一个简单的内核，仅迭代组件（L0、L1...），然后调用 cuBLAS 来处理矩阵乘法？

另请注意，我必须使用 ComponentArray，因为这必须全部发生在 ODE 求解器中。而且，高性能对我来说至关重要。

Answer 1

我认为这里没有发生任何实际的标量索引。

我现在无法立即访问 GPU 机器，但是如果将该循环更改为

会发生什么

for k in names
    getproperty(c, k) .= getproperty(a, k) * getproperty(b, k)
end

我希望这可以消除标量索引警告。

对于示例中的小数组，内核启动开销很可能占主导地位，这就是计算速度慢的原因。我猜你的实际工作量要大得多？

由于循环每次迭代中的计算是相互独立的，因此您可以按照此处的说明异步安排它们：https://cuda.juliagpu.org/stable/usage/multitasking/