IRs, dialects, and lowering
Trace one tensor operation through high-level, structured, GPU, and target representations
Before you begin
Follow a linked prerequisite when you cannot explain it from memory. Background skills without links are stated as abilities rather than hidden terminology.
- Tensor shapes and datatypes
- Reading assignments and function-call results
Compile y = GELU(x @ w + bias) and ask the compiler to dump its program after several passes. One dump can contain a high-level matrix operation and broadcast addition. A later dump can contain tiled loops, explicit memory spaces, and a GPU launch. The final generated file can contain PTX or another target instruction form. Comparing the dumps shows where a mathematical operation became a hardware schedule.
Each compiler dump is an intermediate representation, abbreviated IR: a compiler-owned program form designed for analysis and rewriting. An MLIR dialect is a defined family of operations, types, and attributes inside the Multi-Level Intermediate Representation infrastructure. For example, a tensor-oriented dialect can preserve matrix semantics, while the GPU dialect can express thread blocks, address spaces, and launches. The lesson follows one value across those forms before discussing lowering in general.
A compiler needs several ways to describe the same program because no single representation makes every optimization visible. At a high tensor level, matrix multiplication is one operation with algebraic meaning. At a loop level, tiles, iteration order, and memory accesses become explicit. Near the target, address spaces, warps, vector instructions, and native operations matter. Lowering is the controlled loss of abstraction while preserving semantics.
MLIR calls related operation and type systems dialects, allowing several abstraction levels to coexist and transform into one another. OpenXLA and other compiler stacks use comparable staged ideas. The vocabulary matters less than the principle: knowledge should remain explicit until the pass that can use it, then be converted into a form the next decisions require.
Follow one matrix multiplication through the levels. At the tensor level it is a dot operation with shape and datatype semantics. At a tiled level it becomes loops over output and K tiles. At a GPU level those tiles map to blocks, warps, memory spaces, and synchronization. Near the target they become address calculations and matrix instructions. Each stage answers questions that were unnecessary at the previous stage.
What an intermediate representation contains
Begin with a three-line high-level dump for y = GELU(x @ w + bias): `%0 = matmul %x, %w : tensor<128x4096xbf16>`, `%1 = add %0, %bias : tensor<128x4096xbf16>`, and `%2 = gelu %1 : tensor<128x4096xbf16>`. The words `matmul`, `add`, and `gelu` name operations: actions represented in the compiler program. `%0`, `%1`, and `%2` name the results produced by those operations. The type after the colon says that each result is a 128 by 4096 tensor containing BF16 values.
Read the first data dependency directly from the printed names. The matrix operation defines `%0`; the addition reads `%0` and defines `%1`; GELU reads `%1` and defines `%2`. This convention is called static single assignment, or SSA, because each result name is defined by exactly one operation. A result can have several readers, but a later operation does not overwrite the result under the same name. The explicit producer and readers let the compiler ask whether an intermediate has one consumer, whether two operations can be fused, and when the intermediate stops being needed.
An operation can also carry attributes: fixed facts written on the operation rather than supplied as ordinary tensor inputs. A transpose operation can carry a permutation such as `[1, 0]`; a matrix operation can carry precision or contracting-dimension information. Later representations can attach a memory space to say whether a buffer resides in device HBM or on-chip shared memory. Operations containing a loop or conditional can own a structured region containing other operations. Attributes, memory spaces, and structured regions are introduced when the running program needs them rather than as an unrelated vocabulary list.
The printed form is designed for compiler work, not only for display. A verifier checks that `%0` exists before the addition uses it and that tensor shapes and datatypes satisfy each operation's rules. An analysis can compute which operations read `%0`, whether two buffers can refer to the same storage, or when a result is no longer live. A rewrite pass uses those facts to replace one valid fragment with another, such as incorporating `%1` and `%2` into a supported matrix epilogue.
Why several abstraction levels are necessary
At a high level, a matrix multiplication is one operation. The compiler knows the complete algebraic structure and can combine the operation with a bias or choose a library call. The representation does not yet state which thread loads A[17,9]. The high-level omission is useful because early passes can reason about the complete operation without reconstructing the operation from scalar loops.
A structured tensor or loop level introduces iteration spaces and tiles while preserving relationships among dimensions. A GPU level introduces blocks, threads, shared memory, and barriers. A target level introduces concrete instructions and address spaces. No single level makes operator algebra, loop scheduling, and instruction encoding equally easy to analyze.
Dialects and mixed representations
In MLIR, a dialect defines related operations, types, attributes, and verification rules. A module can contain operations from several dialects when their interfaces are compatible. Compatible interfaces permit progressive conversion instead of requiring the complete program to move from one monolithic IR to another in one pass.
High-level tensor operations can coexist with lower-level arithmetic or control while a pass converts only the operations it understands. Interfaces let generic passes ask questions across dialects. The design supports compiler stacks for different frameworks and targets without requiring every frontend to implement every low-level optimization independently.
Lowering and legalization
Lowering replaces an operation with a more explicit implementation. One dot operation could lower to a vendor-library call, a generated tiled kernel, or scalar loops. The selected path depends on target capabilities, shape, layout, and cost. Lowering therefore includes an implementation decision; it is not a line-by-line change of syntax.
A conversion target defines which operations are legal after a stage. Rewrite patterns eliminate illegal operations while preserving semantics. Type conversion and temporary bridge operations may be necessary while converted and unconverted regions coexist. Verification after each stage catches structural failures close to the pass that introduced them.
Bufferization is a specific lowering decision that converts tensor semantics into explicit buffer, or memref, semantics. Tensor SSA values describe results without requiring each result to own a new allocation. Bufferization analyzes uses, aliases, and lifetimes to decide whether an existing buffer can hold a result or whether a copy or allocation is necessary. Poor decisions at this boundary can create extra memory traffic even when later kernel scheduling is efficient.
Information can be lost too early
Decomposing a matrix multiplication into generic scalar loops can erase the fact that the loop nest is a matrix operation. A later pass may be unable to recover that structure reliably, so it cannot select a tensor-core kernel or fuse a supported epilogue. Lowering should retain useful semantics until the passes that consume them have run.
Inspect IR at meaningful boundaries. If an expected fusion is absent, find the first stage where producer and consumer are separated. If a transpose remains, find whether layout assignment lacked a constraint. If a dynamic dimension blocks tiling, identify the missing bound or divisibility fact. The first changed stage is more actionable than the final assembly alone.
Worked example: Following a fused linear activation through IR
Source computes y = gelu(xW + b). The expectation is one matrix kernel with a fused epilogue.
- At graph IR, verify that matmul, bias broadcast, and GELU appear in one captured region with compatible users and no side effects.
- Record the tensor shapes, layouts, and users at this level. If the matmul output has another consumer, an epilogue fusion can duplicate work or require materialization.
- After fusion or library selection, inspect whether bias and activation are represented as an epilogue or remain separate tensor operations.
- At GPU or target IR, confirm the matrix path and where bias and activation execute. An unsupported GELU form may force materialization despite high-level fusion.
- Compare generated launches and HBM traffic. Confirm that the IR transformations explain the timeline and profiler measurements.
- If fusion is lost, identify the first IR dump where the operations separate and inspect the pass decision or legality condition at that stage.
Conclusion: IR inspection locates the exact stage where the intended fusion was preserved or lost, turning ‘the compiler failed’ into a narrow transformation question.
Common errors
- Reading only the final assembly. Final assembly shows what exists but not which high-level opportunity disappeared earlier.
- Lowering aggressively by hand. Once semantics such as matrix structure or broadcasting are erased, later passes have less knowledge.
- Assuming IR syntax is the goal. The useful skill is tracing invariants and decisions across representations.
- Treating every printed conversion as runtime data movement. Some operations only change the compiler's view of a value; verify generated memory operations before counting bytes.
Reference summary
An intermediate representation makes program semantics explicit enough for analysis and transformation. MLIR allows multiple dialects—families of operations, types, and attributes—to coexist. High-level tensor operations preserve algebra and shape; structured loop/tile forms expose schedules; GPU dialects express launches and address spaces; LLVM/NVVM-level forms approach target instructions. Lowering progressively replaces high-level operations while preserving defined meaning.
%0 = stablehlo.dot_general %a, %b ...
↓ algebraic simplification / fusion / layout
%1 = linalg.matmul ins(%a, %b) outs(%c)
↓ tiling / mapping
gpu.launch blocks(...) threads(...) { ... }
↓ NVVM / LLVM / PTX
- Analysis asks facts about the program: shapes, aliases, side effects, liveness, or dependencies.
- Transformation rewrites while preserving the declared semantics.
- Legalization defines which operations may remain at the next level.
- Bufferization decides how tensor values map to storage and lifetimes.
- Target lowering selects hardware-specific instructions, layouts, and runtimes.
Knowledge check
Why can early lowering of an attention operation to PTX prevent later optimizations?
Show the answer guide
- A high-level attention operation preserves semantic facts: the relationship among Q, K, and V, the softmax reduction axes, masks, precision rules, and opportunities to tile or fuse the complete operation.
- Lowering immediately to PTX replaces that operation with target-specific instructions, addresses, and synchronization. Later passes can no longer safely infer the full attention semantics from the low-level instruction stream.
- The lost knowledge can prevent online-softmax transformation, fusion across attention boundaries, layout propagation, alternative tiling, memory planning that removes the score matrix, or retargeting to a different accelerator.
- Lowering should occur only after transformations that require the high-level meaning have run; target details should be introduced progressively through intermediate representations.
Compare early and staged lowering of attention
Use an MLIR-, XLA-, or torch.compile-based example for scaled dot-product attention with Q, K, and V shaped [2,8,1024,64]. Capture a high-level graph or IR, an optimized tensor or loop IR, and the final GPU or PTX form. Create or identify a path that treats an equivalent low-level attention call as opaque before high-level fusion.
Evidence to keep: Submit the source program, all three IR dumps, an annotated operation lineage showing where softmax, masking, and score materialization remain visible, plus launch-count and memory-traffic measurements for the staged and opaque paths. State the exact optimization that the opaque early-lowered form prevents.
With staged representations visible, we can study the main optimization decisions: fusion, layout propagation, schedule selection, and memory planning.
Sources and further reading
The lesson explanations synthesize primary papers and official references. Follow the originals for proofs, exact APIs, architecture details, and version-sensitive behavior.