WebGPU 通用计算实战：从 Compute Shader 到浏览器端本地 AI 推理

// matmul.wgsl —— 分块矩阵乘法 C = A * B
// 每个 workgroup 计算一个 TILE×TILE 的输出块
const TILE : u32 = 16u;

struct Params {
  M : u32,
  N : u32,
  K : u32,
};

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

var<workgroup> tileA : array<f32, 256>;  // TILE * TILE = 256
var<workgroup> tileB : array<f32, 256>;

@compute @workgroup_size(16, 16)
fn main(@builtin(workgroup_id) wg  : vec2<u32>,
        @builtin(local_invocation_id) lid : vec2<u32>) {
  let row = wg.y * TILE + lid.y;
  let col = wg.x * TILE + lid.x;
  var sum : f32 = 0.0;

  let tiles = (params.K + TILE - 1u) / TILE;
  for (var t : u32 = 0u; t < tiles; t = t + 1u) {
    // 协同加载分块到共享内存
    let aIdx = row * params.K + (t * TILE + lid.x);
    let bIdx = (t * TILE + lid.y) * params.N + col;
    tileA[lid.y * TILE + lid.x] = select(0.0, A[aIdx], row < params.M && (t*TILE+lid.x) < params.K);
    tileB[lid.y * TILE + lid.x] = select(0.0, B[bIdx], col < params.N && (t*TILE+lid.y) < params.K);
    workgroupBarrier();  // 等所有人加载完再算

    for (var k : u32 = 0u; k < TILE; k = k + 1u) {
      sum = sum + tileA[lid.y * TILE + k] * tileB[k * TILE + lid.x];
    }
    workgroupBarrier();  // 算完再加载下一块
  }

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

const adapter = await navigator.gpu.requestAdapter();
const device  = await adapter.requestDevice();
const queue   = device.queue;

// 假设 A 是 M×K，B 是 K×N
const M = 512, N = 512, K = 512;
const bytes = (n) => n * 4;

const bufA = device.createBuffer({ size: bytes(M*K), usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST });
const bufB = device.createBuffer({ size: bytes(K*N), usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST });
const bufC = device.createBuffer({ size: bytes(M*N), usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC });
const bufP = device.createBuffer({ size: 12,         usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST });

queue.writeBuffer(bufA, 0, float32A);
queue.writeBuffer(bufB, 0, float32B);
queue.writeBuffer(bufP, 0, new Uint32Array([M, N, K]));

const module = device.createShaderModule({ code: wgslSource });
const pipeline = device.createComputePipeline({
  layout: "auto",
  compute: { module, entryPoint: "main" },
});

const bindGroup = device.createBindGroup({
  layout: pipeline.getBindGroupLayout(0),
  entries: [
    { binding: 0, resource: { buffer: bufP } },
    { binding: 1, resource: { buffer: bufA } },
    { binding: 2, resource: { buffer: bufB } },
    { binding: 3, resource: { buffer: bufC } },
  ],
});

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));  // 注意是 workgroup 数量
pass.end();
queue.submit([encoder.finish()]);

// 读回结果
const out = device.createBuffer({ size: bytes(M*N), usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ });
// ... 一次 copyBufferToBuffer + mapAsync 即可拿到 C

import { pipeline, env } from "https://cdn.jsdelivr.net/npm/@xenova/transformers";

// 强制走 WebGPU 后端
env.backends.onnx.ep = "webgpu";   // 而非 wasm / cpu

const generator = await pipeline("text-generation", "Xenova/Qwen2.5-0.5B-Instruct");
const out = await generator("用一句话解释 WebGPU 的优势：", { max_new_tokens: 64 });
console.log(out[0].generated_text);