Understanding MLIR Rewrite Patterns

Photo by Ahmad Dirini on Unsplash

MLIR (Multi-Level Intermediate Representation) is a flexible compiler framework that allows us to define and transform operations through Rewrite Patterns. These patterns enable us to automatically replace one operation with another, which is useful for optimization, lowering, or dialect conversion.

Why Use Rewrite Patterns?

• Optimization: Convert expensive operations into more efficient ones.
• Dialect Conversion: Convert operations from one dialect to another.
• Canonicalization: Replace redundant patterns with cleaner representations.

In this blog, we will learn how to define and apply a simple MLIR rewrite pattern that converts one operation into another.

Step 1: Define MLIR Operations

Let’s define two simple operations in a custom dialect (MyDialect).

• mydialect.const: Represents a constant value.
• mydialect.const2: Another constant operation with a different name 😅.

We will write a rewrite pattern to convert const into const2.

Define const and const2 in TableGen (MyDialectOps.td)

def MyDialect_Const : MyDialectOp<"const"> {
  let summary = "Returns a constant";
  let arguments = (ins AnyAttr : $op1);
  let results = (outs AnyType : $res);
}

def MyDialect_Const2 : MyDialectOp<"const2"> {
  let summary = "Returns a constant 2";
  let arguments = (ins AnyAttr : $op);
  let results = (outs AnyType : $res);
}

Here,

• op1 and op are attributes (constant values).
• res is the result type of the operation.

Step 2: Define the Rewrite Pattern

We now define a pattern that transforms const into const2.

def const1toconst2 : Pat<(MyDialect_Const $a), (MyDialect_Const2 $a)>;

This pattern tells MLIR:

• Match mydialect.const operation.
• Replace it with mydialect.const2.

Step 3: Test MLIR Input Before and After Rewrite

Before Applying Rewrite

module {
  func.func @main() {
    %0 = "mydialect.const"() <{op1 = 33 : i32}> : () -> i32
    return
  }
}

Here, we have a function using "mydialect.const" with an integer value of 33.

After Applying Rewrite

module {
  func.func @main() {
    %0 = "mydialect.const2"() <{op = 33 : i32}> : () -> i32
    return
  }
}

The const operation has been replaced with const2, and the attribute has been renamed from op1 to opas per the const2 definition in the td file.

Step 4: Apply the Rewrite Pattern in C++

/ Create a pattern rewrite set
    mlir::RewritePatternSet patterns(&context);
    populateWithGenerated(patterns);  // Automatically adds all generated patterns

    // Apply the patterns
    if (failed(applyPatternsAndFoldGreedily(module.get(), std::move(patterns)))) {
        llvm::errs() << "Pattern application failed.\n";
        return 1;
    }

Conclusion

MLIR Rewrite Patterns are a powerful way to automate code transformations in a compiler pipeline. In this blog, we: ✅ Defined MLIR operations (const, const2) ✅ Created a simple rewrite pattern (const1toconst2) ✅ Applied the transformation using C++

Now you have a basic understanding of how to write and apply MLIR rewrite patterns! 🚀

The complete code is available in our GitHub repository: ***Blog CodeBase Repo.***