PyTorch→CUDA LLM Compiler

With the constraints of 5,000 lines of Python, 2 weeks, and no external libraries, the author built an LLM compiler from scratch and laid out a pipeline that lowers PyTorch graphs through Torch IR → Tensor IR → Loop IR → Tile IR → Kernel IR → CUDA.

The core primitives of Tensor IR are three: Elementwise, Reduction, and IndexMap. Computations like add, mul, exp, rsqrt, axis reductions like sum, max, prod, and layout transformations like reshape, transpose, slice, unsqueeze, concat are all reduced to these three.

As an example, RMSNorm unfolds from a single PyTorch op into a Tensor IR of 15 nodes, and Linear decomposes into a broadcast mul and sum. The intermediate M×K×N shape is not actually materialized, and is later merged into a single kernel during the fusion stage.

Consecutive IndexMaps like transpose and slice are collapsed into one by composing the coordinate mappings. They are then lifted into Loop IR, adjacent loops are fused to reduce global memory round-trips, and the final CUDA kernel is emitted.

Part 1 covers capture, decomposition, IndexMap composition, loopification, and fusion, while the upcoming Part 2 will explain Tile IR, matmul optimization, and code generation details.