Fusion Planner

The fusion planner operates on an intermediate representation (IR) of the neural network, typically a dataflow graph where nodes are operators (e.g., Conv, MatMul, Add) and edges are data tensors. This representation must explicitly capture data dependencies and tensor shapes. The planner analyzes this graph to identify fusion groups—subgraphs of connected operators that are candidates for combination. Frameworks like MLIR or XLA HLO provide the structured IR necessary for this analysis.

This component applies a set of fusion heuristics to identify profitable fusion opportunities. It uses pattern matching to recognize common, beneficial subgraphs. For example:

• Elementwise Fusion: Combining pointwise ops like Add, ReLU, Sigmoid.
• Vertical Fusion: Merging a producer (e.g., Conv) with its direct consumer (e.g., BatchNorm).
• Horizontal Fusion: Combining independent ops that share an input. Canonical patterns like Conv-BN-ReLU are hard-coded rules. More advanced planners use a pattern library that can be extended for new operator types.

The cost model is a predictive function that estimates the execution time, memory usage, or hardware utilization of a potential fused kernel versus the unfused baseline. It evaluates fusion profitability. Key inputs include:

• Operator characteristics: FLOPs, memory bandwidth requirements.
• Data access patterns: Stride, contiguity, reuse distance.
• Hardware parameters: Cache sizes, register file limits, compute unit throughput. The model must balance memory-bound fusion (reducing DRAM traffic) against compute-bound fusion (increasing arithmetic intensity). An inaccurate cost model can lead to performance-degrading fusions.

Given the exponential number of possible fusion groupings, the planner employs a search algorithm to find a near-optimal plan. Common strategies include:

• Greedy Algorithms: Iteratively fuse the most profitable pair until no beneficial fusions remain.
• Dynamic Programming: For sequences like chains of elementwise ops, using optimal substructure.
• Graph Partitioning: Treating fusion as a problem of partitioning the computational graph into clusters, constrained by fusion legality (e.g., no cycles introduced). The output is a fusion plan—a mapping of each original operator to a specific fused kernel.

Once a fusion plan is decided, the planner passes the specification of each fusion group to a kernel code generator. This backend component is responsible for producing the low-level, hardware-specific code (e.g., CUDA, Metal, Vulkan) for the fused kernel. The planner must provide a precise description of the fused operation's logic, including:

• Fused computation schedule.
• Memory access patterns.
• Expected tensor shapes and data types. This handoff is crucial in compilers like TVM or Torch Inductor, where the planner and code generator are separate subsystems.

Modern fusion planners are not purely static; they incorporate profile-guided optimization. They may:

• Execute micro-benchmarks on the target hardware to calibrate the cost model.
• Use just-in-time (JIT) fusion to make final decisions at runtime based on actual input shapes and device properties.
• Cache profiling results for common operator patterns to accelerate future planning. This feedback loop allows the planner to adapt to different GPUs, CPUs, or specialized accelerators (NPUs), ensuring the fusion plan is tailored to the specific execution environment's characteristics.

The core intelligence of any fusion planner is its cost model and search algorithm.

• Analytical Cost Models: Estimate execution time based on hardware parameters (memory bandwidth, compute throughput) and operation characteristics (FLOP count, bytes accessed). They quickly evaluate the profitability of a candidate fusion.
• ML-Based Cost Models: Used in systems like TVM, these are neural networks trained on profiling data to predict kernel runtime more accurately, especially for complex, fused kernels.
• Search Algorithms:
  • Greedy Algorithms: Traverse the graph, fusing nodes if a local cost model predicts improvement. Fast but can be suboptimal.
  • Dynamic Programming: Explores all possible fusion partitions of a sequence (e.g., a chain of ops) to find a globally optimal plan.
  • Evolutionary/ML-Guided Search: Explores a broader space of possibilities, as used in auto-schedulers.

The low-level implementation of a fusion plan. A fused kernel is a single, unified GPU or accelerator kernel that executes the combined logic of multiple primitive operations. This eliminates intermediate memory stores/loads and amortizes kernel launch overhead, the latency cost of submitting work to the GPU. It is the final, compiled output of a fusion planner's decisions.

The graph-level optimization that a fusion planner performs. It involves merging adjacent computational operators (nodes) in a neural network's computational graph into a single, compound operation. The planner's goal is to identify which fusion groups—sets of candidate operators—to merge to minimize data movement between memory hierarchies (DRAM to GPU memory).

The automated process of identifying and merging subgraphs within a computational graph. A fusion planner executes graph fusion, often using pattern matching to recognize canonical sequences like Conv-BN-ReLU. It explores the search space of possible fusions, evaluating each using a cost model for fusion to determine fusion profitability—the net performance gain from merging.

The specialized compiler or compiler pass that houses the fusion planner. Examples include:

• Fusion in XLA: Google's Accelerated Linear Algebra compiler, used by TensorFlow/JAX.
• Fusion in TVM: Apache TVM's stack, which uses auto-scheduling.
• Fusion in MLIR: Uses dialects like Linalg within the Multi-Level IR. These compilers provide the infrastructure (IR, cost models) the planner requires.

The two primary execution strategies for a fusion planner:

• Just-In-Time Fusion: Fusion decisions and kernel generation occur dynamically at runtime. This allows optimization based on actual input shapes (e.g., torch.compile).
• Ahead-of-Time Fusion: Fusion is performed statically during compilation, producing a pre-optimized executable. This trades runtime flexibility for reduced latency and deployment footprint.

Core tactical approaches a planner uses:

• Vertical Fusion: Merging a producer operator with its direct consumer.
• Horizontal Fusion: Merging independent, parallel operators.
• Elementwise Fusion: Combining pointwise ops (ReLU, Add).
• Memory-Bound vs Compute-Bound Fusion: Optimizing for data movement reduction or arithmetic intensity. Canonical patterns include Fused Multi-Head Attention (e.g., FlashAttention) and Fused Conv-BN-ReLU.