Fusion in XLA

Vertical fusion merges a sequence of dependent operations where the output of one operator is the immediate input to the next. This is the most common and profitable fusion pattern.

• Mechanism: Chains operations like MatMul -> BiasAdd -> ReLU into a single kernel.
• Primary Benefit: Eliminates the write and subsequent read of the intermediate tensor to slow global memory (e.g., HBM).
• Example: Fusing a convolution with its following batch normalization and activation function is a classic vertical fusion that can yield >2x speedup by keeping data in registers or shared memory.

Horizontal fusion combines multiple independent operations that consume the same input tensor or operate in parallel within the dataflow graph.

• Mechanism: Executes several element-wise operations on the same data in a single kernel pass.
• Primary Benefit: Amortizes kernel launch overhead and improves memory bandwidth utilization by reading the input data once for multiple computations.
• Example: Applying both a Sin and a Cos function to the same tensor can be fused horizontally. The compiler must ensure the fused operations have compatible shapes and can be mapped to the same parallel execution grid.

This is a specific, highly profitable case of vertical/horizontal fusion targeting elementwise operations.

• Definition: An operation is elementwise if it computes each output element independently using only the corresponding input elements (e.g., Add, ReLU, Tanh).
• Mechanism: XLA fuses long chains of pointwise ops into a single kernel. The combined operation performs all per-element computations before writing the final result to memory.
• Benefit: Dramatically reduces memory-bound bottlenecks. For example, a graph with ten sequential pointwise ops would perform ten global memory reads and writes without fusion, but only one read and one write with fusion.

XLA performs fusion on its High-Level Operations (HLO) IR, which is a compiler-friendly representation of the computation.

• Process: The input graph (e.g., from TensorFlow or JAX) is lowered to HLO. The HLO fusion pass analyzes the dataflow and applies fusion heuristics.
• Cost Model: The compiler uses a cost model to estimate the profitability of fusing a candidate group, considering factors like reduced memory bytes accessed versus potential increase in kernel register pressure.
• Output: The pass outputs a new HLO graph where HloFusionInstruction nodes replace the original subgraphs of primitive operations.

XLA does not fuse all possible operations; it uses a fusion planner to make greedy or cost-model-driven decisions.

• Key Consideration - Fusion Profitability: Not all fusion is beneficial. Fusing a very large, compute-intensive operation (like a convolution) with a tiny one may not justify the added complexity and could hurt occupancy.
• Heuristics: The planner uses rules such as:
  • Fuse producers into consumers unless the producer has multiple users.
  • Prefer to fuse operations that are "expensive" to materialize (large outputs).
  • Avoid fusing operations that would create a kernel with excessive register usage.
• Goal: To generate a fusion plan that minimizes total execution time, not just kernel count.

XLA recognizes and aggressively fuses common, performance-critical subgraph patterns found in neural networks.

• Conv-BN-ReLU: The quintessential fused layer in CNNs. XLA fuses the Convolution, Batch Normalization (scale, shift, mean, variance), and ReLU activation. The BN parameters are compiled directly into the kernel, and the activation is applied before writing to global memory.
• LayerNorm + GeLU: A common pattern in transformers. The elementwise GeLU activation is fused with the preceding LayerNorm operation.
• Impact: These hand-optimized equivalent patterns are generated automatically, providing performance comparable to manually written fused kernels without requiring user intervention.

This is the quintessential fused pattern for convolutional neural networks. XLA fuses a Convolution layer, an Add operation for the bias term, and a non-linear activation function (e.g., ReLU, Sigmoid) into one kernel.

• Eliminates two intermediate tensor writes and reads.
• Keeps data in fast GPU registers or shared memory between operations.
• Example: Conv2D + BiasAdd + ReLU becomes a single __nv_fused_conv2d_bias_relu kernel call.

XLA aggressively fuses Batch Normalization with surrounding operations. The canonical pattern is Conv-BatchNorm-Activation. During inference, the batch norm's mean, variance, scale, and offset are statically folded into the preceding convolution's weights and bias, creating a mathematically equivalent but faster single operation.

