Kernel Fusion

Each kernel launch on an NPU incurs significant scheduling overhead from the host driver. Fusing multiple operations into a single kernel eliminates these repeated launch costs. This is critical for models with many small, sequential layers where launch latency can dominate execution time. For example, fusing a convolution, bias add, and ReLU activation transforms three distinct dispatches into one.

The primary performance gain comes from eliminating intermediate store/load cycles to high-latency global memory (HBM or DRAM). Without fusion, each kernel writes its output to global memory, which the next kernel must then read. Fusion keeps intermediate tensors in faster memory hierarchies:

• Registers: For thread-local data.
• Shared Memory: For data shared within a thread block. This can reduce global memory traffic by 2-3x for common operator chains, moving the workload from being memory-bound toward being compute-bound.

Fused kernels exhibit superior temporal and spatial locality. Data fetched into caches or shared memory is reused immediately by subsequent fused operations before being evicted. This maximizes the utility of every memory transaction. It also enables more effective use of memory coalescing, as data access patterns can be optimized across the entire fused operation sequence rather than per-kernel.

A single, larger fused kernel provides the compiler with a broader view of the computation graph. This allows for more aggressive instruction scheduling and software pipelining across the original kernel boundaries. The compiler can interleave instructions from different logical operations, hiding arithmetic and memory latencies and better utilizing the NPU's functional units.

When operations are separate, each kernel must allocate registers for its full working set. Fusion allows live ranges of intermediate variables to be contained within the fused kernel, often reusing the same physical registers across different stages of the computation. This reduces the total register footprint, minimizing costly register spilling to local memory, which is a major performance penalty on accelerators.

Kernel fusion creates new opportunities for downstream compiler passes:

• Constant Propagation & Folding: Constants from earlier ops can be folded into later operations.
• Common Subexpression Elimination (CSE): Redundant calculations across the original kernel boundaries can be identified and removed.
• Dead Code Elimination (DCE): Unused outputs from intermediate stages are never materialized.
• Advanced Tiling: The entire fused operation can be tiled as a unit, optimizing data movement for the combined working set.

This pattern fuses two or more elementwise operations that are applied independently to each element of a tensor. The compiler creates a single kernel where each thread reads an input element, performs the chained operations in registers, and writes the final result.

• Example: A sequence like output = relu(bias_add(convolution(input, weights))) where the bias addition and ReLU are fused with the preceding convolution.
• Benefit: Eliminates the need to write the intermediate bias_add result to global memory and read it back for the relu operation, drastically reducing memory traffic.

This pattern fuses a reduction operation (e.g., sum, max) with a preceding producer kernel. Instead of writing a full tensor to memory only to immediately read it back and reduce it, the reduction is performed on-the-fly as the data is generated.

• Example: Computing the L2 norm of a vector. A naive approach generates the vector of squared values, writes it to memory, then launches a second kernel to sum them. A fused kernel squares each element and immediately adds it to an accumulator in shared memory or registers.
• Benefit: Transforms a memory-bound reduction into a compute-bound operation, avoiding the round-trip to global memory for the intermediate tensor.

Also known as vertical fusion, this is the most general pattern. It merges two separate kernels where the output of the first (producer) is the direct input to the second (consumer). The compiler stitches their computation graphs into one.

• Key Challenge: The loops and data access patterns of the two kernels may differ. The compiler must perform loop alignment and potentially loop fission on the fused kernel to correctly match the iteration spaces.
• Benefit: The intermediate tensor resides entirely in fast memory (registers or shared memory), bypassing slow global memory. This is critical for deep networks with many sequential layers.

This pattern fuses two or more independent kernels that operate on different input data but have the same loop structure. The compiler creates a single kernel where each thread or thread block computes the results for multiple, unrelated operations.

• Example: Applying separate batch normalization and dropout layers to the same input tensor in parallel. A fused kernel would read the input once, compute both normalized and dropped-out values, and write two output tensors.
• Benefit: Improves arithmetic intensity and kernel occupancy by amortizing the overhead of memory loads and kernel launch across multiple operations. It better saturates the compute units.

A specialized pattern for operations like convolutions or pooling that use a sliding window over an input. Fusion occurs when multiple stencil-based layers are applied sequentially (e.g., Conv → Conv).

