MLIR has become the de facto infrastructure for building machine learning compilers, not because it was designed specifically for ML, but because its dialect system solves the composability problems that make ML compiler stacks so difficult to maintain. The project, part of the LLVM ecosystem, provides a way to layer abstractions from framework-level operations down to hardware instructions without collapsing into a single monolithic IR.
Why ML compilers need dialects instead of just operations
A compiler-based approach turns a combinatorial problem into a linear one. If features A, B, and C each work independently, their combinations should also work — a property that hand-tuned libraries struggle to achieve. But ML presents a massive semantic gap: source programs operate on high-dimensional tensors in Python frameworks, while target hardware executes low-dimensional vector or scalar instructions on accelerators. The breadth of frameworks (TensorFlow, PyTorch, TFLite) and hardware targets (CPU, GPU, specialized ASICs) makes a single operation set insufficient.
MLIR addresses this with dialects: namespaces containing coherent sets of operations, types, and attributes that model a specific abstraction level. A compiler pieces together dialects like building blocks, adding custom extensions where needed. Two characteristics make this work: operations can nest internal structure through regions, and dialects define their own types to signal abstraction levels.
Regions let operations carry structure
Traditional IRs offer only basic blocks and functions as structural containers. MLIR operations can nest regions — subgraphs of operations — enabling structured operations that carry payload computations. The linalg.generic op exemplifies this: it captures a perfect loop nest implicitly through indexing maps, while its region holds the elementwise computation. Functions and modules are themselves payload-carrying operations under the hood. Regions act as natural boundaries for pattern application, simplifying transformations that would otherwise require analyzing explicit loop structures.
Types map to abstraction levels
Operations compute on values; types represent the abstraction level. An addition on tensors, buffers, or scalars is structurally similar but operates at fundamentally different levels: tensors map to framework abstractions, buffers to memory hierarchies, scalars to chip registers. Each dialect defines its own types, and MLIR's type conversion infrastructure lets rules compose across dialects. Converting between dialects with different types means converting abstraction levels — the core of compiler lowering.
High-level model dialects form the hourglass top
Framework-specific dialects (tf, tfl, torch) import source models faithfully. They live in framework repositories because of tight coupling. The compiler then funnels these into common model dialects — mhlo (descended from XLA, bridging TensorFlow) and tosa (a specification with precise numeric guarantees, increasingly multi-framework). This hourglass structure keeps the middle layers manageable.
The linalg dialect drives middle-level transformations
At the middle level, dialects enjoy design flexibility because they do not model external entities. The linalg dialect is central: its operations capture perfect loop nests implicitly via indexing maps, with regions holding payload computation. This makes fusion as simple as composing inverse indexing maps — inverse(producerIndexMap).compose(consumerIndexMap) — instead of analyzing induction variable ranges. Named ops like linalg.matmul and linalg.conv* are syntactic sugar over linalg.generic, all sharing consistent structure so transformations work uniformly.
Both mhlo and tosa lower to linalg while staying in tensor abstractions. The implicit loop nest enables tiling and fusion critical for tiled architectures. A TOSA conv2d converts to linalg ops, then tiles into nested scf.for loops with tensor.extract_slice operations feeding smaller linalg.conv_2d_nhwc_hwcf and linalg.generic ops.
Bufferization and distribution bridge to hardware
Tensors are immutable SSA values, making analysis simple. Eventually they must become mutable buffers (memref type) — a process called bufferization. MLIR pushes this late, often after vectorization, to keep transformations mechanical. After bufferization, distribution maps problem tiles to hardware tiles (CPU threads, GPU workgroups, workitems). The tiled convolution becomes nested loops over memref.subview operations feeding linalg ops with ins/outs operands.
The flow uses tensor, memref, arith, and math dialects for tensor manipulation, buffer access, arithmetic, and math intrinsics respectively. Lower-level dialects (vector, GPU, SPIR-V, LLVM) handle the final descent to hardware instructions, though those belong to separate compilation stages.
What this means for compiler engineers
The dialect hierarchy is not a fixed pipeline but a composable toolkit. Teams building ML compilers can adopt the dialects matching their source frameworks and hardware targets, extend them where needed, and rely on MLIR's conversion infrastructure to connect the layers. The implicit structure in linalg ops and the type-driven lowering strategy reduce the engineering burden of maintaining transformation passes across the full stack. As tosa adoption grows and bufferization matures, the middle layers are converging toward a shared foundation that multiple frontends and backends can reuse.