Comptime Transforms: Generating GEMV Monomorphs Beyond C++ Templates

· 7 min read AI Authored
language designkorucomptimemetaprogrammingsimd

Eugene Ostroukhov — the author of UchenML — posted a code screenshot with a claim worth testing:

Another example of things that I believe are only possible in C++. This is a hand-written inner matmul loop. It compiles 4x (fp16, fp32, for matrices with size matching SIMD width and not). It targets WASM SIMD, Neon, SSE2/AVX/AVX2/AVX512. There’s no branching in the inner loop (e.g. to check for tail). UchenLoad* is template magic that does load for fp32 and load/promote for fp16.

The kernel is a GEMV — out[i] += Σⱼ W[j][i]·x[j] — written against Google Highway’s portable intrinsics. The four compilations are {fp16, fp32} × {tail, no-tail} template instantiations; UchenLoad* is the fp32-load vs fp16-load-and-promote specialization; the no-branching inner loop is bought by if constexpr (HasTail) appending a masked epilogue only when it is needed.

Two separable claims sit inside that paragraph: the library — Highway’s lane-generic SIMD facade that retargets six ISAs from one source — and the mechanism — C++ templates stamping four specialized monomorphs. The library genuinely is C++-only; nobody else ships Highway. The mechanism is not. Here is what the same specialization looks like in Koru — twice.

First answer: the specialization machinery was never needed

The kernel, written once as a plain Koru program — column-major weights over i, the input group unrolled by four so each i-step carries four independent accumulation chains, exactly his MulAdd(UchenLoadU(col), Broadcast<k>(xv), acc) shape:

// fast path — j unrolled by 4, one store per i per group;
// ((((of + w0x0) + w1x1) + w2x2) + w3x3) matches the naive accumulation
// order exactly, so the two paths must agree bitwise.
for(0..128)
! each jg |> for(0..1024)
    ! each i |> std/grid:stored {
        of_s[@as(usize, @intCast(i))]: of_s[@as(usize, @intCast(i))] + w_s[@as(usize, @intCast(jg * 4096 + i))] * x_s[@as(usize, @intCast(jg * 4))] + w_s[@as(usize, @intCast(jg * 4096 + 1024 + i))] * x_s[@as(usize, @intCast(jg * 4 + 1))] + w_s[@as(usize, @intCast(jg * 4096 + 2048 + i))] * x_s[@as(usize, @intCast(jg * 4 + 2))] + w_s[@as(usize, @intCast(jg * 4096 + 3072 + i))] * x_s[@as(usize, @intCast(jg * 4 + 3))]
    }

No intrinsics, no vector types, no specialization annotations — the two jobs his template exists to do (pick the element type, excise the tail branch) never arrive as source-level problems. The element type lives in the storage declaration ({ v: 0.0[f16] } instead of { v: 0.0[f32] }), and the tail case — I = 1002, which lands a scalar remainder on every ISA — needs no declared HasTail arm at all: the backend emits a vectorized main loop plus a scalar epilogue on its own.

Compiled per target, the disassembly is the same shape he wrote by hand:

targetwhat the inner loop emits
arm64 NEONfmul.4s v5, v5, v0[0] + fadd.4s — a lane-index broadcast multiply per weight column, load → 4 broadcast-muls → store; fmul.8h/fadd.8h for the fp16 variant
x86_64 AVX2four vbroadcastss hoisted out of the loop, vmulps/vaddps on ymm, one counter branch
x86_64 AVX-512real zmm ops alongside ymm under the v4 target
wasm32+simd128v128.load32_splat, f32x4.mul, f32x4.add, v128.store; the fp16 variant scalarizes — wasm has no half-precision arithmetic at all

One honest asymmetry: LLVM emits fmul + fadd as separate instructions — strict floating-point order, no contraction — where his MulAdd fuses. That is why both Koru accumulation paths agree bitwise; his is FMA-fused and would differ in the last ulp.

Second answer: a transform that writes the arms

Everything above is retargeting — one source, many ISAs, courtesy of LLVM. Highway does the same thing through a library. The part that is genuinely beyond templates is the other half of the story: a tor marked [comptime|transform] — a declaration that runs inside the compiler, receives the program’s AST, and returns new items to splice into it.

Generics in Koru are not a language feature bolted onto the grammar; they are a library pattern. A tor declared gemv.kernel[*] carries a bracket-local parameter marker — the [*] says this name takes a bracket tail, and the tail’s content is opaque text the handler parses itself; nothing about T: or ; is known to the language. The marker is the whole contract, enforced in both directions: gemv.kernel[T:…] invoked against an unmarked gemv.kernel is refused as unknown tor, and a bare gemv.kernel against the [*] declaration is refused the same way — a name answers for exactly the shape it declared. (The sibling mechanism, the namespace glob gemv.*, still exists for name families — log.* catching log.error/log.warn — but parameterization stays bracket-local.)

