The Subtree Knows Where to Run

· 14 min read

The Subtree Knows Where to Run

This is a roadmap post. Some of what follows ships today and links to the post that proves it. Some of it is a reach we can see clearly from where we’re standing. We mark the line between them explicitly, because that line is the interesting part: the reach is short.

The question nobody answers at compile time

You have a compute kernel. Should it run on the CPU or the GPU?

For small inputs the CPU wins, decisively — launching a GPU kernel and copying data across the bus is fixed overhead that dwarfs the actual work. For large inputs the GPU wins, once there’s enough parallel work to amortize that cost. The crossover is a threshold, and it depends on the data and the machine.

Almost every system answers this at runtime — a check, a branch, a dispatch. Not because runtime is better, but because deciding it at compile time is ergonomically brutal on most platforms. You’d need to profile representative inputs, model the target machine, orchestrate a two-phase build, and manage stale profiles. Most teams look at that and take the runtime overhead instead. It’s a rational trade — the cost of the good path is just too high.

The claim of this post is narrow and, we think, true: in Koru that cost is mostly already paid. The pieces that make compile-time placement expensive elsewhere are, here, existing language features with passing tests. What’s missing is smaller than what’s present.

Let’s walk the stack.

Piece 1 — Variants already select a target

Koru has proc variants: multiple implementations of one event, same contract, chosen at compile time. The variant-selection post’s own example includes a GPU one:

~event compute { x: i32 }
| result i32

~proc compute|naive { /* debuggable */ }
~proc compute|fast  { /* Gauss's formula */ }
~proc compute|gpu   { /* GPU kernel */ }

And build-time selection already exists — a comptime event maps events to variants per build config, and the call site doesn’t change:

~[release]std/build:variants { "compute": "fast", "graphics:blur": "gpu" }

So the selection mechanism for “run this one on the GPU” is not hypothetical. It’s the same mechanism that picks fast over naive. The target is just another variant tag, and gpu is already in the namespace the emitter knows.

Piece 2 — The kernel scope is already a structural, closed region

The catch with ~proc compute|gpu today is that the body is hand-written GLSL — the human writes the shader; Koru orchestrates it. That’s the same escape hatch as |zig or |js. It’s useful, but it is not “compile a Koru subtree to the GPU automatically.”

The automatic version needs a subtree the compiler can retarget — one where the compiler, not a human, decides the iteration. Koru already has exactly that shape in its computation kernels. A kernel describes a relationship between data elements, not a loop:

std/kernel:pairwise {
    bodies.vel       -= d * mag * bodies.other.mass
    bodies.other.vel += d * mag * bodies.mass
}

kernel:pairwise and kernel:self are compile-time transforms that synthesize the iteration. And crucially, the kernel scope is closed: a validator rejects any non-kernel construct inside it — no arbitrary if, no ad-hoc loops, no branch constructors. That validator is a wall. Inside the wall, the language is a small, compiler-owned sub-grammar. Outside, it’s the full language.

That wall is the whole game. Inside it, retargeting the emitted iteration from a CPU loop to a GPU kernel is “teach the transform a second output,” not “rewrite anything.” The kernel source doesn’t change. Nobody writes a shader.

Piece 3 — “Can it run there?” is a capability, and capabilities are checkable

That validator — the wall — is a hand-rolled denylist today. Koru is already generalizing it into something declarative: the trellis, a set of pattern laws over the shape of the program’s flow tree, checked at compile time.

~std/trellis:define("kernel-shape")
| `std/kernel:init/kernel/.*` _ => ok
| `.*std/kernel:pairwise`     _ => error "pairwise found outside kernel init"

~std/trellis:enforce("kernel-shape")

The trellis’s own design notes name the destination out loud: it’s “the categorization substrate for variant capability-float (gpu-legal, simd-fusable, …).” A gpu-legal law is just a trellis whose error arms reject the shapes a GPU can’t run.

But the real endpoint is bigger than pattern-matching shapes. It’s a capability lattice — annotations that compose: an event marked gpu-legal requires its callees to be gpu-legal; allocation isn’t; unbounded recursion isn’t. The compiler derives a subtree’s capability from its parts, the same way purity already propagates. Trellis is rung one of that ladder — the checker you can write today with zero new machinery. The inference is the reach.

The clean division that falls out: capability decides what’s eligible to cross the wall; the backend decides how to lower it. Neither subsumes the other.

Piece 4 — MLIR is a sibling backend, not a new final IR

Koru compiles to Zig today. It also, already, compiles to JavaScript — two live backends off one AST. So a third target isn’t a rewrite of the pipeline; it’s another leaf in a fan-out that already has two.

For GPU that leaf is MLIR. Two things worth stating plainly, because they’re the load-bearing facts:

  • MLIR works on sub-trees. GPU compilation in MLIR is built around outlining: you lift a kernel region into a small module, lower it to a device binary, embed the blob, and leave the rest of the program alone. Partial-in, blob-out is the design center — which maps exactly onto “take a subtree.” You do not move the whole program.
  • MLIR is a sibling to Zig, not a replacement. Host code keeps lowering through Zig. Only the gpu-legal subtrees cross to MLIR. This is heterogeneous compilation, and it’s how every GPU system works — CUDA, SYCL, Mojo, JAX. Nobody recompiles their host language to the device. The host stays host.

So the picture is: Koru → Zig for the host, Koru → MLIR → SPIR-V/PTX for the kernels, glued at the seam. Additive, not a reroute.

Piece 5 — Taps are the instrument the whole thing needs