• Training vs. Inference: Folding is primarily an inference-time optimization.
• Extended Patterns: Can include preceding Conv or MatMul and following Activation.

The fundamental building block of transformer and fully-connected networks. XLA fuses a Matrix Multiplication (MatMul), the addition of a bias vector, and frequently an activation like ReLU or Gaussian Error Linear Unit (GeLU).

• Critical for LLMs: This fusion is a major performance factor in the dense layers of transformers.
• Reduces Memory Bandwidth: The large output of the MatMul is consumed immediately without a round-trip to global GPU memory.

XLA fuses sequences of pointwise (elementwise) operations that apply a function independently to each tensor element. Common examples include:

• Add -> Tanh

• Multiply -> Add (e.g., a residual connection adjustment)

• Long chains of Unary Ops (e.g., Cast -> Log -> Neg)

• Key Benefit: Amortizes kernel launch overhead across many light operations.

• Horizontal Fusion: Can also fuse independent elementwise ops on the same input.

XLA fuses a reduction operation (e.g., Sum, Max) with a preceding elementwise operation that prepares the data. Instead of writing the full intermediate tensor, the fused kernel performs the elementwise op and immediately accumulates the result.

• Pattern: Broadcast -> Elementwise -> Reduce.
• Example: Calculating Softmax involves a Exp, Reduce(Sum) across a dimension, and a Divide. Key parts of this chain are fused.
• Impact: Dramatically reduces memory footprint for large reduction dimensions.

XLA doesn't fuse blindly. Its fusion planner uses a cost model to assess fusion profitability. It analyzes:

• Data Dependencies: Only fuses producers with direct consumers (vertical fusion).
• Memory vs. Compute Bound: Prioritizes fusing memory-bound ops to reduce data movement.
• Kernel Launch Overhead: Weighs savings against potential downsides like increased register pressure or decreased parallelism.

The planner identifies subgraphs matching canonical patterns and decides if generating a custom fused kernel will be faster than executing the ops separately.

Kernel fusion is the low-level compiler technique of combining multiple GPU or accelerator kernels into a single, unified kernel. This eliminates the kernel launch overhead and temporary memory writes/reads between operations. It is the fundamental hardware-level mechanism that enables the performance gains from higher-level operator fusion.

• Primary Benefit: Reduces latency from repeated kernel launch setup and teardown.
• Key Challenge: Requires careful management of on-chip resources like registers and shared memory.
• Example: Fusing an element-wise tanh operation with a preceding matrix multiplication into one CUDA kernel.

Operator fusion is a graph-level optimization that merges adjacent nodes (operators) in a neural network's computational graph. XLA performs this before generating low-level kernels. The goal is to minimize intermediate tensor materialization in high-bandwidth memory (HBM).

• Vertical Fusion: Chains a producer operator with its consumer (e.g., a matmul feeding directly into a bias_add).
• Horizontal Fusion: Combines independent operators that share an input or have similar shapes.
• Profitability Analysis: The compiler uses a cost model to decide if fusion reduces memory traffic more than it increases register pressure.

Graph fusion is the automated process of identifying and merging subgraphs within a full computational graph. In XLA, this involves pattern matching to find known profitable sequences (like Conv2D -> BatchNorm -> ReLU). The compiler partitions the graph into fusion groups, where all internal edges represent fast, on-chip communication.

• Fusion Planner: The component that explores the search space of possible fusions.
• Fusion Group: The final set of operators scheduled for a single kernel.
• Scope: Works across both linear chains and more complex, branching subgraphs.

Fusion heuristics are the rule-based or ML-driven policies a compiler uses to decide what to fuse. A cost model for fusion predicts the execution time of a potential fusion group, balancing:

• Reduced Global Memory Access: The primary benefit.
• Increased Register Pressure: Can limit parallelism if too high.
• Kernel Launch Overhead Savings: More significant for many small ops.
• Potential for Parallelism Loss: Fusing independent ops may reduce concurrent execution.

XLA's fusion planner uses such models to maximize fusion profitability.