突破 WebGPU 算力瓶颈：现代 GPU 架构下的并行前缀和（Prefix Sum）极致优化指南

并行前缀和（Prefix Sum，又称 Scan）是并行计算中最基础且最重要的算法骨架之一。从物理引擎（如粒子系统、流体模拟）、GPU 排序（如 Radix Sort），到光线追踪（BVH 树构建）及无损数据压缩，Scan 算法的吞吐量直接决定了上层应用的性能上限。

在 WebGPU 环境下，虽然 WGSL（WebGPU Shading Language）屏蔽了底层平台的差异，但底层的硬件逻辑（如 VRAM 带宽限制、LDS 银行冲突、Warp/Wavefront 调度）依然存在。本文将深入现代 GPU 架构，探讨如何在 WebGPU 中榨干硬件性能，实现超大规模数据下 Prefix Sum 的极致优化。

1. 瓶颈分析：GPU 硬件架构视角

在编写任何优化代码之前，必须先明确硬件的底层约束。

• 内存带宽受限（Memory Bound）：Scan 算法的算术强度（Arithmetic Intensity）极低，每个输入元素通常只需要几次加法操作。这意味着算法的瓶颈不在 ALU（算术逻辑单元），而在 VRAM（显存）带宽。任何多余的全局内存（Global Memory）读写都会带来灾难性的吞吐量下滑。
• 共享内存冲突（LDS Bank Conflicts）：在 GPU 局部共享内存（WebGPU 中的 workgroup 存储）中，内存被划分为多个等宽的 Bank（通常为 32 个，每个 Bank 宽 4 字节）。如果一个工作组（Workgroup）内的多个线程同时访问同一个 Bank 的不同地址，就会引发串行化延迟（Bank Conflict）。
• 指令发射与分支分化（Branch Divergence）：传统的 Blelloch 或 Kogge-Stone 算法在树状结构的后期阶段，会有大量线程处于闲置状态（使用 if 分支过滤）。这会导致严重的 SIMD 效率低下。

2. 局部优化：工作组级 Scan 的极致避坑

在单个 Workgroup 内部，我们通常使用共享内存来缓存数据并进行局部扫面。最常用的是 Kogge-Stone（步骤少，工作量稍大）和 Blelloch（工作量低，但步骤较多且容易产生 Bank 冲突）算法。

这里我们以 Blelloch Scan 为例，展示如何消除 Bank 冲突。

消除 Bank 冲突（Bank Conflict Free Padding）

当线程通过 stride 访问局部内存时，如果步长是 2 的幂（2, 4, 8, 16...），就会产生严重的 Bank 冲突。解决方法是插入空填充（Padding）。通过对索引进行位移映射，使原本映射到相同 Bank 的数据错开。

// 128 线程的工作组，需要扫描 256 个元素
const WORKGROUP_SIZE = 128u;
const SHARED_MEM_SIZE = 256u;

// 通过 Padding 避免 Bank 冲突的索引转换函数
fn lds_index(idx: u32) -> u32 {
    // 每 32 个元素插入一个 4 字节的空隙
    return idx + (idx >> 5u); 
}

var<workgroup> temp: array<u32, SHARED_MEM_SIZE + (SHARED_MEM_SIZE >> 5u)>;

局部扫描完整 WGSL 实现

结合 Padding 技术，以下是一个高效的工作组级 Exclusive Scan 实现：

