Intermediate Representation (IR)

An IR's abstraction level defines its distance from source code or hardware. High-Level IRs (HIR) preserve rich semantic information (e.g., loops, function calls) for architecture-agnostic optimizations. Mid-Level IRs (MIR) are often control-flow graphs suitable for analysis and transformation. Low-Level IRs (LIR) resemble assembly, exposing hardware details like registers and instruction selection. Compilers like LLVM use a tiered system, progressively lowering from a HIR to a LIR.

Most modern ML IRs represent computation as a Directed Acyclic Graph (DAG). Nodes represent operations (ops) or tensors, and edges represent data dependencies. This structure enables:

• Topological sorting for correct execution order.
• Easy analysis of data flow and operator lifetimes for static memory planning.
• Efficient application of transformations like operator fusion and dead code elimination.
• Frameworks like TensorFlow (GraphDef), PyTorch (TorchScript), and ONNX use this graph paradigm.

A core purpose of an IR is to decouple model definition from hardware execution. A platform-independent IR (e.g., ONNX, MLIR's dialect-ir) allows:

• A single model definition to target CPUs, GPUs, NPUs, and microcontrollers.
• Hardware vendors to develop specialized compiler backends that lower the IR to their native instructions.
• Enables retargeting without modifying the original model, a key tenet of frameworks like Apache TVM and Google's XLA.

Advanced IRs provide mechanisms to model diverse computational domains. MLIR pioneered the dialect system, where the IR can be extended with custom operations, types, and attributes. This allows:

• A Neural Network dialect for high-level ML ops.
• An LLVM dialect for low-level CPU/GPU code.
• Domain-specific dialects for linear algebra, affine loops (polyhedral model), or quantum computing.
• Dialects can be progressively lowered and legalized to a final target, enabling complex, multi-stage compilation.

The IR is the substrate for compiler passes—transformations that analyze and rewrite the graph. Common passes include:

• Canonicalization: Simplifies patterns to a standard form.
• Constant Folding: Pre-computes operations on constant values.
• Common Subexpression Elimination (CSE): Removes redundant calculations.
• Algebraic Simplification: Applies mathematical identities (e.g., x*1 => x). Passes can be greedy, pattern-based, or guided by a cost model, and are often applied in a specific pipeline.

This axis defines when graph structure is finalized. Static IR (e.g., TensorFlow 1.x, ONNX) fixes the graph topology and tensor shapes ahead-of-time (AOT), enabling aggressive optimization. Dynamic IR (e.g., PyTorch's eager mode) is built at runtime, offering flexibility. Modern systems use tracing or scripting to capture dynamic execution into a static graph, or employ just-in-time (JIT) compilation to specialize a dynamic graph based on runtime shapes (shape inference).

A High-Level Graph IR represents a neural network as a computational graph of abstract operators, independent of hardware specifics. This is the initial representation after importing a model from frameworks like PyTorch or TensorFlow.

• Purpose: Enables hardware-agnostic graph-level optimizations like constant folding, dead code elimination, and operator fusion.
• Examples: ONNX (Open Neural Network Exchange) Graph, TensorFlow GraphDef, PyTorch FX Graph, and Relay IR (from Apache TVM).
• Characteristics: Nodes represent coarse-grained operations (e.g., Conv2D, MatMul); tensors have abstract data types and symbolic shapes.

A Lowered IR is the result of graph lowering, where high-level operators are decomposed into primitive, target-specific operations. It bridges the gap between abstract computation and hardware execution.

• Purpose: Exposes low-level details for backend-specific optimizations like loop tiling, vectorization, and efficient memory layout planning.
• Examples: MLIR's Affine Dialect for loop nests, Halide IR for image processing pipelines, TensorFlow XLA's HLO IR after target-specific lowering, and LLVM IR for CPU targets.
• Characteristics: Includes memory buffers, explicit loops, and address calculations. It is often linear or in Static Single Assignment (SSA) form.

A Hardware-Specific IR is a final, vendor-optimized representation that maps directly to the execution model of a particular accelerator, such as a GPU, NPU, or DSP.

• Purpose: Encodes instructions, memory hierarchies, and parallelism models (e.g., warps, workgroups) unique to the target silicon. Enables kernel auto-tuning and optimal operator dispatch.
• Examples: NVIDIA PTX (Parallel Thread Execution), CUDA Graphs, Google TPU's HLO (after further TPU lowering), ARM Compute Library's Graph IR, and proprietary NPU compiler IRs (e.g., Qualcomm SNPE, Apple Core ML).
• Characteristics: Contains explicit parallel constructs, synchronization primitives, and hardware intrinsic functions.

Specialized IRs encode optimizations for model compression, representing networks after quantization or pruning. They are crucial for efficient on-device inference.

• Purpose: Explicitly represents low-precision data types (e.g., int8, uint4) and sparse tensor encoding formats (e.g., CSR, Block-Sparse) to enable integer-only kernels and skip zero computations.
• Examples: TFLite's FlatBuffer schema with quantized operators, ONNX with QuantizeLinear/DequantizeLinear nodes, MLIR's TOSA dialect with quantization attributes, and frameworks using the Sparse Tensor Core extension in LLVM.
• Optimization: Quantization folding is a key pass that merges fake quantization nodes with linear ops in these IRs.

A Polyhedral IR uses mathematical polyhedra to model loop nests and their data dependencies, enabling powerful automatic transformations for compute-intensive, regular kernels like convolutions and matrix multiplications.

• Purpose: Provides a unified model for complex loop optimizations such as loop tiling, fusion, skewing, and interchange while preserving correctness.
• Examples: MLIR's Affine Dialect, the ISL (Integer Set Library) framework, and the core IR of the Polyhedral Model used in compilers like Pluto.
• Use Case: Critical for optimizing linear algebra operations on CPUs and GPUs, guided by a cost model and the Roofline model to achieve peak performance.

A Runtime IR or JIT IR is an intermediate representation that is constructed, optimized, and compiled during program execution. It enables dynamic optimizations based on runtime information.

• Purpose: Allows for Just-In-Time (JIT) compilation where optimizations like shape specialization, kernel selection, and graph rewriting can use concrete tensor shapes and values known only at runtime.
• Examples: PyTorch's TorchScript IR (for JIT), TensorFlow's XLA JIT compilation path, Apache TVM's runtime module, and Triton's JIT IR for GPU programming.
• Advantage: Balances optimization potential with flexibility, avoiding the need for Ahead-Of-Time (AOT) compilation to handle all possible input variations.

The process of progressively transforming a high-level, hardware-agnostic IR into a lower-level, target-specific IR or machine code. This involves mapping abstract operations to concrete hardware instructions and memory layouts.

• High-Level IR (HIR): Captures model architecture and operator semantics.
• Mid-Level IR (MIR): Platform-independent, optimized for analysis and transformations.
• Low-Level IR (LIR): Close to machine code, includes register allocation and instruction selection.
• Target: Enables a single model definition (e.g., PyTorch) to run on diverse backends like CPU, GPU, or NPU.

A compiler optimization that combines multiple sequential operations in a computational graph into a single, compound kernel. This reduces intermediate memory writes and kernel launch overhead.

• Fusion Patterns: Common fusions include element-wise operations (e.g., ReLU, Add) fused into a preceding convolution or matrix multiplication.
• Memory Bandwidth: Primary benefit is reduced DRAM traffic by keeping intermediate results in faster cache or registers.
• Kernel Launch Overhead: Eliminates the latency of scheduling multiple small GPU or NPU kernels.
• Example: Fusing a Conv2D, BatchNorm, and ReLU layer into one ConvBNReLU kernel.

A compile-time optimization that evaluates and replaces subgraphs comprised entirely of constants with their precomputed result tensor. This eliminates runtime computation.

• Static Computation: Moves work from runtime to compile-time.
• Graph Simplification: Removes nodes, making subsequent optimizations like dead code elimination more effective.
• Common Targets: Fixed tensor manipulations, shape calculations, and static control flow conditions.
• Example: The expression add(constant(5), constant(3)) is replaced with constant(8) in the IR.

An optimization pass that identifies and removes operations whose outputs do not contribute to the final output of the computational graph. This reduces both computation and memory usage.

• Liveness Analysis: Determines if a tensor's value is ever read by a node that reaches the graph's output.
• Side Effects: Operations with side effects (e.g., printing, variable updates) are typically not eliminated.
• Prunes Unused Branches: Can remove entire subgraphs from unused conditional branches.
• Impact: Critical for cleaning graphs after other transformations like constant folding.

An optimization that identifies redundant calculations of identical expressions and replaces duplicate computations with a reference to a single cached result.

• Expression Matching: The compiler must prove two subgraphs are functionally identical.
• Memory vs. Compute Trade-off: Stores the result of the first computation, trading memory for reduced FLOPs.
• Scope: Can be applied locally within a basic block or globally across the entire graph.
• Example: If sin(x) is computed in two separate branches of a conditional, CSE computes it once and reuses the value.

The process of propagating and determining the dimensions (shape) and data types (dtype) of all tensors in a computational graph based on the known properties of input nodes.

• Forward Propagation: Rules for each operator define output shape given input shapes (e.g., Conv2D output channels).
• Prerequisite for Optimization: Enables static memory planning, efficient kernel selection, and validation.
• Dynamic Shapes: Some frameworks support partial shape inference where dimensions can be dynamic (symbolic).
• Error Detection: Catches shape mismatches early (e.g., attempting to multiply incompatible matrices).