Pattern Rewriting and Lowering

Static intermediate representations provide the structure for analysis, but the primary value of a compiler lies in its ability to transform code. Within MLIR, transformations are not monolithic passes but rather collections of modular, granular rewrite rules. Whether optimizing a matrix multiplication within the linalg dialect or lowering a tensor operation to explicit loops in the scf (Structured Control Flow) dialect, the mechanism remains consistent. The two primary frameworks for mutation are the Greedy Pattern Rewrite Driver, typically used for optimization, and the Dialect Conversion Framework, used for lowering between abstraction levels.

The Pattern Match and Rewrite Approach

At the core of MLIR transformations is the Pattern class. Unlike traditional compilers that might iterate over a basic block and manually manipulate instructions, MLIR abstracts this into a strictly defined match-and-rewrite sequence. This approach decouples the definition of a transformation from the strategy used to apply it.

A rewrite pattern consists of two distinct phases:

1. Match: The infrastructure inspects an operation to determine if it satisfies specific constraints (e.g., is this a linalg.matmul with constant weights?).
2. Rewrite: If the match succeeds, the rewriter generates new operations and replaces the original root operation.

For C++ engineers, this is implemented by deriving from OpRewritePattern<OpType>. The complexity arises not in the logic itself but in maintaining the integrity of the underlying IR, such as use-def chains and block arguments.

The lifecycle of a single operation within the greedy pattern rewrite driver involves iterative matching and worklist management until convergence is reached.

Declarative Rewrite Rules (DRR)

While C++ offers full control, MLIR encourages defining patterns declaratively using TableGen. This is known as Declarative Rewrite Rules (DRR). DRR allows compiler engineers to express the source DAG and the result DAG succinctly. The TableGen backend automatically generates the C++ boilerplate, ensuring that pattern application is safe and verifying that types match constraints automatically.

For example, folding a transpose into a convolution involves purely structural matching. In C++, this requires checking attributes, operand types, and result users. In DRR, it is expressed as a direct mapping: (transpose (conv $a, $b)) -> (conv_transposed $a, $b).

Canonicalization and Optimization

The most common application of pattern rewriting is canonicalization. This involves iterative optimizations like constant folding, dead code elimination, and algebraic simplification. MLIR applies these patterns using a greedy strategy. The driver applies patterns repeatedly until the IR reaches a fixed point where no further patterns match.

This "optimization-by-convergence" strategy is powerful but requires careful design of the patterns. If Pattern A transforms $X \to Y$ and Pattern B transforms $Y \to X$, the compiler will enter an infinite loop. Developers must ensure that patterns monotonically decrease the complexity or cost of the IR, or distinct them into separate passes.

The Dialect Conversion Framework

While the greedy driver is sufficient for optimization within the same abstraction level, it struggles with lowering, the process of converting high-level dialects (like tensor) to low-level dialects (like memref or llvm). Lowering often involves changing types (Type Conversion) and ensuring that the resulting IR is strictly legal for a specific target.

The Dialect Conversion Framework solves this by introducing the concept of a Conversion Target. The target defines what is "legal."

Defining Legality

Legality is not binary; it is granular. An operation can be:

• Legal: The operation is supported by the target and should be preserved.
• Illegal: The operation must be converted. If it remains after the pass, the compilation fails.
• Dynamic: The operation is legal only if specific conditions are met (e.g., a convolution is legal only if the stride is 1).

The framework traverses the IR. When it encounters an illegal operation, it searches for a registered ConversionPattern that can lower it to a sequence of legal operations.

The interaction between legality checks and type conversion drives the lowering process. The framework automatically injects cast operations when bridging differing type systems.

Type Conversion and Materialization

A critical challenge in lowering is that operands often change types. When lowering linalg on tensors to linalg on buffers, a tensor<4x4xf32> becomes a memref<4x4xf32>.

If an operation is converted but its user has not been converted yet, the IR becomes invalid because the producer yields a memref while the consumer expects a tensor. The Dialect Conversion framework handles this by allowing the user to define source and target materializations. The framework temporarily inserts "cast" operations (unrealized conversion casts) to glue the new and old type systems together. Once the conversion is complete, if all casts cancel out, they are removed. If casts remain, the conversion is deemed partial or incomplete.

Implementing a Lowering Pass

To implement a lowering pass, such as converting a custom MyDialect to the standard LLVM dialect, you must coordinate the components discussed above.

1. The Type Converter: Define how your high-level types map to LLVM types. For instance, mapping a custom ComplexType to a struct of two floats !llvm.struct<(f32, f32)>.
2. The Conversion Target: explicitly mark MyDialect operations as illegal and LLVM dialect operations as legal.
3. Patterns: Implement ConversionPattern classes. In the matchAndRewrite method, you use the TypeConverter to get the expected LLVM types for the operands, create the new LLVM instructions, and replace the root op.

Consider a scenario where we lower a high-level AddOp operating on tensors to a loop nest. The pattern does not merely replace one op with another; it generates a scf.for loop structure. Inside the loop body, it inserts the scalar addition. The rewriter provided by the conversion framework tracks these insertions, ensuring that if the overall conversion fails (perhaps due to a missing pattern for a different op), all these changes are rolled back atomically.

Pattern Granularity and Fusion

A significant architectural decision in MLIR is the preference for many small, composable patterns over large, complex ones. In previous compiler generations, a "Fusion Pass" might be a 5000-line logic block that handles every permutation of Convolution, ReLU, and BatchNorm.

In MLIR, this is decomposed:

1. A pattern to fuse Conv + Bias.
2. A pattern to fuse Conv + ReLU.
3. A pattern to fuse Conv + BatchNorm.

The greedy driver applies these opportunistically. If Conv + Bias + ReLU appears, the first pattern runs, producing ConvBias. Then, the second pattern recognizes ConvBias (which implements the same interface as Conv) and fuses the ReLU. This composability reduces maintenance burden and allows new operators to benefit from existing optimizations simply by implementing the correct interfaces or traits.

By mastering pattern rewriting and the conversion framework, you gain control over the entire compilation pipeline, enabling the translation of abstract mathematical models into highly optimized machine code tailored for specific hardware accelerators.