• Mechanism: The compiler can fuse the stencil computations so that intermediate pixel values, once loaded into shared memory or registers for the first convolution, are immediately reused for the subsequent convolution before being evicted.
• Benefit: Dramatically improves data reuse and reduces the total volume of data fetched from global memory, which is the primary bottleneck for stencil operations.

This pattern fuses operations that involve data-dependent control flow, such as a thresholding operation (e.g., ReLU) fused with its producer. The conditional logic is evaluated inline within the fused kernel.

• Implementation: The compiler embeds the conditional branch (e.g., max(0, x)) directly into the fused kernel's instruction stream. On architectures with warp divergence, careful design is needed to maintain efficiency.
• Benefit: Removes the overhead of a separate kernel launch and intermediate storage for what is often a very cheap operation, keeping the data pathway entirely in registers.

Loop fusion is a compiler optimization that merges two or more adjacent loops iterating over the same range into a single loop. This reduces loop overhead (e.g., increment and branch instructions) and, crucially, improves data locality by keeping intermediate results in faster memory hierarchies like registers or caches, rather than writing them to and reading them from slower global memory. It is a prerequisite optimization often performed before kernel fusion on the resulting compound loop body.

• Example: Fusing a loop that computes element-wise activation with a subsequent loop that performs a reduction.
• Contrast with Kernel Fusion: Loop fusion operates on the abstract syntax tree or intermediate representation; kernel fusion specifically creates a single, launchable GPU/NPU kernel.

Operation fusion is the process of combining multiple primitive tensor operations—such as a convolution, followed by a bias addition, followed by a ReLU activation—into a single, compound kernel. This is the most direct antecedent to kernel fusion in deep learning frameworks. It eliminates the scheduling overhead of launching separate kernels and, most importantly, avoids the intermediate memory traffic of storing the output of one operation to global memory only to immediately reload it as the input to the next.

• Framework Implementation: Performed by graph-level compilers like XLA (for TensorFlow/JAX), TVM, or PyTorch's TorchInductor.
• Key Benefit: Transforms memory-bound chains of operations into a more compute-bound single kernel, moving closer to the hardware's peak performance.

Kernel tiling (or loop tiling) is a loop transformation that partitions a large iteration space into smaller, regular blocks or tiles. The primary goal is to fit the working set of data for a tile into a faster, but limited, memory hierarchy (e.g., an NPU's shared memory or a CPU's L1 cache). This dramatically reduces accesses to slower global memory.

• Mechanism: Outer loops step through tiles, while inner loops process within a tile.
• Synergy with Fusion: Tiling is often applied after fusion to manage the memory footprint of the newly created, larger fused kernel. The optimal tile size is a critical parameter found via auto-tuning.

Common subexpression elimination is a compiler optimization that identifies redundant computations of identical expressions within a scope. It calculates the expression once, stores the result in a temporary register, and reuses that value, eliminating duplicate work. This is a fundamental optimization that becomes increasingly powerful after kernel fusion, as the fused kernel often contains repeated address calculations or shared intermediate values from the originally separate operations.

• Impact: Reduces arithmetic operations and register pressure.
• Example: In a fused (a+b)*c + (a+b)*d kernel, CSE would compute (a+b) once.

Dead code elimination is a compiler pass that removes code which does not affect the final output of the program. This includes computations whose results are never used and statements that are unreachable. DCE is essential after aggressive optimizations like fusion and tiling, as these transformations can create dead code (e.g., partial writes to temporary buffers that are no longer needed, or legacy loop control logic). Removing it reduces binary size, register pressure, and execution time.

• Triggered by: Constant propagation, fusion, and other transformations that make code redundant.

Profile-guided optimization is a two-stage compilation technique. First, the program is instrumented and run on representative training data to collect execution profiles (e.g., which branches are taken, how often loops iterate). Second, the compiler uses this real-world data to make better optimization decisions. For kernel fusion, PGO can inform:

• Fusion Candidates: Identifying which sequences of kernels are most frequently executed together.
• Inlining Decisions: Determining whether to inline a small function into a hot kernel.
• Branch Prediction: Optimizing the layout of conditionals within a fused kernel. This data-driven approach leads to more stable performance gains in production.