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. All shapes are rank 0. We allow rank 1 shapes here, but ALU ops must be rank 0.
UOp(op, src:tuple[UOp, ...], arg)
Recursive properties introduced:
dtype(void,bool,i32,f32)
shape:tuple[int, ...]
bounds:tuple[int, int]
Ops introduced:
CONST -- type determined by the arg
lambda calculus:
PARAM(ParamArg) -- arg input to the call
CALL(CallArg) -- src[0] is the function body, src[1:] are the params...
lil movement:
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
SHRINK(offset, shape) -- offset can be symbolic, shape must be integer
ALU:
CAST/RECIPROCAL/TRUNC
ADD/MUL/MAX/XOR
CMPLT -- returns bool
SHL/FLOORDIV/FLOORMOD -- i32 only
WHERE
The graph rewrite engine should be capable of "running" these programs via const folding also rendering them to C. We support dtype "broadcasting" here.
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)
versions:dict[UOp, tuple[UOp, ...]] <-- a list of UOps that store to a buffer
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, we load/store from/to GLOBAL buffer
AFTER -- AFTER(buf, <store>)
This can render (loopless) C functions that act on memory and immediates. Loads and stores must be fully ordered for programs to be valid. Loads and stores are bounds checked to never be out of range.
Lectures:
- wed - toposort and rendering, C refresher
- fri - memory management and runtime
- mon - type systems for shape, dtype, and addrspace
loops · gemms · convs
Recursive properties introduced:
ranges:dict[UOp, None]
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. Note, basic rangeify is allowed here now.
broadcasting · reduction · rangeify
We still keep the rank < 2 restriction on shapes, but ALU can be rank 1 now.
Ops introduced:
REDUCE (axis count / op)
STAGE -- temp buffer
EXPAND
Rangeify removes shapes from all ops. We have to confront questions about materialization here.
movement ops
Now we introduce movement ops and loops, allowing rank >= 2 shapes for the first time. Note how you can derive gemms and convs from the basic movement ops.
Ops introduced:
RESHAPE/PERMUTE
STACK supports arbitrary inputs, SHRINK supports STACKs in the params.
PAD and FLIP are composite ops here. This compiler is now capable of compiling any model, splitting it into kernels to avoid repeated compute, and compiling it to C.
Autodiff only works on pure shape graphs, not STORE, CALL, AFTER, RANGE, END.
Here you will run Qwen, you'll also learn about Taylor expansion for SIN/LOG2/EXP2/SQRT.
autodiff · threads · FLOPS counting · fast GEMMs
axistype(LOOP,GLOBAL)
Ops introduced:
GRADIENT/DETACH
You'll add GRADIENT/DETACH ops and train MNIST. You'll extend your C program to support threads and begin considering GPUs.
parallel · GPUs · hardware accelerators · tensor cores
axistype(LOOP,GLOBAL,LOCAL)
addrspace(GLOBAL,LOCAL,ALU)
Here we add kernels, GPUs, axis mappings, and tensor cores. This can produce decent torch-competitive CUDA/HIP code now.
upcasting · 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.
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 |
please make it open source, I want to learn thru this material