Quick start

This example compiles Python kernel source to a WAVE binary and runs it through the emulator. It works without a GPU. Build the toolchain and install the Python SDK using the installation instructions.

Vector addition

Save this as vector_add.py:

import wave_gpu
from wave_gpu.runtime import compile_kernel, launch_emulator

source = """
@kernel
def vector_add(a: f32[:], b: f32[:], out: f32[:], n: u32):
    gid = thread_id()
    if gid < n:
        out[gid] = a[gid] + b[gid]
"""

a = wave_gpu.array([1.0, -2.5, 3.25, 4.0])
b = wave_gpu.array([5.0, 6.5, -7.25, 8.0])
out = wave_gpu.zeros(4)
binary = compile_kernel(source, "python")
launch_emulator(binary, [a, b, out], [4], (1, 1, 1), (4, 1, 1))
print(out.to_list())

Run it:

python3 vector_add.py

The result is [6.0, 4.0, -4.0, 12.0]. The launch uses one workgroup with four threads. The scalar argument 4 sets n; each thread adds one pair of elements.

SDK guides

The Python guide describes compilation and emulator execution. The Rust guide also covers launch through the vendor runtimes. See the C++ guide and TypeScript guide for their APIs.

For direct assembly, continue to Writing kernels. For hardware requirements, see Supported GPUs.