GPU compute from Python and JavaScript in the browser

What the GPU Lab is

A vendored JavaScript package, gpu-lab version 0.1.0, with four interchangeable backends: webgpu.mjs (WGSL compute shaders with a workgroup size of 64), webgl2.mjs (GLSL fragment shaders over float textures), wasm.mjs (C kernels compiled with Clang to kernels.wasm) and cpu.mjs (a JavaScript reference). auto tries WebGPU, then WebGL2, then compiled WASM, then JavaScript, and records every fallback it attempted. Hardware paths reject recognized software renderers, because a missing GPU is not a successful GPU check.

There is no native GPU support. zipp py on a graph program uses the local CPU reference evaluator and reports backend: "cpu-python"; it never starts CUDA, PyTorch or a native GPU API.

Operations and limits

Graphs are plain data (version: 1) over a fixed set of float32 operations: input, full, scalar or shape-matched add, sub and mul, relu, positive masks, transpose, matmul, whole-tensor sum, and a toroidal Game of life step. The runtime does not turn arbitrary Python or JavaScript into shaders.

From Python

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())   # callback receives [0, 60, 120]

Graph.submit(callback, on_error=None, **outputs) records the graph as plain data and posts it as a host request; the callback receives a dict with backend, the named outputs (each with shape, dtype and a flat data list) and the host's stats. Tensors support +, -, *, @, .relu(), .sum(), .transpose(), .positive() and .life(). The host adapter (createPythonGPUAdapter(engine, compute, { allowExecute: true, maxPending: 16 })) drains the request, validates, executes, and answers with pythonCall("__zipp_py_deliver", ...). Execution is granted explicitly, an explicit backend choice is never downgraded to CPU, and a late result never reaches a disposed engine. Run natively with zipp py, the same program evaluates on the CPU reference and reports cpu-python.

From JavaScript, without the VM

The runtime is also usable from plain browser JavaScript with no Zipp engine involved:

import { createRuntime } from "/gpu-lab/src/runtime.mjs"

const runtime = await createRuntime({ backend: "webgl2" })
const result = await runtime.execute({
  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 }],
})
console.log(result.backend, result.outputs.values.data)   // "webgl2" [0, 60, 120]
runtime.dispose()

The playground currently wires zipp_gpu for Python guests only; a JavaScript-guest bridge exists for custom embedders but is not yet enabled in the stock playground.

torch.compile on the GPU

The Torch subset can route a model through the GPU Lab. torch.compile(model)(x) returns a pending result whose .submit(callback, on_error=None) delivers the tensor asynchronously; that method is a Zipp extension, because browser GPU completion is asynchronous and PyTorch's synchronous API cannot be reproduced. Supported inference ops are elementwise add, sub and mul, ReLU, matmul, whole-tensor sum and mean, transpose, square, scalar division, and nn.Linear / nn.ReLU / nn.Sequential combinations.

torch.compile(step, training=True) captures one training step: exactly one zero_grad(), a scalar loss.backward() and optimizer.step() with plain SGD (learning rate, weight decay, maximize and parameter groups; momentum, Nesterov, Adam and higher-order gradients are rejected). Parameters are uploaded and read back on every call; there is no resident model state on the device. A five-step fixture matches CPU PyTorch 2.11.