跪拜 Guibai
← Back to the summary

Offloading JavaScript Computation to the GPU via Web Workers and WebGPU

In the previous article, we analyzed in detail what Web Workers are all about and made it clear that almost all computations on a webpage happen on the CPU.

See the previous article: # A bit dry, frontend architecture basics: What exactly is a Web Worker? Is it the same concept as threads in Java?

But we all know that most AI computations now—whether inference or training—actually happen on the GPU.

Or graphics rendering; this whole series of computations happens on the GPU.

That's also why Nvidia's stock is screaming.

Computation

If we make complex JavaScript computations run not on the CPU, but on the GPU,

can we break through the bottleneck of CPU thread limits and achieve parallel computation across many threads?

The answer is: yes!

Let's talk about Web Workers first. The child thread opened by a Worker runs JavaScript code, and the threads it opens are CPU threads.

So it's "only natural" that these computations happen on the CPU, and the Worker itself does not possess GPU computing capability.

But we can take a different approach, referencing the logic of how large models run on GPUs. (I'll dig a hole here and come back later to explain why large models run on GPUs)

That is, let the CPU act as the "commander" throughout the entire task, with the GPU merely acting as the "soldier" doing the work.

At this point, the Worker no longer undertakes the computation task itself but becomes a "dispatcher" that commands the GPU to work.

Through Worker instructions, it calls the WebGPU API and submits computation commands to the GPU.

Here, the real computation logic actually happens on the GPU side; the Worker here only issues commands and receives results.

Let's look at a simple Demo:

// calc.worker.gpu.js
self.onmessage = async (e) => {
    // Access navigator.gpu through the Worker
    const adapter = await navigator.gpu.requestAdapter()
    const device = await adapter.requestDevice()
}

The GPU, as an independent hardware device, is completely "decoupled" from the CPU side in terms of undertaking computation tasks.

Once the CPU side finishes dispatching tasks, it can actually continue running other computations, and finally just wait to receive the results.

Problems

But there is a rather troublesome point to note here:

If computation happens on the CPU, the results are stored in CPU memory.

But for computation placed on the GPU, the results are stored in GPU video memory.

The CPU cannot directly read results from GPU video memory; it needs to first copy the results back to its "own" memory.

Also, note that I'm talking about WebGPU, not WebGL.

There is a fundamental difference between the two. WebGL's design goal is graphics rendering; it lacks general-purpose compute shaders, so it almost doesn't support complex computations.

WebGPU, on the other hand, has a compute pipeline, which can perform large-scale general-purpose computation.

According to official explanations, we can roughly consider WebGPU as version 2.0 of WebGL. (I won't expand on this here; those interested can check out the WebGPU section on MDN.)

Based on this theory, we have obtained a method to break through the CPU bottleneck.

image.png

The overall process is:

  1. The main thread creates a Worker and sends input data.
  2. Inside the Worker: get GPU adapter → create device → create buffer → write WGSL compute shader → build compute pipeline.
  3. Write input data to GPU buffer.
  4. Submit compute commands to GPU for execution (asynchronous).
  5. GPU computation finishes, map and copy results from GPU video memory to CPU memory buffer.
  6. The Worker sends the results back to the main thread via postMessage.

Now that we understand the theory, let's look at the code:

// Main thread
const gouWoker = new Woker('calc.worker.gpu.js')

gpuWoker.onmessage = (e) => {
    console.log('Computation result:', e.data)
}

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

Create a calc.worker.gpu.js file

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;
}
`;

// Initialize GPU device
async function initGPU() {
  const adapter = await navigator.gpu.requestAdapter();
  if (!adapter) throw new Error("Browser does not support WebGPU");
  device = await adapter.requestDevice();

  // Create compute pipeline
  const shaderModule = device.createShaderModule({ code: shaderCode });
  computePipeline = device.createComputePipeline({
    layout: "auto",
    compute: {
      module: shaderModule,
      entryPoint: "main"
    }
  });
}

// Execute computation
async function runCompute(inputFloatArr) {
  if (!device) await initGPU();
  const elementCount = inputFloatArr.length;
  const byteSize = inputFloatArr.byteLength;

  // Create storage buffer
  const storageBuffer = device.createBuffer({
    size: byteSize,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST
  });

  // Create readback buffer: GPU → CPU, MAP_READ flag indicates it can be mapped to CPU memory for reading
  const readbackBuffer = device.createBuffer({
    size: byteSize,
    usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ
  });

  // Write input data to GPU storage buffer
  device.queue.writeBuffer(storageBuffer, 0, inputFloatArr);

  // Create bind group, bind buffer to shader
  const bindGroup = device.createBindGroup({
    layout: computePipeline.getBindGroupLayout(0),
    entries: [
      { binding:0, resource: { buffer: storageBuffer } }
    ]
  });

  // Submit compute task to GPU
  const commandEncoder = device.createCommandEncoder();
  const passEncoder = commandEncoder.beginComputePass();
  passEncoder.setPipeline(computePipeline);
  passEncoder.setBindGroup(0, bindGroup);
  // workgroup_count: groups of 8, calculate total batches
  passEncoder.dispatchWorkgroups(Math.ceil(elementCount / 8));
  passEncoder.end();

  // After GPU computation completes, copy results to readable readback buffer
  commandEncoder.copyBufferToBuffer(storageBuffer, 0, readbackBuffer, 0, byteSize);
  const commands = commandEncoder.finish();
  device.queue.submit([commands]);

  // ========== Core steps to get GPU computation results ==========
  // mapAsync: Asynchronously map GPU video memory to CPU memory, this step is async! Must await
  await readbackBuffer.mapAsync(GPUMapMode.READ);
  // Get CPU-side ArrayBuffer view
  const gpuResultBuffer = readbackBuffer.getMappedRange();
  // Copy a copy of the data (mapped memory cannot be directly postMessage'd, it becomes invalid after unmap)
  const result = new Float32Array(gpuResultBuffer);

  // Unmap, release resources, must call
  readbackBuffer.unmap();

  // Destroy buffers to release video memory
  storageBuffer.destroy();
  readbackBuffer.destroy();

  return result;
}

// Worker message listener
self.onmessage = async (e) => {
  if(e.data.type === 'runCompute') {
    try {
      const res = await runCompute(e.data.inputData);
      // Send GPU computation results back to main thread
      self.postMessage({ result: res });
    } catch(err) {
      console.error(err);
    }
  }
};

There are several core points in the entire execution process:

  1. Computation results cannot be directly read by the CPU; a readback buffer needs to be created.
  2. When using getMappedRange() to get the ArrayBuffer, you must copy a new array. Because after calling unmap(), the mapped memory will be immediately reclaimed and destroyed.
  3. The process of using mapAsync() to map GPU video memory to CPU memory is asynchronous. If the GPU computation hasn't finished, it will hang or throw an error. So you must wait for all execution to complete, and only mapAsync after submit.
  4. Definitely! Definitely! Definitely do not forget unmap(). Without unmapping, subsequent GPU operations will directly freeze.

Doesn't it feel very complex? That's right, it is complex.

This set of operations is basically never used in regular JavaScript computation tasks. The above operations almost exclusively occur in one scenario: browser-side AI inference.

That is, WebLLM. I haven't seen this solution applied in any other scenario so far.

Limitations

Can this set of operations really save a significant amount of computation time at the business level?

Not necessarily!

First, copying data from CPU to GPU to dispatch computation tasks is itself time-consuming.

Copying computation results from GPU back to CPU is the same.

Moreover, Worker cold start itself is time-consuming, WebGPU initialization is also time-consuming, and creating computation tasks inside the GPU also takes time.

Most critically, WebGPU technology currently has limited support. Running it requires a local server, i.e., localhost.

Or HTTPS; HTTP won't work!

image.png

This is also why this solution is only suitable for WebLLM: CPU computation struggles, but GPU achieves large-scale parallel computation.

In this scenario, the data copying between CPU and GPU is no longer the "bulk" of the time consumed.

Ps: If you rattle off this whole set of concepts in an interview, I believe it could impress some interviewers.