Fusion in MLIR

MLIR represents computations within specialized dialects, which are intermediate representations tailored for specific domains. Key dialects for fusion include:

• Linalg Dialect: Provides structured, loop-free operations (like linalg.generic) perfect for representing and transforming dense linear algebra, enabling pattern matching for fusions like Conv-BN-ReLU.
• Affine Dialect: Represents nested loops with affine constraints, allowing for precise analysis and fusion of loop nests (loop fusion).
• Tensor Dialect: Handles tensor types and operations, serving as a bridge between high-level frameworks and lower-level computational dialects. This separation allows fusion logic to be applied at the most appropriate abstraction level before lowering to hardware-specific forms like LLVM or GPU dialects.

MLIR uses a Declarative Rewrite Rule (DRR) system to define fusion transformations. Engineers specify source and target patterns using a high-level domain-specific language, and MLIR's framework automatically applies them. For example:

• A rule can match a linalg.matmul followed by a linalg.elemwise_add.
• It can rewrite this pattern into a single linalg.generic operation that performs the fused computation. This declarative approach makes fusion transformations composable, maintainable, and less error-prone than manual C++ rewrite logic, while enabling aggressive combining of elementwise and reduction operations.

A core utility within the Linalg dialect is the linalg.fuse transformation. It performs tile-and-fuse by:

1. Tiling: Decomposing a producer operation (e.g., a convolution) into smaller tiles.
2. Fusion: Fusing the computation of a consumer operation (e.g., ReLU) directly into the loop that produces each tile of the producer. This creates locality of reference, where the consumer uses the producer's output tile while it is still in fast cache or registers, eliminating a full round-trip to slow memory. The transformation is parametric in tile sizes, allowing tuning for different cache hierarchies.

MLIR does not perform fusion in one step. It uses progressive lowering, where high-level, platform-agnostic fused operations are gradually lowered through multiple intermediate representations. Critical to this process are interfaces (like TilingInterface, FusionInterface):

• Operations that implement the TilingInterface can be split into tiles.
• Operations with FusionInterface can be analyzed for data dependencies and legality of fusion. This allows generic fusion passes to operate on any operation that implements these interfaces, making the infrastructure extensible to new dialects and operations beyond Linalg.

Not all fusion is beneficial. MLIR's fusion passes integrate cost models and profitability heuristics to decide when to fuse. They analyze:

• Memory Traffic: Estimating the reduction in DRAM bandwidth from keeping intermediate results local.
• Arithmetic Intensity: Assessing whether fusion creates a more balanced compute-to-memory ratio.
• Kernel Launch Overhead: Evaluating the savings from reducing the number of discrete GPU kernels.
• Parallelism Constraints: Ensuring fusion does not serialize independent operations or exceed hardware resource limits (e.g., GPU register pressure). Decisions can be static (Ahead-of-Time) or informed by runtime profiling data.

MLIR's Affine and SCF (Structured Control Flow) dialects provide the representation to analyze and fuse loops. This is a foundational optimization for memory-bound kernels.

• Mechanism: Identifies parallel affine.for or scf.for loops with compatible iteration spaces and data dependencies.
• Benefit: Fuses multiple passes over data into one, keeping intermediate values in fast cache or registers.
• Example: Fusing an elementwise tanh operation with a subsequent elementwise add that uses the same tensor, eliminating a full write/read cycle to main memory.

The Linalg dialect represents tensor contractions and generic linear algebra operations in a declarative, loop-free form. Fusion is a key step in its compilation flow.

• Mechanism: After applying tiling (splitting work into blocks), the compiler fuses the producer of a tile with its consumer.
• Benefit: Enables producer-consumer fusion where the producer's output is consumed immediately within the same kernel, crucial for complex operators like linalg.generic (elementwise ops) feeding into linalg.matmul.
• Result: Creates a single, efficient kernel for a sub-computation like a tiled and fused matrix multiplication chain.

A common and highly profitable fusion target in MLIR is combining chains of pointwise operations.

• Representation: Often expressed as linalg.generic ops or sequences of arith and math dialect operations.
• Process: The compiler pattern-matches a sequence of independent per-element ops (e.g., add → relu → cast) and generates a single loop body performing all computations.
• Impact: Dramatically reduces kernel launch overhead and intermediate memory allocation for activation functions and data type conversions in neural networks.

This is the classic fusion case where an operation's output is an immediate input to another. MLIR's use-def chains and SSA form make these dependencies explicit for the compiler.

• Scope: Applied within a basic block or across logically connected operations in the IR.
• Example: Fusing a linalg.conv_2d operation with its immediately following linalg.elemwise_binary (add bias) operation. The bias addition is pulled into the convolution's output computation.
• MLIR Advantage: Dialects like Linalg maintain structured semantics, allowing the compiler to reason about and safely fuse these operations without breaking correctness.

Horizontal fusion combines independent operations that could be executed in parallel, often to improve resource utilization and amortize memory access costs.

• Mechanism: Identifies operations that share similar input structures or iteration patterns but compute different outputs.
• Use Case: Fusing two independent linalg.generic operations that both read from the same input tensor but apply different transformations (e.g., computing both mean and variance in a single pass).
• Benefit: Increases arithmetic intensity and improves cache line utilization by reading the input data once for multiple computations.