Here’s where Koru has something most toolchains don’t. Deciding placement well means measuring — and the hardest part of compile-time specialization elsewhere is building an instrument that can observe the right thing at the right seam without wrecking the code or the measurement.

Koru’s event taps are that instrument, and they’re a first-class, zero-cost-when-off, compile-time feature. A universal tap cross-cuts every transition in the program:

~tap(* -> *)
| Profile p |> record(p.source, p.timestamp_ns)

And taps come in three tiers — a cost/detail dial that happens to match the instrumentation lifecycle exactly:

TierSizeCarries
Transition12 bytessource, dest, branch (enums) — the cheap firehose: which seams are hot
Profile32+ bytes+ timestamp — timing and timelines
Auditvariable+ the serialized payload — the actual data at the seam

That bottom tier is the one that closes the loop. An Audit tap sees the real payload flowing through a seam — not just that a transition happened, but the actual sizes and shapes:

~tap(kernel_seam -> *)
| Audit a |> variant_picker:observe(a)   // a.payload holds the real values

Which means the “representative inputs” problem — the thing that makes autotuning a chore — dissolves. You don’t guess which sizes matter. You observe the distribution the code actually sees, on the actual target, at zero cost when the instrument is off.

The composition — where it all lands

Put the pieces together and the runtime-vs-compile-time question stops being a trade you’re forced into and becomes one you win:

  1. A kernel subtree clears the capability check (gpu-legal) — it can run on the GPU. Static.
  2. Taps observe the real per-seam data distribution on the target, keyed to a machine fingerprint. This is profile-guided optimization — except the profile carries payload semantics (the actual n), not just edge counts.
  3. The fingerprint does double duty: it’s the profile’s key and its staleness detector — a new GPU or driver changes the fingerprint, which is the signal to recalibrate.
  4. The compiler reads the shape of each seam’s distribution and picks per seam:
    • unimodal (always large, or always small) → statically select one leaf and fuse across the seam, keeping data resident on the device. No branch.
    • bimodal (sometimes tiny, sometimes huge) → keep a runtime branch — but only here, where the data proves it earns its keep.
  5. Anything not known at build time falls back to runtime dispatch — the portable path, used exactly where specialization isn’t available.

The prize in step 4 is fusion. A runtime branch structurally can’t fuse across a seam — both paths must survive, so data can’t be pinned to one side. Static selection can, and in GPU work the launch-and-transfer overhead usually dominates the compute, so fusing adjacent kernels and keeping data resident is often a bigger win than the per-kernel CPU/GPU choice itself.

And because Koru’s compile time is programmable, the classic three-tool PGO ceremony collapses into one comptime pass with two modes: no profile yet → emit a benchmark harness that measures exactly the seams this binary contains; profile present → select, fuse, specialize. Run once to get the harness, run it on the box to get the fingerprinted profile, compile again to get the specialized binary. No external autotuner — the compiler already knows every seam, because it generated them.

What’s in the tree today, and what’s the reach

The honest ledger — and the reason this is a roadmap and not a fantasy:

In the tree today (each with a passing test or a shipping post):

  • Proc variants, including a |gpu variant and build-time variant selection.
  • The kernel scope as a structural, compiler-generated region with a validator wall that rejects non-kernel constructs.
  • Trellis: declarative shape laws over the flow tree, enforce (halt) and check (comptime dispatch).
  • Taps: universal, zero-cost-when-off, three tiers up to Audit — which exposes the serialized payload. Payload access is tested.
  • Two live structural backends off one AST (Zig and JavaScript), proving the fan-out is real.
  • A GLSL → SPIR-V → Vulkan-wrapper pipeline behind the |gpu variant (compile-only spike today; it does not yet dispatch).
  • A |mlir AOT backend for CPU. A |mlir variant’s body is lowered (mlir-optmlir-translate), compiled (clang -c), and linked into the binary — no runtime machinery, the object is built at compile time. A scalar kernel genuinely emits MLIR, lowers, links, and runs: compute|mlir(21)42, pinned by a passing test. Scalar-in/out and CPU-only so far, and the toolchain path is hardcoded — a proof-of-concept knot the deps/requires family will untie. But the pipe is real: Koru source → emitted MLIR → lowered → linked → ran.

The reach (the slice list):

  • A variant-picker instrumentation library: a tap consumer that histograms parameter sizes per seam and writes a fingerprinted profile. New consumer on an existing surface — application code, not language machinery.
  • A machine fingerprint probe (write-once in the stdlib, reused everywhere).
  • A second output leaf on the kernel transform: have kernel:self/pairwise generate MLIR where they generate a Zig loop today. The MLIR backend now exists to emit into (above) — so the reach is the compiler writing the kernel, not a hand-written |mlir body.
  • MLIR/generated forms for the handful of dense-buffer primitives a kernel calls.
  • Capability inference: generalize the validator wall from a hand-rolled denylist into a composable gpu-legal annotation lattice. Trellis is rung one.
  • The comptime PGO pass: benchmark-harness mode ⇄ select-and-fuse mode, reading the profile.

Six slices. None of them is “invent a systems language that lowers to MLIR.” Every one is a bounded reach from something that already exists and already has a test.

Why this is on-thesis

Koru’s whole move is to take things that are normally runtime or manual and press them into the compile-time type system. Resource safety: runtime discipline → phantom obligations. Control flow: runtime branching → the type system. This is that same move, pointed at performance placement. The CPU/GPU decision — deferred to a runtime check everywhere else because compile-time is too painful — moves to compile time, and the thing that makes it not painful is machinery Koru already built for its own reasons.

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

That’s the road. Most of it is already paved.