Kernel Fusion

The primary bottleneck for edge AI is memory bandwidth, not compute. Each standalone kernel (e.g., Conv → BatchNorm → ReLU) must:

• Read inputs from global memory (slow, high energy).
• Write intermediate results back to global memory (slow, high energy).
• Be launched by the host CPU, incurring kernel launch overhead.

Fusion eliminates these intermediate global memory round-trips by keeping data in fast registers or shared memory between fused operations.

The compiler identifies fusible patterns within the computational graph. Common patterns include:

• Elementwise Ops: Consecutive ReLU, Sigmoid, Add.
• Producer-Consumer: Conv → BatchNorm → Activation.
• Reduction Ops: Softmax (which involves Exp, Sum, Div).

Once identified, the compiler replaces the matched subgraph with a single, custom fused kernel. This is a key pass in graph optimization frameworks like XLA, TVM, and MLIR.

The compiler generates a new, monolithic kernel that performs the computation of all fused operations. This involves:

• Fused Load/Store: Input tensors are loaded once; intermediate results are passed via registers.
• Loop Fusion: Nested loops from separate operations are combined, improving data locality.
• Intermediate Computation: Values are computed on-the-fly and consumed immediately, never written to main memory.

For example, a fused Conv-BN-ReLU kernel computes the convolution output, applies batch norm scaling/shifting, and clips with ReLU—all within a single loop nest before storing the final result.

The efficiency of a fused kernel depends on the target hardware's architecture. Compilers like TVM use auto-tuning to search for optimal parameters:

• Thread Block Size: How work is partitioned across GPU/NPU cores.
• Tile Sizes: For managing data movement between memory hierarchies.
• Vectorization Width: To maximize SIMD unit utilization.

A kernel fused for an Arm CPU with NEON instructions will have a different optimal configuration than one for an NVIDIA GPU or a Qualcomm Hexagon NPU.

Kernel fusion directly targets the key metrics for edge deployment:

• Latency Reduction: Can reduce end-to-end inference time by 20-50% for models with many sequential small ops by eliminating kernel launch overhead and memory stalls.
• Power Efficiency: Dramatically reduces energy consumption per inference by minimizing costly DRAM accesses. Power is often proportional to data movement.
• Deterministic Execution: Fewer kernel launches and memory allocations lead to more predictable, real-time-friendly inference loops.

Fusion is not always applicable or beneficial:

• Register Pressure: Over-fusing can require more registers than available, causing register spilling to slower memory, which hurts performance.
• Parallelism Reduction: Fusing independent operations that could run in parallel may reduce throughput on highly parallel hardware.
• Compiler Complexity: Requires deep hardware knowledge and sophisticated pattern matching. Not all operation sequences are legally fusible (e.g., ops with data-dependent shapes).
• Portability: A hand-tuned fused kernel for one accelerator may not perform well on another, challenging cross-platform deployment.

XLA's fusion pass aggressively combines sequences of element-wise, reduction, and broadcast operations. It uses a cost model to decide fusion profitability, aiming to keep fused computations within fast on-chip memory (e.g., GPU shared memory or CPU L1 cache).

• Fusion Types: Horizontal (parallel ops) and vertical (producer-consumer ops).
• Key Benefit: Dramatically reduces reads/writes to high-latency global memory by keeping intermediate results in registers or cache.
• Example: A pattern like Softmax(Add(Broadcast, Tensor)) is fused into a single GPU kernel.

TVM implements fusion via its Relay IR and the FuseOps pass. It partitions the computational graph into fusible subgraphs, then generates a single fused kernel for each partition using the TVM code generator (TIR).

• Strategy: Uses a rule-based or cost-model-driven approach to group operators.
• Hardware-Aware: Fusion decisions account for target hardware constraints like register count and memory hierarchy.
• Result: Produces tailored, high-performance kernels for CPUs, GPUs, and custom accelerators from a fused subgraph.

Compilers built on MLIR use progressive lowering through multiple abstraction layers (e.g., linalg, vector, gpu dialects) to enable fusion. Transformations are expressed as portable, composable rewrite rules.

