Fusion in TVM

Fusion begins in TVM's high-level Relay Intermediate Representation (IR). The compiler performs graph-level pattern matching to identify subgraphs of operators that are profitable to combine. Common patterns like Conv2D + BatchNorm + ReLU are detected and replaced with a single, compound CallNode representing the fused operator. This high-level fusion reduces the number of operators passed to the low-level scheduler and defines the fusion boundaries.

Each fused operator group is lowered to Tensor Expressions (TE), a domain-specific language for describing tensor computations. The key to TVM's flexibility is its scheduling primitives applied to these TEs:

• compute_inline: Fuses a producer operation directly into the consumer's loop nest.
• compute_at: Computes a producer tensor within a specific loop iteration of the consumer.
• cache_read / cache_write: Explicitly manages data movement, creating fusion opportunities by keeping intermediate results in fast memory (e.g., shared memory on GPUs). These primitives give fine-grained control over loop fusion, tiling, and memory hierarchy mapping.

TVM's Ansor auto-scheduler automatically generates high-performance schedules, including fusion decisions, without manual tuning. It:

• Explores a search space of possible loop transformations, including various fusion strategies (e.g., inline, compute_at).
• Uses a cost model to predict the performance of different fused kernel implementations.
• Learns over time which fusion patterns are most effective for the target hardware (CPU/GPU). This automates the complex trade-off between reduced memory traffic (from fusion) and potential loss of parallelism.

After scheduling, TVM generates highly optimized, target-specific low-level code for the fused kernel. It leverages:

• LLVM for CPUs and ARM architectures.
• CUDA/NVPTX or ROCm for GPUs.
• OpenCL for diverse accelerators. The code generator applies final low-level optimizations like vectorization and register allocation to the fused loop structure. The output is a minimal kernel launch overhead, as many primitive ops are executed within a single kernel.

A primary benefit of TVM's fusion is optimized memory access patterns. By fusing operators:

• Intermediate tensors are kept in registers or thread-local memory, eliminating costly writes and reads to global GPU memory (DRAM).
• Memory accesses are coalesced across fused operations, improving bandwidth utilization.
• Data reuse is maximized within cache hierarchies (L1/L2 on CPU, shared memory on GPU). This transforms memory-bound operations (like element-wise activations) into compute-bound kernels when fused with compute-intensive ops like convolutions.

TVM does not fuse blindly; it uses heuristics and models to assess fusion profitability. Key considerations include:

• Operator Dependencies: Only fusible within data-flow paths.
• Memory Footprint: Fusing very large operators can exceed register pressure or shared memory limits, causing spills.
• Parallelism: Fusion can sometimes reduce opportunities for parallel execution.
• Hardware Characteristics: The optimal fusion strategy differs for a multi-core CPU versus a massively parallel GPU. TVM's cost models encode these hardware traits to make profitable decisions.

The combination of multiple pointwise operations that apply an independent function to each element of a tensor. This is one of the most fundamental and profitable fusion patterns.

• Examples: A sequence like Add → ReLU → Sigmoid.
• Mechanism: TVM's scheduler can fuse these operations by generating a single loop that applies all functions in one pass over the data.
• Benefit: Eliminates multiple passes over the same memory, drastically reducing memory bandwidth pressure and kernel launch overhead.

The fusion of a producer operator with its immediate consumer operator along the dataflow edge of the graph. This chains sequentially dependent operations.

• Example: Fusing a MatMul with a subsequent BiasAdd.
• TVM Implementation: TVM's compute_at and fuse primitives allow the computation of the consumer to be scheduled "inside" the loop nest of the producer.
• Benefit: The intermediate result from the producer is kept in fast registers or cache for immediate consumption, avoiding a costly write and read to global memory.

The fusion of multiple independent operators that share a common input or operate in parallel on different tensors. This pattern increases thread utilization.

• Example: Applying separate ReLU and Tanh activations to two different tensors produced from the same layer.
• TVM Implementation: Achieved by computing the independent operations within the same loop nest or kernel, often using TVM's compute_inline or parallel scheduling constructs.
• Benefit: Amortizes kernel launch overhead across multiple operations and can improve cache locality if the inputs are related.

Fusing an operation that performs a reduction (like sum or max) with preceding elementwise operations. This is critical for patterns like batch normalization.

