深入 WebGPU Compute Shader：如何实现高性能 2D FFT 图像频域处理

在传统的 Web 图像处理中，空间域滤波（如高斯模糊、边缘检测）通常通过卷积核（Convolution Kernel）实现。当卷积核尺寸较大（例如 $51x51$ 以上）时，空间域卷积的计算复杂度会急剧上升至 $O(N^2 \cdot K^2)$。

利用傅里叶变换（FFT），我们可以将图像转换到频域，将复杂的空间域卷积简化为频域的点乘（Hadamard 积），计算复杂度直接降至 $O(N^2 \log N)$。

随着 WebGPU 标准的落地，借助 Compute Shader（计算着色器）和 Workgroup Shared Memory（工作组共享内存），我们可以在浏览器端实现工业级的实时 2D FFT 图像处理。本文将深入探讨如何在 WebGPU 中构建高性能的 2D FFT 渲染管线。

一、 算法选型：为什么不要在 GPU 上生搬硬套 Cooley-Tukey？

经典教科书中的 Cooley-Tukey FFT 算法采用“蝶形算子”进行分治，但它有一个致命的缺点：需要进行比特反转（Bit-Reversal）置换。

在 CPU 上，比特反转的开销尚可接受；但在 GPU 上，这种非连续、跨步长的内存访问会导致极严重的非合并内存访问（Uncoalesced Memory Access），极大地压低显存带宽利用率。

推荐方案：Stockham 自动排序算法

在 GPU 编程中，Stockham 算法是更佳的选择。它的核心优势在于：

1. 无需显式的比特反转置换。
2. 每一轮迭代（Pass）的输入和输出都是自然顺序。
3. 极其适合利用 GPU 的 Workgroup Shared Memory（L1 缓存级别） 进行乒乓交替读写（Ping-pong buffering）。

二、 核心架构：行列法（Row-Column Method）实现 2D FFT

对一副 $W \times H$ 的图像进行 2D FFT，通常采用一维分步法：

+------------+     1D FFT (按行)      +------------+       转置 (Transpose)      +------------+
|  原始图像  | ---------------------> |  行变换结果 | -------------------------> |  转置后图像 |
+------------+                        +------------+                             +------------+
                                                                                       |
+------------+       转置 (Transpose)     +------------+     1D FFT (按行)             |
|  2D 频域   | <--------------------- |  列变换结果 | <------------------------+
+------------+                        +------------+

整个 Pipeline 由以下 Compute Pass 组成：

1. Row FFT Pass：对每一行数据进行 1D FFT。
2. Transpose Pass：将矩阵行列转置（使列数据在内存中连续，便于下一步处理）。
3. Column FFT Pass（实际上是对转置后的行再次进行 1D FFT）。
4. Transpose Pass：再次转置还原。

三、 手撕 WGSL：基于共享内存的一维 FFT 实现

以下是基于 WebGPU 的 WGSL 实现的一维基-2（Radix-2）Stockham FFT 核心 Compute Shader。该着色器利用了工作组共享内存来消除全局内存（VRAM）的频繁读写。

// 假设处理的行长度为 512
const N: u32 = 512u;
const LOG2_N: u32 = 9u;

// 绑定资源
@group(0) @binding(0) var<storage, read> input_buf : array<vec2<f32>>;
@group(0) @binding(1) var<storage, read_write> output_buf : array<vec2<f32>>;

// 声明工作组共享内存，用于存储中间乒乓数据
var<workgroup> shared_ping : array<vec2<f32>, N>;
var<workgroup> shared_pong : array<vec2<f32>, N>;

// 复数乘法
fn complex_mul(a: vec2<f32>, b: vec2<f32>) -> vec2<f32> {
    return vec2<f32>(
        a.x * b.x - a.y * b.y,
        a.x * b.y + a.y * b.x
    );
}

// 计算旋转因子 (Twiddle Factor) W_N^k
fn get_twiddle(k: u32, n: u32) -> vec2<f32> {
    let theta = -2.0 * 3.14159265359 * f32(k) / f32(n);
    return vec2<f32>(cos(theta), sin(theta));
}

