前端转AI要了解的技术，其他人不用看。前端架构基础之：让Js的计算运行在GPU上

上一篇文章中咱们详细分析了 Web Worker 到底是怎么回事儿，也清楚了网页上几乎所有的计算都是发生在CPU上的。

详见上篇文章：# 有点干，前端架构基础之：Web Worker到底是什么？它和Java中的线程是一个道理吗？

但是我们都知道，现在 AI 的运算其实大多都是发生在GPU上的，无论是推理，还是训练。

又或者是图形渲染，这一系列的计算都发生在GPU上。

这也是为什么英伟达股票嗷嗷叫的原因。

计算

如果，我们让Js的复杂计算不运行在 CPU 上，而是运行在 GPU 上。

是否能够打破 CPU 线程上限的瓶颈，从而实现N多个线程的并行计算？

答案是：可以！

咱们先说 Web Worker，Worker 新开的子线程里面跑的是 Js 代码，开辟的线程也是 CPU 的线程。

所以这些计算都发生在 CPU 是"理所当然"的，而且 Worker 本身并不具备 GPU 运算能力。

但是我们可以另辟蹊径，参考大模型运行在 GPU 上的逻辑。~~（这里挖个坑，回头讲讲为什么大模型运行在GPU上）~~

也就是说让 CPU 作为"指挥官"存在于整个任务中，GPU 只是作为干活的"小兵"存在。

这时候 Woker 其实就不再承担计算任务了，而是变成了"指挥" GPU 干活的"调度员"。

通过 Woker 指令调用 WebGPU API，向 GPU 提交计算指令。

这时候，真正的计算逻辑其实发生在 GPU 端， Woker 在这里只是下发指令/接收结果。

简单看一下 Demo：

js 复制代码

// calc.worker.gpu.js
self.onmessage = async (e) => {
    // 通过 Worker 访问 navigator.gpu
    const adapter = await navigator.gpu.requestAdapter()
    const device = await adapter.requestDevice()
}

GPU 作为一个独立的硬件设备，承担计算任务和 CPU 这边其实是彻底"切割"开的。

当 CPU 这边将任务下发完以后，其实就可以继续向下跑其他计算了，最后等着接受一下结果就好。

问题

但是这里需要注意一个比较麻烦的点：

如果是在 CPU 中运算，运算结果是存放在 CPU 内存中的。

而放在 GPU 中的运算，运算结果是存放在 GPU 显存中的。

CPU 无法直接读取 GPU 显存中的结果，它需要先将结果拷贝回"自己"的内存中。

另外这里要看到我说的是 WebGPU ，而不是 WebGL。

这两个存在本质上的区别，WebGL 的设计目标是图形渲染，没有通用计算着色器，所以几乎不支持做复杂运算。

而 WebGPU 拥有compute pipeline计算管线，可以进行大规模通用计算。

根据官方的解释，我们大概可以将 WebGPU 看作是 WebGL 的2.0版本。（这里不做展开了，感兴趣的可以去看看 MDN 中 WebGPU 篇）

根据这个理论，我们已经得到了能够突破 CPU 瓶颈的方法。

整体流程是：

主线程创建 Worker，下发输入数据。
在 Worker 内部：获取 GPU 适配器 → 创建设备 → 创建缓冲区 → 编写 WGSL 计算着色器 → 构建计算管线。
把输入数据写入 GPU 缓冲区。
提交计算命令给 GPU 执行（异步）。
GPU 运算结束，把结果从 GPU 显存映射拷贝到 CPU 内存缓冲区。
Worker 通过 postMessage 将结果发回主线程。

理论搞明白了，咱们看一下代码：

js 复制代码

// 主线程
const gouWoker = new Woker('calc.worker.gpu.js')

gpuWoker.onmessage = (e) => {
    console.log('运算结果：', e.data)
}

gpuWoker.postMessage({
    type: 'runCompute',
    inputData: new Float32Array([1, 2, 3, 4, 5, 6, 7])
})

创建一个 calc.worker.gpu.js 文件

js 复制代码

let device = null;
let computePipeline = null;

const shaderCode = `
@group(0) @binding(0) var<storage, read_write> data: array<f32>;

@compute @workgroup_size(8)
fn main(@builtin(global_invocation_id) id: vec3u) {
  let index = id.x;
  data[index] = data[index] * 2.0;
}
`;

