Koru now runs on the GPU

· 5 min read

Koru now runs on the GPU

The roadmap post went up yesterday, and it ended on a promise:

The subtree already knows its shape. It already knows what it’s allowed to do. Soon it will know where it runs.

Soon turned out to be today.

What landed

The same Koru kernel — the exact source, one parameter added — now runs on the GPU of the machine that compiled it. No hand-written shader: the compiler writes the kernel, lowers it to a device binary, and emits the host code that dispatches it.

import std/kernel
import std/io

std/kernel:shape(Body) {
    mass: f32,
}

std/kernel:init(Body) {
    { mass: 1.0 },
    { mass: 2.0 },
    { mass: 3.0 },
}
| kernel k |> std/kernel:self|mlir[gpu] { k.mass *= 2.0 }
| computed c |> std/io:print.blk {
        mass[0]={{ c[0].mass:f }} mass[1]={{ c[1].mass:f }} mass[2]={{ c[2].mass:f }}
    }

That’s the whole program — a pure-Koru .k file, no ~ in sight. The kernel that runs on the GPU is written entirely in Koru; nobody writes a shader.

The only new thing on the page is [gpu]. Drop it and the identical kernel runs as a CPU loop; add it and the identical kernel runs on the GPU. That is the roadmap’s whole claim made concrete: the target is a parameter of the lowering, not a rewrite.

What that one parameter sets in motion, all at compile time:

  1. The kernel:self transform generates a gpu.module — block-per-element indexing, a StorageBuffer memref, the SPIR-V entry-point ABI — instead of the CPU scf.for loop.
  2. The build lowers it (convert-gpu-to-spirv), serializes a .spv blob, and runs spirv-val as a build gate: an invalid blob is a loud compile failure, never a shipped artifact.
  3. The compiler emits Vulkan host glue and links it — upload the buffer, build a pipeline from the blob, dispatch, read back — and the buffer comes home mutated.

Run it on an Apple M2 Pro (Vulkan via MoltenVK) and it prints mass[0]=2 mass[1]=4 mass[2]=6. Those numbers were computed on the GPU.

The tests

These are the ground-truth regression tests behind the claim — the source, the expected output, and the pass/fail state, live:

How it got built — both halves proven

The arc was split so each rung could be proven before the next was wired:

  • Reach SPIR-V. Teach the kernel transform to emit the GPU form and route it through the SPIR-V pipeline to a validated blob. At this rung the blob exists and validates; nothing dispatches it yet.
  • Run it. Hand-spike the raw Vulkan dispatch first — standalone, the exact shape the compiler would emit — until it computed {1,2,3} → {2,4,6} on the actual device. Then teach the compiler to emit that proven glue and link libvulkan.

No runtime engine, no JIT. The .spv blob is produced at build time like any object file, and the host side is ordinary Vulkan calls the compiler writes and links. It is AOT the whole way down — the same posture as the MLIR CPU backend: partial-in, blob-out, glued at the seam, host stays host.

The honest ledger

The roadmap drew the line between shipped and reach explicitly, so this post owes the same:

  • This is a correctness-and-capability milestone, not a speed claim. The input is three floats. It proves the pipe — Koru source → generated MLIR → SPIR-V → dispatch → result — end to end on real hardware. It says nothing yet about throughput; a three-element kernel is exactly where the CPU wins. Placement economics — the crossover, fusion, keeping data resident — is still the reach the roadmap describes.
  • f32 only, on Apple. An f64 shape under [gpu] is a compile-time wall, not a crash (test 390_104).
  • Selection is explicit today — the variant and the target, by hand at the call site (|mlir[gpu]). That’s the scaffolding. The roadmap’s selection layer moves the choice off the call site: a build config picks the target per build profile — a gpu build and a cpu build off the same source, the way build:variants already selects fast over naive — and eventually the compiler infers it globally, from capability and profile. The call-site knob is a labelled placeholder that dies the day that lands.
  • The toolchain paths are pinned. The MoltenVK and MLIR prefixes are hardcoded today — the same proof-of-concept knot the deps/requires work will untie.

None of that changes the load-bearing fact: a subtree the compiler owns crossed from the CPU to the GPU, automatically, and ran.

Why it’s on-thesis

Koru’s move is to take what is normally runtime or manual and press it into compile time. The compiler already writes the kernel. Now it also decides — for this one explicit parameter, later by inference — where that kernel runs, and emits everything needed to run it there. Same source, two machines, chosen by the build.

The subtree knew its shape. It knew what it was allowed to do. Now it knows where it runs.