@compute @workgroup_size(256) // 每个工作组 256 个线程，处理 512 个点（每个线程处理两个点）
fn main(
    @builtin(global_invocation_id) global_id : vec3<u32>,
    @builtin(local_invocation_id) local_id : vec3<u32>,
    @builtin(workgroup_id) workgroup_id : vec3<u32>
) {
    let row_offset = workgroup_id.x * N;
    let tid = local_id.x;

    // 1. 从全局内存加载数据到共享内存
    shared_ping[tid] = input_buf[row_offset + tid];
    shared_ping[tid + N / 2u] = input_buf[row_offset + tid + N / 2u];
    
    workgroupBarrier();

    var use_ping = true;

    // 2. Stockham 蝶形变换循环
    for (var stage = 0u; stage < LOG2_N; stage = stage + 1u) {
        let sub_size = 1u << stage;          // 当前子 DFT 的大小
        let double_sub_size = sub_size << 1u; // 合并后的大 DFT 大小

        // 计算当前线程要处理的元素索引
        let thread_block = tid / sub_size;
        let thread_offset = tid % sub_size;
        let read_idx0 = thread_block * sub_size + thread_offset;
        let read_idx1 = read_idx0 + N / 2u;

        let write_idx0 = thread_block * double_sub_size + thread_offset;
        let write_idx1 = write_idx0 + sub_size;

        // 计算旋转因子
        let twiddle = get_twiddle(thread_offset, double_sub_size);

        if (use_ping) {
            let u = shared_ping[read_idx0];
            let v = shared_ping[read_idx1];
            let t = complex_mul(twiddle, v);

            shared_pong[write_idx0] = u + t;
            shared_pong[write_idx1] = u - t;
        } else {
            let u = shared_pong[read_idx0];
            let v = shared_pong[read_idx1];
            let t = complex_mul(twiddle, v);

            shared_ping[write_idx0] = u + t;
            shared_ping[write_idx1] = u - t;
        }

        // 同步工作组内的所有线程，确保当前 Stage 写入完毕
        workgroupBarrier();
        use_ping = !use_ping;
    }

    // 3. 写回全局内存
    if (use_ping) {
        output_buf[row_offset + tid] = shared_ping[tid];
        output_buf[row_offset + tid + N / 2u] = shared_ping[tid + N / 2u];
    } else {
        output_buf[row_offset + tid] = shared_pong[tid];
        output_buf[row_offset + tid + N / 2u] = shared_pong[tid + N / 2u];
    }
}

四、 压榨性能：那些不可忽视的 WebGPU 优化细节

1. 规避 Shared Memory Bank Conflict（银行冲突）

在上面的 WGSL 示例中，如果多个线程同时访问共享内存的同一个 Bank，硬件会使访问串行化，导致性能急剧下降。

• 解决思路：在高级实现中，可以通过对共享内存的数组大小进行微调（Padding，如声明 array<vec2<f32>, N + 1> 避开偶数对齐），或者精心设计索引映射来打散访问路径。

2. 高性能矩阵转置（Transpose Pass）

转置操作是 2D FFT 的中转站。直接写回全局内存会导致非连续写入（Stride Write）。

• 优化策略：使用 $16 \times 16$ 或 $32 \times 32$ 的 二维共享内存块（Tile）。
• 线程组先将一个正方形 Tile 读入共享内存，在共享内存内部完成局部转置，然后**连续（Coalesced）**地写入目标全局内存。这样可以保证 VRAM 的读取和写入全部都是合并访问。

// 典型的高效 Tile 转置核心逻辑片段
var<workgroup> tile: array<array<vec2<f32>, 16>, 16>;
// ... 读入 tile ...
workgroupBarrier();
// ... 转置写出到 output，保证写出地址在显存中连续 ...

3. 预计算旋转因子（Twiddle Factors Lookup Table）

在 Compute Shader 中实时计算 sin() 和 cos() 是非常昂贵的。

• 优化方式：在初始化阶段，使用 JS/TS 预先计算好一维 FFT 所需的所有 Twiddle 因子，并存入一个只读的 GPUBuffer 或纹理中。在 Shader 中直接根据索引读取（Texture Fetch 或 Buffer Load），可以节省大量的 ALU 计算单元。

4. 开启 fp16 支持（shader-f16）

对于图像处理而言，32位单精度浮点数（f32）的精度往往是冗余的。如果你使用的是 Chrome 121+ 等支持 shader-f16 扩展的环境：

• 可以申请 f16 类型的 WebGPU 缓冲区。
• 这不仅可以将 VRAM 带宽占用直接减半，还能在支持半精度计算的 GPU 上获得双倍的算力吞吐（如 Nvidia Tensor Cores 或 AMD Vega+）。

五、 图像频域滤波的应用场景

完成 2D FFT 转换后，图像就变成了频域复数矩阵。此时，我们可以进行各种有趣的滤波操作，最后通过 2D IFFT（逆傅里叶变换，只需将旋转因子取共轭并除以 $N \times M$）还原回空间域。

// 频域理想低通滤波器的 WGSL 伪代码
let center = vec2<f32>(f32(width) / 2.0, f32(height) / 2.0);
let dist = distance(current_coords, center);
if (dist > cutoff_radius) {
    freq_data[idx] = vec2<f32>(0.0, 0.0); // 过滤高频
}

结语

在 WebGPU 中实现高性能 2D FFT，代表了前端复杂计算从“能跑”向“工业级性能”的跃迁。通过合理运用 Stockham 算法、Workgroup 共享内存 和 Tile 矩阵转置，你可以在现代浏览器中轻松实现 60 帧以上的实时超大图像频域处理。这不仅适用于传统的图像特效，也是在 Web 端部署轻量级深度学习（如 CNN 频域加速）和复杂音频谱图分析的重要技术基石。