The Problem: Why LLVM IR Wasn't Enough

LLVM is effective for what it does: take moderately low-level IR and turn it into optimized machine code for a dozen architectures. But it has a fundamental limitation, there's only one level of IR, and it sits at roughly the abstraction level of C.

If you're building a compiler for a high-level language (say, a tensor processing DSL), you have to lower your high-level concepts to LLVM IR early. Once you do, information is lost. The LLVM optimizer sees loops and pointer arithmetic, not matrix multiplications and convolutions. It can't fuse two matrix operations that it doesn't know are matrix operations.

This is the "abstraction gap" that plagued ML compilers for years:

High-level DSL: matmul(A, B) + transpose(C)
 │
 ┌───────────────┤ (information loss here)
 ▼ │
LLVM IR: nested loops, gep instructions, loads/stores
 → LLVM can optimize the loops but can't see the linear algebra

Every ML framework solved this differently: TensorFlow had XLA, PyTorch had TorchScript then Torch-MLIR, TVM had its own IR, Halide had its own IR. Each one reinvented type systems, SSA construction, pass infrastructure, and canonicalization. Massive duplication.

MLIR's Core Insight: Dialects All the Way Down

MLIR's key insight is that you don't need one IR, you need a framework for building many IRs that can coexist and interoperate. That framework is the dialect system.

A dialect is a namespace containing:

  • Operations (the "instructions" of your IR)
  • Types (the type system for your domain)
  • Attributes (compile-time constants and metadata)
  • Interfaces (traits that operations can implement for generic transformations)

The magic is that operations from different dialects can exist in the same IR, and you can progressively lower from high-level dialects to low-level ones:

// High-level: linalg dialect (linear algebra operations)
%result = linalg.matmul ins(%A, %B : tensor<64x128xf32>, tensor<128x256xf32>)
 outs(%C : tensor<64x256xf32>) -> tensor<64x256xf32>

// Mid-level: affine dialect (polyhedral loop nests)
affine.for %i = 0 to 64 {
 affine.for %j = 0 to 256 {
 affine.for %k = 0 to 128 {
 %a = affine.load %A[%i, %k] : memref<64x128xf32>
 %b = affine.load %B[%k, %j] : memref<128x256xf32>
 %c = affine.load %C[%i, %j] : memref<64x256xf32>
 %prod = arith.mulf %a, %b : f32
 %sum = arith.addf %c, %prod : f32
 affine.store %sum, %C[%i, %j] : memref<64x256xf32>
 }
 }
}

// Low-level: llvm dialect (LLVM IR operations)
%ptr = llvm.getelementptr %base[%idx] : (!llvm.ptr, i64) -> !llvm.ptr
%val = llvm.load %ptr : !llvm.ptr -> f32

At each level, you can apply domain-specific optimizations that would be impossible at lower levels. The linalg level knows about tiling and fusion strategies for matrix operations. The affine level can apply polyhedral transformations. The llvm level does register allocation and instruction selection.

Building a Custom Dialect: A Toy Example

The following example builds a simple dialect for a hypothetical image processing DSL. It defines operations for loading images, applying filters, and compositing.

Step 1: Define the Dialect in TableGen

// ImageOps.td
include "mlir/IR/OpBase.td"
include "mlir/IR/AttrTypeBase.td"
include "mlir/Interfaces/SideEffectInterfaces.td"

// Define the dialect
def ImageDialect : Dialect {
 let name = "image";
 let cppNamespace = "::image";
 let summary = "A dialect for image processing operations";
 let description = [{
 The image dialect provides operations for loading, transforming,
 and compositing images with type-safe channel information.
 }];
}

// Define a custom type for images with channel info
def Image_ImageType : TypeDef<ImageDialect, "Image"> {
 let mnemonic = "img";
 let parameters = (ins
 "int64_t":$width,
 "int64_t":$height,
 "int64_t":$channels // 1=gray, 3=RGB, 4=RGBA
 );
 let assemblyFormat = "`<` $width `,` $height `,` $channels `>`";
}

// Load operation
def Image_LoadOp : Op<ImageDialect, "load", [Pure]> {
 let summary = "Load an image from a path";
 let arguments = (ins StrAttr:$path);
 let results = (outs Image_ImageType:$result);
 let assemblyFormat = "$path attr-dict `:` type($result)";
}