@compute @workgroup_size(WORKGROUP_SIZE)
fn main(
    @builtin(local_invocation_id) local_id: vec3<u32>,
    @builtin(workgroup_id) group_id: vec3<u32>
) {
    let tid = local_id.x;
    
    // 1. 协作式加载：每个线程加载 2 个元素到工作组内存（LDS）
    let idx_a = tid * 2u;
    let idx_b = tid * 2u + 1u;
    
    let lds_a = lds_index(idx_a);
    let lds_b = lds_index(idx_b);

    // 假设 input_buffer 是我们的全局数据源
    temp[lds_a] = input_buffer[group_id.x * SHARED_MEM_SIZE + idx_a];
    temp[lds_b] = input_buffer[group_id.x * SHARED_MEM_SIZE + idx_b];

    // 2. Up-Sweep (Reduction) 阶段
    var offset = 1u;
    for (var d = SHARED_MEM_SIZE >> 1u; d > 0u; d >>= 1u) {
        workgroupBarrier();
        if (tid < d) {
            let ai = lds_index(offset * (2u * tid + 1u) - 1u);
            let bi = lds_index(offset * (2u * tid + 2u) - 1u);
            temp[bi] += temp[ai];
        }
        offset <<= 1u;
    }

    // 3. 清零最后一个元素以实现 Exclusive Scan
    if (tid == 0u) {
        temp[lds_index(SHARED_MEM_SIZE - 1u)] = 0u;
    }

    // 4. Down-Sweep 阶段
    for (var d = 1u; d < SHARED_MEM_SIZE; d <<= 1u) {
        offset >>= 1u;
        workgroupBarrier();
        if (tid < d) {
            let ai = lds_index(offset * (2u * tid + 1u) - 1u);
            let bi = lds_index(offset * (2u * tid + 2u) - 1u);
            
            let t = temp[ai];
            temp[ai] = temp[bi];
            temp[bi] += t;
        }
    }
    workgroupBarrier();

    // 5. 写回全局显存
    output_buffer[group_id.x * SHARED_MEM_SIZE + idx_a] = temp[lds_a];
    output_buffer[group_id.x * SHARED_MEM_SIZE + idx_b] = temp[lds_b];
}

3. 全局优化：处理超大数据的架构方案

当输入数据量极大（例如 $10^7$ 级），单个工作组无法一次性处理时，必须在多个工作组之间建立协同机制。目前主要有两种流派：

方案 A：多级 Reduce-Scan-Scan（经典三阶段法）

这是最通用且兼容性最好的方案。将全局数据划分为多个 Block：

1. Reduce 阶段：每个 Workgroup 计算自己 Block 的局部和，并将该 Block Sum 写入临时的辅助 Buffer。
2. Scan 阶段：对辅助 Buffer（Block Sums）进行单工作组 Scan，得到每个 Block 的全局偏移地址。
3. Scan & Add 阶段：再次启动 Compute Pipeline，每个 Workgroup 扫描自己的数据，并在最后加上第二阶段计算出的该 Block 全局偏移量。

优点：无跨工作组同步依赖，完全符合 WebGPU 标准管线，兼容各种中低端 GPU。
缺点：需要多次 Pipeline 切换，数据不得不反复读写 VRAM，在显存带宽限制严重的设备上开销明显。

方案 B：单通道扫描（Single-Pass Scan / Decoupled Look-back）

这是现代 GPU（如 CUDA CUB 库）中主流的高效扫描方案。它通过 Cooperative Groups（协作组） 和原子操作，在单次 Pipeline Dispatch 内完成所有数据的 Scan。

在 WebGPU 中，由于缺乏严格的跨工作组硬同步屏障（Barrier），实现 Decoupled Look-back 需要极其精妙的设计：

1. 动态块分配（Dynamic Block Indexing）：使用一个全局原子变量 atomicAdd 来动态获取当前工作组要处理的 Block ID，防止因 GPU 调度顺序不确定导致的死锁。
2. 状态标志位（Status Flags）：为每个 Block 分配一个状态字：
   • X（无有效数据）
   • A（该 Block 的局部 Reduce 和已计算完毕）
   • P（该 Block 到起点的全局 Prefix Sum 已计算完毕）
3. 回看机制（Look-back）：第 $i$ 个 Block 率先完成局部 Reduction 后，向前的 Block（$i-1, i-2, \dots$）进行回看。如果前驱 Block 的状态是 P，直接累加并更新自己的状态为 P；如果是 A，则累加该局部和并继续向前看，直到遇到 P 为止。

WebGPU 实现难点：
WGSL 自 WebGPU 1.0 起支持 atomic 类型和原子操作，但必须要非常小心地处理内存可见性。目前由于缺少 atomic_thread_fence 的细粒度控制，单通道 Scan 在部分非统一内存架构（NUMA）GPU 上可能面临正确性挑战，通常需要依赖 atomicExchange 或 atomicCompareExchangeWeak 配合 workgroupBarrier() 进行硬件状态推进。

