<!-- Canonical: https://lucioduran.com/blog/llvm-mlir-compiler-infrastructure-custom-dialects
Language: en
Content-SHA256: 937d0e41a8afc754f68c747ea0aeca70090b1b8cc1a6d1837dded647693673db
-->

Compilers

# LLVM MLIR: Building Custom Compiler Dialects for Domain-Specific Optimization

2026-01-2711 min read

Lucio Durán
Engineering Manager & AI Solutions Architect

Also available in: [Español](<https://lucioduran.com/es/blog/llvm-mlir-compiler-infrastructure-custom-dialects>)

## The Problem: Why LLVM IR Wasn't Enough

[LLVM](<https://llvm.org/>) 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:

```mlir
// 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

```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:

```cpp
// 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:

```cpp
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:

```mlir
// 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
# 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

```bash
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

```bash
# 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.

## Frequently asked questions

### What is MLIR?

MLIR (Multi-Level Intermediate Representation) is a compiler infrastructure within the LLVM project that provides a flexible, extensible framework for building domain-specific compilers using composable 'dialects', collections of types, operations, and transformation rules.

### How is MLIR different from LLVM IR?

LLVM IR is a single, fixed-level representation designed for general-purpose CPU code generation. MLIR supports multiple abstraction levels simultaneously, letting you represent high-level domain concepts (tensor operations, hardware intrinsics) alongside low-level code.

### What are MLIR dialects?

Dialects are modular collections of operations, types, and attributes in MLIR. Each dialect defines its own vocabulary, for example, the 'linalg' dialect has operations for linear algebra, while the 'gpu' dialect has operations for GPU kernels. You can mix dialects freely in the same IR.

### Who uses MLIR in production?

Google (TensorFlow/XLA, IREE), AMD (ROCm compiler stack), Intel (oneAPI), Modular (Mojo language), and several ML hardware startups use MLIR as their compiler backbone.

## Tips

1. **Tip:** When defining custom MLIR ops, always use TableGen (` .td ` files) instead of hand-writing C++ op classes. TableGen auto-generates verifiers, builders, parsers, and printers. Writing ops by hand is a debugging nightmare.
2. **Tip:** MLIR's ` -mlir-print-ir-after-all ` flag dumps the IR after every pass. Pipe it through ` diff ` to see exactly what each transformation does. This is the single most useful debugging technique for pass development.
3. **Tip:** Start with the ` func ` and ` arith ` dialects when learning MLIR. They're simple enough to understand in an afternoon but demonstrate all the core concepts, regions, blocks, SSA values, and type systems.
4. **Tip:** Use ` mlir-opt --canonicalize ` as your first optimization pass. Canonicalization patterns (constant folding, dead code elimination, identity removal) are registered per-op and run to fixpoint. It catches 80% of obvious redundancies.

- llvm
- mlir
- compilers
- dialects
- optimization
- machine-learning
- dsl
- intermediate-representation

[Back to Blog](<https://lucioduran.com/blog>)