Operator Fusion

TensorFlow Lite's converter performs graph-level operator fusion, combining sequences like Conv2D + BiasAdd + Activation into a single node. Its XNNPACK delegate then implements these fused patterns as highly optimized kernels for CPU targets (Arm Cortex-A/Cortex-M). This two-stage approach separates hardware-agnostic graph optimizations from hardware-specific kernel implementations.

• Common Fusions: FULLY_CONNECTED + RELU, DEPTHWISE_CONV_2D + RELU6.
• Target: Mobile CPUs and microcontrollers via TFLite Micro.

Apache TVM employs operator fusion as a core optimization within its Relay intermediate representation (IR). Its BYOC (Bring Your Own Codegen) framework allows vendors to define custom fusion patterns for their hardware. TVM uses a cost model to decide fusion strategies, balancing the benefits of reduced memory access against potential register pressure and lower parallelism.

• Approach: Pattern matching on computational graph followed by kernel code generation.
• Flexibility: Enables hardware-specific fusion rules for NPUs, GPUs, and custom accelerators.

Compilers built on MLIR (Multi-Level Intermediate Representation) perform fusion through progressive lowering and dialect-specific optimization passes. IREE, for GPU/CPU targets, fuses operations within its Linalg and Vector dialects. The XLA/MLIR-HLO dialect includes a dedicated fusion pass that groups operations to be emitted as a single loop nest, maximizing data locality.

• Mechanism: Transformations within specialized MLIR dialects (linalg, tosa, mhlo).
• Benefit: Enables aggressive fusion across complex operation graphs before hardware-specific codegen.

TensorRT applies kernel fusion and layer fusion during graph optimization to minimize data transfers between GPU kernels. It combines adjacent layers (e.g., convolution, bias, and activation) into a single CUDA kernel. This reduces launch overhead and leverages on-chip memory, which is critical for latency-sensitive deployments.

• Primary Target: NVIDIA GPUs for data center and edge (Jetson).
• Key Fusion: Horizontal fusion of independent operations and vertical fusion of dependent operations.

These libraries provide pre-fused, hand-optimized kernels for Arm processors. The Arm Compute Library (for Cortex-A) offers kernels like NEGEMMConvolutionLayer that fuse convolution, bias, and activation. CMSIS-NN (for Cortex-M) provides fused functions such as arm_convolve_HWC_q7_fast() which integrate convolution, ReLU, and downscaling for microcontrollers.

• Methodology: Manual kernel creation, not automated graph fusion.
• Advantage: Provides highly efficient, deterministic kernels for constrained cores.

ONNX Runtime performs operator fusion during graph optimization sessions. Its efficiency depends on the selected Execution Provider (EP). The CPU EP fuses supported patterns, while hardware-specific EPs (e.g., for NVIDIA TensorRT, Intel OpenVINO) may apply more aggressive, hardware-tailored fusion during graph translation to their native format.

• Cross-Platform: Fusion logic is part of the common graph optimizer.
• Provider-Specific: Final fused kernel implementation is delegated to the EP (e.g., TensorRT EP does GPU fusion).

Kernel fusion is the low-level implementation of operator fusion, where the compute kernels for consecutive layers are merged. This reduces:

• Kernel launch overhead from the runtime scheduler.
• Global memory accesses by keeping intermediate tensors in faster cache or registers.
• Instruction count by eliminating redundant loops and memory loads/stores. For example, fusing a Convolution, Batch Normalization, and ReLU activation into a single kernel is a common optimization in frameworks like TensorFlow Lite for Microcontrollers and Apache TVM.

Graph optimization is the broader compiler phase where operator fusion occurs. A model's computational graph is analyzed and transformed before code generation. Key passes include:

• Constant folding: Pre-computing operations on static tensors.
• Dead code elimination: Removing unused operations or branches.
• Operator fusion: Merging compatible sequential nodes.
• Layout transformation: Optimizing tensor data formats for the target hardware. Tools like MLIR (Multi-Level Intermediate Representation) and XLA (Accelerated Linear Algebra) apply these transformations to generate highly efficient code for diverse backends.

Hardware-aware compilation tailors the fused operator kernels to the specific capabilities of the target silicon. This involves:

• Selecting optimal data types (e.g., int8 vs. float32) based on the processor's ALU support.
• Utilizing specialized instructions like Arm CMSIS-NN kernels for Cortex-M or DSP extensions for signal processing.
• Tiling strategies that match the cache hierarchy of the accelerator (e.g., a microNPU like the Ethos-U55). The compiler must understand hardware constraints to generate the most efficient fused kernel, rather than a generic one.

Apache TVM is an open-source compiler stack for machine learning that excels at operator fusion and hardware-aware optimizations. Its architecture is central to understanding modern ML compilation:

• Uses a graph-level intermediate representation (Relay) to apply high-level optimizations like fusion.
• Employs a tensor expression language (TE) to describe kernel computations.
• Features an auto-scheduler (Ansor) that can automatically generate and optimize fused kernels for novel operators or hardware targets. TVM demonstrates how fusion is not a manual process but an automated search for the optimal schedule of operations on a given device.