// 初始化GPU设备
async function initGPU() {
  const adapter = await navigator.gpu.requestAdapter();
  if (!adapter) throw new Error("浏览器不支持WebGPU");
  device = await adapter.requestDevice();

  // 创建计算管线
  const shaderModule = device.createShaderModule({ code: shaderCode });
  computePipeline = device.createComputePipeline({
    layout: "auto",
    compute: {
      module: shaderModule,
      entryPoint: "main"
    }
  });
}

// 执行运算
async function runCompute(inputFloatArr) {
  if (!device) await initGPU();
  const elementCount = inputFloatArr.length;
  const byteSize = inputFloatArr.byteLength;

  // 创建存储缓冲区
  const storageBuffer = device.createBuffer({
    size: byteSize,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST
  });

  // 创建回读缓冲区：GPU → CPU，MAP_READ 标记代表可以映射到CPU内存读取
  const readbackBuffer = device.createBuffer({
    size: byteSize,
    usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ
  });

  // 将输入数据写入GPU存储缓冲区
  device.queue.writeBuffer(storageBuffer, 0, inputFloatArr);

  // 创建绑定组，把缓冲区绑定到着色器
  const bindGroup = device.createBindGroup({
    layout: computePipeline.getBindGroupLayout(0),
    entries: [
      { binding:0, resource: { buffer: storageBuffer } }
    ]
  });

  // 提交计算任务给GPU
  const commandEncoder = device.createCommandEncoder();
  const passEncoder = commandEncoder.beginComputePass();
  passEncoder.setPipeline(computePipeline);
  passEncoder.setBindGroup(0, bindGroup);
  // workgroup_count：8个一组，计算总批次
  passEncoder.dispatchWorkgroups(Math.ceil(elementCount / 8));
  passEncoder.end();

  // GPU计算完成后，把计算结果拷贝到可读回读缓冲区
  commandEncoder.copyBufferToBuffer(storageBuffer, 0, readbackBuffer, 0, byteSize);
  const commands = commandEncoder.finish();
  device.queue.submit([commands]);

  // ========== 获取GPU计算结果核心步骤 ==========
  // mapAsync：异步把GPU显存映射到CPU内存，这一步是异步！必须await
  await readbackBuffer.mapAsync(GPUMapMode.READ);
  // 获取CPU端ArrayBuffer视图
  const gpuResultBuffer = readbackBuffer.getMappedRange();
  // 拷贝一份数据（map出来的内存不能直接postMessage，unmap之后会失效）
  const result = new Float32Array(gpuResultBuffer);

  // 解除映射，释放资源，必须调用
  readbackBuffer.unmap();

  // 销毁buffer释放显存
  storageBuffer.destroy();
  readbackBuffer.destroy();

  return result;
}

// worker消息监听
self.onmessage = async (e) => {
  if(e.data.type === 'runCompute') {
    try {
      const res = await runCompute(e.data.inputData);
      // 将GPU计算结果发送回主线程
      self.postMessage({ result: res });
    } catch(err) {
      console.error(err);
    }
  }
};

整个执行过程中有几个核心点：

计算结果是无法直接被CPU读取到的，需要创建一个回读缓冲区。
通过 getMappedRange() 获取 ArrayBuffer，一定要拷贝一份新数组 出来。因为调用unmap()之后映射内存会立刻被回收销毁。
使用 mapAsync() 将GPU显存映射到CPU内存的过程是异步的，如果计算在GPU没跑完会卡住或者报错。所以必须等全部执行完成，submit之后再 mapAsync。
一定！一定！一定不要忘记 unmap ()，不解除映射，后续 GPU 操作会直接卡死。

是不是感觉非常复杂，复杂就对了。

这套操作其实在常规Js计算任务中根本不会拿来用，上述操作基本都发生在浏览器端 AI 推理这一种场景下。

也就是WebLLM，我到目前没见过其他场景上这套方案。

限制

这套操作在业务层面真的能够节约大量的计算时间吗？

其实不见得！

首先，从 CPU 拷贝数据到 GPU 下发计算任务本身很耗时。

从 GPU 拷贝计算结果到 CPU 也一样。

而且 Worker 冷启动本身也耗时，WebGPU 初始化也耗时，GPU 内创建计算任务也会耗时。

最关键的是 WebGPU 这个技术目前为止支持度有限，运行起来需要本地启动，也就是 localhost。

或者是 HTTPS，HPPT不行！

这也是为什么这套方案仅适合 WebLLM 的原因，CPU 运算吃力，但 GPU 实现大规模并行计算。

这种场景下 CPU 与 GPU 之间的数据拷贝反而不是耗时的"大头"了。

Ps: 如果面试的时候叽里呱啦来这一套，我相信能够打动部分面试官。