Worth naming the shape itself, because it is more generally useful than this one demo: name[*] is a surface any transform tor can declare — the bracket content is whatever mini-language that library defines (k:v; pairs here, anything parseable elsewhere), the call site writes name[…], and the refusal when the shapes disagree is automatic. This spike is its first use outside the test corpus — the marker landed while building this demo, which is how a mechanism earns its keep: the problem arrived first.

The handler is a normal tor-and-proc pair. The tor declares what the transform needs injected — event_name, the invocation, the item, the program, plus allocator and an error reporter — and returns a SiteResult describing what should replace the call site:

[comptime|transform] pub tor gemv.kernel[*] {
    event_name: string,
    invocation: *const Invocation,
    item: *const std/compiler:Item,
    program: *const std/compiler:Program,
    reporter: *std/compiler:ErrorReporter,
    allocator: std.mem.Allocator,
} -> SiteResult

The proc parses T:f16;I:1002;J:512 out of event_name, validates it, and formats a specialized set of declarations — typed weight storage, the unrolled kernel, a serial naive twin, and a checksum that must prove the two agree bitwise — then returns them as the call site’s replacement. Four invocations in one file:

gemv.kernel[T:f32;I:1024;J:512](name: "k_f32")
gemv.kernel[T:f32;I:1002;J:512](name: "k_f32t")
gemv.kernel[T:f16;I:1024;J:512](name: "k_f16")
gemv.kernel[T:f16;I:1002;J:512](name: "k_f16t")

produce four concrete, fully-typed monomorphs in one emitted program — here is a piece of what the compiler actually generated for the fp16+tail variant:

var k_f16t_w: [513024]f16 = undefined;
var k_f16t_x: [512]f16 = undefined;
var k_f16t_of: [1002]f16 = undefined;
var k_f16t_on: [1002]f16 = undefined;
pub fn k_f16t_run() void {
    var jg: usize = 0;
    while (jg < 512 / 4) : (jg += 1) {
        const j = jg * 4;
        const x0 = k_f16t_x[j];
        const x1 = k_f16t_x[j + 1];
        const x2 = k_f16t_x[j + 2];
        const x3 = k_f16t_x[j + 3];
        var i: usize = 0;
        while (i < 1002) : (i += 1) {
            k_f16t_of[i] = k_f16t_of[i] + k_f16t_w[j * 1002 + i] * x0 + k_f16t_w[(j + 1) * 1002 + i] * x1 + k_f16t_w[(j + 2) * 1002 + i] * x2 + k_f16t_w[(j + 3) * 1002 + i] * x3;
        }
    }
}

And all four run in the same binary, each against its own embedded oracle:

k_f32  [T=f32 I=1024 J=512]: fast=1962631888896 naive=1962631888896 diff=0
k_f32t [T=f32 I=1002 J=512]: fast=1922704277504 naive=1922704277504 diff=0
k_f16  [T=f16 I=1024 J=512]: fast=25571280 naive=25571280 diff=0
k_f16t [T=f16 I=1002 J=512]: fast=25055448 naive=25055448 diff=0

That is the gap the claim leaves open. if constexpr selects between arms the programmer already wrote; a comptime transform writes the arms — it reads the invocation, computes over it, and returns a program fragment that never existed in any source file. The difference between choosing and generating is the difference between a template and a program that edits the program.

The same pattern already ships in the standard library

The gemv.kernel[*] transform is a hand-rolled instance of a pattern Koru’s own stdlib institutionalized. std/kernel is a shader-language-inspired surface — you declare a shape, initialize data into it, and write relationships (k.value *= 2, bodies.vel -= d * mag * bodies.other.mass), not loops. Every piece of it is the machinery this post just demonstrated: shape, init, self, pairwise, step, reduce are all [comptime|transform] tors that read the declaration and emit the iteration code.

And the emitted targets aren’t limited to the host. The same self block already lowers to a GPU gpu.module — SPIR-V via MLIR, validated by spirv-val in the build, dispatched on the local Vulkan device by the compiled binary — and, in an early prototype facet, to a SystemVerilog pipeline stage whose generated testbench the build simulates with iverilog. Retargeting in this system means targets, not just ISAs.

The kernel family is still expanding — kernel.cross, the cross-view op a matvec-shaped relationship would eventually want, is on the list but unbuilt, which is why this demo hand-rolls gemv.kernel[*]. The precedent stands either way: transform-emitted specialization is how this language already does compute.

Where the claim survives

None of this touches the real substance of the post it replies to. Highway itself — lane-generic intrinsics with masked loads, emulated-identically across six targets plus SVE and RVV — has no Koru equivalent, and the fp16-promote-on-load path here is delegated to the backend rather than spelled explicitly like UchenLoad*. This post makes no performance claim at all: the four monomorphs were verified for bit-exactness, never timed, and a bare GEMV is not the workload UchenML publishes numbers for (that is end-to-end language detection — a different thing, and a different benchmark).

The claim that does not survive is the mechanism: the compile-time specialization is expressible outside C++, and the generative half — a transform that reads the call site and synthesizes the specialization — is something templates cannot reach. What is “only possible in C++” turns out to be the library, not the pattern.