// Blur operation
def Image_BlurOp : Op<ImageDialect, "blur", [Pure]> {
 let summary = "Apply Gaussian blur";
 let arguments = (ins
 Image_ImageType:$input,
 F32Attr:$sigma
 );
 let results = (outs Image_ImageType:$result);
 let assemblyFormat = "$input `,` $sigma attr-dict `:` type($input) `->` type($result)";

 // Verify input and output have same dimensions
 let verifier = [{
 auto inType = getInput().getType().cast<ImageType>();
 auto outType = getResult().getType().cast<ImageType>();
 if (inType.getWidth() != outType.getWidth() ||
 inType.getHeight() != outType.getHeight())
 return emitOpError("input and output dimensions must match");
 return success();
 }];
}

// Composite operation
def Image_CompositeOp : Op<ImageDialect, "composite", [Pure]> {
 let summary = "Composite two images with blend mode";
 let arguments = (ins
 Image_ImageType:$foreground,
 Image_ImageType:$background,
 StrAttr:$blend_mode // "over", "multiply", "screen"
 );
 let results = (outs Image_ImageType:$result);
 let assemblyFormat = [{
 $foreground `,` $background `,` $blend_mode attr-dict `:`
 type($foreground) `,` type($background) `->` type($result)
 }];
}

Step 2: Write a Domain-Specific Optimization Pass

Now let's write a pass that fuses consecutive blur operations into a single blur with combined sigma:

// FuseBlursPass.cpp
#include "mlir/Pass/Pass.h"
#include "ImageOps.h"
#include <cmath>

namespace {
struct FuseBlursPass
 : public mlir::PassWrapper<FuseBlursPass,
 mlir::OperationPass<mlir::func::FuncOp>> {

 void runOnOperation() override {
 getOperation().walk([](image::BlurOp outerBlur) {
 // Check if the input to this blur is another blur
 auto innerBlur = outerBlur.getInput().getDefiningOp<image::BlurOp>();
 if (!innerBlur)
 return;

 // Check that the inner blur has no other users
 if (!innerBlur.getResult().hasOneUse())
 return;

 // Fuse: sigma_combined = sqrt(sigma1^2 + sigma2^2)
 float sigma1 = innerBlur.getSigma().convertToFloat();
 float sigma2 = outerBlur.getSigma().convertToFloat();
 float fusedSigma = std::sqrt(sigma1 * sigma1 + sigma2 * sigma2);

 // Replace the outer blur with a single blur on the inner's input
 mlir::OpBuilder builder(outerBlur);
 auto fusedBlur = builder.create<image::BlurOp>(
 outerBlur.getLoc(),
 outerBlur.getType(),
 innerBlur.getInput(),
 builder.getF32FloatAttr(fusedSigma));

 outerBlur.replaceAllUsesWith(fusedBlur.getResult());
 outerBlur.erase();
 if (innerBlur.getResult().use_empty())
 innerBlur.erase();
 });
 }

 mlir::StringRef getArgument() const final { return "image-fuse-blurs"; }
 mlir::StringRef getDescription() const final {
 return "Fuse consecutive blur operations";
 }
};
} // namespace

// Register the pass
std::unique_ptr<mlir::Pass> createFuseBlursPass() {
 return std::make_unique<FuseBlursPass>();
}

Step 3: Lower to Standard Dialects

Eventually, we need to lower our image operations to something LLVM can compile. Here's a sketch of lowering image.blur to loop nests:

struct BlurOpLowering : public mlir::OpConversionPattern<image::BlurOp> {
 using OpConversionPattern::OpConversionPattern;

