GPU: WebGL and WebGPU

For raw GPU work, Python gets the full WebGL 2 and WebGPU interfaces, the same as JavaScript. You write your own shaders and talk to the graphics card directly. Shaders are written as Python strings: GLSL for WebGL, WGSL for WebGPU.

Most programs never need this: 2D functions and 3D functions do the GPU work for you. Use this page when you want your own shaders, or are bringing a WebGL or WebGPU example across from JavaScript.

What differs from ordinary drawing

• Ask the canvas for a GPU context. gl = canvas.getContext("webgl2"), or canvas.getContext("webgpu"). It takes the place of ctx for that canvas: one canvas draws with either the 2D functions or the GPU, not both.
• Typed arrays come from the page: js.Float32Array([0, 0.8, -0.8, -0.7]), js.Uint16Array([...]).
• Constants are read from the context, as in JavaScript: gl.ARRAY_BUFFER, gl.TRIANGLES, gl.FLOAT.
• Reading a value back needs await, for example ok = await gl.getProgramParameter(prog, gl.LINK_STATUS) or await gl.getShaderInfoLog(shader). Each read waits a frame, so check once at the start, not every frame.
• WebGPU set-up calls return promises. await them and you get back the object to use: adapter = await js.navigator.gpu.requestAdapter(), then device = await adapter.requestDevice().
• Keyword arguments become the descriptor objects WebGPU takes: device.createShaderModule(code=code) passes {code: code}. Nested parts are Python dictionaries and lists.
• Draw each frame in draw(), as with everything else.

WebGL 2: a shader-coloured triangle

A triangle whose colours move, worked out per pixel by a fragment shader.

gl = canvas.getContext("webgl2")

VS = """#version 300 es
in vec2 pos; out vec2 v;
void main() { v = pos; gl_Position = vec4(pos, 0.0, 1.0); }"""

FS = """#version 300 es
precision mediump float; in vec2 v; uniform float t; out vec4 col;
void main() { col = vec4(0.5 + 0.5 * sin(t + v.x * 3.0), 0.5 + 0.5 * v.y, 0.8, 1.0); }"""

def shader(kind, src):
    s = gl.createShader(kind)
    gl.shaderSource(s, src)
    gl.compileShader(s)
    return s

prog = gl.createProgram()
gl.attachShader(prog, shader(gl.VERTEX_SHADER, VS))
gl.attachShader(prog, shader(gl.FRAGMENT_SHADER, FS))
gl.linkProgram(prog)

buf = gl.createBuffer()
gl.bindBuffer(gl.ARRAY_BUFFER, buf)
gl.bufferData(gl.ARRAY_BUFFER, js.Float32Array([0, 0.8, -0.8, -0.7, 0.8, -0.7]), gl.STATIC_DRAW)
loc = gl.getAttribLocation(prog, "pos")
gl.enableVertexAttribArray(loc)
gl.vertexAttribPointer(loc, 2, gl.FLOAT, False, 0, 0)
t_loc = gl.getUniformLocation(prog, "t")
t = 0.0

ok = await gl.getProgramParameter(prog, gl.LINK_STATUS)
print("linked:", ok)

def draw():
    global t
    t += 0.05
    gl.viewport(0, 0, canvas.width, canvas.height)
    gl.clearColor(0.05, 0.05, 0.1, 1)
    gl.clear(gl.COLOR_BUFFER_BIT)
    gl.useProgram(prog)
    gl.uniform1f(t_loc, t)
    gl.drawArrays(gl.TRIANGLES, 0, 3)

Tap to enlarge

How it fits together:

• VS and FS are the two shaders, as Python strings. The vertex shader places each corner; the fragment shader colours each pixel.
• shader() compiles one; the program joins both. print("linked:", ok) shows whether they linked. If not, print(await gl.getProgramInfoLog(prog)) says why.
• The buffer holds the three corners as a Float32Array. vertexAttribPointer tells the shader to read them two numbers at a time.
• Each frame, draw() sets the size, clears, sends the time t to the shader with uniform1f, and draws.
• canvas.width and canvas.height are the canvas's size in real pixels, which is what viewport wants.

WebGPU: a triangle

WebGPU is the newer GPU interface. It needs a recent Safari on Mac, iPhone and iPad, or a recent desktop Chrome or Edge.

gpu = js.navigator.gpu
adapter = await gpu.requestAdapter()
device = await adapter.requestDevice()
context = canvas.getContext("webgpu")
fmt = await gpu.getPreferredCanvasFormat()
context.configure(device=device, format=fmt, alphaMode="opaque")

code = """
@vertex fn vs(@builtin(vertex_index) i: u32) -> @builtin(position) vec4f {
  var p = array<vec2f, 3>(vec2f(0.0, 0.7), vec2f(-0.7, -0.6), vec2f(0.7, -0.6));
  return vec4f(p[i], 0.0, 1.0);
}
@fragment fn fs() -> @location(0) vec4f { return vec4f(1.0, 0.6, 0.2, 1.0); }
"""
module = device.createShaderModule(code=code)
pipeline = device.createRenderPipeline(
    layout="auto",
    vertex={"module": module, "entryPoint": "vs"},
    fragment={"module": module, "entryPoint": "fs", "targets": [{"format": fmt}]},
)

def draw():
    encoder = device.createCommandEncoder()
    view = context.getCurrentTexture().createView()
    rp = encoder.beginRenderPass(colorAttachments=[{"view": view, "loadOp": "clear", "storeOp": "store",
                                                    "clearValue": {"r": 0.05, "g": 0.05, "b": 0.1, "a": 1}}])
    rp.setPipeline(pipeline)
    rp.draw(3)
    rp.end()
    device.queue.submit([encoder.finish()])

Tap to enlarge

How it fits together:

• requestAdapter() finds the graphics card and requestDevice() opens it. Both are awaited.
• context.configure(...) connects the canvas to the device, in the colour format the screen prefers.
• code is one WGSL string holding both shaders: vs places the three corners, fs colours them orange.
• The pipeline joins the shaders into one drawing recipe. Its descriptor is written as keyword arguments with dictionaries inside.
• Each frame, draw() records a render pass that clears the canvas and draws three corners, then submits it to the GPU.

Related

• 3D functions
• Bringing a JavaScript example across
• Reading a value back