如何在Julia中将嵌套向量传入GPU内核？

我尝试将嵌套Vector传入GPU内核，但运行失败，以下是最小复现代码：

using CUDA

function kernel(data::CuDeviceVector{CuDeviceVector{Float64, 1}, 1}) 
    # 暂时无操作
end

# 创建包含10个浮点数的向量
floats = Vector{Float64}()
for i in 1:10
    push!(floats, Float64(1))
end

# 将向量包装到另一个向量中
outer = Vector{CuArray{Float64}}()
push!(outer, floats |> CuArray)

# 上传到GPU
final = outer |> CuArray

# 调用内核
@cuda threads=1 kernel(final)

错误发生在final = outer |> CuArray这一行，报错信息：

ERROR: LoadError: CuArray only supports element types that are stored inline

请问如何在GPU中使用嵌套Vector？

解决方案 #

报错的核心原因是：CuArray只能存储可内联（inline）的基础元素类型，而CuArray本身是包含设备指针、尺寸等元数据的复杂类型，无法直接作为另一个CuArray的元素存储。下面提供两种可行的实现方案：

方式一：扁平数组+索引表 #

把所有嵌套子数组合并成一个大的扁平数组，同时维护一个索引数组记录每个子数组的起始位置和长度，适合子数组长度不固定的场景：

using CUDA

# 准备设备上的嵌套子数组
sub_arrays = [CuArray(rand(Float64, 10)), CuArray(rand(Float64, 5)), CuArray(rand(Float64, 8))]

# 展平所有子数组为单个设备数组
flat_data = vcat(sub_arrays...)
# 生成索引表：每个元素存储子数组的起始下标和长度
indices = CUDA.@sync CuArray([(sum(length.(sub_arrays[1:i-1]))+1, length(sub_arrays[i])) for i in 1:length(sub_arrays)])

function kernel(flat_data::CuDeviceVector{Float64,1}, indices::CuDeviceVector{Tuple{Int,Int},1})
    idx = threadIdx().x
    if idx <= length(indices)
        start, len = indices[idx]
        # 通过视图获取子数组
        sub_vec = @view flat_data[start:start+len-1]
        # 这里添加子数组的处理逻辑
        @cuprintln("子数组长度：", len, "，第一个元素：", sub_vec[1])
    end
end

@cuda threads=length(indices) kernel(flat_data, indices)

方式二：设备指针数组 #

创建存储设备指针的数组，在内核中将指针转换为CuDeviceVector使用，适合需要保留原嵌套结构的场景：

using CUDA

# 准备设备上的嵌套子数组
sub_arrays = [CuArray(rand(Float64, 10)), CuArray(rand(Float64, 5)), CuArray(rand(Float64, 8))]

# 存储每个子数组的设备指针
ptr_array = CuArray([pointer(arr) for arr in sub_arrays])
# 存储每个子数组的长度
lengths = CuArray([length(arr) for arr in sub_arrays])

function kernel(ptr_array::CuDeviceVector{Ptr{Float64},1}, lengths::CuDeviceVector{Int,1})
    idx = threadIdx().x
    if idx <= length(ptr_array)
        ptr = ptr_array[idx]
        len = lengths[idx]
        # 将指针转换为设备向量
        sub_vec = CuDeviceVector(len, ptr)
        # 这里添加子数组的处理逻辑
        @cuprintln("子数组长度：", len, "，第一个元素：", sub_vec[1])
    end
end

@cuda threads=length(ptr_array) kernel(ptr_array, lengths)

额外提示 #

如果所有子数组长度完全相同，直接使用二维CuArray（比如CuArray{Float64,2}）是最高效的方式，内核中通过行/列索引即可访问对应子数组。

内容的提问来源于stack exchange，提问作者Bobface