A hands-on one semester course where students build their own compiler from scratch, starting from elementwise programs and ending with training SOTA LLMs on GPUs. This course aggressively builds on the previous week, and is an exercise in slop management. If you let any slop in early, it will compound and you will not finish the class.
Course description: This course covers the design and implementation of a modern machine learning compiler, and examines the interaction between IR design, hardware capabilities, and the structure of machine learning programs. Topics covered include term rewriting, code generation, movement operators, kernel fusion, memory hierarchies, GPU architecture, automatic differentiation, and flash attention. It is a project course, providing experience with performance-oriented programming, managing a codebase that grows all semester, and working in 1 or 2 person teams, culminating in a compiler capable of training modern LLMs.
Prerequisites: This course has no formal prerequisites, but it is designed for someone who is deeply familiar with the material in 15-213 or equivalent, highly skilled in the language they choose to implement their compiler in, capable of managing programming projects that are long in timeline, and comfortable with C and some GPU programming. There will be frequent in-person quizzes and lectures will not be posted online, so only sign up if you plan to attend class.
This syllabus will outline the spec, but not the implementation. The implementation and choice of programming language is up to the student. There is no starter code. We will establish a universal, simple, interchange format for UOps.
# gemm_256x256.uir
%0 = buffer 65536 : dtype=f32 slot=0
%1 = buffer 65536 : dtype=f32 slot=1
%2 = reshape %0, (256, 256, 1)
%3 = reshape %1, (256, 1, 256)
%4 = mul %2, %3
%5 = reduce %4 : op=add pop=1
- 55% -- compiler (weeks 1-10)
- 5% -- 10 tests -- due every monday (weeks 1-10), judged on correctness by openly accessible binary oracle
- 20% -- correctness -- due every wednesday, passing all student tests + teacher tests + basic performance gates
- 15% -- style (did you write slop? can I read your code? can it be simpler?)
- 15% -- performance (weeks 6-10, graded on a curve)
- 25% -- final project (weeks 11-14)
- 20% -- in person quizzes at the start of class (will be at least 20 of them)
No midterm or final. < 70 = F, 70-80 = C, 80-90 = B, 90+ = A.
uops · elementwise ops · symbolic · rewrites · renderer
Students become familiar with the UOp and write a compiler capable of compiling and simplifying elementwise programs. We introduce the graph structure and rewriting. The UOp is hash consed.
UOp(op, src:tuple[UOp, ...], arg)
Recursive properties introduced:
dtype(void,bool,i32,f32)
shape:tuple[int, ...]
Ops introduced:
PARAM(ParamArg) -- arg input to the call
CALL(CallArg) -- src[0] is the function body, src[1:] are the params...
CONST -- type determined by the arg
STACK(src, ...) -- new leading axis, all src has same shape. until week 4, inputs can only be STACK/CONST.
INDEX -- shape from (n,) -> (), pop axes off the left with srcs
CAST/RECIP
ADD/MUL/CMPLT
FLOORDIV/FLOORMOD
WHERE
The graph rewrite engine should be capable of "running" these programs via const folding.
Lectures:
- wed - Course intro, turing completeness, the UOp, the format, the first assignment
- fri - ALU stuff, bottom up greedy rewrite engines, const folding
- mon - call/param stuff, symbolic simplification, computer algebra systems
Compilers must be handed in by Wednesday every week before 10 AM class.
Students write a renderer and a runtime capable of compiling and executing C programs.
Recursive properties introduced:
addrspace(GLOBAL,ALU)
Ops introduced:
BUFFER(ParamArg) -- global persistent buffer (never in call)
ALLOC(ParamArg) -- call scope storage, example "float a[10]"
LOAD/STORE -- load has 1 src, store has 2
AFTER -- AFTER(buf, <store>)
This can render (loopless) C functions that act on memory and immediates.
Lectures:
- wed - toposort and rendering, C refresher
- fri - memory management and runtime
- mon - type systems for shape, dtype, and addrspace
loops · gemms · convs
Ops introduced:
RANGE -- src is the end of the range (3 means 0,1,2) and arg is the range name
END -- srcs are (passthrough, range)
Should be simple, and we can write gemms and convs in this language.
movement ops · broadcasting · reduction · rangeify
Now we introduce movement ops and loops. Note how you can derive gemms and convs from the basic movement ops.
Ops introduced:
STACK -- now allows arbitrary inputs
RESHAPE/EXPAND/PERMUTE/FLIP/PAD/SHRINK
REDUCE (axis count / op)
Rangeify removes shapes from all ops.
This compiler is now capable of compiling any model, splitting it into kernels to avoid repeated compute, and compiling it to C.
Ops introduced:
STAGE -- a temp buffer you can index into
Autodiff only works on pure shape graphs, not STORE, CALL, AFTER, RANGE, END. Here you will train your first MNIST.
FLOPS counting · upcasting · fast GEMMs · flash attention
Now things get fast. Still only on CPU, but we can now produce SOTA-competitive C code for GEMMs, reduces, and convs.
Ranges can have axis type LOOP, UPCAST, or THREAD.
parallel · GPUs · hardware accelerators · tensor cores
Here we add kernels, GPUs, axis mappings, and tensor cores. This can produce decent torch-competitive CUDA/HIP code now.
real models · symbolic
Here we implement a real LLM and train models.
Use your compiler to implement a paper, port it to strange hardware, mostly anything
| Weeks | Topics | Milestone |
|---|---|---|
| 1–2 | UOps, elementwise ops, symbolic, rewrites | Compile & simplify elementwise programs to C |
| 3–4 | Loops, movement ops, reduction, rangeify | Compile any model to C (slow, but correct) |
| 5–6 | Memory hierarchies, upcasting, GEMMs, convs | SOTA-competitive CPU code |
| 7–8 | Call, GPUs, accelerators, tensor cores | Torch-competitive CUDA/HIP code |
| 9–10 | Real models, autodiff | Train a real LLM |
| 11+ | Student projects | Choose a project and extend your compiler |