WebGPU 计算着色器矩阵乘法优化：从 Naive 到 Tiled 深度实战

矩阵乘法（GEMM）是科学计算、机器学习推理、物理仿真和图形渲染的核心算子。如何在 WebGPU 的计算着色器中高效实现矩阵乘法，直接关系到浏览器端 AI 推理、GPGPU 计算乃至跨平台图形应用的实际性能上限。本文将从零构建一个完整的 WebGPU GEMM 管线，逐步剖析从朴素实现到分块（Tiled）优化、Workgroup Memory（工作组内存）利用的全链路工程实践。

WebGPU 计算管线基础架构

WebGPU 的计算着色器运行在 GPU 通用计算单元上，其执行模型与 CUDA/OpenCL 类似但又受 Web 平台沙箱约束。一个完整的计算管线由以下核心组件构成：

• Device：GPU 设备抽象，负责资源分配
• Pipeline：计算管线状态对象，绑定着色器与绑定组
• Bind Group：资源绑定（缓冲区、采样器等）
• Workgroup：线程组，是 GPU 调度的基本单位

计算着色器使用 WGSL（WebGPU Shading Language）编写，语法接近 Rust，具有强类型和内存安全保证。

Naive 矩阵乘法：迈出第一步

首先实现一个最直接的矩阵乘法版本，逐元素计算 C = A × B。

JavaScript 端管线配置

async function initGPUMatMul() {
  const adapter = await navigator.gpu.requestAdapter();
  const device = await adapter.requestDevice();

  // 矩阵维度：M×K 乘以 K×N = M×N
  const M = 1024, K = 1024, N = 1024;
  const bufferSizeA = M * K * 4; // f32 = 4 bytes
  const bufferSizeB = K * N * 4;
  const bufferSizeC = M * N * 4;

  // 创建缓冲区
  const bufferA = device.createBuffer({
    size: bufferSizeA,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST
  });
  const bufferB = device.createBuffer({
    size: bufferSizeB,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST
  });
  const bufferC = device.createBuffer({
    size: bufferSizeC,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC
  });

  // 创建Bind Group Layout
  const bindGroupLayout = device.createBindGroupLayout({
    entries: [
      { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } },
      { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } },
      { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } }
    ]
  });

  const pipeline = device.createComputePipeline({
    layout: device.createPipelineLayout({ bindGroupLayouts: [bindGroupLayout] }),
    compute: {
      module: device.createShaderModule({ code: naiveShaderCode }),
      entryPoint: 'main'
    }
  });

  const bindGroup = device.createBindGroup({
    layout: bindGroupLayout,
    entries: [
      { binding: 0, resource: { buffer: bufferA } },
      { binding: 1, resource: { buffer: bufferB } },
      { binding: 2, resource: { buffer: bufferC } }
    ]
  });

  // 提交计算命令
  const encoder = device.createCommandEncoder();
  const pass = encoder.beginComputePass();
  pass.setPipeline(pipeline);
  pass.setBindGroup(0, bindGroup);
  pass.dispatchWorkgroups(Math.ceil(N / 16), Math.ceil(M / 16));
  pass.end();
  device.queue.submit([encoder.finish()]);
}

WGSL Naive 着色器

@group(0) @binding(0) var<storage, read> A: array<f32>;
@group(0) @binding(1) var<storage, read> B: array<f32>;
@group(0) @binding(2) var<storage, read_write> C: array<f32>;

@compute @workgroup_size(16, 16)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    let row = global_id.x;
    let col = global_id.y;
    let M = 1024u;
    let K = 1024u;
    let N = 1024u;

    if (row >= M || col >= N) { return; }

    var sum: f32 = 0.0;
    for (var k: u32 = 0u; k < K; k = k + 1u) {
        sum = sum + A[row * K + k] * B[k * N + col];
    }
    C[row * N + col] = sum;
}

这个朴素实现的问题极为严重：每个输出元素需要 K 次全局内存读取，而 GPU 的全局内存延迟通常在 400-800 个时钟周期。当 M=N=K=1024 时，全局内存访问量约为 2×1024³ = 2 GFLOPs 对应的输入读取，但有效计算同样约为 2 GFLOPs，算术强度仅为 1:1，远未达到现代 GPU 的计算与带宽平衡点。

分块矩阵乘法：挖掘数据重用

分块（Tiled）矩阵乘法的核心思想是将大矩阵拆分为小块，让每个 Workgroup 负责计算 C 矩阵的一个 Tile。关键在于：A 矩阵的一个 Tile 被多个 B Tile 引用，反之亦然，从而大幅减少全局内存访问次数。

算法原理

假设 TILE_SIZE = 16，每个 Workgroup 计算 C 矩阵中一个 16×16 的区域。分块后，每个输出元素的计算迭代 K/TILE_SIZE 次，每次加载 A 的一个 16×? 行和 B 的一个 ?×16 列到 Workgroup Memory 中。

