docs / GPU Programming / Cookbook

Cookbook

Each recipe below is verified by the compiler test suite. Copy one into a .mi file and run it with miri run yourfile.mi.

Vector add

The simplest kernel: two device buffers added element-wise into a third, then read back.

use system.gpu
use system.collections.array

const N = 4

gpu let a = [1.0, 2.0, 3.0, 4.0]
gpu let b = [5.0, 6.0, 7.0, 8.0]
gpu var dst = [0.0, 0.0, 0.0, 0.0]

forall i in 0..N
    dst[i] = a[i] + b[i]

let host = dst
println(f'{host[0]} {host[1]} {host[2]} {host[3]}')

Output: 6.0 8.0 10.0 12.0.

Buffer-reuse pipeline

Two kernels run over the same gpu var with no readback between them, so the device buffer is uploaded once and reused. The telemetry counters prove the cost model: one upload, two launches, one readback.

use system.gpu

const N = 8

gpu_reset_telemetry()
gpu var data = [0, 0, 0, 0, 0, 0, 0, 0]

forall i in 0..N
    data[i] = i + 8

forall i in 0..N
    data[i] = data[i] + 8

let host = data
println(f'{host[7]} {gpu_uploads()} {gpu_launches()} {gpu_readbacks()} {gpu_fences()}')

Output: 23 1 2 1 1data[7] = 7 + 8 = 15, then 15 + 8 = 23. Then 1 upload, 2 launches, 1 readback, 1 fence. No cross-residency assignment sits between the two forall blocks, which is the visible marker that the buffer is shared.

SAXPY (fused multiply-add)

The coefficient is a host scalar, captured into the kernel as a read-only uniform.

use system.gpu

gpu let x = [1.0, 2.0, 3.0, 4.0]
gpu let y = [5.0, 6.0, 7.0, 8.0]
gpu var dst = [0.0, 0.0, 0.0, 0.0]

let a = 2.0
forall i in 0..4
    dst[i] = a * x[i] + y[i]

let host = dst
println(f'{host[0]} {host[1]} {host[2]} {host[3]}')

Output: 7.0 10.0 13.0 16.0.

Matrix multiply (naive)

A 2×2 product mapped as one thread per output cell, with the dot product computed by an in-kernel while loop. This is the correctness illustration — every thread reads global memory directly.

use system.gpu
use system.collections.array

gpu let a = [1.0, 2.0, 3.0, 4.0]
gpu let b = [5.0, 6.0, 7.0, 8.0]
gpu var c = Array<f32, 4>()

forall idx in 0..4
    let row = idx / 2
    let col = idx - row * 2
    var sum = 0.0
    var k = 0
    while k < 2
        sum = sum + a[row * 2 + k] * b[k * 2 + col]
        k = k + 1
    c[idx] = sum

let host = c
println(f'{host[0]} {host[1]} {host[2]} {host[3]}')

Output: 19.0 22.0 43.0 50.0. For the optimized counterpart — shared-memory tiles, barriers, explicit workgroup shape — see Shared Memory & Barriers.

Parallel reduction

On-device tree reduction returning a scalar — one launch, one 4-byte readback.

use system.gpu

gpu let data = [1, 2, 3, 4, 5, 6, 7, 8]
let sum = data.reduce(0, fn(a int, b int) int: a + b)
println(f'sum = {sum}')

Output: sum = 36. The result is a gpu-resident scalar. The host let binding is the fence. See On-Device Reduction for keeping it on the device.

Atomic histogram

Concurrent scatter writes through Atomic<u32> buckets.

use system.gpu
use system.gpu.atomic

fn main()
    gpu var hist = Array<Atomic<u32>, 256>()
    forall i in 0..147456
        atomic_add(hist, i % 256, 1 as u32)
    let host = hist
    let bucket_0 = host[0]
    let bucket_1 = host[1]
    println(f"bucket_0={bucket_0} bucket_1={bucket_1}")

Output: bucket_0=576 bucket_1=576.

Warp shuffle reduction

Register-only tree reduction over one subgroup — no shared memory, no barrier.

use system.gpu
use system.collections.array

gpu fn warp_reduce_sum(input Array<int, 32>, dst out Array<int, 1>)
    let lane = kernel.warp.lane_id
    var v = input[lane]

    v = v + kernel.warp.shuffle_down(v, 16)
    v = v + kernel.warp.shuffle_down(v, 8)
    v = v + kernel.warp.shuffle_down(v, 4)
    v = v + kernel.warp.shuffle_down(v, 2)
    v = v + kernel.warp.shuffle_down(v, 1)

    if lane == 0
        dst[0] = v

fn main()
    gpu let input = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32]
    gpu var dst = Array<int,1>()
    warp_reduce_sum(input, dst).launch(Dim3(1, 1, 1), Dim3(32, 1, 1))
    let result = dst
    println(f'{result[0]}')

Output: 528.