← Back to projects
Summary

This was a passion project. I'd just built a general-purpose language compiler in my compilers course (CMPUT 415), and I wanted to take that same toolkit (grammar, AST, type checking, and lowering through MLIR/LLVM) into the machine-learning compiler world to understand how systems like Triton and XLA actually work. So I built one end to end, and it's a real pipeline, not a toy. You write a tile-level kernel with a separate schedule block; it parses (ANTLR4), type-checks with clang-style diagnostics, lowers into a custom MLIR dialect, and then splits into two backends off that one dialect: a CPU path that emits real x86-64 AVX2 assembly, and a GPU path that lowers the same source to Ampere tensor-core ops (mma.sync, cp.async, ldmatrix). It's an IR/codegen-level project (I haven't run kernels on hardware or benchmarked them yet), but every stage of the compiler is built and tested, and the two-backends-from-one-dialect design is the part I'm proudest of.

Stack

What it's built with.

Compiler Frontend

  • ANTLR4
  • AST Design
  • Type Systems
  • DSL Design

MLIR & Codegen

  • MLIR Dialect Authoring (ODS/TableGen)
  • linalg / memref / scf
  • gpu / nvgpu (tensor cores)
  • LLVM IR → x86 AVX2

Build & Tooling

  • C++17
  • CMake
  • lit / FileCheck
Details

How it works.

What it does

One tile-level source language, two very different targets. A kernel is written once, algorithm plus a separate schedule, and the compiler carries it through a custom MLIR dialect down to either x86-64 AVX2 assembly or NVIDIA Ampere tensor-core ops. The frontend (grammar, AST, type checker) and the dialect are the fixed core; each backend is a set of pattern-rewrite lowerings on top.

I built it to learn this layer by doing rather than reading. Modern ML compilers lean on the same staged-lowering ideas I'd just used on a general-purpose language in CMPUT 415, grammar, AST, a type checker, IR, and passes, only pointed at tensors and hardware instead of a stack machine. Retargeting that mental model from "compile a language" to "compile a matmul onto AVX2 and tensor cores" is exactly what I wanted out of it. It's IR/codegen-level: the CPU path produces real assembly on disk, the GPU path produces correct tensor-core IR, and hardware execution and benchmarking are scoped as the next step.

The interesting part: two backends, one dialect

"Vectorize and fuse" on CPU and "manage the memory hierarchy and hit tensor cores" on GPU are genuinely different problems, but because both consume the same `tile` dialect, adding the second backend was a set of new lowering patterns, not a rewrite. That is the whole argument for an MLIR-style design, and getting to feel it firsthand (instead of taking it on faith from a paper) was the point.

On CPU the schedule drives the optimizations: vector widths clamp to the target ISA, a matmul and its elementwise epilogue fuse into one op, and the outer loop nest can promote to a parallel form. On GPU the tile's memory-space attribute drives codegen: global→shared becomes `cp.async`, shared→register becomes `ldmatrix`, and the matmul body becomes a cooperative `mma.sync`. Same frontend contract, completely different backend.

Highlights

The things I'm proudest of.

  • ▹A real language frontend: an ANTLR4 grammar for a tile-level kernel DSL with first-class tile types (shape + layout + memory space), load/store/matmul/elementwise/reduction ops, and a separate Halide/Triton-style schedule block (tile, bind, vectorize, unroll, parallel, pipeline, prefetch, fuse, swizzle), algorithm and schedule are decoupled, exactly like the systems it learns from.
  • ▹A type checker with clang-style diagnostics (`file:line:col: error:`, notes, source-line carets) that catches shape mismatches, dangling iteration-variable references, invalid bind dimensions, and out-of-range pipeline depths at compile time.
  • ▹A custom MLIR dialect authored in ODS/TableGen with real op verifiers, the shared middle layer both backends lower from. Getting this dialect right (encoding memory hierarchy, tiling, and schedule intent) is what makes the whole two-target design work.
  • ▹CPU backend, all the way to metal: `tile` → linalg + memref + scf → LLVM IR → x86-64 AVX2 assembly via llc, emitting real `.s` on disk. Includes vertical fusion (matmul + elementwise epilogue), ISA-aware vector-width clamping (AVX2/AVX-512), algebraic canonicalization, and parallel-loop promotion.
  • ▹GPU backend to tensor cores: the same source lowers to `gpu` + `nvgpu` for NVIDIA sm_80, `nvgpu.mma.sync` (16×8×16, fp16 in / fp32 accum), `cp.async` global-to-shared staging, and `ldmatrix` shared-to-register loads. This path is complete at the IR level; the NVVM→PTX step and running on real silicon are the honest next milestones.
  • ▹Backed by a lit/FileCheck test suite across the parser, type checker, dialect verifiers, and every lowering pass. No JIT, execution, or perf benchmarking yet, that's the deliberate boundary of the project, and I'd rather state it than imply it.