Workflow: ML Compute Pipeline

Fresh 🌱

Why WebGPU for ML

• 3x+ faster ML inference vs WebGL
• Native compute shaders — no fragment shader hacks
• Storage buffers up to 128MB+ — handles large tensors
• Atomic operations for reduction passes
• Subgroup extensions for warp-level parallelism (Chrome 134+)

Pipeline Architecture

TensorFlow.js WebGPU Backend

import * as tf from '@tensorflow/tfjs';
import '@tensorflow/tfjs-backend-webgpu';

// Enable WebGPU backend
await tf.setBackend('webgpu');
await tf.ready();

console.log('Backend:', tf.getBackend()); // 'webgpu'

// Now all tf operations run on GPU
const model = await tf.loadGraphModel('model/model.json');
const inputTensor = tf.browser.fromPixels(imageElement);
const predictions = model.predict(inputTensor.expandDims(0));

Manual Compute Shader for Inference

// Example: ReLU activation layer
const reluShader = `
  @group(0) @binding(0) var<storage, read>       input  : array<f32>;
  @group(0) @binding(1) var<storage, read_write> output : array<f32>;

  @compute @workgroup_size(256)
  fn main(@builtin(global_invocation_id) id : vec3u) {
    let i = id.x;
    if (i >= arrayLength(&input)) { return; }
    output[i] = max(0.0, input[i]);
  }
`;

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

function runReLU(inputBuffer, outputBuffer, size) {
  const bindGroup = device.createBindGroup({
    layout: pipeline.getBindGroupLayout(0),
    entries: [
      { binding: 0, resource: { buffer: inputBuffer } },
      { binding: 1, resource: { buffer: outputBuffer } }
    ]
  });

  const enc = device.createCommandEncoder();
  const pass = enc.beginComputePass();
  pass.setPipeline(pipeline);
  pass.setBindGroup(0, bindGroup);
  pass.dispatchWorkgroups(Math.ceil(size / 256));
  pass.end();
  device.queue.submit([enc.finish()]);
}

Performance Measurement (with Dev Features)

Enable chrome://flags/#enable-webgpu-developer-features for nanosecond precision:

// Timestamp queries for GPU timing
const querySet = device.createQuerySet({
  type: 'timestamp',
  count: 2
});

const resolveBuffer = device.createBuffer({
  size: 16, // 2 x 8 bytes (u64)
  usage: GPUBufferUsage.QUERY_RESOLVE | GPUBufferUsage.COPY_SRC
});

const enc = device.createCommandEncoder();
const pass = enc.beginComputePass({
  timestampWrites: {
    querySet,
    beginningOfPassWriteIndex: 0,
    endOfPassWriteIndex: 1
  }
});
// ... dispatch compute work ...
pass.end();
enc.resolveQuerySet(querySet, 0, 2, resolveBuffer, 0);
device.queue.submit([enc.finish()]);

// Read timing results
const stagingBuf = device.createBuffer({
  size: 16,
  usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST
});
// copyBufferToBuffer + mapAsync to read nanosecond timestamps