《Born》第6章：WebGPU 后端——零 CGO 的 GPU 加速

WebGPU 是 Born 最激进的架构选择。它让 Go 程序在不写一行 C/C++、不依赖 CUDA Toolkit 的情况下，直接调用 GPU 进行深度学习计算。

为什么选 WebGPU

架构：Queue → CommandEncoder → ComputePass

CPU (Go)                      GPU
  │                            │
  ├─> 创建 Buffer ────────────>│
  ├─> 创建 ShaderModule ──────>│
  ├─> 创建 ComputePipeline ───>│
  │                            │
  ├─> 写入数据 to Buffer ─────>│
  ├─> 编码 ComputePass ───────>│
  │   (bindGroup, dispatch)    │
  ├─> Submit Queue ───────────>│  执行 shader
  │                            │
  ├─> Read Buffer <────────────┤  读取结果

Born 的 WebGPU 后端封装了这整个流程。用户只需要：

backend := webgpu.New()
x := tensor.Randn[float32](tensor.Shape{1000, 1000}, backend)
y := x.MatMul(x.T()) // 自动在 GPU 上执行

WGSL：WebGPU 的着色语言

Born 内嵌了 53 个 WGSL compute shader，覆盖所有核心算子：

// add.wgsl
@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> result: array<f32>;

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    let idx = gid.x;
    if (idx < arrayLength(&result)) {
        result[idx] = a[idx] + b[idx];
    }
}

每个 shader 在程序启动时预编译，后续调用直接复用 pipeline。

延迟求值

WebGPU 后端的杀手锏是延迟求值：

x := tensor.Randn[float32](shape, webgpuBackend)
y := x.Add(x)      // 不执行，加入队列
z := y.MulScalar(2) // 不执行，加入队列
w := z.ReLU()      // 不执行，加入队列

result := w.Data()  // 此时批量提交所有命令

好处：

• 隐藏 CPU↔GPU 通信延迟
• 允许后端做算子融合优化
• 减少 GPU 提交开销

性能基准

📘 《Born》连载技术书，第 6/22 章。