Graph Fusion

The compiler first traverses the computational graph to identify subgraphs of operators that are profitable to fuse. This is often done via pattern matching against known, high-benefit sequences (e.g., Conv-BatchNorm-ReLU). The process analyzes data dependencies and tensor shapes to find fusible groups, known as fusion groups.

A cost model evaluates whether fusing a candidate group will yield a net performance gain. It estimates:

• Reduction in kernel launch overhead and global memory traffic.
• Increase in arithmetic intensity and data locality.
• Potential downsides like increased register pressure or reduced parallelism. The fusion planner uses this model to construct an optimal fusion plan, prioritizing memory-bound operations where fusion benefits are highest.

Once a fusion group is finalized, the compiler generates a single fused kernel. This involves:

• Lowering the fused operator subgraph to a unified intermediate representation.
• Applying loop fusion and scheduling transformations to optimize for cache hierarchy (e.g., L1, shared memory).
• Emitting target-specific code (e.g., CUDA, Metal) for the new compound operation. Compilers like XLA, TVM, and MLIR specialize in this automated kernel generation.

Fusion strategies are categorized by the dataflow relationship between operators:

• Vertical Fusion: Merges a producer operator with its consumer (e.g., a matrix multiplication followed by a bias add). This reduces intermediate results written to slow memory.
• Horizontal Fusion: Merges independent operators that consume the same input or operate in parallel. This amortizes kernel launch overhead and can improve memory coalescing. Compilers use heuristics to apply the most beneficial strategy per subgraph.

Fusion can occur at different stages in the compilation pipeline:

• Ahead-of-Time (AOT) Fusion: The graph is analyzed, fused, and compiled to a static executable before deployment. This minimizes runtime overhead and is used for fixed graph deployments.
• Just-in-Time (JIT) Fusion: Performed dynamically at runtime (e.g., via torch.compile). This allows fusion decisions to adapt to dynamic input shapes and specific hardware contexts, offering flexibility for dynamic graphs.

Certain fused operator patterns are so ubiquitous they have dedicated, hand-tuned implementations:

• Fused Conv-BN-ReLU: The standard building block of CNNs, combining convolution, batch normalization, and activation.
• Fused Multi-Head Attention: A key optimization in transformers, famously implemented by FlashAttention. It fuses the entire attention mechanism (QKV projection, scoring, softmax, aggregation) into one IO-aware kernel to minimize HBM accesses.
• Elementwise Fusion: Chains pointwise ops (e.g., Add, Sigmoid, Mul) that operate independently on each tensor element.

Operator fusion is the foundational graph-level optimization that merges adjacent computational nodes in a neural network's computational graph into a single, compound operation. Its primary goal is to minimize the overhead of writing and reading intermediate results to and from slow global memory (e.g., GPU HBM).

• Mechanism: The compiler identifies a chain of operations (e.g., matmul -> add -> relu) and replaces them with a single, custom operator.
• Benefit: Eliminates multiple kernel launch latencies and reduces memory bandwidth pressure, which is often the bottleneck for inference.
• Example: Fusing a convolution, batch normalization, and ReLU activation into one kernel is a standard optimization in frameworks like TensorFlow (via XLA) and PyTorch (via torch.compile).

Kernel fusion is the low-level implementation of operator fusion, where multiple GPU or accelerator kernels are combined into a single, unified kernel. This occurs after the graph-level decision to fuse operators has been made.

• Core Challenge: Writing an efficient fused kernel requires expert-level knowledge of hardware memory hierarchies (registers, shared memory, L1/L2 cache) and parallelism models (warps, thread blocks).
• Compiler Role: Modern ML compilers like TVM and MLIR can automatically generate high-performance fused kernels, moving this task from manual coding to automated scheduling.
• Impact: Reduces kernel launch overhead, a significant latency source when executing thousands of small, sequential operations.

A fusion compiler is a specialized compiler or compiler pass responsible for automatically performing graph and kernel fusion. These are critical infrastructure components in modern ML frameworks.

• XLA (Accelerated Linear Algebra): Google's compiler for TensorFlow and JAX. It performs ahead-of-time (AOT) or just-in-time (JIT) fusion by aggressively identifying fusible subgraphs and generating efficient code for TPUs, GPUs, and CPUs.
• Apache TVM: An open-source compiler stack that uses its scheduling language (AutoTVM, Ansor) to explore the search space of possible kernel fusions and optimizations for diverse hardware backends.
• MLIR (Multi-Level Intermediate Representation): Provides reusable infrastructure for building compilers. Its Linalg and Affine dialects offer powerful, pattern-based transformations for representing and performing fusion on computational graphs.

Fusion heuristics and cost models are the decision-making engines that determine which operator groups are profitable to fuse. A compiler cannot fuse all possible operators, as excessive fusion can lead to register spilling or reduced parallelism.

• Heuristics: Rule-based logic (e.g., "always fuse elementwise operations with their data-dependent producer").
• Cost Model: A predictive model that estimates the execution time or memory traffic of fused vs. unfused versions. It considers:
  • Data dependency patterns.
  • Operator arithmetic intensity (compute-bound vs. memory-bound).
  • Hardware-specific parameters (memory bandwidth, cache sizes).
• Fusion Planner: The component that uses the cost model to construct an optimal fusion plan, often by exploring a constrained search space.

These are two fundamental patterns for grouping operators within a computational graph.

• Vertical Fusion (Producer-Consumer): Merges operators that are sequentially dependent. This is the most common pattern, chaining operations to keep intermediate data in fast registers or cache.
  • Example: Fusing a layernorm operation with a subsequent linear projection.
• Horizontal Fusion (Parallel/Sibling): Merges independent operators that consume the same input tensor or operate in parallel. This increases workload per kernel and amortizes launch overhead.
  • Example: Fusing the separate q, k, v linear projections at the start of a multi-head attention block into a single, batched matrix multiplication.
• Hybrid Approaches: Advanced compilers perform both types, often within the same subgraph, to maximize data locality and hardware utilization.

These are hand-optimized, fused implementations of common neural network patterns. They serve as both performance targets for compilers and as manually crafted kernels for maximum speed.

• Fused Conv-BN-ReLU: The quintessential CNN fusion, combining Convolution, Batch Normalization, and ReLU Activation. It folds the BN parameters into the conv weights at runtime, making the normalization and activation effectively free.
• Fused Multi-Head Attention / FlashAttention: A revolutionary fused kernel that computes the entire attention mechanism in a single pass while minimizing reads/writes to high-bandwidth memory (HBM). FlashAttention uses IO-aware algorithms to recompute attention scores on-chip (SRAM), enabling training and inference with extremely long sequences.
• Fused LayerNorm + GeGLU/SiLU: Common in transformer FFN blocks, fusing the normalization and gated activation functions reduces memory traffic for a key, memory-bound operation.