py2wgsl — Python → WGSL compiler

Compiles a restricted subset of Python into complete WGSL compute shader modules that run on WebGPU. Kernels are ordinary Python functions parsed with the ast module — the Python is never executed; it is translated.

Built for warppool, the distributed WebGPU compute pool, and usable anywhere WGSL runs.

Install

pip install py2wgsl        # zero runtime dependencies
pip install py2wgsl[gpu]   # + wgpu, to execute kernels from Python

Usage

from py2wgsl import Array, f32, compile_kernel, global_id, array_length

def scale_add(scale: f32, offset: f32, values: Array[f32], result: Array[f32]):
    i = global_id()
    if i < array_length(result):
        result[i] = values[i] * scale + offset

kernel = compile_kernel(scale_add, workgroup_size=64)
print(kernel.wgsl)         # full WGSL module, ready for createShaderModule()
print(kernel.buffers)      # binding layout metadata for setting up bind groups

There is also compile_kernel_source(source_str) for compiling from a source string, and python -m py2wgsl.demo prints two example kernels (including a Monte Carlo pi kernel like this project's hand-written shader).

Parameter → binding mapping

Buffers the kernel writes to (via buf[i] = ...) get read_write access; the rest are read. The entry point is always fn main with @builtin(global_invocation_id).

Supported subset

• Assignments (first assignment declares a typed var), annotated assignments, augmented assignments
• if / elif / else, while, for i in range(...) with constant step, break, continue, bare return
• Arithmetic, comparison, boolean, and bitwise operators; a if c else b (compiles to select)
• Math builtins bare or as math.*: sqrt, sin, cos, floor, min, max, clamp, pow, abs, ... and math.pi / math.e / math.tau
• Casts: f32(x), i32(x), u32(x)
• GPU builtins: global_id(), local_id(), workgroup_id(), array_length(buf), barrier(), storage_barrier()

Semantics: WGSL wins over Python

These are deliberate; the compiler errors on the ambiguous cases it can catch:

• Integers are 32-bit and wrap; they are not Python bignums
• / on two integers is a compile error — use // or cast with f32()
• // on negative integers truncates toward zero (Python floors)
• No implicit int/float conversions — cast explicitly (bare literals are fine: x * 2 works when x: f32, matching WGSL abstract literals)
• Conditions must be bool — no truthiness
• Variables first assigned inside a block are scoped to that block
• No atomics, vectors, textures, or user-defined function calls (yet)

Verified on GPU

The generated WGSL is validated by tests/test_py2wgsl.py (structure and error handling, no GPU needed) and has been executed end-to-end on WebGPU via wgpu-py (Metal): elementwise results match a Python reference and the compiled Monte Carlo kernel estimates pi correctly.