MLIR (Multi-Level Intermediate…

Home› Knowledge Base› MLIR (Multi-Level Intermediate Representation)

MLIR (Multi-Level Intermediate Representation) is an extensible compiler infrastructure for building domain-specific compilers — developed by Google and now part of LLVM, MLIR enables ML frameworks to define custom optimizations and target diverse hardware through a flexible, composable IR system.

What Is MLIR?

Definition: Framework for building and composing compiler IRs.
Origin: Google, now LLVM project.
Purpose: Simplify compiler construction for ML and beyond.
Key Feature: Multiple abstraction levels in one framework.

Why MLIR Matters

Fragmentation: Each framework had its own compiler stack.
Reuse: Share optimizations across frameworks/targets.
Flexibility: Custom dialects for domain-specific needs.
Hardware Diversity: Single path to many accelerators.
Performance: Systematic optimization opportunities.

MLIR Architecture

Dialect System:

<svg viewBox="0 0 760 470" xmlns="http://www.w3.org/2000/svg" font-family="-apple-system,Segoe UI,Roboto,sans-serif">
  <rect x="0" y="0" width="760" height="470" fill="#0d1117"/>
  <text x="380" y="28" fill="#e6edf3" font-size="21" font-weight="700" text-anchor="middle">MLIR — Multi-Level Intermediate Representation</text>
  <text x="380" y="48" fill="#8b98a5" font-size="12" text-anchor="middle">extensible compiler infrastructure with multiple abstraction levels (dialects) in one framework</text>

  <!-- Main: dialect stack diagram -->
  <rect x="30" y="65" width="700" height="270" rx="6" fill="#080d14" stroke="#233043" stroke-width="1"/>
  <text x="380" y="86" fill="#e6edf3" font-size="11" font-weight="600" text-anchor="middle">MLIR Dialect Lowering Stack</text>

  <!-- High-level dialects (top) -->
  <rect x="80" y="100" width="140" height="38" rx="4" fill="#0b1220" stroke="#a78bfa" stroke-width="1.2"/>
  <text x="150" y="118" fill="#c4b5fd" font-size="9" font-weight="600" text-anchor="middle">torch / stablehlo</text>
  <text x="150" y="132" fill="#6b7684" font-size="7.5" text-anchor="middle">ML framework graphs</text>

  <rect x="250" y="100" width="140" height="38" rx="4" fill="#0b1220" stroke="#a78bfa" stroke-width="1.2"/>
  <text x="320" y="118" fill="#c4b5fd" font-size="9" font-weight="600" text-anchor="middle">linalg / tensor</text>
  <text x="320" y="132" fill="#6b7684" font-size="7.5" text-anchor="middle">structured ops (matmul)</text>

  <rect x="420" y="100" width="140" height="38" rx="4" fill="#0b1220" stroke="#a78bfa" stroke-width="1.2"/>
  <text x="490" y="118" fill="#c4b5fd" font-size="9" font-weight="600" text-anchor="middle">tosa / mhlo</text>
  <text x="490" y="132" fill="#6b7684" font-size="7.5" text-anchor="middle">hardware-agnostic ops</text>

  <rect x="590" y="100" width="110" height="38" rx="4" fill="#0b1220" stroke="#a78bfa" stroke-width="1.2"/>
  <text x="645" y="118" fill="#c4b5fd" font-size="9" font-weight="600" text-anchor="middle">arith / math</text>
  <text x="645" y="132" fill="#6b7684" font-size="7.5" text-anchor="middle">scalar ops</text>

  <text x="690" y="108" fill="#c4b5fd" font-size="7">high-level</text>

  <!-- Arrows down -->
  <path d="M380,140 L380,155" fill="none" stroke="#3a4453" stroke-width="1"/>
  <polygon points="377,153 380,159 383,153" fill="#3a4453"/>

  <!-- Mid-level dialects -->
  <rect x="120" y="162" width="130" height="38" rx="4" fill="#0b1220" stroke="#34d399" stroke-width="1"/>
  <text x="185" y="180" fill="#6ee7b7" font-size="9" font-weight="600" text-anchor="middle">scf / affine</text>
  <text x="185" y="194" fill="#6b7684" font-size="7.5" text-anchor="middle">loops, tile, fuse</text>

  <rect x="270" y="162" width="130" height="38" rx="4" fill="#0b1220" stroke="#34d399" stroke-width="1"/>
  <text x="335" y="180" fill="#6ee7b7" font-size="9" font-weight="600" text-anchor="middle">memref / bufferize</text>
  <text x="335" y="194" fill="#6b7684" font-size="7.5" text-anchor="middle">memory allocation</text>

  <rect x="420" y="162" width="130" height="38" rx="4" fill="#0b1220" stroke="#34d399" stroke-width="1"/>
  <text x="485" y="180" fill="#6ee7b7" font-size="9" font-weight="600" text-anchor="middle">vector</text>
  <text x="485" y="194" fill="#6b7684" font-size="7.5" text-anchor="middle">SIMD vectorization</text>

  <rect x="570" y="162" width="120" height="38" rx="4" fill="#0b1220" stroke="#34d399" stroke-width="1"/>
  <text x="630" y="180" fill="#6ee7b7" font-size="9" font-weight="600" text-anchor="middle">gpu dialect</text>
  <text x="630" y="194" fill="#6b7684" font-size="7.5" text-anchor="middle">kernel launch</text>

  <text x="690" y="172" fill="#6ee7b7" font-size="7">mid-level</text>

  <!-- Arrows down -->
  <path d="M380,202 L380,217" fill="none" stroke="#3a4453" stroke-width="1"/>
  <polygon points="377,215 380,221 383,215" fill="#3a4453"/>

  <!-- Low-level dialects -->
  <rect x="150" y="224" width="130" height="38" rx="4" fill="#0b1220" stroke="#f59e0b" stroke-width="1"/>
  <text x="215" y="242" fill="#fbbf24" font-size="9" font-weight="600" text-anchor="middle">llvm dialect</text>
  <text x="215" y="256" fill="#6b7684" font-size="7.5" text-anchor="middle">LLVM IR equivalent</text>

  <rect x="300" y="224" width="130" height="38" rx="4" fill="#0b1220" stroke="#f59e0b" stroke-width="1"/>
  <text x="365" y="242" fill="#fbbf24" font-size="9" font-weight="600" text-anchor="middle">nvvm / rocdl</text>
  <text x="365" y="256" fill="#6b7684" font-size="7.5" text-anchor="middle">GPU backend IR</text>

  <rect x="450" y="224" width="130" height="38" rx="4" fill="#0b1220" stroke="#f59e0b" stroke-width="1"/>
  <text x="515" y="242" fill="#fbbf24" font-size="9" font-weight="600" text-anchor="middle">spirv</text>
  <text x="515" y="256" fill="#6b7684" font-size="7.5" text-anchor="middle">Vulkan compute</text>

  <text x="690" y="240" fill="#fbbf24" font-size="7">low-level</text>

  <!-- Arrows down to hardware -->
  <path d="M380,264 L380,279" fill="none" stroke="#3a4453" stroke-width="1"/>
  <polygon points="377,277 380,283 383,277" fill="#3a4453"/>

  <!-- Hardware targets -->
  <rect x="100" y="286" width="80" height="30" rx="3" fill="#1a0a0a" stroke="#f87171" stroke-width="0.8"/>
  <text x="140" y="305" fill="#fca5a5" font-size="8" text-anchor="middle">x86 / ARM</text>

  <rect x="200" y="286" width="80" height="30" rx="3" fill="#1a0a0a" stroke="#f87171" stroke-width="0.8"/>
  <text x="240" y="305" fill="#fca5a5" font-size="8" text-anchor="middle">NVIDIA GPU</text>

  <rect x="300" y="286" width="80" height="30" rx="3" fill="#1a0a0a" stroke="#f87171" stroke-width="0.8"/>
  <text x="340" y="305" fill="#fca5a5" font-size="8" text-anchor="middle">AMD GPU</text>

  <rect x="400" y="286" width="80" height="30" rx="3" fill="#1a0a0a" stroke="#f87171" stroke-width="0.8"/>
  <text x="440" y="305" fill="#fca5a5" font-size="8" text-anchor="middle">TPU</text>

  <rect x="500" y="286" width="80" height="30" rx="3" fill="#1a0a0a" stroke="#f87171" stroke-width="0.8"/>
  <text x="540" y="305" fill="#fca5a5" font-size="8" text-anchor="middle">custom ASIC</text>

  <rect x="600" y="286" width="80" height="30" rx="3" fill="#1a0a0a" stroke="#f87171" stroke-width="0.8"/>
  <text x="640" y="305" fill="#fca5a5" font-size="8" text-anchor="middle">FPGA</text>

  <!-- Bottom: key insight -->
  <rect x="30" y="348" width="700" height="82" rx="6" fill="#080d14" stroke="#233043" stroke-width="1"/>
  <text x="380" y="368" fill="#e6edf3" font-size="10" font-weight="600" text-anchor="middle">Why MLIR Matters for AI Hardware</text>
  <text x="50" y="390" fill="#8b98a5" font-size="8.5">• One framework replaces N separate compiler stacks (TVM, XLA, Triton internals all converging on MLIR)</text>
  <text x="50" y="406" fill="#8b98a5" font-size="8.5">• Custom dialects let chip startups add hardware-specific ops without rewriting the entire compiler</text>
  <text x="50" y="422" fill="#8b98a5" font-size="8.5">• Progressive lowering: optimize at each level, verify correctness, then lower further</text>

  <text x="380" y="452" fill="#6b7684" font-size="11" text-anchor="middle">MLIR is the LLVM of AI — every new accelerator compiler is being built on it (2024+).</text>
</svg>

Key Dialects:

Dialect      | Purpose
-------------|----------------------------------
tf           | TensorFlow operations
torch        | PyTorch operations
stablehlo    | Stable HLO (cross-framework)
linalg       | Generic linear algebra
tensor       | Tensor operations
scf          | Structured control flow
memref       | Memory references
arith        | Arithmetic operations
gpu          | GPU abstractions
llvm         | LLVM IR target

How MLIR Works

Example Lowering:

Input (PyTorch):
  y = torch.matmul(A, B)

↓ torch dialect
  %y = torch.matmul %A, %B

↓ linalg dialect
  %y = linalg.matmul ins(%A, %B) outs(%C)

↓ scf/memref
  scf.for %i = 0 to %M {
    scf.for %j = 0 to %N {
      scf.for %k = 0 to %K {
        %a = memref.load %A[%i, %k]
        %b = memref.load %B[%k, %j]
        %c = memref.load %C[%i, %j]
        %prod = arith.mulf %a, %b
        %sum = arith.addf %c, %prod
        memref.store %sum, %C[%i, %j]
      }
    }
  }

↓ Target (LLVM or GPU)

MLIR in ML Ecosystem

Framework Integration:

Framework        | MLIR Usage
-----------------|----------------------------------
TensorFlow       | XLA uses MLIR (StableHLO)
PyTorch          | torch-mlir, torch.compile
JAX              | JAX → StableHLO → MLIR
IREE             | End-to-end MLIR compiler
OpenXLA          | Cross-framework compilation

torch-mlir Example:

import torch
import torch_mlir

class MyModel(torch.nn.Module):
    def forward(self, x, y):
        return torch.matmul(x, y)

model = MyModel()
example_inputs = (torch.randn(4, 8), torch.randn(8, 16))

# Export to MLIR
mlir_module = torch_mlir.compile(
    model,
    example_inputs,
    output_type="stablehlo"
)

print(mlir_module)

Advantages of MLIR

For Compiler Developers:

Benefit              | Description
---------------------|----------------------------------
Reusable passes      | Share optimizations across dialects
Type system          | Rich, extensible type support
Verification         | Built-in IR validation
Debugging            | Great tooling (mlir-opt, etc.)
Documentation        | Operation definitions are docs

For Hardware Vendors:

Benefit              | Description
---------------------|----------------------------------
Single entry point   | Support TF, PyTorch, JAX via MLIR
Focus on backend     | Framework integration handled
Community            | Leverage ecosystem work
Portability          | Standard representation

Common Passes

Pass                   | Purpose
-----------------------|----------------------------------
Canonicalization       | Simplify patterns
CSE                    | Common subexpression elimination
Inlining               | Inline function calls
Loop fusion            | Combine loops
Tiling                 | Partition for parallelism
Bufferization          | Convert tensors to memrefs

MLIR is the foundation of modern ML compiler stacks — by providing a flexible, extensible framework for building domain-specific compilers, it enables the systematic optimization needed to extract maximum performance from diverse AI hardware.

mlircompilerintermediatedialectloweringxla

Explore 500+ Semiconductor & AI Topics

From EUV lithography to CUDA optimization — search the full knowledge base or chat with our AI assistant.

🔍 Search Topics 💬 Ask CFSGPT 📚 Browse All