• Linalg Fusion: Fuses producer-consumer linalg operations (like matmul + add) by tiling and promoting results to shared memory in a single kernel.
• Tile and Fuse: A common pattern that combines loop tiling for data locality with operation fusion.
• Advantage: The modular IR allows fusion strategies to be reused across different compiler backends and hardware targets.

Inductor, PyTorch's JIT compiler, performs fusion on a lowered operator graph (FX Graph). It identifies fusible pointwise, reduction, and memory-bound operations, then generates fused kernels via Triton (for GPUs) or C++/OpenMP (for CPUs).

• Automatic Discovery: The compiler automatically patterns match common sequences like linear + relu + dropout.
• Kernel Generation: Uses the Triton language to generate efficient GPU kernels for the fused subgraph.
• Goal: Minimize kernel launches and global memory accesses in eager-mode PyTorch models.

ONNX Runtime performs graph optimizations, including fusion, during model load. Providers like CUDA, TensorRT, and OpenVINO have dedicated fusion passes that map subgraphs to pre-fused, highly optimized kernels in their respective libraries.

• Graph Rewriting: Identifies patterns like Conv + BatchNormalization + Activation and replaces them with a single fused node.
• Provider-Specific Kernels: Leverages vendor-optimized fused kernels (e.g., cuDNN's CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_GEMM for Conv-Bias-Activation).
• Benefit: Provides portable optimization across diverse hardware via a unified API.

TFLite uses operator fusion as a primary graph optimization to reduce latency and binary size for mobile and edge devices. It fuses common patterns during model conversion (.tflite file generation).

• Common Fusions: FullyConnected + Activation, DepthwiseConv2D + Activation, BatchNormalization folding into previous convolution.
• Delegation: For accelerators (e.g., NPUs), fused subgraphs are often delegated as a single unit to proprietary driver kernels.
• Impact: Critical for performance on resource-constrained devices where kernel launch overhead is proportionally significant.

A high-level graph optimization that merges multiple sequential neural network operations (e.g., Conv2D, BatchNorm, ReLU) into a single, compound operator before low-level code generation. This creates the opportunity for kernel fusion by reducing the granularity of the computational graph. It is a prerequisite optimization that defines the fused operation's semantics for the subsequent compiler backend.

A compiler pass that transforms a neural network's computational graph by applying semantic-preserving transformations to improve execution efficiency. Key techniques include:

• Operator Fusion and Constant Folding
• Dead Code Elimination to remove unused operations
• Common Subexpression Elimination to reuse computed values These high-level optimizations restructure the model for more efficient low-level code generation, directly enabling techniques like kernel fusion.

A compiler optimization that transforms scalar operations to execute simultaneously on multiple data points using Single Instruction, Multiple Data (SIMD) or vector processor instructions. While kernel fusion reduces kernel launch overhead and global memory accesses, vectorization maximizes data parallelism within a kernel. Modern compilers apply both: a fused kernel's internal loops are heavily vectorized to fully utilize the processor's vector units.

A compiler optimization that partitions large tensors into smaller, fixed-size blocks (tiles) to fit within the processor's fast, hierarchical memory (e.g., L1/L2 cache). This is critical for operations like matrix multiplication and convolution. Kernel fusion often creates new opportunities for tiling by combining multiple loops, allowing for more efficient data reuse across the fused operations and significantly reducing trips to slow main memory (DRAM).

A compiler optimization that pre-allocates and reuses memory buffers for all intermediate tensors at compile time. This eliminates dynamic memory allocation overhead during inference. Kernel fusion synergizes powerfully with this technique: by fusing operations, intermediate tensors that were previously written to and read from global memory become short-lived, temporary values that can be held in registers or a small, statically allocated scratchpad, drastically reducing the total memory footprint.

An automated compiler process that searches a vast space of possible low-level code implementations, parameters, and schedules to find the optimal version for a specific workload and hardware target. This is essential for implementing efficient kernel fusion, as the compiler must evaluate different fusion strategies, tile sizes, and vectorization factors. Auto-tuning uses empirical performance measurements (e.g., via Profile-Guided Optimization) to select the best-fused kernel configuration.