Version: 1.1.0

GPU programming

GPU kernels are Mojo functions that run on an accelerator. You use the same structs, parameters, and traits as CPU code, plus a grid of threads.

Essentials

• A kernel reads its thread index from the GPU index helpers.
• Parameters carry tile size and dtype so the compiler specializes the kernel.
• The MAX library supplies device buffers and launch helpers.
• You can drop to vendor assembly such as PTX when a hardware instruction has no wrapper yet.
• Start from a correct CPU version, then move the inner loop.

Example

def vector_add(a: UnsafePointer[Float32], b: UnsafePointer[Float32], out: UnsafePointer[Float32], n: Int):
    var i = global_idx.x
    if i < n:
        out[i] = a[i] + b[i]

Keep in mind

The puzzles linked from the docs home are the shortest path to the launch and memory rules. See GPU puzzles.