Photo by Rene Böhmer on Unsplash

Understanding MLIR Passes Through a Simple Dialect Transformation

In this blog, we’ll walk through the process of creating a custom dialect in MLIR, lowering its operations to another dialect (in this case, arith)

This tutorial assumes a basic understanding of MLIR concepts. We’ll cover:

1. Creating a Custom Dialect
2. Adding an Operation to the Dialect
3. Writing a Pass to Lower the Operation
4. Driving the Pass in an MLIR Program

1. Creating a Custom Dialect

A dialect in MLIR defines a collection of operations, types, and attributes. Here’s how we define our simple MyDialect:

def MyDialect : Dialect {
    let summary = "A sample dialect for understanding";
    let name = "mydialect";
    let cppNamespace = "mlir::mydialect";
}

Base Operation Class

class MyDialectOp<string mnemonic> : Op<MyDialect, mnemonic>
{
  let summary = "Operation Class";
}

2. Operation Definition

The Const operation doesn't take any parameters or return values in this example, making it simple for educational purposes. Its declaration is written in TableGen and compiled into C++ headers:

def MyDialect_Const : MyDialectOp<"const"> {
    let summary = "Returns a constant";
}

Pass Definition

def convertmydialect2arith : Pass<"convert-mydialect-to-arith", "mlir::ModuleOp"> {
    let summary = "Convert our dialect operations to arith dialect";
    let dependentDialects = ["mlir::mydialect::MyDialect"];
}

3. Writing a Pass to Lower the Operation

Passes are the heart of MLIR transformations. Our pass, convertmydialect2arith, will replace each instance of Const with an equivalent arith.constant operation.

struct convertmydialect2arith : public mlir::mydialect::impl::convertmydialect2arithBase<convertmydialect2arith> {
    void runOnOperation() override {
        auto mod = getOperation();

        mod->walk([&](mlir::mydialect::Const constOpIter) {
            // Use the OpBuilder to create the replacement operation
            OpBuilder b(constOpIter);

            auto newOp = b.create<mlir::arith::ConstantOp>(
                b.getUnknownLoc(),               // Location
                b.getI32Type(),                  // Type
                b.getI32IntegerAttr(33)          // Value (constant 33)
            );

            // Replace and erase the original operation
            constOpIter->replaceAllUsesWith(newOp);
            constOpIter->erase();
        });
    }
};

4. Driving the Pass

Now, let’s create a program that:

• Instantiates the dialect and adds operations.
• Runs the pass to lower Const to arith.constant.

int main() {
    MLIRContext context;
    context.getOrLoadDialect<mydialect::MyDialect>();
    context.getOrLoadDialect<func::FuncDialect>();
    context.getOrLoadDialect<arith::ArithDialect>();

    // Create a module and define a function
    OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
    OpBuilder builder(&context);

    auto funcType = builder.getFunctionType({}, {});
    auto funcOp = builder.create<func::FuncOp>(builder.getUnknownLoc(), "main", funcType);
    module->push_back(funcOp);

    // Add a `Const` operation
    Block *entryBlock = funcOp.addEntryBlock();
    builder.setInsertionPointToStart(entryBlock);
    builder.create<mydialect::Const>(builder.getUnknownLoc());
    builder.create<func::ReturnOp>(builder.getUnknownLoc());

    module->dump();  // Dump IR before the pass

    // Add and run the pass
    PassManager pm(&context);
    pm.addPass(createconvertmydialect2arith());
    if (failed(pm.run(*module))) {
        llvm::errs() << "Failed to run passes\n";
        return 1;
    }

    module->dump();  // Dump IR after the pass
    return 0;
}

Running the Example

Once compiled, the program will:

1. Dump the IR before the pass, showing the Const operation.
2. Transform Const into arith.constant using the custom pass.
3. Dump the IR after the pass.

Before Custom Pass:

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

After Custom Pass:

module {
  func @main() {
    %0 = arith.constant 33 : i32
    return
  }
}

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