Unlock 20‑30× GPU Speed: WebGPU in Three.js, Babylon.js, and TensorFlow.js

WebGPU is a groundbreaking web graphics technology that, despite 0% browser support and an unfinished standard, is already supported by major 3D and game frameworks such as Three.js and Babylon.js, and is used by TensorFlow.js to accelerate mobile deep learning with 20–30× performance gains over WebGL.

1. Using WebGPU in Three.js

With Three.js wrappers you can directly generate WebGPU calls. First import the required libraries:

import * as THREE from 'three';
import * as Nodes from 'three-nodes/Nodes.js';
import { add, mul } from 'three-nodes/ShaderNode.js';
import WebGPU from './jsm/capabilities/WebGPU.js';
import WebGPURenderer from './jsm/renderers/webgpu/WebGPURenderer.js';
...

Then the code resembles ordinary WebGL code, but you use WebGPURenderer instead of the WebGL renderer:

renderer = new WebGPURenderer();
renderer.setPixelRatio(window.devicePixelRatio);
renderer.setSize(window.innerWidth, window.innerHeight);
container.appendChild(renderer.domElement);
...

If the wrapper is insufficient, you can extend it with WGSL (WebGPU Shading Language), similar to GLSL or HLSL:

material = new Nodes.MeshBasicNodeMaterial();
material.colorNode = desaturateWGSLNode.call({ color: new Nodes.TextureNode(texture) });
materials.push(material);

const getWGSLTextureSample = new Nodes.FunctionNode(`
  fn getWGSLTextureSample(tex: texture_2d<f32>, tex_sampler: sampler, uv:vec2<f32>) -> vec4<f32> {
    return textureSample(tex, tex_sampler, uv) * vec4<f32>(0.0, 1.0, 0.0, 1.0);
  }
`);

A complete example renders a dancing figure in about 100 lines of code (image omitted).

2. Using WebGPU in Babylon.js

Babylon.js follows a similar pattern. The main difference is that it checks support for specific features, such as compute shaders, rather than a blanket WebGPU availability.

const supportCS = engine.getCaps().supportComputeShaders;

If you switch the context to WebGL2, the code changes accordingly, and Babylon.js also provides a playground for vertex and pixel shaders.

3. Accelerating Deep Learning with WebGPU

TensorFlow.js adopts WebGPU for inference because compute shaders are crucial for deep‑learning workloads. The accelerated code is largely implemented in WGSL and compiled to GPU instructions.

getUserCode(): string {
  const rank = this.xShape.length;
  const type = getCoordsDataType(rank);
  const start = this.xShape.map((_, i) => `uniforms.pad${i}[0]`).join(',');
  const end = this.xShape.map((_, i) => `uniforms.pad${i}[0] + uniforms.xShape${rank > 1 ? `[${i}]` : ''}`).join(',');
  const userCode = `
    if (index < uniforms.size) {
      let start = ${type}(${start});
      let end = ${type}(${end});
      let outC = getCoordsFromIndex(index);
      if (any(outC < start) || any(outC >= end)) {
        setOutputAtIndex(index, uniforms.constantValue);
      } else {
        let coords = outC - start;
        setOutputAtIndex(index, getX(${rank > 1 ? '[coords[0], coords[1], coords[2], coords[3]]' : 'coords'}));
      }
    }
  `;
  return userCode;
}

4. Writing WebGPU Code Without a Framework

WebGPU programming follows the modern Vulkan‑style pipeline. You create a GPUCanvasContext, request an adapter and device, write WGSL shaders, build a render pipeline, encode commands, and submit them.

async function testGPU() {
  const canvas = document.getElementById('webgpu');
  const gpuContext = canvas.getContext('webgpu');
  const adapter = await navigator.gpu.requestAdapter();
  const device = await adapter.requestDevice();
  const presentationFormat = gpuContext.getPreferredFormat(adapter);
  gpuContext.configure({ device, format: presentationFormat });

  const triangleVertWGSL = `
    @stage(vertex)
    fn main(@builtin(vertex_index) VertexIndex : u32) -> @builtin(position) vec4<f32> {
      var pos = array<vec2<f32>, 3>(
        vec2<f32>(0.0, 0.5),
        vec2<f32>(-0.5, -0.5),
        vec2<f32>(0.5, -0.5));
      return vec4<f32>(pos[VertexIndex], 0.0, 1.0);
    }
  `;
  const redFragWGSL = `
    @stage(fragment)
    fn main() -> @location(0) vec4<f32> {
      return vec4<f32>(1.0, 0.0, 0.0, 1.0);
    }
  `;
  const pipeline = device.createRenderPipeline({
    vertex: { module: device.createShaderModule({ code: triangleVertWGSL }), entryPoint: 'main' },
    fragment: { module: device.createShaderModule({ code: redFragWGSL }), entryPoint: 'main', targets: [{ format: presentationFormat }] },
    primitive: { topology: 'triangle-list' },
  });

  const commandEncoder = device.createCommandEncoder();
  const textureView = gpuContext.getCurrentTexture().createView();
  const renderPassDescriptor = {
    colorAttachments: [{ view: textureView, loadValue: { r: 1, g: 1, b: 1, a: 1 }, storeOp: 'store' }],
  };
  const passEncoder = commandEncoder.beginRenderPass(renderPassDescriptor);
  passEncoder.setPipeline(pipeline);
  passEncoder.draw(3, 1, 0, 0);
  passEncoder.end();
  device.queue.submit([commandEncoder.finish()]);
}

testGPU();

Because browsers have not yet fully supported WebGPU, you need a cutting‑edge build such as Chrome Canary with the "enable‑unsafe‑webgpu" flag turned on.

Conclusion

Compared with OpenGL ES‑based WebGL 1.0, WebGPU is closer to Vulkan and can better exploit modern GPU capabilities. Frameworks like Three.js and Babylon.js already demonstrate its rendering potential, while TensorFlow.js shows its compute advantage. Although still experimental, WebGPU’s underlying Vulkan and Metal technologies are mature, and once widely available it will empower metaverse‑type applications and future front‑end development.