算法复杂度分析：
- 全局内存访问：O(M×N×K / TILE_SIZE)
- 相比朴素版本，内存流量降低 TILE_SIZE 倍
- 当 TILE_SIZE = 16 时，算术强度从 1 提升到 8

WGSL Tiled 着色器

const TILE_SIZE: u32 = 16u;

@group(0) @binding(0) var<storage, read> A: array<f32>;
@group(0) @binding(1) var<storage, read> B: array<f32>;
@group(0) @binding(2) var<storage, read_write> C: array<f32>;

// Workgroup Memory：各 Workgroup 私有的高速共享内存
var<workgroup> tileA: array<f32, 256>; // 16×16
var<workgroup> tileB: array<f32, 256>; // 16×16

@compute @workgroup_size(16, 16)
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 M = 1024u;
    let K = 1024u;
    let N = 1024u;

    let row = workgroup_id.y * TILE_SIZE + local_id.y;
    let col = workgroup_id.x * TILE_SIZE + local_id.x;

    var sum: f32 = 0.0;
    let numTiles = (K + TILE_SIZE - 1u) / TILE_SIZE;

    for (var t: u32 = 0u; t < numTiles; t = t + 1u) {
        // 协作加载：每个线程加载一个元素到 Workgroup Memory
        let aCol = t * TILE_SIZE + local_id.x;
        let aRow = row;

        if (aRow < M && aCol < K) {
            tileA[local_id.y * TILE_SIZE + local_id.x] = A[aRow * K + aCol];
        } else {
            tileA[local_id.y * TILE_SIZE + local_id.x] = 0.0;
        }

        let bRow = t * TILE_SIZE + local_id.y;
        let bCol = col;

        if (bRow < K && bCol < N) {
            tileB[local_id.y * TILE_SIZE + local_id.x] = B[bRow * N + bCol];
        } else {
            tileB[local_id.y * TILE_SIZE + local_id.x] = 0.0;
        }

        // 同步：确保所有线程完成加载
        workgroupBarrier();

        // 计算当前 Tile 的部分积
        for (var k: u32 = 0u; k < TILE_SIZE; k = k + 1u) {
            sum = sum + tileA[local_id.y * TILE_SIZE + k] * tileB[k * TILE_SIZE + local_id.x];
        }

        // 同步：确保所有线程完成计算后再加载下一 Tile
        workgroupBarrier();
    }

    if (row < M && col < N) {
        C[row * N + col] = sum;
    }
}

这个 Tiled 实现的关键优化在于：

1. Workgroup Memory 复用：tileA 和 tileB 使用 GPU 的 shared memory（通常 32-64KB/SM），延迟仅为 1-3 个时钟周期
2. 协作加载（Cooperative Loading）：16×16 = 256 个线程并行加载，理论带宽利用率接近峰值
3. 双 Barrier 隔离：workgroupBarrier() 确保内存可见性，避免数据竞争

进一步优化：向量化与寄存器分块

在 Tiled 基础上还可以进一步优化，利用 WGSL 的 vec4 类型和寄存器分块策略。

const TILE_SIZE: u32 = 16u;
const VEC_SIZE: u32 = 4u;

@group(0) @binding(0) var<storage, read> A: array<f32>;
@group(0) @binding(1) var<storage, read> B: array<f32>;
@group(0) @binding(2) var<storage, read_write> C: array<f32>;

var<workgroup> tileA: array<f32, 256>;
var<workgroup> tileB: array<f32, 256>;

@compute @workgroup_size(16, 16)
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 M = 1024u;
    let K = 1024u;
    let N = 1024u;

    let localRow = local_id.y;
    let localCol = local_id.x;
    let globalRow = workgroup_id.y * TILE_SIZE + localRow;
    let globalCol = workgroup_id.x * TILE_SIZE + localCol;

    // 每个线程计算 2×2 寄存器分块（Register Blocking）
    var sum00: f32 = 0.0;
    var sum01: f32 = 0.0;
    var sum10: f32 = 0.0;
    var sum11: f32 = 0.0;

    let numTiles = K / TILE_SIZE;

    for (var t: u32 = 0u; t < numTiles; t = t + 1u) {
        // 协作加载 4 个元素（vec4 利用内存合并）
        let loadIdx = localRow * (TILE_SIZE / VEC_SIZE) + localCol;
        let aRowBase = globalRow;
        let bColBase = globalCol;

        // 批量加载
        let aCol = t * TILE_SIZE + localCol * VEC_SIZE;
        for (var c: u32 = 0u; c < VEC_SIZE; c = c + 1u) {
            let idx = localRow * TILE_SIZE + localCol * VEC_SIZE + c;
            if (aRowBase < M && aCol + c < K) {
                tileA[idx] = A[aRowBase * K + aCol + c];
            } else {
                tileA[idx] = 0.0;
            }

            let bRow = t * TILE_SIZE + localRow;
            let bCol = globalCol * VEC_SIZE + c;
            if (bRow < K && bCol < N) {
                tileB[idx] = B[bRow * N + bCol];
            } else {
                tileB[idx] = 0.0;
            }
        }

        workgroupBarrier();

        // 寄存器分块计算
        for (var k: u32 = 0u; k < TILE_SIZE; k = k + 1u) {
            let aVal0 = tileA[localRow * TILE_SIZE + k];
            let aVal1 = tileA[(localRow + 8u) * TILE_SIZE + k];
            let bVal0 = tileB[k * TILE_SIZE + localCol];
            let bVal1 = tileB[k * TILE_SIZE + (localCol + 8u)];

            sum00 += aVal0 * bVal0;
            sum01 += aVal0 * bVal1;
            sum10 += aVal1 * bVal0;
            sum11 += aVal1 * bVal1;
        }

        workgroupBarrier();
    }

    // 写回结果
    if (globalRow < M && globalCol < N) {
        C[globalRow * N + globalCol] = sum00;
        C[(globalRow + 8u) * N + globalCol] = sum10;
        C[globalRow * N + (globalCol + 8u)] = sum01;
        C[(globalRow + 8u) * N + (globalCol + 8u)] = sum11;
    }
}

