Fusion Heuristics

The primary constraint for fusion is the dataflow graph. Heuristics analyze producer-consumer relationships to identify valid fusion groups.

• Vertical Fusion: Merges sequentially dependent operators (e.g., a Convolution followed by BatchNorm).
• Horizontal Fusion: Merges independent operators that share a common input or operate in parallel.
• Fusion Group Discovery: Algorithms traverse the graph to find connected subgraphs where intermediate tensors are only consumed within the group, making them candidates for elimination.

A core goal of fusion is to reduce memory bandwidth pressure. Heuristics evaluate the memory-bound vs. compute-bound nature of operations.

• Fusion for Cache: Prefers fusing operations where intermediate results can stay in fast cache (L1, shared memory), avoiding costly trips to global GPU memory.
• Kernel Launch Overhead: Combining ops amortizes the fixed cost of launching a GPU kernel.
• Elementwise Ops: Operations like ReLU or Add are ideal fusion targets as they are memory-intensive; fusing them with a compute-heavy op (like a MatMul) hides their memory latency.

Heuristics classify operators by their computational profile to balance the fused kernel's workload.

• Compute-Bound Operations: Dense matrix multiplications (MatMul, Conv) have high arithmetic intensity. Fusing lighter ops with them improves overall utilization.
• Memory-Bound Operations: Pointwise activations (Sigmoid, Tanh) and reductions are limited by memory speed. They benefit greatly from fusion.
• Canonical Patterns: Compilers use pattern matching to recognize and fuse known profitable sequences like Conv-BN-ReLU or Linear-GELU automatically.

Profitability is hardware-dependent. Heuristics use a cost model to estimate execution time on the target accelerator (e.g., NVIDIA GPU, Google TPU).

• Register Pressure: Fusing too many ops can exceed the GPU's register file, causing register spilling to slower memory and degrading performance.
• Thread Block Scheduling: The fused kernel must efficiently map to the GPU's execution units (SMs).
• Compiler Backends: Heuristics differ between XLA, TVM, and torch.compile's Inductor, as each has unique optimization passes and hardware targets.

The final decision is a trade-off analysis. A fusion planner evaluates potential groups against a cost model.

• Benefit: Reduced global memory accesses, lower launch overhead, improved cache locality.
• Cost: Potential for increased register usage, reduced parallelism, and more complex kernel code.
• Search Space: For complex graphs, the planner may explore multiple fusion plans to find a near-optimal solution, as exhaustive search is often intractable.

The timing of the fusion decision introduces different heuristic considerations.

