%uML Compiler from Scratch
From compiler foundations to frontier systems

Building an ML Compiler from Scratch

The book follows one small compiler, uopc, as it is built week by week: a graph IR and a rewrite engine first, then C code generation, loops, kernel fusion, automatic differentiation, CPU and GPU optimization, and training a GPT. Continue through compiler analysis, GPU engineering, dynamic execution, numerical methods, and distributed systems, with runnable labs and a frontier-lab preparation track.

Start reading Jump to chapter 1 PDF · 815 pages
Chapters
37
Diagrams
46
Words
269k
Compiler
2,803 lines
gemm.uirinput
# C = A @ B for 4x4 matrices, from movement ops only
%0 = buffer 16 : dtype=f32 slot=0
%1 = buffer 16 : dtype=f32 slot=1
%2 = reshape %0, (4, 1, 4)
%3 = reshape %1, (4, 4)
%4 = permute %3, (1, 0)
%5 = reshape %4, (1, 4, 4)
%6 = mul %2, %5
%7 = reduce %6 : op=add pop=1
parse → rangeify → simplify → linearize → render
gemm.cgenerated by uopc
void gemm(float* restrict data0, float* restrict data1, float* restrict data2) {
  for (int ridx0 = 0; ridx0 < 4; ridx0++) {
    int alu1 = (ridx0*4);
    for (int ridx1 = 0; ridx1 < 4; ridx1++) {
      float acc0;
      acc0 = 0.0f;
      for (int ridx2 = 0; ridx2 < 4; ridx2++) {
        float val2 = acc0;
        float val3 = data0[(alu1+ridx2)];
        float val4 = data1[((ridx2*4)+ridx1)];
        acc0 = (val2+(val3*val4));
      }
      float val5 = acc0;
      data2[(alu1+ridx1)] = val5;
    }
  }
}
MNIST, week 5
94%
test accuracy, 15 ms per step
GEMM 1024³, week 6
321
GFLOP/s, from 5 naive (numpy: 449)
Flash attention, week 6
3 ways
unfused, fused, online softmax
GPT, weeks 9–10
2.11
val loss on Tiny Shakespeare, from 3.83

Contents

Front matter
  1. Preface: How to Read This Book
Front matter
  1. Foundations: From Programming to Tensor Compilers
Part I

The IR

  1. Week 1The UOp: A Graph IR
  2. Week 1Term Rewriting and the Rewrite Engine
  3. Week 1Constant folding and symbolic algebra
Part II

From graphs to machine code

  1. Week 2Memory, rendering C, and the runtime
  2. Week 3Loops and the linearizer
  3. Week 4Movement ops and rangeify
Part III

A real compiler

  1. Week 5Kernel fusion and fission
  2. Week 5Automatic differentiation
  3. Week 5A Tensor library and your first MNIST
Part IV

Making it fast

  1. Week 6Fast CPU kernels
  2. Week 6Fusion, recompute, and flash attention
Part V

Accelerators

  1. Week 7GPU architecture and mapping kernels
  2. Week 8Shared memory and tensor cores
Part VI

Real models

  1. Weeks 9–10Training a GPT
  2. Weeks 9–14Symbolic Shapes, Inference, and Where to Go Next
Appendices
  1. Appendix AUOp and .uir Reference
  2. Appendix BTesting, Debugging, and Slop Management
  3. Appendix CGlossary and Quiz Bank
  4. Appendix DComplete source code
Part VII

Advanced engineering

  1. Advanced trackThe Syllabus Contract and Completion Gates
  2. Advanced trackControl Flow, SSA, Analysis, and Lowering
  3. Advanced trackGPU Engineering: Correctness, Throughput, and Attention
  4. Advanced trackSymbolic Shapes, Guards, Mutation, and Graph Capture
  5. Advanced trackMemory Planning, Precision, and Distributed Execution
  6. CapstoneFrom Course Completion to Independent Engineering
Part VIII

Frontier systems and research

  1. Advanced trackNumerical Analysis for Compiler and Kernel Engineers
  2. Advanced trackScheduling, Autotuning, and Backend Construction
  3. Advanced trackIrregular Workloads and Heterogeneous Accelerators
  4. Frontier preparationA Frontier-Lab Preparation Map
Part IX

Implementing deeper compilers

  1. Compiler internalsBuilding a Source Compiler: Contracts, CFGs, and SSA Construction
  2. Compiler internalsOptimization with Executable-Path Reasoning
  3. Compiler internalsExecuting SSA, Destroying Phis, and Allocating Registers
  4. Tensor compilerA Dynamic Tensor Compiler with Fusion and Planned Memory
  5. Training compilerDifferentiating and Training Through the Dynamic Compiler
  6. GPU implementationRegister-Tiled GEMM on a Real GPU

Run it yourself

Python 3.11+, numpy and gcc with OpenMP. An NVIDIA GPU is optional; GPU kernels are also checked on the CPU through an emulator.

git clone https://github.com/StarAtNyte/ml-compiler-from-scratch.git
cd ml-compiler-from-scratch
python -m venv .venv && .venv/bin/pip install -e ".[dev]"
.venv/bin/python -m pytest -q              # 102 tests
.venv/bin/python examples/mnist_mlp.py     # train MNIST with your compiler
.venv/bin/python examples/gpt.py           # train a GPT, then sample from it