MLIR Architecture: The Unified Intermediate Representation of AI Compilers

Summary: Triton compiles "kernel code"; XLA compiles "whole-graph execution plans"—and their shared foundation is MLIR, LLVM's compiler infrastructure for "multi-level intermediate representations". This post clarifies MLIR's core abstractions (Dialect/Operation/SSA), the lowering pipeline, and why AI compilers all choose to "stand on MLIR".

0. Experiment Goals

• Understand MLIR's three core abstractions: Dialect, Operation, SSA values
• Read the IR lowering pipeline: high-level dialects → mid-level dialects → LLVM IR
• Clarify MLIR's relationship with Triton and XLA (who stands on whom)
• Read mlir-opt IR output (a basic skill when debugging compilers)

1. Why MLIR: One IR Is Not Enough

How it works:

• LLVM IR is "close to the machine": good for register allocation and instruction selection, but it cannot express "this is a matmul" or "this is a convolution"—AI ops get shredded into countless scalar instructions in LLVM IR, losing all high-level optimization information.
• MLIR's core idea: not "one IR" but an IR framework—you define your own dialects, each with its own op set, and lower between dialects level by level.
• The benefit: optimizations happen at the most suitable level—graph fusion at high dialect layers (tosa/linalg), vectorization mid-level, codegen at the bottom (LLVM).

AI model (high-level semantics)
   │ dialects: tosa / linalg (tensor semantics; op fusion/shape inference)
   ▼
mid-level dialects (vectorization, tiling)
   │ dialects: arith / vector / scf
   ▼
LLVM IR (close to machine)
   │
   ▼
machine code

2. Core Abstractions: Dialect / Operation / SSA

// A sample MLIR IR (tosa + arith dialects mixed)
// %0, %1 are SSA values (immutable, explicit dataflow)
func.func @add_relu(%arg0: tensor<4x4xf32>, %arg1: tensor<4x4xf32>) -> tensor<4x4xf32> {
  %0 = tosa.add %arg0, %arg1 : tensor<4x4xf32>     // tensor add (high-level semantics)
  %1 = tosa.relu %0 : tensor<4x4xf32>              // ReLU
  return %1 : tensor<4x4xf32>
}

Step-by-step:

• Dialect: a set of ops—tosa (tensor ops), arith (scalar arithmetic), scf (structured control flow), vector (vectorization), linalg (linear algebra)—each dialect is a "semantic layer".
• Operation: the basic unit of IR, shaped dialect.op operands : types. Ops can nest (regions/blocks).
• SSA values: each value assigned once and immutable—dataflow is explicitly visible, making analysis (liveness, dependencies) and transformation (rewriting) simple.

3. The Lowering Pipeline: How High-Level Becomes Machine Code

# Run a lowering pipeline with mlir-opt (example: tosa → linalg → LLVM)
mlir-opt --tosa-to-linalg --linalg-bufferize --convert-linalg-to-llvm input.mlir

# Observe the IR change at each step (essential for compiler debugging)
mlir-opt --tosa-to-linalg input.mlir   # watch tosa dialect become linalg

How it works:

• Pass: MLIR optimization/lowering happens in passes; one pass rewrites a dialect's ops into another dialect (or optimizes within the same dialect).
• Pipeline: passes run in sequence: tosa-to-linalg (semantic lowering) → linalg-bufferize (tensors → buffers) → vectorization → LLVM.
• Extensibility: new hardware integration = new dialect + new passes, reusing the entire LLVM backend—the fundamental reason MLIR became the infrastructure of AI compilers.

4. The MLIR Ecosystem: Who Stands on It

Compiler/framework	Relationship with MLIR	Notes
XLA (TPU/GPU)	graph IR (HLO) expressed within the MLIR framework	Google family
Triton	frontend dialect lowered into MLIR (tt dialect)	the bottom layer of KB Triton
torch-mlir	lowers torch FX graphs into MLIR	another branch of the PyTorch compile stack
IREE	MLIR-based general compile/runtime for edge	mobile deployment

Engineering notes:

• Learning MLIR ≠ learning one compiler: it is "compiler infrastructure"; once you understand the Dialect/pass model, XLA/Triton/torch-mlir architectures become instantly readable.
• When tuning torch.compile or Triton and hitting confusing errors, open the compiler's MLIR dump (TORCH_LOGS or Triton's IR output)—locating the problem at the IR level is faster than guessing from source.
• Domestic hardware (Ascend/Cambricon) compile stacks generally extend MLIR dialects—understanding MLIR means understanding domestic compiler architecture (echoing Domestic Accelerator Porting).

5. Engineering Takeaways

1. MLIR is an IR framework, not a single IR: dialect layering; optimize at the most suitable level.
2. Three abstractions: Dialect (semantic layer) / Operation (basic unit) / SSA (explicit dataflow).
3. Lowering = pass pipelines: high-level dialects lower to LLVM IR level by level, observable at each step.
4. Ecosystem map: XLA/Triton/torch-mlir/IREE all stand on MLIR.
5. Debugging skill: look at the IR dump on compile errors; observe step by step with mlir-opt.
6. Domestic compile stacks are MLIR-family too: understanding MLIR = understanding their architecture.
7. No need for mastery, just literacy: framework engineers only need to read IR and run passes to validate.

6. Further Reading

• MLIR official docs (dialects and pass system)
• MLIR paper: Scaling Compiler Optimization with Extensible Infrastructure
• KB: 08. AI Compiler · Triton Compiler · PyTorch Compile Stack · Domestic Accelerator Porting