• Example: Fusing the mean (sum/N) and variance calculations in a BatchNorm layer with the subsequent normalization and scaling operations.
• TVM Mechanism: TVM's scheduler can fuse the reduction loops with the computation loops, often using rfactor to manage parallel reduction strategies.
• Benefit: Intermediate reduction values are computed and consumed on-chip, avoiding multiple sweeps through large tensors in global memory.

The canonical deep learning fusion pattern that combines a Convolution, Batch Normalization, and ReLU activation into a single kernel. This is a prime target for TVM's graph-level pattern matcher.

• Pattern Recognition: TVM identifies the subgraph conv2d → batch_norm → relu and marks it as a fusion group.
• Kernel Generation: TVM's schedule fuses the loops, applying the BN scaling/shifting and the ReLU clamp directly to the convolution's output before it is written to memory.
• Impact: This fusion can provide a >2x speedup for CNN inference by eliminating two intermediate tensor writes and reads.

A critical pattern in Transformer feed-forward networks, fusing the matrix multiplication, bias addition, and non-linear GELU activation. This is a key optimization for LLM inference.

• Pattern: matmul → add (bias) → gelu.
• TVM Approach: TVM's Ansor auto-scheduler or manual TE scheduling can generate a single kernel where the GELU approximation is applied elementwise to the MatMul result in registers.
• Benefit: Significantly reduces latency in the transformer block, which is often memory-bound. This fusion is analogous to the FusedMultiplyAdd + activation fusion in libraries like cuDNN and oneDNN.

Kernel fusion is the foundational low-level compiler optimization that combines multiple GPU or accelerator kernels into a single, unified kernel. This directly reduces kernel launch overhead—the latency of submitting work to the GPU—and improves data locality by keeping intermediate results in fast registers or shared memory instead of writing them to slow global memory. It is the mechanical implementation of the fusion concept.

Operator fusion is the graph-level optimization that identifies and merges adjacent computational nodes (operators) in a neural network's computational graph. While kernel fusion works at the hardware instruction level, operator fusion works at the framework abstraction level (e.g., combining a Conv2d, BatchNorm, and ReLU). TVM's fusion capabilities bridge these levels by transforming fused operator patterns into optimized fused kernels.

A fusion compiler is a specialized compiler or compiler pass designed to automate fusion optimizations. Key examples include:

• XLA (Accelerated Linear Algebra): Google's compiler for TensorFlow/JAX, known for aggressive fusion.
• MLIR (Multi-Level Intermediate Representation): Provides dialects (like Linalg) and transformation infrastructure for fusion.
• TVM's Relay and TIR: TVM uses its Relay frontend for graph-level fusion and its TensorIR (TIR) for scheduling and kernel generation. These compilers use fusion heuristics and cost models to decide fusion profitability.

Canonical fusion patterns are pre-defined, high-value operator sequences that compilers target for fusion. Common examples include:

• Fused Conv-BN-ReLU: The quintessential CNN layer fusion, combining Convolution, Batch Normalization, and activation.
• Fused Multi-Head Attention: A complex fusion that combines projection, scoring, softmax, and aggregation into one kernel, as seen in FlashAttention.
• Elementwise Fusion: Merging pointwise ops like Add, Mul, Sigmoid that operate independently on each tensor element. TVM's pattern matching infrastructure automatically identifies and applies these patterns.

Fusion decisions are guided by the nature of the workload and hardware constraints:

• Vertical Fusion: Merging a producer operator with its consumer (sequential dependency).
• Horizontal Fusion: Merging independent operators that share an input or run in parallel.
• Memory-Bound Fusion: Focused on reducing data movement between GPU memory hierarchies.
• Compute-Bound Fusion: Aimed at increasing arithmetic intensity to saturate GPU compute units. TVM's fusion-aware scheduling and fusion planner evaluate these strategies to generate an optimal execution plan.

Fusion can be applied at different stages of the compilation lifecycle:

• Ahead-of-Time Fusion (AOT Fusion): Fusion and kernel generation occur statically before deployment, producing a portable, optimized binary. Ideal for fixed model architectures and target hardware.
• Just-In-Time Fusion (JIT Fusion): Fusion decisions are made dynamically at runtime (e.g., on first model execution). This allows optimization based on actual input shapes and hardware context, as used by torch.compile in PyTorch. TVM supports both modes, allowing engineers to choose the appropriate trade-off between compilation time and runtime flexibility.