如何在 Julia 中将矩阵数组存储在 CUDA 数组中

Question

假设我有一个数组

A

，它在每个索引/位置

A[x]

处存储一个2×2矩阵

x

。在我的代码中，我想要执行基本操作，例如逐点乘法和加法，但也对数组进行一个元素的循环移位，然后与同类数组进行逐点乘法。

我想利用我的 GPU 来完成这些任务。我熟悉如何为特定任务编写内核函数，但我不确定如何分配数组以充分利用我的计算资源。

我已经明白了

A=[CUDA.zeros(Float32, 2, 2) for x in 1:100]

可以解决，但我不太确定这是否会将整个数组结构保存在 GPU 上（是吗？）。由于我无法分配具有可变元素的 CUDA 数组，如

A=CuArray([CUDA.zeros(Float32, 2, 2) for x in 1:100])

导致

CuArray 仅支持内联分配的元素类型。
CuArray{Float32, 2, CUDA.Mem.DeviceBuffer} 是可变类型

所以我想知道在 Julia 中是否有一种自然的方法可以做到这一点。另外，有没有办法将矩阵保存为指定大小（但可变）的数组（因此是大小数组的 CUDA 数组）？

我知道我可以将相同的信息存储在数组中

A=CUDA.zeros(Float32, 100, 4) 

但是，我不确定如何利用

LinearAlgebra

库进行矩阵运算。

编辑： 我尝试遵循评论中的建议，并尝试过

function kernel_add!(C, A, B, N)
    idx = threadIdx().x + (blockIdx().x - 1) * blockDim().x
    if idx <= N
        C[:, :, idx] = A[:, :, idx] + B[:, :, idx]
    end
    return
end

在数组上执行

N = 100
A = CUDA.ones(Float32, 2, 2, N);
B = CUDA.zeros(Float32, 2, 2, N);
C = CUDA.zeros(Float32, 2, 2, N);

与

@cuda threads = 512 blocks = cld(N, 512) kernel_add!(C, A, B, N)

但是，我无法使其正常工作。在核函数中使用数组切片有问题吗？这是错误消息：

InvalidIRError：为 kernel_add 编译 MethodInstance！(::CuDeviceArray{Float32, 3, 1}, ::CuDeviceArray{Float32, 3, 1}, ::CuDeviceArray{Float32, 3, 1}, ::Int64) 导致无效的 LLVM红外 原因：不支持通过文字指针进行调用（调用 ijl_alloc_string）

Answer 1

您还可以创建一个三阶张量 CUDA.zeros(2,2,100)，并且可能使用 Tullio.jl 进行将使用 GPU 的复杂操作