Elementwise Fusion

Elementwise fusion is strictly limited to pointwise or elementwise operations. These are mathematical functions applied independently to each element of a tensor, with no cross-element dependencies. Common examples include:

• Activation functions: ReLU, Sigmoid, Tanh, GELU
• Arithmetic operations: Add, Multiply, Subtract (with broadcasting)
• Simple transformations: Absolute value, Negation, Scaling Because each output element depends solely on the corresponding input element(s) at the same position, the computations are embarrassingly parallel and perfectly suited for fusion into a single, coalesced loop.

The primary performance benefit of elementwise fusion is the drastic reduction of memory bandwidth pressure. Without fusion, each intermediate result is written to global GPU memory (DRAM) by one kernel, only to be immediately read back by the next. Fusion eliminates these intermediate tensor stores and loads.

Key Mechanism: The fused kernel computes the entire chain of operations for each element, holding the intermediate scalar value in a thread register or shared memory, before writing only the final result to DRAM. This transforms a memory-bound sequence of light computations into a more compute-efficient kernel.

Each GPU kernel launch incurs a fixed scheduling and dispatch overhead. For a sequence of small, elementwise ops, this launch latency can dominate total execution time. Elementwise fusion amortizes this overhead across all fused operations.

Performance Impact: Launching one fused kernel instead of N separate kernels can provide significant speedups for short operator chains, especially for small batch sizes or tensor shapes where individual kernel runtime is low. This is a classic technique to improve latency for inference workloads.

Elementwise fusion is typically performed automatically by deep learning compilers during graph lowering or kernel generation. The compiler identifies fusible subgraphs within the computational graph.

Key Compiler Components:

• Pattern Matcher: Identifies chains of elementwise ops (e.g., Add -> ReLU -> Sigmoid).
• Fusion Planner/Heuristics: Decides if fusion is profitable based on data dependencies and a cost model.
• Kernel Code Generator: Emits a single CUDA/HIP/MLIR kernel implementing the fused operation sequence. Major frameworks leverage this in backends like XLA, TorchInductor (via torch.compile), TVM, and MLIR.

Elementwise fusion manifests in two primary dataflow patterns:

Vertical Fusion (Chaining): Merges a producer elementwise op with its consumer elementwise op. This is the most common form, creating longer chains (e.g., LayerNorm -> SiLU). It directly reduces intermediate memory traffic.

Horizontal Fusion (Side-by-Side): Merges independent elementwise ops that consume the same input tensor. For example, computing both the ReLU and Sigmoid of a tensor in one pass. This improves compute density and amortizes memory reads of the common input.

Not all elementwise operator sequences are profitable to fuse. Compilers use heuristics to avoid negative fusion.

Constraints and Trade-offs:

• Register Pressure: A long fusion chain may require many live variables, exhausting GPU registers and causing spilling to slower memory.
• Parallelism Reduction: Excessive fusion can create a single, large kernel that limits opportunities for concurrent kernel execution.
• Divergent Control Flow: Combining ops with very different branching patterns (e.g., complex conditionals) can lead to warp divergence and underutilization. The compiler's cost model must balance reduced memory traffic against these potential downsides.

Fusion can be applied at different stages of the compilation pipeline, with distinct trade-offs.

• Just-In-Time (JIT) Fusion:
  • Performed at runtime (e.g., PyTorch's torch.compile, JAX).
  • Adapts to dynamic input shapes and graph variations.
  • Incurs one-time compilation overhead.
• Ahead-of-Time (AOT) Fusion:
  • Performed during offline compilation (e.g., TVM, XLA AOT).
  • Eliminates runtime compilation cost.
  • Produces a static, deployable binary optimized for a fixed graph.

Kernel fusion is a low-level compiler optimization that combines multiple GPU or accelerator kernels into a single, unified kernel. This eliminates the kernel launch overhead and associated synchronization points between successive operations. It is the fundamental hardware-level technique that enables higher-level operator fusion.

• Primary Goal: Reduce latency from repeated kernel submission and global memory access.
• Scope: Works at the level of CUDA/GPU kernels or other accelerator-specific code.

Operator fusion is a graph-level optimization that merges adjacent computational nodes (operators) in a neural network's computational graph. It creates a single, compound operation, which is then implemented via a fused kernel. This is the direct parent concept of elementwise fusion.

• Graph Transformation: Occurs during graph compilation (e.g., in XLA, TorchScript, ONNX Runtime).
• Targets Memory Bandwidth: Primarily reduces intermediate tensor writes to and reads from high-latency global memory (HBM/DRAM).

These are two fundamental patterns for grouping operations:

• Vertical Fusion (Producer-Consumer): Merges a sequence of dependent operations. A classic example is fusing a matrix multiplication with a following bias addition and ReLU activation. Elementwise fusion is often vertical.
• Horizontal Fusion (Parallel/Sibling): Merges independent operations that consume the same input or operate in parallel. An example is applying multiple, different pointwise activations to the same tensor in separate branches, then combining them.

Graph fusion is the automated process of identifying subgraphs (called fusion groups) within a computational graph that are profitable to fuse. A fusion planner uses fusion heuristics and a cost model to make these decisions.

• Fusion Group: A set of operators marked for combined execution.
• Profitability Analysis: The compiler must decide if fusion reduces runtime, considering trade-offs like increased register pressure or reduced parallelism.

Specialized compilers that perform aggressive fusion as a core optimization:

• XLA (Accelerated Linear Algebra): Google's compiler for TensorFlow/JAX. Known for aggressive horizontal and vertical fusion, often fusing dozens of operations into a single kernel.
• TVM (Tensor Virtual Machine): Uses its scheduling language (AutoTVM, Ansor) to explicitly define and automate fusion strategies for diverse hardware backends.
• MLIR (Multi-Level IR): Provides intermediate representations (e.g., Linalg dialect) and transformation infrastructure to represent and perform fusion at multiple abstraction levels.

Pre-defined, hand-optimized kernels for common operator sequences. These serve as templates for compilers and are highly performance-critical:

• Fused Conv-BN-ReLU: The standard fused block in convolutional networks (Convolution + Batch Normalization + ReLU).
• Fused Multi-Head Attention: A single kernel for the entire attention mechanism, as popularized by FlashAttention.
• Fused LSTM/GRU Gates: Combining the multiple matrix operations within a recurrent cell.
• Elementwise Chains: Sequences like Add -> Sigmoid -> Multiply (e.g., in a gating mechanism).