How AI Compilers Work: Front-End Fusion, Back-End Kernels, AutoTuning

Why can't a PyTorch model run directly on a GPU? Between them sits the AI compiler — it translates "the framework's computational graph" into "efficient kernels on hardware." This article explains the three-layer structure, operator fusion, and autotuning, linking to MLIR Architecture and Triton in Practice.

1. Why an AI Compiler

Between framework and hardware lies a semantic gap:

PyTorch (high-level graph)        GPU instructions (low-level kernels)
      │                                    ▲
      └────────  AI compiler  ─────────────┘
        graph opt → IR → kernel gen/tuning

Translating op-by-op results in: each operator launched separately, intermediate tensors repeatedly hitting VRAM, terrible bandwidth utilization (see Operator & Runtime Ledger). The compiler exists to optimize globally while translating.

2. Three-Layer Structure

Layer	Job	Typical techniques
Front-end	graph optimization	operator fusion, constant folding, dead-code elimination
Middle	unified IR	MLIR, TOSA, ONNX
Back-end	kernel generation & scheduling	codegen, autotuning, memory planning

Takeaway: front-end optimizations are hardware-agnostic (benefit any card); back-end is hardware-bound (each vendor implements its own). This is the core motivation for MLIR's layered design — write the front-end once, implement back-ends separately.

3. Operator Fusion: The Most Important Front-End Optimization

The gain is already quantified in Operator Ledger (2 writes + 2 reads saved). By type:

Fusion type	Example	Gain
Adjacent fusion	Conv+BN, GELU+Add	one fewer intermediate in VRAM
Cross-structure fusion	full attention fusion (FlashAttention)	I/O complexity down a notch
Compute-comm fusion	gradient sharding & comm overlap	comm hidden inside compute

Takeaway: fusion is "putting the I/O ledger into the op boundary" — for bandwidth-bound ops (most), fusing one intermediate saves 30-50%.

4. Back-End Autotuning: Three Routes

Kernel parameters (block size, loop unrolling, SMEM partition) explode combinatorially; manual tuning is infeasible:

Route	Principle	Example
Hand templates + heuristics	preset templates, rule-based selection	cuDNN/cuBLAS heuristics
Auto search	compile-time enumeration + measured scoring	TVM AutoTVM, Triton autotune
ML-driven	learn a kernel performance model	XLA, some research systems

Takeaway: production prefers "templates + heuristics" (cuDNN's algorithm selector is most stable); research/new chips prefer "auto search." Triton's @triton.autotune is the practical entry (see Triton Learning Path).

5. PyTorch 2.0 Compile Stack: torch.compile

torch.compile(model)
   ├── FX graph capture (Python layer → computational graph)
   ├── TorchInductor (IR optimization + codegen)
   └── generates Triton / C++ kernels

• Graph capture solves "dynamic Python is hard to compile" — freeze it into a graph first;
• TorchInductor generates Triton kernels by default, balancing performance and readability;
• Gain: end-to-end training speedup of 10-40% (compile + fusion + kernel optimization), echoing PyTorch Compile Stack.

6. Checklist

1. If training/inference performance is stuck, check whether ops "land in VRAM every step" — enable torch.compile or write fused kernels;
2. On new chips/new ops, try Triton autotune first; freeze template params for production;
3. For cross-hardware deployment, use MLIR/TOSA as the middle layer — don't write one front-end per hardware;
4. Profile before/after each compile optimization (see Acceptance & Baselines).

---

Organized and written by this site for learning purposes; specific versions and measurements are subject to on-site conditions.