This course will focus on modern compiler design in the era of AI programming frameworks (e.g., PyTorch, JAX) and AI hardware accelerators. AI models are evolving and continue to rely on high-performance computing. Specialized hardware for AI is often hard to program to exploit peak performance. AI models are also evolving in ways that are coupled with hardware strengths. This course will teach how to build effective compiler systems for the next generation of programming frameworks and chips in a layered way to improve hardware usability and deliver high performance as automatically as possible.
As background, we will cover modern compiler design principles and concepts (such as SSA) that we have learnt over the past few decades while building infrastructure like LLVM – the lessons learnt and the mistakes made. The course will then cover the MLIR compiler infrastructure in some depth, focusing on how MLIR can be used to build the next generation of compiler systems, in particular, high-level compilers (similar to Torch Inductor or XLA) and mid-level compilers like Triton. The course will also cover effective optimization for parallelism and locality, as well as mapping to specialized matrix/tensor units commonly found in AI chips. Tiling and cross-operator fusion are examples of transformations that will be covered to emphasize where compilers excel over hand-written libraries and kernels, while also showing how the right intermediate-representation abstractions should be chosen to perform such optimizations effectively.