Fusion-Aware Scheduling

The primary objective is to restructure computation to keep intermediate tensor data within fast, on-chip memory hierarchies (e.g., GPU shared memory or L1/L2 caches). By fusing producer and consumer operators, the scheduler eliminates expensive round-trips to global memory (DRAM). This is critical for memory-bound operations where data transfer, not computation, is the bottleneck.

• Example: Fusing an element-wise activation (ReLU) directly after a matrix multiplication allows the activation to be computed on the multiplication's output while it's still in registers or cache, preventing a write and subsequent read from DRAM.

Each independent GPU kernel launch incurs fixed scheduling and dispatch latency. Fusion-aware scheduling aims to reduce the total number of kernel launches by grouping many small, sequential operations into fewer, larger fused kernels. This amortizes the launch overhead over more useful work, improving overall throughput, especially for networks with many light-weight layers.

• Technical Detail: On modern GPUs, kernel launch overhead can be on the order of microseconds. For models with hundreds of small ops (e.g., activations, biases), this overhead can dominate execution time if not fused.

This objective focuses on making better use of the hardware's raw compute power (FLOPS). The scheduler combines compute-bound operations (like convolutions) with adjacent memory-bound operations (like biases or scaling) into a single kernel. This increases the arithmetic intensity (FLOPs per byte of DRAM access) of the fused kernel, moving the workload closer to the hardware's peak computational throughput.

• Result: A fused Conv-BN-ReLU kernel performs the convolution, batch normalization scaling/shifting, and ReLU in one pass, maximizing FLOPs utilization compared to three separate, memory-intensive kernels.

Fusion-aware scheduling uses loop transformations—tiling, unrolling, vectorization—not just for a single operation, but across fused operator boundaries. The scheduler can tile a loop nest that encompasses multiple fused ops, allowing shared data to be staged in a tile and reused across all operations within that tile. This is more effective than optimizing each op in isolation.

• Mechanism: When ops A and B are fused, the scheduler can create a single, complex loop structure that loads input data, computes A, immediately consumes its output for B, and then writes the final result, all within a tailored tile size that fits the cache.

There is a fundamental tension between fusion and parallelism. Fusing too many operators can create monolithic kernels that limit parallel execution opportunities (e.g., by reducing independent work for GPU thread blocks). A key objective of the scheduler is to model this trade-off.

• The Scheduler's Decision: It uses a cost model to evaluate whether fusing a group of ops will improve performance more than executing them in parallel. It must consider register pressure, shared memory usage, and available parallelism on the target hardware (SM count on GPU).

The scheduler's decisions are tightly coupled to the target accelerator's architecture. Its objective is to produce a fusion plan that allows the downstream fusion compiler (e.g., in XLA, TVM, or MLIR) to generate an optimal fused kernel for that specific hardware.

• Hardware Awareness: The schedule for an NVIDIA A100 (with Tensor Cores) will differ from one for an AMD MI250X or an Intel CPU. The scheduler considers factors like memory bandwidth, cache sizes, vector unit width, and supported instruction sets (e.g., DP4A for INT8) when making fusion decisions to enable the most efficient final kernel code.

The graph-level optimization that merges adjacent computational operators in a neural network's computational graph into a single, compound operation. This is the primary goal that fusion-aware scheduling aims to enable.

• Reduces intermediate memory transfers by keeping data in registers or cache.
• Creates fusion groups that become candidates for a single kernel.
• A prerequisite optimization that scheduling decisions must preserve and exploit.

The low-level implementation technique that combines the logic of multiple primitive operations into a single, unified GPU or accelerator kernel. This is the executable result of successful operator fusion.

• Eliminates kernel launch overhead by reducing total launch count.
• Improves data locality by avoiding global memory round-trips for intermediate results.
• Can be hand-written (e.g., Fused Conv-BN-ReLU) or compiler-generated.

A fundamental compiler transformation that merges multiple adjacent loops iterating over the same range into a single loop. This is a classic technique often applied within a generated fused kernel.

• Reduces loop overhead from branch prediction and increment operations.
• Dramatically improves cache utilization by processing data while it is hot.
• A key intra-kernel optimization enabled by a scheduler that fuses operators with compatible iteration spaces.

The automated process of identifying and merging subgraphs of operators within a computational graph. This uses pattern matching and cost models to create efficient fused kernels.

• Pattern Matching for Fusion identifies canonical sequences like Linear -> GELU.
• A Fusion Planner explores the search space of possible fusions.
• Fusion Profitability analysis determines if fusion yields a net performance gain, considering register pressure and parallelism.

A specialized compiler or compiler pass responsible for performing fusion optimizations. These are the tools that implement fusion-aware scheduling.

• XLA (Accelerated Linear Algebra): Google's compiler for TensorFlow/JAX, known for aggressive fusion.
• TVM (Apache TVM): Uses a scheduling language and auto-scheduling to generate fused kernels.
• MLIR (Multi-Level IR): Provides dialects (e.g., Linalg, Affine) and infrastructure to represent and transform computations, enabling fusion across abstraction levels.

The two primary strategies for when fusion decisions and kernel generation occur.

• Just-In-Time Fusion (JIT Fusion): Decisions are made dynamically at runtime (e.g., via torch.compile). Adapts to specific input shapes but incurs compilation overhead.
• Ahead-of-Time Fusion (AOT Fusion): The entire model is statically compiled and fused before deployment. Produces a portable, optimized executable with zero runtime compilation cost.
• The choice influences the scheduler's flexibility and the deployment artifact.