A single-instruction GPU virtual machine and toolchain
Here's a dumb idea that turned out to work. A GPU virtual machine with exactly one instruction, so that divergence stops being something I ever have to think about.
Sometimes you have to run a program you didn't write and won't see until runtime, a node graph an artist wired together in an editor, a scrap of user-generated content that amounts to a small expression. You don't want to run it once, you want it per element, the same program over a million different inputs at once on the GPU. Ideally you'd hand it to the shader compiler, but in most places this comes up you can't JIT a shader at runtime, so the program has to live as data and something on the GPU has to interpret it.
The obvious way to write that interpreter is an opcode table. You fetch an instruction, look at its tag, and jump to whatever implements add or mul or sin. On a CPU this is essentially free and nobody thinks about it, but on a GPU it's poison. A warp runs its lanes in lockstep and diverges whenever a value that differs between lanes reaches control flow, and an opcode interpreter lets that happen everywhere. Some handlers branch internally on per-lane values, not the ones backed by a hardware instruction but the builtins a compiler emulates in software, which is more of them than you would think, and any complex operation you write yourself. The interpreted program's own branches are worse, since they make the program counter data-dependent, and once the per-lane counters disagree every fetch after that disagrees too. The shader has usually made data-dependent decisions before the loop even starts, and you often want different lanes running different programs by indexing the stream by a thread id, where they are not even fetching the same opcode.
A single instruction removes all of it since the program counter only increments the same for every lane. A decision is a conditional write through the guard rather than a branch, so a computed value never reaches it, and lanes on different programs still run the same instruction every step and differ only in operands. So even per-lane streams never diverge either. The example here runs one program across every lane, but the design is what keeps all of the above divergence-free. Nothing a lane computes ever reaches control flow. You can run a different program per lane if you want to.
What that instruction has to be
What would an instruction like that even look like? It sounds like a nebulous idea, and not one that should be possible, that a single operation could stand in for any program at all. It turns out it can, and the requirements are surprisingly minimal. It has to do arithmetic, because in the end these programs are just math over their inputs. And it has to be able to make a decision, because a program that cannot branch isn't worth much. The constraint that matters is that the decision cannot become a real branch, otherwise you're right back to the divergence problem the whole idea is meant to avoid, so it has to choose between values without ever jumping.
A fused multiply-add covers the arithmetic, and a guard placed on top of it covers the decision. I call the combination an FMAG, which stands for fused multiply-add if greater (than zero) / guard. The idea of a machine with a single instruction already has a name, the one instruction set computer or OISC, though it's usually treated as a curiosity rather than something you would actually build on. This is one of the few legitimate uses of it I'm aware of.
What's in here
This repository is the VM and the toolchain that feeds it, a small C library. You write programs either in the shading language or directly against the SSA CFG IR builder, since the whole point is to JIT them as data at runtime, and the library optimizes and lowers whichever you use into a flat FMAG stream. A reference VM runs that stream so you can test it on the CPU. There's also a command line driver along with, somewhat absurdly for a machine with a single instruction, a text assembler and matching disassembler for the masochists who can think in these FMAG instructions (so far I'm the only one xD)
Roughly, the path a program takes through it is as follows.
- Shading language (
tools/sl.c). An example frontend, here to show how you might use the builder rather than something you have to use, so it lives intoolsand is built entirely on the public library rather than being part of it. A small typed language off32throughf32x4with functions, locals,if/else, calls and the usual operators, compiling to ordinary IR for the passes below to chew on. - IR builder (
fmag_ir.h) is an SSA CFG. You build functions, blocks, constants and FMAG operations and wire them up with calls, branches and returns, with the common operations (add,mul,select,absand so on) provided as sugar over plain FMAGs. - Optimization and lowering. This is where most of the work happens, and honestly most of the entire repository, since it turns out there is a surprising amount of optimization to do around a single instruction. The full set of passes, in pipeline order, each detailed in OPTIMIZATIONS.md, is:
- Inlining, every call cloned in leaf first until the module is one function, recursion rejected.
- Ternary recognition,
if/elsediamonds whose arms either return or rejoin a merge are turned straight into selects with the merge PHIs rewritten as selects, and forwarding and empty blocks spliced out. - If-conversion, whatever structured control flow is left after that is predicated down into a single block.
- Constant folding, literal arithmetic plus the algebraic identities like
x*1,x*0, a select whose two arms are the same value, and a guard already known true or false, all backed by a shared constant pool so identical literals are interned. - Guard stripping, peeling
(x > 0) ? 1 : 0sign booleans and their1 - ...complements so a comparison can serve as a guard directly. - Multiply-add fusion, folding a producer
a*1 + 0into the operand or addend slot of whatever consumes it. - Commutative canonicalization, moving a lone constant factor into
band folding two constant factors together. - Common subexpression elimination.
- Dead code elimination.
- Min and max orientation, flipping the arms so the kept value ends up in the lower numbered register and saves a move.
- Contraction and reassociation, enabled only when the builder is created in fast math mode, merging constant factors and addends and dropping
x * 0. - Register allocation during lowering, biased toward the result registers, recycling each register at its last use, reusing the dying destination as the fallback, and shuffling the results into the output registers at the end with a scratch register to break any cycle.
- Assembler and disassembler (
fmag_as.h). A plain text form, one%rd = fmag %ra, b, c, guardper line with a.def name, inputs, outputsdirective, printed back out by the disassembler. - VM (
fmag_vm.h) is the CPU interpreter used for testing. It runs a compiled stream over a register file and leaves the outputs in the low registers. - CLI (
tools/cli.c). Thecldriver compiles a shading language file (-c), assembles a text program (-a), and runs (-r), disassembles (-d) and inspects (-i).fmagfiles. - Superoptimizer (
tools/superopt.c). Finds the shortest sequence of FMAGs that computes a given function and proves it against a reference, exhaustively for one input or by sampling for two. It runs either a breadth-first search over register machine states for the true minimum on short programs, or a STOKE style stochastic search for the longer ones. Several of the sequences the standard library and optimizer use came out of it. - Bindings (
odin/fmag.glsl) Bindings to use this in Odin. Which is the programming language we use at JangaFX.
You can build the lib sources directly, but if you'd rather not, unified.py amalgamates the whole library into a single fmag.h in the stb style, so you just drop that one header into your project and define FMAG_IMPLEMENTATION in one translation unit to pull the implementation in.
The instruction
As stated, the entire machine is built around one operation:
fmag(a, b, c, d, e) = e > 0 ? fma(a, b, c) : d
When the guard
eis positive you get the FMA. Otherwise you just getdback.
The reason a single instruction can cover this much ground comes down to what the two halves give you. An FMA can evaluate any polynomial on its own, since chaining it together is just Horner's rule, and once you have polynomials a couple of Newton iterations give you the rest of a math library: division, square roots, reciprocals, all of it. The guard is the other half of the instruction, a conditional that is essentially free, because (e > 0) ? x : y compiles to a select rather than a jump.
On the GPU the guard is meant to behave as a select, which is the same thing you would get from mix(y, x, e > 0) or a cndmask. This means a branch buried inside one of these programs only costs a conditional move, and the warp stays together. That single property is what makes the whole scheme viable here. The hardware never has to mask off lanes or serialize paths, because there are no paths to serialize.
Arithmetic
When you pin the guard to a positive constant, the conditional cancels out and what remains is ordinary arithmetic. Every common math operation falls out naturally
fma(a, b, c) = a*b + c = fmag(a, b, c, 0, 1) # 1 > 0, always the product
add(a, b) = a + b = fma(a, 1, b) # a*1 + b
sub(a, b) = a - b = fma(b, -1, a) # b*-1 + a
mul(a, b) = a * b = fma(a, b, 0) # a*b + 0
neg(a) = -a = fma(a, -1, 0)
When the guard is something you actually computed instead of a constant, the same instruction picks between two values
select(t, f, g) = (g > 0) ? t : f = fmag(t, 1, 0, f, g)
relu(x) = max(x, 0) = fmag(x, 1, 0, 0, x)
max(a, b) = (a-b > 0) ? a : b = fmag(a, 1, 0, b, sub(a, b))
min(a, b) = (a-b > 0) ? b : a = fmag(b, 1, 0, a, sub(a, b))
abs(x) = (x > 0) ? x : -x = fmag(x, 1, 0, neg(x), x)
That is already most of a math library out of a single instruction. The comparisons and selects give you masks, clamping, and the piecewise behavior range reduction depends on, and the multiply-add gives you Horner and Newton. Between the two halves, very little arithmetic is left uncovered.
The encoding
An instruction is four 32-bit words, sized to fit into a uvec4 that can be sent to a shader.
packet
0-31: "b operand"
32-63: "c operand"
64-95: "guard operand"
96-111: "d destination register"
112-127: "a multiply register"
b, c and the guard are NaN-boxed. The payload is a register index when the value is a quiet NaN, and the float literal itself when it is not. The one restriction this imposes is that NaN and infinity cannot be used as literals, because the tag would become ambiguous. Though in practice nobody was going to write a literal NaN anyway.
The last word holds two bare register indices, the destination d and the first multiply operand a, 16 bits each, with no room left for a NaN tag. The destination also serves as the fallback, which is a deliberate design choice rather than a limitation of the encoding. The instruction reads r[d], possibly overwrites it with the FMA result, and on the false arm leaves it unchanged, so the result and the fallback share the same slot. In effect an FMAG is a conditional write to a register. This dual role is what keeps the encoding at exactly 16 bytes per instruction, which matters when you are streaming a program through a buffer.
Picking which slot is a register
The destination always has to be a register, since the result has to be written somewhere. The first multiply operand a is forced to be a register too, and the reasoning there is straightforward. There is no point in supporting two literal factors in a*b. If both were constants the product would just be constant folded ahead of time, so at most one factor is ever a literal and a can always be the register without giving anything up. That leaves b, c and the guard free to be a register or an immediate, and funnelling a into the remaining register-only slot is also what makes the whole 16 byte encoding fit together cleanly.
The interpreter
One of the more satisfying properties of this design is how small the interpreter turns out to be. The entire execution engine is only a few lines of GLSL. It decodes the three NaN-boxed operands with isnan, splits the last word into its two register indices, reads r[a] directly since a is always a register, and uses a branchless mix.
float r[16]; // Size this to the register count FMAG reports, see below vec3 loadops(uvec3 o) { vec3 f = uintBitsToFloat(o); uvec3 i = o & 15u; // This should be masked to the register size return mix(f, vec3(r[i.x], r[i.y], r[i.z]), isnan(f)); } void exec(uvec4 op) { vec3 v = loadops(op.xyz); uint d = op.w & 15u; // This should be masked to the register size uint a = (op.w >> 16) & 15u; // This should be masked to the register size r[d] = mix(r[d], fma(r[a], v.x, v.y), v.z > 0.0); } void main() { for (uint i = 0u; i < u_length; ++i) exec(b_code[i]); // inputs were loaded into r[0..], outputs are read back out of r[...] }
Note
That mix is the boolean select overload, not a lerp. It throws away the operand it did not pick.
The register file r has to be sized to the number of scalar registers a program needs, which is the regs field FMAG writes into the header (every vector is just that many scalar slots). Since r is a local array it weighs directly on register pressure inside the shader, so you do not want it oversized. The IR optimizer and the lowering passes work hard to minimize the count, and it is usually small, but you still have to commit to a size at compile time in the shader. One approach is to specialize the interpreter for a few powers of two and pick the smallest that fits at runtime. Another is to settle on one large size and accept the pressure. Sixteen is a decent middle ground for most programs.
The inner loop on RDNA
To see whether all of this actually delivers on its promise, I ran the above interpreter through RGA for RDNA3 (gfx1100) to inspect the ISA the GPU would execute. The inner loop is about as tight as I could have hoped for. There is a single loop branch and nothing else. No divergence, no complex control flow and the register reads and writes come out as plain indexed moves.
_L1: s_buffer_load_b128 s[12:15], s[4:7], s3 ; fetch instruction uniformly s_and_b32 s10, s15, 15 ; d = da & 15 s_bfe_u32 s0, s15, 0x40010 ; a = (da >> 16) & 15 s_and_b32 s1, s12, 15 ; index of b s_and_b32 s11, s13, 15 ; index of c s_and_b32 s2, s14, 15 ; index of guard s_mov_b32 m0, ... / v_movrels_b32 v.., v0 ; r[idx] reads v_cmp_u_f32 ... ; isnan on b, c, guard v_cndmask_b32 ... ; operand select v_fmac_f32 v19, v17, v18 ; r[a]*b + c v_cmp_lt_f32 vcc, 0, v20 ; guard > 0 v_cndmask_b32 v16, v16, v19, vcc ; guard select v_movreld_b32 v0, v16 ; r[d] = result s_cbranch_scc1 _L1 ; loop
The instruction fetch is a scalar load, s_buffer_load_b128, because the program counter is the same across the warp, so the entire warp pulls a single instruction at a time from the same address. This doesn't need to be the case, and I emplore you to do per-lane programs if desired. In either case, there is very little actual math in the body here, three v_cmp_u for the isnan tests on b, c and the guard, three v_cndmask for the operand selection, one v_fmac for the fused multiply-add itself, and the guard's compare and select. Everything else is moving values in and out of the register file with v_movrels and v_movreld. The loop is dominated by data movement rather than computation, which is exactly what you want from an interpreter whose entire job is to shuttle values through a single repeated operation and the body is branchless aside from the loop back-edge. The body is meant to be partially unrolled in the shader, doing several instructions per iteration, which amortizes the scalar fetch and the back-edge across the batch and lets the indexed register moves of one instruction overlap the arithmetic of the next. Since none of it can diverge, an unrolled stretch just runs straight through.
A note on how this was built
I prototyped the idea first in Python, then in Odin, to convince myself it actually held together before committing to it. What you see here is the serious implementation in C. I used Claude heavily throughout, both to prototype and experiment with the idea and to write a lot of the documentation you are reading.