Pattern Matching for Fusion

The compiler scans the computational graph's dataflow structure to identify connected subgraphs that match predefined templates. This involves:

• Graph isomorphism checking to match operator types and connectivity.
• Canonicalization of the graph to ensure patterns are recognized regardless of minor syntactic differences in the original model definition.
• Templated patterns like Conv → BatchNorm → ReLU or Linear → GELU are stored in a pattern library for efficient matching.

Not all matched patterns are profitable to fuse. A cost model evaluates the potential performance gain by estimating:

• Reduction in kernel launch overhead from fewer GPU dispatches.
• Memory bandwidth savings from eliminating intermediate tensor writes to global memory.
• Potential downsides like increased register pressure or reduced occupancy in the fused kernel. The fusion planner uses this analysis to accept or reject fusion candidates, ensuring net performance improvement.

Pattern matching distinguishes between two primary fusion strategies based on data dependencies:

• Vertical Fusion (Producer-Consumer): Merges sequentially dependent operators (e.g., MatMul → Add). This is the most common target, reducing intermediate memory traffic.
• Horizontal Fusion (Sibling Operations): Merges independent operators that consume the same input tensor (e.g., two parallel SiLU activations). This is less common but reduces kernel launch overhead. Pattern matchers are tuned to identify profitable opportunities for both types.

Efficient pattern matching relies on libraries of hand-crafted or learned fusion templates optimized for specific domains:

• Vision Networks: Canonical patterns include Conv-BN-ReLU, DepthwiseConv-PointwiseConv (MobileNet blocks).
• Transformer Models: Patterns like fused Multi-Head Attention (FlashAttention) or LayerNorm-Linear.
• Elementwise Chains: Sequences of pointwise ops like Add → Tanh → Multiply. These libraries are often embedded within compilers like XLA, TVM, or PyTorch's Inductor.

Pattern matching operates on a compiler's Intermediate Representation (IR), which provides a normalized view of the computation. Key integrations include:

• MLIR Dialects: Using the Linalg or TOSA dialects to represent and match patterns in a hardware-agnostic way.
• Graph Rewriting: Once a pattern is matched, the compiler's IR is rewritten, replacing the matched subgraph with a single, abstract fused operator node.
• Lowering Path: This fused node is later lowered to a target-specific fused kernel during code generation.

The timing of pattern matching is a critical system design choice:

• Just-In-Time (JIT) Fusion: Patterns are matched and kernels are generated at runtime (e.g., via torch.compile). This allows optimization based on dynamic input shapes but adds compilation overhead.
• Ahead-of-Time (AOT) Fusion: Patterns are matched during an offline compilation step (e.g., with XLA AOT). This eliminates runtime overhead, producing a static, pre-fused executable ideal for deployment. The pattern matching algorithm is fundamentally similar, but the integration into the compilation pipeline differs.

This is the quintessential fusion pattern in convolutional neural networks (CNNs). It combines a Convolution layer, a Batch Normalization layer, and a ReLU activation into a single kernel.

• Mechanism: The compiler fuses the linear convolution, the affine transformation of batch norm (scale and shift), and the non-linear ReLU clamp into one operation.
• Benefit: Eliminates two intermediate tensor writes/reads to global memory, drastically reducing memory bandwidth pressure.
• Example: Found in nearly all modern CNN architectures like ResNet, EfficientNet, and MobileNet.

This pattern fuses the entire multi-head attention mechanism—query/key/value projection, attention scoring, softmax, and output projection—into a single, highly optimized kernel.

• Core Innovation: Pioneered by FlashAttention, which uses an IO-aware algorithm to keep intermediate attention matrices in fast SRAM, avoiding costly writes to high-bandwidth memory (HBM).
• Impact: Enables training and inference of transformers with extremely long sequences (e.g., 128K+ tokens) by reducing memory complexity from quadratic to linear in sequence length.
• Implementation: Critical in production systems using models like GPT-4, Llama, and Claude.

This pattern targets sequences of pointwise operations where each operator applies a function independently to every element of a tensor.

• Common Operations: Includes Add, Mul, ReLU, Sigmoid, Tanh, Gelu, and Swish.
• Fusion Logic: The compiler identifies a chain of these operations and generates a single kernel that applies the composed mathematical function per element.
• Performance Gain: Amortizes kernel launch overhead and maximizes arithmetic intensity by keeping data in registers or cache across the entire sequence.

