In short: An AI compiler such as MLIR, XLA, or TVM takes a high-level model graph and progressively lowers it, fuses operations, and schedules memory and parallelism so a GPU, TPU, or custom accelerator runs the same math with far fewer kernel launches and memory round-trips.
When you write y = relu(conv2d(x, w) + b) in PyTorch or JAX, nothing on the chip understands that line. Between your Python and the silicon sits an AI compiler whose job is to turn a mathematical description of a neural network into efficient machine kernels. Getting this translation right routinely changes inference latency by 2–5x on the exact same hardware, which is why every serious accelerator team ships a compiler, not just a driver.
What "lowering" actually means
A model starts life as a graph of coarse operators — matmul, convolution, softmax, layernorm. The compiler does not emit assembly for these in one leap. Instead it lowers the program through a series of intermediate representations (IRs), each closer to hardware. MLIR formalizes this with dialects: the same program can exist simultaneously as tosa or stablehlo (framework-level math), then linalg (loop-nest algebra over tensors), then affine and scf (explicit loops and control flow), then vector and llvm or an accelerator dialect. Each step is a rewrite that preserves semantics but exposes more structure to optimize.
⭐ The power of MLIR is that one infrastructure hosts many dialects. A TPU backend, a GPU backend, and an in-house NPU can share the same linalg and tiling passes, then diverge only at the final hardware dialect — instead of each team writing a compiler from scratch.
Fusion: the single biggest win
Most accelerators are starved for bandwidth, not compute. An A100 does ~19.5 TFLOP/s of FP32 but moves only ~1.5–2 TB/s from HBM. If every elementwise op writes its full result back to memory, you pay that bandwidth repeatedly. Operator fusion collapses a chain like add → relu → multiply into one kernel that keeps intermediate values in registers or shared memory, touching DRAM only for the true inputs and final output.
Consider a bias + gelu over a 32M-element activation tensor (FP16). Unfused, you write and re-read ~128 MB of traffic; fused, you drop roughly half of it. XLA and TVM both do this automatically, and XLA's fusion is a large part of why the same JAX model can outrun a naive PyTorch eager loop by 2–4x before any manual tuning.
⭐ Fusion is not free: fuse too aggressively and a kernel's register pressure spills to local memory, erasing the gain. Compilers model this with cost heuristics, which is why the scheduling search below matters as much as the fusion rule itself.
Scheduling and tiling for the memory hierarchy
A matrix multiply is trivial math but a hard scheduling problem. The compiler must decide tile sizes for the L2/shared-memory/register levels, loop order, vectorization width, and how to map work onto thread blocks and tensor cores. TVM pioneered separating the algorithm (what to compute) from the schedule (how), then searching the schedule space automatically with its Ansor/AutoTVM autotuners. On an unfamiliar operator shape, an autotuned schedule often beats a hand-written vendor kernel by 10–30%.
| Tool | Origin | Core idea | Typical target |
|---|---|---|---|
| MLIR | LLVM / Google | Reusable multi-dialect IR infrastructure | Any (CPU, GPU, NPU) |
| XLA | Whole-graph fusion + layout assignment | TPU, GPU | |
| TVM | Apache / UW | Decoupled schedule + autotuning search | GPU, mobile, FPGA |
Where it meets custom silicon
For a bespoke accelerator, the compiler is the product. Teams often emit a custom MLIR dialect for their systolic array, then reuse upstream tiling and bufferization passes. This is also where EDA and compiler worlds meet: the datapath you synthesize with Synopsys or Cadence flows must match the memory assumptions your compiler bakes into its schedules. If you are building both the chip and its toolchain, our sister brand Avecas VLSI design services works on exactly this hardware/compiler co-design boundary.
The same GEMM intuition that drives fusion and tiling underpins almost every accelerator benchmark — if you want to go deeper on the kernel math, see our companion piece on GEMM optimization, and for the hardware side read systolic array dataflow.
Want to lower your own model to a TVM schedule and read the generated PTX line by line? Our CourseTron AI Hardware track walks through MLIR dialects, fusion passes, and autotuning on real GPUs.
Frequently asked questions
Is an AI compiler different from a normal C compiler?
Yes. A C compiler optimizes scalar control flow; an AI compiler reasons about dense tensor operations, data layout, and the memory hierarchy, and it often searches a large space of equivalent schedules rather than applying fixed passes.
Why do XLA and TVM both build on MLIR now?
MLIR provides reusable IR plumbing — pass managers, dialects, bufferization — so projects stop re-implementing the same infrastructure and can share loop-level optimizations across backends.
Does fusion always make things faster?
No. Fusing too many ops raises register pressure and can spill to slow memory. Compilers use cost models to stop fusing before the kernel becomes inefficient.
Can I just rely on vendor kernels like cuBLAS?
For standard shapes, yes. But for unusual sizes, fused epilogues, or custom hardware, an autotuned compiler schedule frequently beats the stock library.