 mlir::LogicalResult matchAndRewrite(
 image::BlurOp op,
 OpAdaptor adaptor,
 mlir::ConversionPatternRewriter &rewriter) const override {

 auto loc = op.getLoc();
 auto imgType = op.getInput().getType().cast<image::ImageType>();
 int64_t width = imgType.getWidth();
 int64_t height = imgType.getHeight();
 int64_t channels = imgType.getChannels();
 float sigma = op.getSigma().convertToFloat();

 // Compute kernel size from sigma
 int kernelSize = static_cast<int>(std::ceil(sigma * 3.0f)) * 2 + 1;
 int halfKernel = kernelSize / 2;

 // Allocate output buffer
 auto memrefType = mlir::MemRefType::get(
 {height, width, channels}, rewriter.getF32Type());
 auto output = rewriter.create<mlir::memref::AllocOp>(loc, memrefType);

 // Generate nested loops for convolution
 // (simplified, real implementation would generate separable passes)
 auto yLoop = rewriter.create<mlir::affine::AffineForOp>(
 loc, 0, height, 1);
 rewriter.setInsertionPointToStart(yLoop.getBody());

 auto xLoop = rewriter.create<mlir::affine::AffineForOp>(
 loc, 0, width, 1);
 rewriter.setInsertionPointToStart(xLoop.getBody());

 // ... kernel computation loop body ...

 rewriter.replaceOp(op, output);
 return mlir::success();
 }
};

The IR In Action

Here's what a complete pipeline looks like in our image dialect:

// Input IR (high-level image operations)
func.func @process_photo(%path: !image.img<1920, 1080, 3>) -> !image.img<1920, 1080, 3> {
 %photo = image.load "/photos/input.jpg" : !image.img<1920, 1080, 3>
 %blurred1 = image.blur %photo, 2.0 : !image.img<1920, 1080, 3> -> !image.img<1920, 1080, 3>
 %blurred2 = image.blur %blurred1, 3.0 : !image.img<1920, 1080, 3> -> !image.img<1920, 1080, 3>
 %overlay = image.load "/photos/overlay.png" : !image.img<1920, 1080, 4>
 %result = image.composite %overlay, %blurred2, "over" :
 !image.img<1920, 1080, 4>, !image.img<1920, 1080, 3> -> !image.img<1920, 1080, 3>
 return %result : !image.img<1920, 1080, 3>
}

// After blur fusion pass (sigma = sqrt(4 + 9) = 3.606):
func.func @process_photo(%path: !image.img<1920, 1080, 3>) -> !image.img<1920, 1080, 3> {
 %photo = image.load "/photos/input.jpg" : !image.img<1920, 1080, 3>
 %blurred = image.blur %photo, 3.606 : !image.img<1920, 1080, 3> -> !image.img<1920, 1080, 3>
 %overlay = image.load "/photos/overlay.png" : !image.img<1920, 1080, 4>
 %result = image.composite %overlay, %blurred, "over" :
 !image.img<1920, 1080, 4>, !image.img<1920, 1080, 3> -> !image.img<1920, 1080, 3>
 return %result : !image.img<1920, 1080, 3>
}

The fusion pass turned two blur operations into one, cutting the convolution work roughly in half. This optimization is trivial to express at the image dialect level, but would be nearly impossible to discover after lowering to loops.

Real-World MLIR: The ML Compiler Landscape

The biggest users of MLIR are ML compilers, and for good reason. ML workloads have a natural multi-level structure:

Framework level: torch.nn.Linear(256, 128)
 │
 ▼ TorchMLIR
Graph level: linalg.matmul + linalg.add (bias)
 │
 ▼ Tiling + Fusion
Tile level: scf.for loops with vectorized inner bodies
 │
 ▼ Hardware lowering
Hardware level: gpu.launch_func / vector.contract / amx.tile_mulf
 │
 ▼ LLVM lowering
Machine code: AVX-512 / CUDA PTX / AMX instructions

Each level has its own optimization opportunities:

  • Graph level: operator fusion (matmul + bias + relu = one kernel), layout optimization (NHWC vs NCHW)
  • Tile level: tile size selection, cache-aware scheduling, vectorization strategy
  • Hardware level: mapping to specific ISA extensions (AMX for Intel, Tensor Cores for NVIDIA)

IREE: Google's ML Runtime Built on MLIR

IREE (Intermediate Representation Execution Environment) is probably the most complete MLIR-based ML compiler in production. It takes TensorFlow, PyTorch, or JAX models, lowers them through MLIR dialects, and generates optimized code for CPUs, GPUs, and accelerators.

The compilation pipeline:

