Julia CUDA - 在没有 CPU 的情况下保存中间内核结果

Question

考虑以下 CUDA 内核，它计算二维矩阵每一行的平均值。

using CUDA

function mean!(x, n, out)
    """out = sum(x, dims=2)"""
    row_idx = (blockIdx().x-1) * blockDim().x + threadIdx().x
    for i = 1:n
        @inbounds out[row_idx] += x[row_idx, i]
    end
    out[row_idx] /= n
    return
end

using Test

nrow, ncol = 1024, 10
x = CuArray{Float64, 2}(rand(nrow, ncol))
y = CuArray{Float64, 1}(zeros(nrow))
@cuda threads=256 blocks=4 row_sum!(x, size(x)[2], y)
@test isapprox(y, sum(x, dims=2))  # test passed

还要考虑下面的CUDA内核

function add!(a, b, c)
    """ c = a .+ b """
    i = (blockIdx().x-1) * blockDim().x + threadIdx().x
    c[i] = a[i] + b[i]
    return
end

a = CuArray{Float64, 1}(zeros(nrow))
b = CuArray{Float64, 1}(ones(nrow))
c = CuArray{Float64, 1}(zeros(nrow))
@cuda threads=256 blocks=4 add!(a, b, c)
@test all(c .== a .+ b)  # test passed

现在，假设我想编写另一个使用 mean!() 的中间结果的内核。例如，

function g(x, y)
    """ mean(x, dims=2) + mean(y, dims=2) """
    xrow, xcol = size(x)
    yrow, ycol = size(y)
    
    mean1 = CuArray{Float64, 1}(undef, xrow)
    @cuda threads=256 blocks=4 mean!(x, xcol, mean1)

    mean2 = CuArray{Float64, 1}(zeros(yrow))
    @cuda threads=256 blocks=4 mean!(y, ycol, mean2)

    out = CuArray{Float64, 1}(zeros(yrow))
    @cuda threads=256 blocks=4 add!(mean1, mean2, out)
    return out
end

（当然，g() 在技术上不是内核，因为它 returns 什么的。）

我的问题是 g() 是否“正确”。特别是 g() 在 GPU/CPU 之间传输数据是否浪费时间？

例如，如果我的理解是正确的，可以优化 g() 的一种方法是像初始化 mean1 一样初始化 mean2。这是因为在构造 mean2 时，我们实际上首先在 CPU 上创建 zeros(yrow)，然后将其传递给 CuArray 构造函数以复制到 GPU。相反，mean1 被构造但未初始化（由于 undef）因此避免了这种额外的传输。

总而言之，我如何 save/use 中间内核结果，同时尽可能避免 CPU/GPU 之间的数据传输？

Answer 1

您可以直接在 GPU 上生成零数组或向量！

尝试：

CUDA.zeros(Float64, nrow)

一些基准：

julia> @btime CUDA.zeros(Float64, 1000,1000)
  12.600 μs (26 allocations: 1.22 KiB)
1000×1000 CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}:
...

julia> @btime CuArray(zeros(1000,1000))
  3.551 ms (8 allocations: 7.63 MiB)
1000×1000 CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}:
...

Julia CUDA - 在没有 CPU 的情况下保存中间内核结果

Julia CUDA - Saving intermediate kernel results without CPU

cuda

julia

gpu