Fusion Compiler

The compiler analyzes the computational graph to identify subgraph patterns that are canonical candidates for fusion. This involves:

• Recognizing common sequences like Conv-BN-ReLU or linear layers followed by activations.
• Using declarative rules or learned fusion heuristics to match operator combinations.
• Distinguishing between vertical fusion (producer-consumer chains) and horizontal fusion (parallel independent operators). This pattern recognition is the first step before any transformation, allowing the compiler to apply known, high-value optimizations.

Not all fusions are beneficial. The compiler employs a cost model for fusion to estimate the performance impact. This analysis weighs:

• Reduced kernel launch overhead and fewer global memory accesses.
• Potential downsides like increased register pressure or reduced parallelism.
• Whether the fused group is memory-bound (benefiting from less data movement) or compute-bound (benefiting from increased arithmetic intensity). The fusion planner uses this model to construct an optimal fusion plan, exploring a search space to maximize net gain.

Once a fusion group is identified and deemed profitable, the compiler generates a single, unified fused kernel. This process involves:

• Loop fusion: Merging the iteration spaces of the constituent operations.
• Intermediate memory elimination: Writing temporary results directly to registers or shared memory instead of off-chip DRAM.
• Scheduling optimization: Making fusion-aware scheduling decisions for memory hierarchy mapping (e.g., fusion for cache) and parallelization. The output is a custom, high-performance kernel (e.g., a fused multi-head attention kernel) that implements the combined logic.

Fusion is deeply tied to the target accelerator's architecture. The compiler must:

• Understand hardware constraints like shared memory size, register count, and warp scheduling.
• Generate code optimized for specific compute capabilities (e.g., NVIDIA's Tensor Cores).
• Choose optimal numerical precision and memory layouts. This is why frameworks have separate compiler backends like XLA for TPUs, TVM for various CPUs/GPUs, and MLIR dialects for vendor-specific operations. The fusion strategy for a mobile GPU will differ from that for a data center accelerator.

Fusion compilers operate under different timing paradigms:

• Ahead-of-Time Fusion (AOT): Fusion decisions and kernel generation happen statically before deployment. This produces a pre-optimized, standalone executable (e.g., via XLA AOT compilation). Benefits include minimal runtime overhead and predictable performance.
• Just-In-Time Fusion (JIT): Fusion occurs dynamically at runtime, often on the first model execution. This allows optimizations based on actual input shapes and dynamic control flow. PyTorch's torch.compile and its Inductor backend use JIT fusion, as do tracing-based frameworks. JIT offers flexibility but incurs a one-time compilation cost.

A fusion compiler does not operate in isolation. It is part of a larger optimization pipeline that includes:

• Constant folding and dead code elimination.
• Layout transformations and memory pooling.
• Operator lowering from high-level ops to primitive hardware instructions.
• Interaction with coarse-grained launch optimizations like CUDA Graph capture. For example, fusion in MLIR occurs within a pipeline of transformations expressed in different dialects (e.g., transforming linalg ops before lowering to LLVM). This integration ensures fusion complements other critical performance optimizations.

Kernel fusion is the low-level compiler optimization technique that merges multiple computational kernels into a single, unified kernel. This reduces kernel launch overhead and improves data locality by keeping intermediate results in fast registers or shared memory instead of writing them to slower global memory. It is the core mechanical action performed by a fusion compiler.

• Primary Goal: Minimize the number of discrete GPU kernel launches.
• Example: Fusing an elementwise tanh operation with a preceding add operation so both computations happen in one GPU kernel pass over the data.

Operator fusion is the graph-level optimization that identifies and merges adjacent nodes (operators) in a neural network's computational graph. It creates a single, compound operation, which is then implemented via kernel fusion. This is a prerequisite step performed by the fusion compiler's graph optimizer.

• Scope: Works on the abstract computational graph before kernel code generation.
• Decision Basis: Analyzes data dependencies and memory access patterns.
• Contrast with Kernel Fusion: Operator fusion decides what to combine; kernel fusion implements how to combine it at the hardware instruction level.

A fusion group (or fusion region) is a subgraph of operators within a larger computational graph that a compiler has identified as candidates for merging into a single kernel. The fusion compiler's planner partitions the graph into these groups based on fusion heuristics and a cost model.

• Formation Rules: Groups are typically formed from operators with producer-consumer dataflow (vertical fusion) or operators that share a common input (horizontal fusion).
• Boundary: The compiler must decide where one fusion group ends and another begins, balancing fusion profitability with hardware constraints like register pressure.

These are concrete compiler frameworks where fusion logic is implemented. A fusion compiler is often a core pass within these systems.

• XLA (Accelerated Linear Algebra): Google's compiler for TensorFlow/JAX. Known for aggressive fusion using its HLO (High-Level Optimizer) IR and target-specific (TPU/GPU) code generation.
• MLIR (Multi-Level Intermediate Representation): Uses dialects like Linalg and transforms to represent and perform fusion across different levels of abstraction, enabling retargetable optimizations.
• TVM (Apache TVM): Employs its scheduling language (te.schedule) and auto-scheduler (Ansor) to explicitly define or automatically discover high-performance fused kernel schedules.

The fusion heuristics and cost model are the decision-making core of a fusion compiler. They determine fusion profitability—whether fusing a group of operators will yield a net performance gain.

• Heuristics: Rule-based logic (e.g., "always fuse elementwise ops with their direct producer").
• Cost Model: A predictive model that estimates execution time, memory traffic, and register usage of fused vs. unfused versions.
• Key Factors: Data dependency, memory-bound vs. compute-bound nature of ops, intermediate tensor size, and target hardware capabilities.

A fused kernel is the final output of the fusion compiler: a single, executable GPU or accelerator kernel that implements the functionality of an entire fusion group. It can be hand-tuned (library kernel) or compiler-generated.

• Characteristics: Has a single kernel launch, uses thread synchronization primitives internally, and manages intermediate data in shared memory or registers.
• Canonical Examples:
  • Fused Conv-BN-ReLU: The standard building block of CNNs.
  • FlashAttention: A famous fused kernel for the entire multi-head attention mechanism, optimizing for IO efficiency.
• Performance Impact: Eliminates the latency of multiple kernel launches and reduces global memory bandwidth pressure.