A fundamental pattern in fully-connected networks and transformer blocks that fuses a matrix multiplication with a bias addition and an activation function.

• Pattern: Y = Activation( X @ W + b )
• Compiler Action: The matrix multiply, vector bias add, and activation are compiled into one kernel, avoiding storing the large intermediate result X @ W.
• Context: This is a primary target for compilers like XLA and Torch Inductor when optimizing feed-forward networks and MLP layers.

This pattern fuses a reduction operation (like sum or mean) with a preceding elementwise operation, or a subsequent operation that uses the reduction's scalar result.

• Example 1: Mean(Square(X - Y)) (fused MSE loss).
• Example 2: Softmax (which involves a reduction for the normalization denominator).
• Benefit: The intermediate tensors from the elementwise op are reduced in-place within fast shared memory, never written to global memory. This is critical for loss functions and normalization layers.

Compilers use a combination of techniques to identify these canonical patterns within a computational graph.

• Graph Rewriting: The compiler traverses the graph, looking for subgraphs that match a predefined template (e.g., Conv -> BN -> ReLU).
• Cost Model Integration: After identifying a candidate pattern, a cost model estimates the performance benefit of fusion versus execution as separate kernels.
• Compiler Backends:
  • XLA: Uses aggressive pattern matching for HLO (High-Level Optimizer) graphs.
  • TVM: Employs its AutoTVM and Ansor schedulers to search for optimal fusion patterns.
  • MLIR: Uses pattern rewriting systems within dialects like Linalg to declaratively define and match fusion patterns.

The graph-level optimization that merges adjacent computational operators in a neural network's computational graph into a single, compound operation. This is the primary goal that pattern matching serves.

• Key Goal: Minimize intermediate memory transfers between global memory and registers.
• Scope: Works on the abstract computational graph before low-level kernel generation.
• Example: Identifying a Linear -> GELU sequence and replacing it with a single FusedLinearGELU operator node.

A lower-level compiler optimization that combines the implementation of multiple primitive operations into a single, unified GPU or accelerator kernel.

• Implementation Level: Occurs after operator fusion, during code generation for a specific hardware target (e.g., CUDA, Metal).
• Primary Benefit: Eliminates kernel launch overhead and enables fine-grained data locality within shared memory or registers.
• Relationship: Pattern matching at the operator level often dictates which operations will be candidates for kernel fusion.

A set of operators within a computational graph that have been identified by a compiler as candidates to be combined and executed by a single, fused kernel.

• Formation: Created by the pattern matching and fusion planning passes of a compiler.
• Constraints: Operators must have compatible data dependencies and memory access patterns.
• Compiler Role: Compilers like XLA, TVM, and PyTorch's Inductor dynamically create fusion groups based on cost models.

Rule-based or cost-model-driven algorithms used by compilers to decide which sets of operators are profitable to fuse.

• Inputs: Analyze data dependency, operator types, tensor shapes, and estimated memory traffic.
• Output: A binary decision on whether to fuse a matched pattern.
• Examples:
  • Always Fuse: Elementwise operations (e.g., consecutive ReLU, Sigmoid).
  • Cost-Based Fuse: A MatMul followed by a light elementwise op may be fused if the reduction in memory bandwidth outweighs increased register pressure.

The overarching process of automatically identifying and merging subgraphs of operators within a computational graph to create efficient fused kernels. Pattern matching is a core sub-task of this process.

• Automation: Modern compilers perform this automatically across the entire graph.
• Phases: 1) Pattern Matching, 2) Profitability Analysis (Heuristics/Cost Model), 3) Group Formation, 4) Kernel Code Generation.
• Frameworks: Central to the optimization pipelines of MLIR (via Linalg/Affine dialects), Apache TVM, and torch.compile.

The final output of the fusion process: a single, hand-written or compiler-generated GPU/accelerator kernel that implements the combined functionality of multiple primitive operations.

• Characteristics: Manually tuned kernels (e.g., FlashAttention, Fused Conv-BN-ReLU) often set the performance gold standard that compilers aim to match via automatic fusion.
• Benefit: Executes the entire fused operation without writing intermediate results to slow global memory.
• Trade-off: Can increase register usage and limit parallelism, making profitability analysis essential.