Input: StableHLO (from JAX/TF) or Torch-MLIR (from PyTorch)
 → IREE Input dialect (normalize all inputs)
 → IREE Flow dialect (partition into dispatches, manage buffers)
 → IREE Stream dialect (async execution scheduling)
 → IREE HAL dialect (hardware abstraction)
 → Target-specific codegen (LLVM CPU, SPIR-V, CUDA, Metal)

Each arrow is a series of MLIR passes that lower and optimize. The whole pipeline is around 80 passes, and every intermediate state is printable, verifiable MLIR.

Mojo: A Language Built on MLIR

Modular's Mojo programming language is built directly on MLIR. This means Mojo's compiler can express high-level abstractions (Python-like syntax) and low-level hardware operations (SIMD intrinsics) in the same IR, and the optimization pipeline can see both simultaneously.

# Mojo code that compiles through MLIR
fn matmul_tiled[M: Int, N: Int, K: Int](
 A: SIMD[DType.float32, M*K],
 B: SIMD[DType.float32, K*N],
) -> SIMD[DType.float32, M*N]:
 # This compiles to MLIR linalg ops, then gets tiled and vectorized
 # through the standard MLIR pipeline
 var C = SIMD[DType.float32, M*N](0)
 for i in range(M):
 for j in range(N):
 for k in range(K):
 C[i*N + j] += A[i*K + k] * B[k*N + j]
 return C

Mojo's compiler sees the SIMD types and loop structure, matches it against tiling patterns, and generates code that uses hardware vector instructions directly. The MLIR infrastructure handles the type checking, SSA construction, and pass scheduling, Mojo's compiler team focuses on the domain-specific patterns.

Getting Started: The Practical Path

For those wanting to experiment with MLIR, here is a recommended path:

1. Build MLIR from Source

git clone https://github.com/llvm/llvm-project.git
cd llvm-project
mkdir build && cd build

cmake -G Ninja ../llvm \
 -DLLVM_ENABLE_PROJECTS=mlir \
 -DLLVM_BUILD_EXAMPLES=ON \
 -DLLVM_TARGETS_TO_BUILD="host" \
 -DCMAKE_BUILD_TYPE=Release \
 -DLLVM_ENABLE_ASSERTIONS=ON

ninja -j$(nproc)

2. Play with mlir-opt

# Write a simple MLIR program
cat > test.mlir << 'EOF'
func.func @add(%arg0: f32, %arg1: f32) -> f32 {
 %0 = arith.addf %arg0, %arg1 : f32
 return %0 : f32
}
EOF

# Run canonicalization
./bin/mlir-opt --canonicalize test.mlir

# Lower to LLVM dialect
./bin/mlir-opt --convert-func-to-llvm --convert-arith-to-llvm test.mlir

# See all available passes
./bin/mlir-opt --help | grep "Passes:"

3. Study the Toy Tutorial

The MLIR repo includes a complete tutorial that builds a compiler for a toy language from scratch. It's in mlir/docs/Tutorials/Toy/ and covers everything from lexing to LLVM codegen through MLIR.

4. Look at Real Dialects

Study the linalg and affine dialect implementations in the MLIR source tree. They're well-documented and demonstrate production patterns for op definitions, transformations, and lowering.

Why This Matters Beyond Compilers

MLIR's dialect system is genuinely useful outside traditional compiler work:

  • Hardware design: CIRCT (Circuit IR Compilers and Tools) uses MLIR dialects to represent hardware description languages, replacing Verilog/VHDL toolchains.
  • Quantum computing: The quake dialect represents quantum circuits for optimization before mapping to quantum hardware.
  • Database engines: Query planners can use custom dialects to represent relational algebra and optimize execution plans.
  • Numerical simulation: The sparse_tensor dialect handles sparse matrix operations with automatic format selection.

The pattern is always the same: you have a domain with specific operations and optimization rules that are invisible at a lower abstraction level. MLIR gives you the infrastructure to express that domain without building everything from scratch.

MLIR isn't just a better LLVM IR. It's a framework for building compilers at any abstraction level, with infrastructure that took decades to develop in LLVM, available out of the box. If you're building anything that transforms structured representations, whether that's ML models, hardware descriptions, or database queries, MLIR is worth your time.