← Tiled Thoughts
·5 min read

How Rewriting works in MLIR

Contents
  1. TL;DR
  2. Part 1: The moving pieces
  3. RewritePatternSet and PatternRewriter
  4. Greedy vs. Conversion
  5. Part 2: Your first greedy rewrite pattern.
  6. Pattern Definition
  7. Part 3 : Running it: Tiny IR + Command
  8. Build and run
  9. Result
  10. Part 4: Dialect Conversion in a nutshell
  11. Core ingredients
  12. Example: Convert toy.addi to arith.addi
  13. Key differences from greedy patterns:
  14. Part 5 : Folding vs. Patterns
  15. Part 6: Match helpers, benefits, and ordering
  16. Part 7: Debugging and guardrails
  17. Part 8: Mini LIT Test Example
  18. Conclusion
  19. Further Reading

An in-depth look at the rewriting mechanisms in MLIR and how they enable powerful optimizations.

If you only remember one thing from this post: rewriting in MLIR is “find a pattern, make a change, repeat until more changes can’t be made”, with two key components:

  1. Greedy pattern application (canonicalization and local clenups), and
  2. Dialect conversion (legalize/convert regions with invariants about the legal forms of ops).

TL;DR

Part 1: The moving pieces

RewritePatternSet and PatternRewriter

Greedy vs. Conversion

Greedy (Canonicalization and Local Rewrites)

Conversion (Dialect Conversion)

Part 2: Your first greedy rewrite pattern.

Let’s fold away arith.addi %x, 0: i32 into just %x. Yeah, it’s trivial, and MLIR’s canonicalization already does this, but it’s a great starting point.

Pattern Definition

#include "mlir/IR/PatternMatch.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Pass/Pass.h"

using namespace mlir;

namespace {
struct FoldAddIWithZeroPattern : OpRewritePattern<arith::AddIOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(arith::AddIOp op,
                                PatternRewriter &rewriter) const override {
    auto isZeroConst = [](Value v) {
      if (auto c = v.getDefiningOp<arith::ConstantOp>()) {
        if (auto intAttr = dyn_cast<IntegerAttr>(c.getValue()))
          return intAttr.getValue().isZero();
      }
      return false;
    };

    Value lhs = op.getLhs();
    Value rhs = op.getRhs();

    if (isZeroConst(lhs)) {
      rewriter.replaceOp(op, rhs);
      return success();
    }
    if (isZeroConst(rhs)) {
      rewriter.replaceOp(op, lhs);
      return success();
    }
    return failure();
  }
};
} // namespace

struct FoldAddIZeroPass
    : public PassWrapper<FoldAddIZeroPass, OperationPass<func::FuncOp>> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(FoldAddIZeroPass)
  StringRef getArgument() const override { return "fold-addi-zero"; }
  void runOnOperation() override {
    MLIRContext *ctx = &getContext();
    RewritePatternSet patterns(ctx);
    patterns.add<FoldAddIWithZeroPattern>(ctx);

    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
      signalPassFailure();
  }
};

std::unique_ptr<Pass> mlir::createFoldAddIZeroPass() {
  return std::make_unique<FoldAddIZeroPass>();
}

Part 3 : Running it: Tiny IR + Command

Given this tiny IR in test.mlir:

module {
  func.func @foo(%x : i32) -> i32 {
    %c0 = arith.constant 0 : i32
    %y  = arith.addi %x, %c0 : i32
    return %y : i32
  }
}

Build and run

mlir-opt test.mlir --pass-pipeline="builtin.module(func.func(fold-addi-zero))"

Tip: Use –mlir-print-ir-after-all/–mlir-print-ir-before-all to see IR after each pass.

Result

The output IR will have the addition folded away:

// -----// IR Dump Before FoldAddIZeroPass (fold-addi-zero) //----- //
func.func @foo(%arg0: i32) -> i32 {
  %c0_i32 = arith.constant 0 : i32
  %0 = arith.addi %arg0, %c0_i32 : i32
  return %0 : i32
}

module {
  func.func @foo(%arg0: i32) -> i32 {
    return %arg0 : i32
  }
}

Part 4: Dialect Conversion in a nutshell

Greedy rewrites are great for local simplifications, but what if you want to convert ops from one dialect to another while ensuring certain invariants?

Core ingredients

Example: Convert toy.addi to arith.addi


struct ToyAddLowering : public OpConversionPattern<Toy::AddIOp> {
  using OpConversionPattern<toy::TransposeOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(Toy::AddIOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // adaptor carries already-converted operands/types if a TypeConverter is
    // used
    auto resTy = adaptor.getLhs().getType(); // i32 (post-conversion if any)
    auto sum = rewriter.create<arith::AddIOp>(
        op.getLoc(), resTy, adaptor.getLhs(), adaptor.getRhs());
    rewriter.replaceOp(op, sum.getResult());
    return success();
  }
};

struct LowerToyPass
    : public PassWrapper<LowerToyPass, OperationPass<ModuleOp>> {
  void runOnOperation() final {
    MLIRContext *ctx = &getContext();

    // 1) What is legal?
    ConversionTarget target(*ctx);
    target.addLegalDialect<arith::ArithDialect>();
    target.addIllegalDialect<Toy::ToyDialect>();

    // 2) (Optional) TypeConverter, if you need to rewrite types.
    TypeConverter typeConverter; // no-op here

    // 3) Patterns
    RewritePatternSet patterns(ctx);
    patterns.add<ToyAddLowering>(typeConverter, ctx);

    // 4) Apply
    if (failed(applyPartialConversion(getOperation(), target,
                                      std::move(patterns))))
      signalPassFailure();
  }
};

Key differences from greedy patterns:

Part 5 : Folding vs. Patterns

Part 6: Match helpers, benefits, and ordering

Part 7: Debugging and guardrails

Part 8: Mini LIT Test Example

// RUN: mlir-opt %s --pass-pipeline="builtin.module(func.func(fold-addi-zero))" | FileCheck %s
module {
  func.func @foo(%x : i32) -> i32 {
    %c0 = arith.constant 0 : i32
    %y  = arith.addi %x, %c0 : i32
    return %y : i32
  }
}

// CHECK-LABEL: func @foo(
// CHECK-NOT: arith.addi
// CHECK: return [[X:%.*]] : i32
// CHECK: }

This test ensures that after running the fold-addi-zero pass, there are no arith.addi operations left in the function.

Conclusion

Further Reading