MLIR (Multi-Level Intermediate Representation)

The fundamental unit of abstraction in MLIR is a dialect, which defines a set of operations, types, and attributes for a specific domain or level of abstraction. This allows MLIR to natively represent high-level frameworks (like TensorFlow or PyTorch), mid-level linear algebra, and low-level hardware instructions within a single IR. Compiler engineers can create custom dialects to model new accelerators or domain-specific languages, enabling progressive lowering from high-level semantics to target-specific code.

MLIR is designed for multi-level IR, allowing a single program to be represented at multiple abstraction levels simultaneously. Compiler passes progressively transform the IR from higher-level, domain-specific dialects (e.g., linalg for linear algebra) to lower-level, hardware-oriented dialects (e.g., llvm or gpu). This enables:

• Modular optimizations: Apply high-level, domain-specific optimizations early, and low-level, hardware-specific optimizations later.
• Reusable transformations: Common passes like loop unrolling or vectorization can be written once and applied to different dialects.
• Clear separation of concerns: Different compiler teams can work on different abstraction layers.

MLIR provides a single framework to manage compilation for heterogeneous hardware. It can represent computation destined for CPUs, GPUs, NPUs, and custom ASICs within the same module. Key mechanisms include:

• Hardware abstraction layers (HAL): Dialects like gpu and spirv model GPU execution and memory.
• Delegation and partitioning: The IR can be split, with subgraphs delegated to specific accelerator backends.
• Target-independent optimizations: Many optimizations are performed before committing to a specific hardware target, maximizing code reuse.

MLIR uses a declarative system for defining compiler transformations. Engineers write rewrite rules in TableGen (a DSL) that specify how to match a subgraph of operations and replace it with a more optimal one. This enables:

• Reliable optimizations: Rules are applied exhaustively until no more matches are found.
• Examples: Fusing a convolution, batch norm, and ReLU into a single operation, or converting a high-level tensor operation into a loop nest.
• Maintainability: The rule-based system makes compiler passes easier to write, test, and understand than imperative C++ transformation code.

MLIR acts as a hub within the compiler ecosystem. It features importers and exporters for many existing formats, allowing it to ingest models from frameworks like TensorFlow (via TensorFlow Graph), PyTorch (via Torch-MLIR), and ONNX. This enables:

• Vendor-neutral optimization pipeline: Models from any frontend can benefit from MLIR's optimization passes.
• Backend code generation: MLIR can lower optimized IR to LLVM IR for CPUs, generate SPIR-V for GPUs, or produce custom C++ for specific accelerators.
• Toolchain reuse: Existing LLVM tools (debuggers, profilers) can be used on code generated via the LLVM dialect.

MLIR's type system and operation semantics support powerful static analysis at compile time. A core capability is shape inference, which propagates tensor dimensions through the computational graph without executing it. This is critical for:

• Memory planning: Determining buffer sizes for static memory allocation on edge devices.
• Optimization validity: Enabling transformations that depend on known tensor shapes (e.g., loop tiling).
• Early error detection: Catching shape mismatches during compilation, not at runtime. MLIR's structure facilitates these analyses across different dialects.

MLIR includes first-class support for generating high-performance code for NVIDIA GPUs. The NVGPU and NVVM dialects allow compilers to represent and optimize GPU-specific concepts at an abstract level before lowering to LLVM's NVPTX backend. This enables:

• Structured op abstractions for GPU kernels (e.g., gpu.launch, warp-level operations).
• Memory hierarchy management explicit modeling of shared, global, and register memory.
• Automatic optimization of memory transfers and kernel scheduling across multiple levels of IR.

Google's Tensor Processing Units (TPUs) are supported through MLIR's TPU dialect and integration within the TensorFlow/XLA stack. MLIR facilitates the mapping of high-level linear algebra operations to the TPU's systolic array architecture. This involves:

• Spatial partitioning of computations across a TPU pod's cores.
• HBM and on-chip memory management for massive tensors.
• Lowering to the TPU's VLIW instruction set for maximum utilization of the matrix multiply units.

AMD's ROCm open software platform uses MLIR for AI compiler stacks. The MLIR AMDGPU dialect provides abstractions for AMD GPU hardware (like CDNA and RDNA architectures). Projects like MLIR-based frontends for frameworks target ROCm, enabling:

• Portable GPU kernel generation from high-level ML frameworks.
• Leveraging AMD's matrix cores (like WMMA operations) through MLIR's vector dialect.
• Integration with the ROCm driver and runtime for final code generation and execution.

For edge AI, MLIR is used to target Arm's Ethos neural processing units and the CMSIS-NN library for Cortex-M CPUs. Specialized MLIR dialects model NPU command streams and micro-architectural constraints, allowing compilers like Apache TVM (which integrates MLIR) to:

• Perform hardware-aware graph partitioning and operator legalization.
• Schedule data flows between CPU and NPU.
• Generate highly optimized, static code that directly controls the accelerator's tensor processing elements.

An intermediate data structure used within a compiler to represent a program's logic between the high-level source code and low-level machine instructions. MLIR is a novel IR designed to be multi-level, meaning it can represent computations at multiple abstraction levels (e.g., linear algebra, loops, hardware-specific instructions) within a single, unified framework. This enables seamless progressive lowering from a high-level model graph down to efficient machine code.

A domain-specific compiler for linear algebra, developed for TensorFlow. XLA performs graph optimizations (like operation fusion) and target-specific code generation to accelerate model execution on CPUs, GPUs, and TPUs. MLIR shares similar goals but is designed as a general compiler infrastructure. The TensorFlow ecosystem is increasingly adopting MLIR to replace and extend parts of XLA, leading to projects like MLIR-HLO (High-Level Operations), which provides MLIR dialects for XLA-style optimizations.

A standardized software interface that abstracts the specifics of underlying hardware accelerators. In MLIR-based compilers like IREE, the HAL is implemented as an MLIR dialect. This HAL dialect allows the compiler to reason about devices, command buffers, and memory allocation in a hardware-agnostic way before final code generation. It enables a single compiled model to target multiple accelerators (e.g., GPU, NPU, DSP) through different backend implementations of the HAL interface.