Julia中基于ComponentArrays和CUDA的异构尺寸矩阵GPU批量乘法

一、核心问题分析 #

直接在CPU循环中处理GPU端ComponentArray的组件，会触发频繁的主机-设备数据交互，不仅速度慢，还会出现标量索引警告。必须让所有计算逻辑在GPU端批量完成，避免CPU作为中间调度层。

二、现有工具的局限性 #

NNlib.jl的batched_mul!仅支持同尺寸矩阵的3D张量批量乘法，若为不同尺寸的组件补零适配，会造成大量内存浪费，不符合ODE求解的高性能需求，因此需要自定义实现。

三、自定义GPU端批量乘法实现 #

不需要从头编写矩阵乘法CUDA核，我们可以利用CUDA.jl对cuBLAS的封装，在GPU核函数中遍历组件，直接调用优化后的矩阵乘操作。

实现步骤 #

1. 预收集组件元信息：在CPU端提前获取每个组件的尺寸和在数组中的线性偏移，避免在GPU核中处理ComponentArray的元数据（ComponentArray的结构信息仅存储在CPU端）。
2. 编写CUDA核函数：让每个线程块负责一个组件的矩阵乘法，通过指针直接访问GPU内存中的组件数据，调用cuBLAS优化的mul!完成计算。
3. 批量启动核函数：将组件元信息转移到GPU，按组件数量启动对应线程块，一次性完成所有组件的乘法。

完整代码示例 #

using ComponentArrays, CUDA, LinearAlgebra

# 生成CPU端测试数据
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))

# 转移数据到GPU
ag = cu(a)
bg = cu(b)
cg = cu(c)

# CPU端收集组件的尺寸与线性偏移信息
component_info = []
for k in valkeys(a)
    a_view = @view(a[k])
    b_view = @view(b[k])
    c_view = @view(c[k])
    push!(component_info, (
        m = size(a_view, 1),
        n = size(b_view, 2),
        k_size = size(a_view, 2),
        a_offset = first(LinearIndices(a)) + (first(LinearIndices(a_view)) - 1),
        b_offset = first(LinearIndices(b)) + (first(LinearIndices(b_view)) - 1),
        c_offset = first(LinearIndices(c)) + (first(LinearIndices(c_view)) - 1)
    ))
end

# 将组件信息转移到GPU
component_info_gpu = cu(component_info)

# 定义CUDA核函数：每个线程块处理一个组件的矩阵乘法
function component_mul_kernel!(c, a, b, info)
    idx = blockIdx().x
    idx > length(info) && return
    
    comp = info[idx]
    # 通过指针直接获取GPU内存中的组件数据
    a_ptr = pointer(a) + (comp.a_offset - 1) * sizeof(eltype(a))
    b_ptr = pointer(b) + (comp.b_offset - 1) * sizeof(eltype(b))
    c_ptr = pointer(c) + (comp.c_offset - 1) * sizeof(eltype(c))
    
    # 包装为设备数组
    A = CuDeviceArray((comp.m, comp.k_size), a_ptr)
    B = CuDeviceArray((comp.k_size, comp.n), b_ptr)
    C = CuDeviceArray((comp.m, comp.n), c_ptr)
    
    # 调用cuBLAS优化的矩阵乘法
    mul!(C, A, B)
    return
end

# 启动核函数：每个组件对应一个线程块
@cuda threads=1 blocks=length(component_info) component_mul_kernel!(cg, ag, bg, component_info_gpu)

# 验证结果（可选）
cg_cpu = Array(cg)
for k in valkeys(a)
    @assert isapprox(cg_cpu[k], a[k] * b[k]; atol=1e-5)
end

四、ODE求解器中的性能优化要点 #

• 预计算组件信息：在ODE初始化阶段一次性计算组件的尺寸和偏移，避免每个求解步长重复计算。
• 避免动态内存分配：核函数中使用CuDeviceArray直接包装指针，不进行额外内存分配，保证计算效率。
• 利用CUDA流：若ODE求解器支持多流，可以将组件乘法放到独立CUDA流中，与其他计算逻辑重叠执行。

五、轻量化替代方案 #

如果组件数量较少，也可以用CUDA.jl的异步并行机制优化CPU循环，虽然是CPU发起调度，但避免了串行等待：

@sync for k in valkeys(a)
    @async begin
        @view(cg[k]) .= @view(ag[k]) * @view(bg[k])
    end
end

这种方式实现简单，性能优于串行循环，但不如自定义核函数高效，适合组件数量不多的场景。

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