A tour of MLIR: The Dialect Stack Everyone Depends On | Aditya Kumar Light Dark System
Post<br>Cancel<br>A tour of MLIR: The Dialect Stack Everyone Depends On<br>Contents A tour of MLIR: The Dialect Stack Everyone Depends On
If you train or serve models, you depend on MLIR whether or not you have ever written a line of it. XLA lowers through it, Triton is built on it, Mojo is MLIR-native, and Torch-MLIR, IREE, and ONNX-MLIR exist to funnel their respective frontends into it. The reason a single piece of infrastructure ended up underneath so many otherwise-competing stacks is worth understanding, because it explains a lot about how modern ML compilers are actually built, and where their seams are.<br>This post is a tour of MLIR: what it is, the dialect idea that makes it different, how a tensor operation is progressively lowered to machine code, and what the infrastructure does and does not give you.<br>What MLIR Actually Is<br>The common misconception is that MLIR is “another IR like LLVM IR.” It is better described as an IR construction kit. LLVM IR is a single, fixed, low-level representation: roughly a typed assembly with SSA values. That is the right abstraction for the last mile to machine code, and the wrong one for a matrix multiply over tensors. Historically, every domain compiler that needed a higher-level representation invented its own from scratch: XLA had HLO, Halide its own IR, TensorFlow its graph, each shipping with a separate pass manager, serialization format, verifier, and pile of bugs.1<br>MLIR’s premise is that those representations have far more in common than not, and that the common parts can be built once and shared: SSA, a CFG of regions and blocks, a pass infrastructure, a pattern rewriter, location tracking, and verification. What differs between domains is then expressed as a dialect.<br>The unit of everything in MLIR is the Operation . An op has operands and results (SSA values), a set of typed attributes (compile-time constants like shapes or strides), and zero or more regions , which themselves contain blocks of further ops. That last property is what makes the IR genuinely multi-level: a single op can carry a whole nested computation, so a high-level linalg.generic and a low-level llvm.add are the same kind of object at different altitudes. Every op belongs to a dialect, which is simply a namespace for a related family of ops, types, and attributes.<br>Stated as a grammar, the relationship is small and recursive. A dialect supplies vocabulary, the op names, types, and attributes, while the shape of an operation is universal. The form below is simplified from MLIR’s textual grammar to show the essential structure; the authoritative productions are in the Language Reference2:<br>; A dialect is a namespace that contributes a family of ops, types, and attributes.<br>dialect ::= (operation-def | type-def | attribute-def)*
; The grammar of an operation is identical across every dialect.<br>operation ::= [ result ("," result)* "=" ] op-name<br>"(" [ operand ("," operand)* ] ")"<br>attr-dict? region*<br>":" type-signature
op-name ::= dialect-name "." mnemonic ; e.g. linalg.matmul; quoted as a string in the generic form<br>attr-dict ::= "{" attr-entry ("," attr-entry)* "}" ; compile-time constants<br>region ::= "{" block+ "}"<br>block ::= operation+ ; ops hold regions hold ops -> it recurses
result ::= ssa-value ; %C<br>operand ::= ssa-value ; %A, %B<br>type ::= "!" dialect-name "." mnemonic | builtin-type ; e.g. !llvm.ptr; builtins: tensor, memref<br>attr-entry ::= name "=" attribute-value ; e.g. 1 : i64, "foo", #dialect.attr<br>Two consequences fall out of this. First, the op name, the types, and the attributes are all namespaced by a dialect (linalg.matmul, tensor), so “adding a dialect” extends the vocabulary without touching the grammar, which is exactly why the surrounding infrastructure can be dialect-agnostic. Second, because an operation may contain a region, and a region contains blocks of further operations, the structure nests without bound, and that recursion is what lets a single op carry an entire computation rather than one instruction.<br>Dialects: One IR, Many Altitudes<br>The defining feature is that dialects coexist. A module mid-compilation routinely holds ops from several dialects at once, and lowering is the gradual replacement of higher-level ops with lower-level ones until only a target dialect remains. The dialects an ML pipeline passes through, from high to low:<br>High level (what you mean): stablehlo and tosa (whole-tensor operator sets), linalg (structured ops over tensors/buffers), tensor (value-semantic tensor manipulation).3Mid level (how it’s structured): memref (buffers with layout/strides), affine and scf (loop nests and structured control flow), vector (SIMD), arith (scalar math).4Low level (where it runs): llvm (translated to LLVM IR for CPUs), gpu plus nvvm/rocdl (GPU targets), spirv (Vulkan/compute).The skill MLIR encodes is choosing when to drop from one altitude to the next. Stay high too long and you cannot...