4. 极致硬件加速：Subgroup（子群）操作

在现代 GPU 架构（NVIDIA Ampere/Ada Lovelace, AMD RDNA, Apple Silicon）中，最核心的物理执行单元是 Warp（NVIDIA）或 Wavefront（AMD）。WebGPU 近期引入了 Subgroups 扩展（在 Chrome 128+ 中已实验性支持）。

Subgroup 操作是提升 Scan 吞吐量的终极武器。 它允许同一个子群内的线程直接通过寄存器进行数据交换（Shuffle Instruction），完全绕过 L1 Cache 和工作组共享内存（LDS），延迟骤降一个数量级。

启用 Subgroups 扩展

enable subgroups; // 必须在 WGSL 头部声明

使用 Subgroup 算子重构局部 Scan

利用 subgroupExclusiveAdd 或 subgroupInclusiveAdd，我们可以在一行代码内完成 Warp 级的扫描，不需要任何 LDS 声明和同步屏障。

@compute @workgroup_size(256)
fn main(
    @builtin(local_invocation_id) local_id: vec3<u32>,
    @builtin(subgroup_invocation_id) sg_sub_id: u32,  // 子群内线程 ID (0~31 或 0~63)
    @builtin(subgroup_id) sg_id: u32                  // 子群 ID
) {
    let global_idx = ...;
    let value = input_buffer[global_idx];

    // 1. 单步子群内 Exclusive Scan（无锁，寄存器级，极快）
    let sg_prefix_sum = subgroupExclusiveAdd(value);

    // 2. 子群间协作：每个子群的最后一个线程将局部和写入共享内存
    var<workgroup> subgroup_sums: array<u32, 8>; // 假设子群大小为 32，256/32 = 8 个子群
    
    let sg_size = subgroupGetSize(); // 动态获取子群大小
    if (sg_sub_id == sg_size - 1u) {
        subgroup_sums[sg_id] = sg_prefix_sum + value; // 保存该子群的总体和
    }
    workgroupBarrier();

    // 3. 由第一个子群对共享内存中的子群和进行第二次扫描
    if (sg_id == 0u) {
        let lds_val = subgroup_sums[sg_sub_id];
        let lds_prefix = subgroupExclusiveAdd(lds_val);
        subgroup_sums[sg_sub_id] = lds_prefix;
    }
    workgroupBarrier();

    // 4. 将子群前缀累加回每个线程
    let final_prefix_sum = sg_prefix_sum + subgroup_sums[sg_id];
    output_buffer[global_idx] = final_prefix_sum;
}

性能对比：使用 Subgroup 重构后，因为消除了大量的 workgroupBarrier() 和显存/LDS 的读写往返，局部 Scan 的吞吐量通常能获得 2x - 4x 的跃升。

5. 向量化（Vectorization）与显存对齐

为了逼近 GPU 物理带宽极限（如 PCIe / HBM / GDDR 的吞吐上限），必须实现显存合并访问（Coalesced Memory Access）。

• 128-bit 宽度向量化：GPU 的内存控制器每次传输通常是 32 或 64 字节。如果一个线程只读写一个 u32（4字节），内存总线效率较低。
• 在加载和写回阶段，建议使用 vec4<u32>。每个线程一次性读取 16 字节，充分填满显存总线，在着色器内部再将 vec4 展开处理。

// 向量化加载示例
let vec_data: vec4<u32> = global_input_vector[offset];
// 在寄存器内进行局部累加
let sum = vec_data.x + vec_data.y + vec_data.z + vec_data.w;

总结：WebGPU Scan 优化路线图

在具体的 WebGPU 项目开发中，建议采用以下进化路径来获取最大吞吐：

在主流显卡（如 RTX 4060 / Apple M3）上，通过 Subgroups + 向量化 + Bank Conflict 消除，WebGPU 可以在毫秒级处理千万级数据的并行前缀和，让基于 Web 的大型通用计算及 AAA 级图形管线的性能瓶颈彻底成为历史。