如何在Vulkan计算着色器中实现并行求和归约算法

Vulkan 并行求和归约实现方案 #

Vulkan GLSL中已经提供了和OpenCL地址空间完全对应的修饰符，你可以直接对应使用：

• OpenCL的__global对应Vulkan GLSL的存储缓冲区，用buffer块声明，绑定到描述符集的对应插槽即可，读写行为和OpenCL完全一致
• OpenCL的__local对应Vulkan GLSL的shared修饰符，声明的变量会分配在工作组本地内存，无需在宿主端显式申请，由驱动自动管理

计算着色器示例代码 #

#version 450
#extension GL_KHR_vulkan_glsl : enable

// 全局输入缓冲区：存储待求和的原始数据，对应OpenCL __global readonly
layout(set = 0, binding = 0) readonly buffer InputBuf {
    float elements[];
} input_buf;

// 全局输出缓冲区：存储每个工作组的归约结果，对应OpenCL __global writeonly
layout(set = 0, binding = 1) writeonly buffer OutputBuf {
    float elements[];
} output_buf;

// 工作组大小设为256，可根据你的硬件情况调整
layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
// 本地内存数组，对应OpenCL __local float local_cache[256]
shared float local_cache[256];

void main() {
    uint local_idx = gl_LocalInvocationID.x;
    uint global_idx = gl_GlobalInvocationID.x;

    // 加载全局内存数据到本地内存
    local_cache[local_idx] = input_buf.elements[global_idx];
    // 等待工作组内所有线程完成加载
    memoryBarrierShared();
    barrier();

    // 对数步长归约
    for (uint step = 128; step > 0; step >>= 1) {
        if (local_idx < step) {
            local_cache[local_idx] += local_cache[local_idx + step];
        }
        // 等待当前步长所有计算完成，避免读写冲突
        memoryBarrierShared();
        barrier();
    }

    // 每个工作组只需要0号线程把最终结果写回全局内存
    if (local_idx == 0) {
        output_buf.elements[gl_WorkGroupID.x] = local_cache[0];
    }
}

宿主端执行流程 #

1. 首先计算第一轮调度的工作组数量：假设输入数据总长度为N，工作组大小为256，则工作组数量为ceil(N / 256)，调度计算管线完成第一轮归约
2. 如果第一轮输出的结果长度（也就是工作组数量）仍大于256，就把上一轮的输出作为下一轮的输入，重复调度计算管线，直到输出结果只剩1个元素，就是最终的求和结果
3. 如果需要适配不同长度的输入，也可以把本地内存声明为不定长数组shared float local_cache[];，在管线创建时指定本地内存的大小即可，不用每次修改着色器代码

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