Ahead-of-Time Fusion (AOT Fusion)

AOT Fusion performs a complete analysis of the model's computational graph during compilation. This allows for global optimizations that are impossible with just-in-time (JIT) methods, such as:

• Aggressive constant folding and propagation.
• Dead code elimination of unused graph branches.
• Optimal fusion planning using a holistic cost model that considers the entire inference path, not just local operator patterns.

The compiler generates custom, fused kernels tailored to the target hardware's architecture (e.g., NVIDIA Ampere GPU, AMD MI300X, AWS Inferentia). This involves:

• Leveraging specialized instructions (e.g., Tensor Cores, Matrix Cores).
• Optimizing for the target's memory hierarchy (register file, shared memory, cache sizes).
• Applying architecture-aware scheduling for thread blocks and warps. The result is a single, monolithic executable that maximizes hardware utilization.

By resolving all optimization decisions at compile time, AOT Fusion removes critical latency sources present in dynamic systems:

• Zero kernel launch overhead for fused operations, as they are part of a pre-compiled sequence.
• No runtime graph interpretation or JIT compilation delays.
• Minimal host-device synchronization. This leads to predictable, low-latency execution, which is essential for real-time inference serving.

The compilation output is a self-contained binary (e.g., a .so library or a .exe). This provides:

• Deterministic performance: The same binary yields identical performance on identical hardware, crucial for service level agreements (SLAs).
• Deployment portability: The compiled artifact can be deployed independently of the training framework (PyTorch, TensorFlow) and its runtime.
• Simplified deployment: The binary contains all fused kernels, eliminating dependency on a just-in-time compiler in the production environment.

The primary constraint of AOT Fusion is its static nature. The compiled graph is fixed for:

• Specific input shapes (batch size, sequence length, image dimensions).
• A single hardware target and software driver version.
• A predefined computational graph. This means it cannot adapt to dynamic control flow or variable input shapes at runtime without recompilation or graph capture techniques, making it less suitable for highly dynamic models.

AOT Fusion differs fundamentally from JIT Fusion in its timing and scope:

• AOT (Ahead-of-Time): Optimization occurs once, offline. It's a deployment-time step. Examples include compilers like Apache TVM, XLA AOT, and MLIR-based compilers.
• JIT (Just-In-Time): Optimization occurs repeatedly, at runtime. It's an execution-time step, often on the first model run. Examples include PyTorch's torch.compile, TensorFlow XLA JIT, and NVIDIA's CUDA Graphs for launch fusion. JIT trades a one-time compilation cost for flexibility with dynamic inputs.

A runtime optimization where fusion decisions and kernel generation occur dynamically during model execution. Unlike AOT Fusion, JIT Fusion adapts to specific input shapes and hardware contexts at the cost of one-time compilation overhead per unique execution graph.

• Key Trade-off: Flexibility vs. startup latency.
• Framework Examples: torch.compile in PyTorch 2.0 uses JIT Fusion via the Inductor backend.

A graph-level optimization that merges adjacent computational operators in a neural network's computational graph into a single, compound operation. This is the high-level transformation for which AOT Fusion provides a static compilation strategy.

• Goal: Minimize intermediate memory transfers between operators.
• Example: Fusing a Convolution, Batch Normalization, and ReLU activation into one Conv-BN-ReLU operator.

The low-level implementation of operator fusion, where the combined functionality of multiple primitive operations is executed by a single, unified GPU or accelerator kernel. AOT Fusion produces these fused kernels.

• Primary Benefit: Eliminates kernel launch overhead and improves data locality.
• Mechanism: Hand-written or compiler-generated kernels that keep intermediate results in registers or shared memory.

Algorithms and predictive models that guide a compiler's fusion decisions by estimating fusion profitability. They determine which operator groups to fuse.

• Factors: Data dependency, memory access patterns, operation type, and hardware characteristics.
• Goal: Avoid negative effects like increased register pressure or reduced parallelism.
• Component: Part of a Fusion Planner that explores a search space of possible fusions.

Two fundamental patterns for combining operators within a computational graph.

• Vertical Fusion: Merges a producer operator with its consumer (sequential dependence). Example: MatMul -> Add (Bias) fusion.
• Horizontal Fusion: Merges multiple independent operators that consume the same input or execute in parallel. Example: Fusing two separate ReLU operations applied to different tensors.

AOT Fusion strategies must account for both patterns.