• Ahead-of-Time (AOT) Fusion: Performed during compilation (e.g., with XLA). Heuristics can be more aggressive, using known tensor shapes and a comprehensive graph view.
• Just-In-Time (JIT) Fusion: Performed at runtime (e.g., PyTorch's torch.compile). Heuristics must be faster and can adapt to dynamic shapes, but have less context for global optimization.
• CUDA Graphs: Represent a launch-time fusion heuristic, capturing a sequence of kernels into a single, replayable unit to eliminate CPU launch overhead.

Google's Accelerated Linear Algebra (XLA) compiler employs a greedy, producer-consumer fusion strategy. It traverses the HLO (High-Level Optimizer) graph in reverse post-order, prioritizing the fusion of operations that:

• Are elementwise, broadcast, or reduction operations.
• Have a producer-consumer data dependency.
• Will not create a cycle in the fused computation graph. Its primary heuristic is to fuse to create larger, more complex operations that can be efficiently lowered to a single LLVM IR kernel, aggressively minimizing intermediate tensor materialization.

Apache TVM uses a cost-model-driven approach via its Ansor auto-scheduler. Instead of hardcoded patterns, it:

• Generates fusion sketches as high-level loop structures for potential fused subgraphs.
• Samples promising programs by applying transformations like loop tiling and vectorization.
• Measures actual hardware performance on the target device to build a dataset.
• Trains a ML-based cost model to predict the runtime of unseen programs, guiding the search for the optimal fused kernel schedule. This makes it highly adaptive to new hardware.

The Multi-Level Intermediate Representation (MLIR) enables fusion through its Linalg dialect, which represents operations declaratively using indexing maps. Key heuristics include:

• Tile and fuse loop transformations that reorder tiling and fusion steps to maximize data locality.
• Use of the Affine dialect to analyze loop nests for fusion legality (no negative distance dependencies).
• Pattern rewriting rules (e.g., linalg.fuse) that merge operations based on their iterator types (parallel, reduction). This provides a structured, multi-level framework for fusion that is easier to reason about and customize than graph-level approaches.

PyTorch's torch.compile backend, Inductor, performs fusion in two primary phases:

1. Graph-Level Fusion: Uses a pattern matcher on the FX graph to identify and replace subgraphs (e.g., sigmoid + mul -> silu) with single, predefined Composite Automatic Registration (CARD) kernels.
2. Loop-Level Fusion: In the subsequent triton or C++ codegen phase, it applies horizontal fusion heuristics to merge independent pointwise operations that share the same loop iteration space into a single kernel, reducing global memory accesses. Its heuristics are tuned for dynamic Python graphs.

NVIDIA's TensorRT applies fusion during the network optimization phase. Its heuristics are highly specialized for NVIDIA GPUs and involve:

• Vertical fusion of adjacent layers (e.g., Convolution + Bias + ReLU) into a single CUDNN or CUBLAS call.
• Horizontal fusion of operations that can be executed in a single kernel to improve occupancy.
• Automatic precision layer fusion, where it may fuse operations before and after a precision conversion (e.g., FP16 -> FP32).
• Heavy use of hand-written, optimized kernel libraries (cuDNN, cuBLAS) for fused patterns, making its heuristics a blend of pattern matching and library dispatch.

Across compilers, fusion decisions are guided by a common set of profitability factors evaluated by a cost model:

• Data Locality: Will fusion keep intermediate tensors in registers or shared memory, avoiding costly DRAM trips?
• Kernel Launch Overhead: How many synchronization points and kernel launches are eliminated?
• Arithmetic Intensity: Does fusion increase the compute-to-memory-access ratio?
• Parallelism Constraints: Does fusion create a sequential bottleneck or reduce occupancy by increasing register pressure?
• Operator Type: Is the pattern elementwise (highly fusible), a reduction, or a complex operation (e.g., convolution)? The cost model estimates speedup, often using roofline model analysis, to accept or reject a fusion candidate.

Graph fusion is the overarching process of automatically identifying and merging subgraphs of operators within a computational graph. It is the practical application of fusion heuristics, which guide this process through pattern matching and cost models to create efficient fused kernels. The goal is to transform a naive graph of discrete operations into an optimized graph with compound nodes.

A fusion group (or fusion region) is a set of operators within a computational graph that have been identified by the heuristics as candidates to be combined and executed by a single, fused kernel. Defining these groups involves analyzing:

• Data dependencies (producer-consumer chains)
• Operation types (e.g., elementwise, reduction)
• Memory access patterns The fusion planner uses the cost model to evaluate the profitability of turning a candidate group into an actual fused kernel.

The cost model for fusion is the predictive engine within the heuristics. It estimates the performance impact—latency, throughput, memory usage—of fusing a specific group of operators. It evaluates factors like:

• Reduced kernel launch overhead
• Improved data locality and cache hit rates
• Increased register pressure or decreased parallelism (potential downsides)
• Memory bandwidth savings from eliminated intermediate stores This model is often based on hardware characteristics (e.g., GPU SM count, cache sizes) and operator attributes.

A fusion compiler is the specialized software that implements fusion heuristics. It is typically a compiler pass within a larger stack. Key examples include:

• XLA (Accelerated Linear Algebra): Google's compiler for TensorFlow/JAX, known for aggressive fusion.
• TVM's Auto-Scheduler: Searches for optimal operator fusion schedules.
• MLIR Dialects (e.g., Linalg): Provide intermediate representations and transformation utilities for fusion. These compilers consume a high-level computational graph and output a lowered, fused execution plan.

Fusion profitability is the core question heuristics must answer: "Will fusing these operators make execution faster?" It is a trade-off analysis. Benefits include amortized kernel launch overhead and reduced global memory traffic. Costs can include increased kernel complexity, higher register usage (leading to reduced occupancy), and lost opportunity for parallel execution. Heuristics must correctly predict this balance to avoid performance regressions.

Pattern matching for fusion is a key heuristic strategy where the compiler identifies known, profitable subgraph patterns. This is a rule-based approach that bypasses exhaustive search for common cases. Canonical patterns include:

• Conv-BN-ReLU: The standard fused block in CNNs.
• Linear-GELU: Common in transformer feed-forward networks.
• Elementwise chains (e.g., Add -> Mul -> Sin). Once a known pattern is matched, the compiler can apply a pre-verified, optimal fusion strategy.