Running Python Compute on WebGPU Through WebAssembly

The constraint that shaped everything

Zipp's browser build is a sandbox first: interpreter-only, no unsafe code, no ambient browser authority, one synchronous host channel and one asynchronous queue, both default-deny. Giving guest Python a WebGPU device handle would have thrown that away. So the engine never sees a GPU. It sees a host request called gpu.execute whose payload is a graph, and it sees a reply. Everything GPU-shaped lives in the host's JavaScript, where the browser's own permission model already applies.

The protocol

A graph is plain data with version: 1, a list of nodes and a list of named outputs. Each node has an integer id, an op, and references to earlier nodes only, so the graph is acyclic by construction. Operations are deliberately few: input and full tensors, scalar or shape-matched add, sub and mul, relu, positive masks, transpose, matmul, whole-tensor sum, and a toroidal Game of life step. Inputs and outputs are finite float32.

What leaves the engine

{
  "version": 1,
  "nodes": [
    { "id": 0, "op": "input", "shape": [3], "data": [-2, 3, 4] },
    { "id": 1, "op": "input", "shape": [3], "data": [10, 20, 30] },
    { "id": 2, "op": "mul", "a": 0, "b": 1 },
    { "id": 3, "op": "relu", "a": 2 }
  ],
  "outputs": [{ "name": "values", "id": 3 }]
}

The host validates before it allocates: at most 512 nodes, 1,048,576 elements per tensor, 32 MiB of summed logical storage, 100 million estimated operations and 16 outputs by default, with a separate 128 MiB cap on live WebGL2 textures that counts RGBA32F padding because one scalar occupies a 16-byte texel. The limits are host policy and can be raised or lowered per runtime; the point is that the guest cannot set them.

Authoring from Python

main.py

from zipp_gpu import Graph

def show(result):
    print(result["backend"], result["outputs"]["values"]["data"])

g = Graph()
a = g.tensor([-2, 3, 4])
b = g.tensor([10, 20, 30])
g.submit(show, values=(a * b).relu())   # [0, 60, 120]

zipp_gpu is a bundled Python module, not native code. Tensor operators append nodes; submit serializes the program and posts it through the engine's host-request queue. Inside the playground the host runs the graph after the current call returns and delivers the callback between VM calls, so the program keeps going meanwhile. Run natively under zipp py, with no host, the same module evaluates the graph on a CPU reference implementation with the same float32 rounding rules and reports cpu-python; the semantics do not change, only who did the arithmetic.

Four backends, one contract

gpu-lab 0.1.0 backends
Backend	How it computes	Notes
`webgpu`	WGSL compute shaders, `@workgroup_size(64)`, storage buffers	Preferred when available; rejects recognized software adapters
`webgl2`	GLSL fragment shaders over float textures, one pass per node	The widely available path; the Life demo on this site runs here
`wasm`	Freestanding C kernels compiled with Clang to `kernels.wasm`	Portable and deterministic; used by the smoke tests
`cpu-js`	Plain JavaScript reference	Defines the expected float32 rounding: every intermediate rounds to float32, sums reduce pairwise

auto tries them in that order and records every failed initialization in runtime.info().fallbackAttempts, so a page can tell a user that the GPU path was unavailable rather than pretending. An explicit backend choice fails rather than silently downgrading, and a missing GPU is never reported as a successful GPU check.

Delivering results back into a sandbox

The host adapter, createPythonGPUAdapter(engine, compute, { allowExecute: true, maxPending: 16 }), drains takeHostRequests(), validates and executes each graph, and answers with pythonCall("__zipp_py_deliver", [id, reply]), where the reply is either {ok: true, value} or {ok: false, error: {code, message}}. The Python callback runs inside that call. Three rules keep it safe: execution is granted explicitly per engine, there is a cap on pending requests, and a late result never reaches an engine that has been disposed, which matters because Workers are terminated at deadlines and tenants are replaced.

torch.compile, with a caveat in the name

The Torch subset uses the same mechanism. torch.compile(model)(x) records the model's forward pass as a graph and returns a pending result; .submit(callback, on_error=None) delivers a regular CPU tensor asynchronously. torch.compile(step, training=True) captures one training step, forward, backward and a plain SGD update, and runs it on the GPU with parameters uploaded and read back on every call. A five-step fixture matches CPU PyTorch 2.11.

What was measured, and what was not

Correctness was measured: the Life rules match an independent PyTorch implementation on 1×1, 2×2, 7×7 and 96×96 grids, with 301 live cells in generation one; the inference and training fixtures match eager PyTorch. Speed was not. The recorded Game of Life demo on this site was captured on an RTX 5090 through WebGL2 and its provenance file says it demonstrates behaviour, not performance. There is no claim anywhere that this path is faster than running the same arithmetic on the CPU, and for graphs this small it usually is not.

Where it goes next

The obvious next steps are persistent device tensors across requests (so a training loop stops re-uploading its parameters), axis reductions and broadcasting, and wiring the JavaScript-guest adapter into the stock playground, which today serves Python guests only. All of them fit the same rule: the engine never touches the GPU, and every byte the guest sends is validated by the host before it costs anything.