Building a GPU compiler from scratch: from PyTorch to CUDA via multi-level IR (Emmy)

Note: the downloaded content is truncated. The text cuts off midway through the autotuning section, so this summary covers only what appears in the download and does not include the conclusions or benchmark results.

The article is a tutorial published in the Kernel Space newsletter, written by Dmitry Trifonov. It follows a PyTorch operation down to CUDA through progressive lowering, using a compiler called Emmy (version 0.3.22). The author recounts that in 2018 he and a friend began building the Apple Neural Engine lowering stack with no prior compiler experience. In his view, building ML compilers still feels like pioneering work, with many frameworks, ideas and tricks. He considers the IR stack and progressive lowering to be among the most settled parts of the field.

Reproducing it requires Python 3.12+ and installing emmy-ml with PyTorch 2.14.1. The IR can be inspected without a GPU. The benchmarks require an NVIDIA GPU and nvcc. The running example is RMSNorm: it computes the mean of the squares of the row, adds epsilon, applies the inverse square root and scales each element by that value and by its learned weight.

The pipeline has six stages, which the article walks through one by one.

Torch IR. The frontend captures the module using PyTorch's export mechanism. At this stage RMSNorm is still a single operation in the torch.fx graph, and the checkpoint weights appear as load("weight").

Tensor IR. It unifies the frontends. Each framework stores its operations with different conventions. For example, a PyTorch dense layer stores the weight transposed relative to that of other frameworks. Tensor IR makes indexing and reductions explicit through index_map calls, which cover transpositions and broadcasts without additional arithmetic. Thus, an nn.Linear and a direct matmul end up describing the same product of shape [4,3,2] and the same sum over the input axis. RMSNorm is broken down here into arithmetic steps: square, sum, division by the count, addition of epsilon, rsqrt and two multiplications.

Loop IR. It converts each node into scalar loops and fuses the regions it can legally merge. If one CUDA kernel were launched per node, intermediate tensors would have to be written to and read back from memory. The exp(-x) example shows two loops fused into one that reads, negates, applies exp and writes, with no intermediate tensor. RMSNorm ends up as a single kernel, and the statistic stays inside the kernel instead of going through global memory.

Normalization. Within Loop IR, it eliminates duplicate work and puts equivalent programs into canonical form. In the (x*x)+(x*x) example, without normalization there are four loads and two multiplications. With it, there is one load and one multiplication. In addition, x+x+y and y+(x+x) produce the same body because addition is commutative. A kernel's identity depends on the normalized body, the typed inputs and outputs, and the compiler's decisions. That is why an efficient schedule can be reused across equivalent forms.

Tile IR. This is the first stage with GPU-aware decisions. It pairs a schedule-independent Fold tree with separate scheduling decisions, following Halide's algorithm/schedule separation. A Fold expresses a map, a reduction or a scan through operands, a per-element function (lift), an optional initial state and combine operation, and an optional observer. In the RMSNorm example, an inner Fold reduces the squares of the row. Another turns the sum into a scale. The outer one produces each output element. The chosen schedule has 256 workers cooperate on the reduction of each row. Emmy can explore other schedules for the same tree before lowering it to Kernel IR and CUDA.

Schedule codec. Matmul serves to show how several decisions interact: how many warps cover the output, how many partial sums each warp keeps in registers, how the reduction is split and how inputs move through shared memory. The codec records each complete combination as a canonical row, with fields such as WORK, TILE, REDUCE, STAGE and RASTER. The article compares two rows for a 128×128 FP16 matmul, both using the same tensor core instruction. A larger register tile reuses each loaded value of A and B more and improves the compute-to-memory-traffic ratio, but it consumes registers and may leave fewer resident warps. The author notes that measurements are needed to know which one wins. The emitted CUDA includes TMA copies with barriers, fragment loads with ldmatrix and the tensor core mma instruction, with double buffering. Emmy walks through the decision points and combines only compatible options. The author suggests that an LLM agent could propose rows through the same interface and that the compiler would check their legality.

Splitting a kernel. Although Loop IR fuses everything that is legal, it is sometimes worth splitting again, and Tile IR has the hardware information to decide. In the example of two chained matmuls, H=AB and O=HC, the fused version recomputes each element of H for each of the 16 output columns. With PLACE=cut, each element of H is computed once. In exchange, a 4 KiB intermediate and a second kernel launch appear, but each matmul can receive a different schedule.

Autotuning. The final section, which is cut off in the download, starts with RMSNorm as a short tuning example. Its schedule space is small, it becomes a single kernel and it can be sampled quickly. The author notes in advance that the gains are usually small and that matmul and attention offer more range but take much longer to tune. The text received does not include the results.