边界处理与动态维度

实际工程中的矩阵乘法必须处理非对齐维度。关键策略包括：

1. 动态维度传入：通过 Uniform Buffer 传递 M、K、N，避免硬编码
2. 边界检查：每个内存访问前验证索引是否越界
3. 零填充（Zero Padding）：将矩阵维度向上对齐到 TILE_SIZE 的倍数

// Uniform 方式传递维度
const uniformData = new Uint32Array([M, K, N]);
const uniformBuffer = device.createBuffer({
  size: uniformData.byteLength,
  usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST
});
device.queue.writeBuffer(uniformBuffer, 0, uniformData);

struct Dimensions {
    M: u32,
    K: u32,
    N: u32,
};
@group(0) @binding(3) var<uniform> dims: Dimensions;

性能基准测试

在 Apple M2 Pro（19 核 GPU）上测试 1024×1024 矩阵乘法：

分析要点：

• Tiled 版本相比 Naive 实现了 6.7 倍加速，主要来自 Workgroup Memory 利用
• 寄存器分块进一步提升了指令级并行（ILM），减少寄存器压力
• Apple MPS 底层使用 AMX 指令集达到了接近理论峰值的性能
• WebGPU 性能约为原生 Metal 的 63%，差距主要来自 API 抽象层和验证开销

工程实践建议

1. TILE_SIZE 选择

TILE_SIZE = 16 在大多数场景下是较好的平衡点。较大的 TILE_SIZE（如 32）可以减少 Tile 循环次数，但会增加 Workgroup 寄存器压力，可能限制 GPU Occupancy（并行 Wavefront 数量）。

2. 避免 Bank Conflict

GPU Workgroup Memory 分为多个 Bank（通常 32 个），当同一 Workgroup 内的线程访问同一 Bank 的不同地址时会发生 Bank Conflict。在矩阵乘法中，将 tileB 转置存储或调整访问步长可有效避免。

3. 使用 timestamp-query 精确测量

WebGPU 支持 timestamp-query 进行 GPU 端计时，精度远高于 JavaScript performance.now()：

const querySet = device.createQuerySet({
  type: 'timestamp',
  count: 2
});
const queryBuffer = device.createBuffer({
  size: 16,
  usage: GPUBufferUsage.QUERY_RESOLVE | GPUBufferUsage.COPY_SRC
});

const pass = encoder.beginComputePass({
  timestampWrites: {
    querySet,
    beginningOfPassWriteIndex: 0,
    endOfPassWriteIndex: 1
  }
});
// ... dispatch ...
 pass.end();
encoder.resolveQuerySet(querySet, 0, 2, queryBuffer, 0);

4. 矩阵布局优化

对于列优先（Column Major）的矩阵（如 BLAS/LAPACK 中使用），调整着色器中的索引计算可以提升内存合并访问效率。WebGPU 本身不强制存储布局，但 GPU 的 L1/L2 缓存对连续访问模式更友好。

结语

WebGPU 计算着色器的矩阵乘法优化是一项系统工程，需要同时考虑算法层面（分块策略）、硬件特性（Workgroup Memory 大小、Wavefront 宽度）和平台约束（Web API 验证开销）。从 Naive 到 Tiled 再到寄存器分块，每一步优化的本质都在于提升算术强度——让 GPU 的计算单元尽可能少等待内存。对于浏览器端 GPGPU 工作负载，这套优化方法论同样适用于粒子系统、图像处理、神经网络推理等多种场景。

随着 WebGPU 标准的成熟和越来越多浏览器（Chrome、Safari、Firefox）完成对 subgroup operations 和 fp16 算术的支持，浏览器端 GPGPU 的性能天花板将继续上移，使得 Web 应用能够承担此前只有原生应用才能胜任的高性能计算任务。