[Mlir-commits] [mlir] [MLIR][Python] Support canonicalization hooks for Python-defined ops (PR #224806)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Sep 18 22:56:30 PDT 2026
https://github.com/PragmaTwice updated https://github.com/llvm/llvm-project/pull/224806
>From b8b75cc93e8c7de45fc16b4d5298c22b3ca9e177 Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Sat, 19 Sep 2026 13:28:07 +0800
Subject: [PATCH 1/2] [MLIR][Python] Support canonicalization hooks for
Python-defined ops
Allow Python-defined operations to provide get_canonicalization_patterns and
register the hook when loading their dialect. Add dynamic operation definition
lookup and hook registration to the C and Python APIs, with callback cleanup
when the hook is replaced or destroyed.
Route dynamic trait attachment through the definition handle and add a test
that canonicalizes a Python-defined operation to its input.
Test: ninja -C build -j 4 MLIRPythonModules
Test: build/bin/llvm-lit -v mlir/test/python/dialects/ext.py
---
mlir/include/mlir-c/ExtensibleDialect.h | 30 +++++++++
mlir/include/mlir/Bindings/Python/IRCore.h | 15 +++++
mlir/lib/Bindings/Python/IRCore.cpp | 39 +++++++++++-
mlir/lib/CAPI/IR/ExtensibleDialect.cpp | 72 +++++++++++++++-------
mlir/python/mlir/dialects/ext.py | 13 ++++
mlir/test/python/dialects/ext.py | 42 +++++++++++++
6 files changed, 186 insertions(+), 25 deletions(-)
diff --git a/mlir/include/mlir-c/ExtensibleDialect.h b/mlir/include/mlir-c/ExtensibleDialect.h
index f787ff8dd20c1..0b56e9624e7c5 100644
--- a/mlir/include/mlir-c/ExtensibleDialect.h
+++ b/mlir/include/mlir-c/ExtensibleDialect.h
@@ -15,6 +15,7 @@
#define MLIR_C_EXTENSIBLEDIALECT_H
#include "mlir-c/IR.h"
+#include "mlir-c/Rewrite.h"
#include "mlir-c/Support.h"
#ifdef __cplusplus
@@ -34,6 +35,7 @@ extern "C" {
DEFINE_C_API_STRUCT(MlirDynamicOpTrait, void);
DEFINE_C_API_STRUCT(MlirDynamicTypeDefinition, void);
DEFINE_C_API_STRUCT(MlirDynamicAttrDefinition, void);
+DEFINE_C_API_STRUCT(MlirDynamicOpDefinition, void);
#undef DEFINE_C_API_STRUCT
@@ -42,6 +44,7 @@ DEFINE_C_API_STRUCT(MlirDynamicAttrDefinition, void);
/// registered.
/// The ownership of the trait will be transferred to the operation name
/// after this call.
+/// DEPRECATED: Use mlirDynamicOpDefinitionAddTrait instead.
MLIR_CAPI_EXPORTED bool
mlirDynamicOpTraitAttach(MlirDynamicOpTrait dynamicOpTrait,
MlirStringRef opName, MlirContext context);
@@ -179,6 +182,33 @@ mlirDynamicAttrDefinitionGetName(MlirDynamicAttrDefinition attrDef);
MLIR_CAPI_EXPORTED MlirDialect
mlirDynamicAttrDefinitionGetDialect(MlirDynamicAttrDefinition attrDef);
+/// Look up a dynamic operation definition by operation name in a given context.
+/// Note that the operation name must be modeled by dynamic dialect and must be
+/// registered.
+MLIR_CAPI_EXPORTED MlirDynamicOpDefinition
+mlirDynamicOpDefinitionLookup(MlirStringRef opName, MlirContext context);
+
+/// The hook function type for getting canonicalization patterns for a given
+/// dynamic operation definition.
+typedef void (*MlirDynamicOpDefinitionGetCanonicalizationPatternsFn)(
+ MlirRewritePatternSet patterns, MlirContext context, void *userData);
+
+/// The hook function type for destructing the userData in callbacks of a
+/// dynamic operation definition.
+typedef void (*MlirDynamicOpDefinitionDestructUserDataFn)(void *userData);
+
+/// Set the hook returning any canonicalization pattern rewrites that the op
+/// supports for a given dynamic operation definition.
+MLIR_CAPI_EXPORTED void mlirDynamicOpDefinitionSetGetCanonicalizationPatternsFn(
+ MlirDynamicOpDefinition opDef,
+ MlirDynamicOpDefinitionGetCanonicalizationPatternsFn fn,
+ MlirDynamicOpDefinitionDestructUserDataFn destructFn, void *userData);
+
+/// Add a dynamic op trait to the given dynamic operation definition.
+MLIR_CAPI_EXPORTED bool
+mlirDynamicOpDefinitionAddTrait(MlirDynamicOpDefinition opDef,
+ MlirDynamicOpTrait dynamicOpTrait);
+
#ifdef __cplusplus
}
#endif
diff --git a/mlir/include/mlir/Bindings/Python/IRCore.h b/mlir/include/mlir/Bindings/Python/IRCore.h
index b28e6ec606d52..6f2d3bd7279b5 100644
--- a/mlir/include/mlir/Bindings/Python/IRCore.h
+++ b/mlir/include/mlir/Bindings/Python/IRCore.h
@@ -1977,6 +1977,21 @@ class MLIR_PYTHON_API_EXPORTED PyOpAdaptor {
PyOpAttributeMap attributes;
};
+class MLIR_PYTHON_API_EXPORTED PyDynamicOpDefinition {
+public:
+ PyDynamicOpDefinition(std::string_view opName,
+ DefaultingPyMlirContext context)
+ : opDef(mlirDynamicOpDefinitionLookup(
+ MlirStringRef{opName.data(), opName.size()},
+ context.get()->get())) {}
+
+ void setGetCanonicalizationPatternsFn(const nanobind::callable &fn);
+ static void bind(nanobind::module_ &m);
+
+private:
+ MlirDynamicOpDefinition opDef;
+};
+
class MLIR_PYTHON_API_EXPORTED PyDynamicOpTrait {
public:
static bool attach(const nanobind::object &opName,
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index b6d92dc8064a2..00a42410bcfed 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -7,6 +7,7 @@
//===----------------------------------------------------------------------===//
// clang-format off
+#include "Rewrite.h"
#include "mlir/Bindings/Python/Globals.h"
#include "mlir/Bindings/Python/IRCore.h"
#include "mlir/Bindings/Python/NanobindUtils.h"
@@ -22,6 +23,7 @@
#include <array>
#include <cassert>
#include <functional>
+#include <nanobind/nanobind.h>
#include <optional>
#include <string>
@@ -2547,6 +2549,37 @@ void PyOpAdaptor::bind(nb::module_ &m) {
"Returns the attributes of the adaptor.");
}
+void PyDynamicOpDefinition::setGetCanonicalizationPatternsFn(
+ const nb::callable &fn) {
+ fn.inc_ref();
+ mlirDynamicOpDefinitionSetGetCanonicalizationPatternsFn(
+ opDef,
+ [](MlirRewritePatternSet patterns, MlirContext context,
+ void *userData) -> void {
+ nb::handle fn(static_cast<PyObject *>(userData));
+ fn(PyRewritePatternSet(patterns),
+ PyMlirContext::forContext(context).get());
+ },
+ [](void *userData) {
+ nb::handle fn(static_cast<PyObject *>(userData));
+ fn.dec_ref();
+ },
+ static_cast<void *>(fn.ptr()));
+}
+
+void PyDynamicOpDefinition::bind(nb::module_ &m) {
+ nb::class_<PyDynamicOpDefinition>(m, "DynamicOpDefinition")
+ .def(nb::init<std::string_view, DefaultingPyMlirContext>(), "op_name"_a,
+ "context"_a = nb::none(),
+ "Lookup a dynamic operation definition with the given operation "
+ "name.")
+ .def("set_get_canonicalization_patterns_fn",
+ &PyDynamicOpDefinition::setGetCanonicalizationPatternsFn,
+ "Sets the function to get canonicalization patterns for this "
+ "dynamic operation definition.",
+ "fn"_a);
+}
+
static MlirLogicalResult verifyTraitByMethod(MlirOperation op, void *userData,
const char *methodName) {
nb::handle targetObj(static_cast<PyObject *>(userData));
@@ -2569,8 +2602,9 @@ static bool attachOpTrait(const nb::object &opName, MlirDynamicOpTrait trait,
throw nb::type_error("the root argument must be a type or a string");
}
- return mlirDynamicOpTraitAttach(
- trait, MlirStringRef{opNameStr.data(), opNameStr.size()}, context.get());
+ MlirDynamicOpDefinition def = mlirDynamicOpDefinitionLookup(
+ MlirStringRef{opNameStr.data(), opNameStr.size()}, context.get());
+ return mlirDynamicOpDefinitionAddTrait(def, trait);
}
bool PyDynamicOpTrait::attach(const nb::object &opName,
@@ -5416,6 +5450,7 @@ void populateIRCore(nb::module_ &m) {
PyAttrBuilderMap::bind(m);
// Extensible Dialect
+ PyDynamicOpDefinition::bind(m);
PyDynamicOpTrait::bind(m);
PyDynamicOpTraits::IsTerminator::bind(m);
PyDynamicOpTraits::NoTerminator::bind(m);
diff --git a/mlir/lib/CAPI/IR/ExtensibleDialect.cpp b/mlir/lib/CAPI/IR/ExtensibleDialect.cpp
index 2b4187da9eabf..7c7fd4cdd83c8 100644
--- a/mlir/lib/CAPI/IR/ExtensibleDialect.cpp
+++ b/mlir/lib/CAPI/IR/ExtensibleDialect.cpp
@@ -8,6 +8,7 @@
#include "mlir-c/ExtensibleDialect.h"
#include "mlir/CAPI/IR.h"
+#include "mlir/CAPI/Rewrite.h"
#include "mlir/CAPI/Support.h"
#include "mlir/IR/ExtensibleDialect.h"
#include "mlir/IR/OperationSupport.h"
@@ -17,25 +18,14 @@ using namespace mlir;
DEFINE_C_API_PTR_METHODS(MlirDynamicOpTrait, DynamicOpTrait)
DEFINE_C_API_PTR_METHODS(MlirDynamicTypeDefinition, DynamicTypeDefinition)
DEFINE_C_API_PTR_METHODS(MlirDynamicAttrDefinition, DynamicAttrDefinition)
+DEFINE_C_API_PTR_METHODS(MlirDynamicOpDefinition, DynamicOpDefinition)
bool mlirDynamicOpTraitAttach(MlirDynamicOpTrait dynamicOpTrait,
MlirStringRef opName, MlirContext context) {
- std::optional<RegisteredOperationName> opNameFound =
- RegisteredOperationName::lookup(unwrap(opName), unwrap(context));
- assert(opNameFound && "operation name must be registered in the context");
-
- // The original getImpl() is protected, so we create a small helper struct
- // here.
- struct RegisteredOperationNameWithImpl : RegisteredOperationName {
- Impl *getImpl() { return RegisteredOperationName::getImpl(); }
- };
- OperationName::Impl *impl =
- static_cast<RegisteredOperationNameWithImpl &>(*opNameFound).getImpl();
+ MlirDynamicOpDefinition dynamicOpDef =
+ mlirDynamicOpDefinitionLookup(opName, context);
- std::unique_ptr<DynamicOpTrait> trait(unwrap(dynamicOpTrait));
- // TODO: we should enable llvm-style RTTI for `OperationName::Impl` and check
- // whether the `impl` is a `DynamicOpDefinition` here.
- return static_cast<DynamicOpDefinition *>(impl)->addTrait(std::move(trait));
+ return mlirDynamicOpDefinitionAddTrait(dynamicOpDef, dynamicOpTrait);
}
MlirDynamicOpTrait mlirDynamicOpTraitIsTerminatorCreate() {
@@ -134,10 +124,6 @@ bool mlirTypeIsADynamicType(MlirType type) {
return llvm::isa<mlir::DynamicType>(unwrap(type));
}
-MlirTypeID mlirDynamicTypeGetTypeID() {
- return wrap(mlir::DynamicType::getTypeID());
-}
-
MlirType mlirDynamicTypeGet(MlirDynamicTypeDefinition typeDef,
MlirAttribute *attrs, intptr_t numAttrs) {
llvm::SmallVector<mlir::Attribute> attributes;
@@ -186,10 +172,6 @@ bool mlirAttributeIsADynamicAttr(MlirAttribute attr) {
return llvm::isa<mlir::DynamicAttr>(unwrap(attr));
}
-MlirTypeID mlirDynamicAttrGetTypeID(void) {
- return wrap(mlir::DynamicAttr::getTypeID());
-}
-
MlirAttribute mlirDynamicAttrGet(MlirDynamicAttrDefinition attrDef,
MlirAttribute *attrs, intptr_t numAttrs) {
llvm::SmallVector<mlir::Attribute> attributes;
@@ -226,3 +208,47 @@ MlirDialect
mlirDynamicAttrDefinitionGetDialect(MlirDynamicAttrDefinition attrDef) {
return wrap(unwrap(attrDef)->getDialect());
}
+
+MlirDynamicOpDefinition mlirDynamicOpDefinitionLookup(MlirStringRef opName,
+ MlirContext context) {
+ std::optional<RegisteredOperationName> opNameFound =
+ RegisteredOperationName::lookup(unwrap(opName), unwrap(context));
+ assert(opNameFound && "operation name must be registered in the context");
+
+ // The original getImpl() is protected, so we create a small helper struct
+ // here.
+ struct RegisteredOperationNameWithImpl : RegisteredOperationName {
+ Impl *getImpl() { return RegisteredOperationName::getImpl(); }
+ };
+ OperationName::Impl *impl =
+ static_cast<RegisteredOperationNameWithImpl &>(*opNameFound).getImpl();
+
+ // TODO: we should enable llvm-style RTTI for `OperationName::Impl` and check
+ // whether the `impl` is a `DynamicOpDefinition` here.
+ return wrap(static_cast<DynamicOpDefinition *>(impl));
+}
+
+void mlirDynamicOpDefinitionSetGetCanonicalizationPatternsFn(
+ MlirDynamicOpDefinition opDef,
+ MlirDynamicOpDefinitionGetCanonicalizationPatternsFn fn,
+ MlirDynamicOpDefinitionDestructUserDataFn destructFn, void *userData) {
+ auto deleter = [destructFn](void *data) {
+ if (destructFn)
+ destructFn(data);
+ };
+
+ auto ownedUserData =
+ std::unique_ptr<void, decltype(deleter)>(userData, deleter);
+
+ unwrap(opDef)->setGetCanonicalizationPatternsFn(
+ [fn, ownedUserData = std::move(ownedUserData)](
+ RewritePatternSet &patterns, MLIRContext *context) {
+ fn(wrap(&patterns), wrap(context), ownedUserData.get());
+ });
+}
+
+bool mlirDynamicOpDefinitionAddTrait(MlirDynamicOpDefinition opDef,
+ MlirDynamicOpTrait dynamicOpTrait) {
+ std::unique_ptr<DynamicOpTrait> trait(unwrap(dynamicOpTrait));
+ return unwrap(opDef)->addTrait(std::move(trait));
+}
diff --git a/mlir/python/mlir/dialects/ext.py b/mlir/python/mlir/dialects/ext.py
index 491008f86023f..3b60bd6921bd9 100644
--- a/mlir/python/mlir/dialects/ext.py
+++ b/mlir/python/mlir/dialects/ext.py
@@ -668,6 +668,18 @@ def _attach_traits(cls) -> None:
):
ir.DynamicOpTrait.attach(cls.OPERATION_NAME, cls)
+ @classmethod
+ def _attach_hooks(cls) -> None:
+ if hasattr(cls, "get_canonicalization_patterns"):
+
+ def get_canonicalization_patterns(patterns, context):
+ with context:
+ (cls.get_canonicalization_patterns)(patterns)
+
+ ir.DynamicOpDefinition(
+ cls.OPERATION_NAME
+ ).set_get_canonicalization_patterns_fn(get_canonicalization_patterns)
+
@classmethod
def _emit_operation(cls) -> None:
ctx = ConstraintLoweringContext()
@@ -992,6 +1004,7 @@ def load(cls) -> None:
for op in cls.operations:
op._attach_traits()
+ op._attach_hooks()
_cext.globals._register_dialect_impl(cls.DIALECT_NAMESPACE, cls, replace=True)
diff --git a/mlir/test/python/dialects/ext.py b/mlir/test/python/dialects/ext.py
index bbaca38dd130d..8ae34829903f5 100644
--- a/mlir/test/python/dialects/ext.py
+++ b/mlir/test/python/dialects/ext.py
@@ -3,6 +3,7 @@
from mlir.ir import *
from mlir.dialects import arith
from mlir.dialects.ext import *
+from mlir.passmanager import PassManager
from mlir.rewrite import *
from mlir import ir
from typing import Any, Optional, Sequence, TypeVar, Union
@@ -1191,6 +1192,47 @@ class RecursivelySpeculatableOp(
)
+# CHECK: TEST: testExtDialectCanonicalizationPatterns
+ at run
+def testExtDialectCanonicalizationPatterns():
+ class TestCanonicalization(Dialect, name="ext_canonicalization"):
+ pass
+
+ class CanonicalizeOp(TestCanonicalization.Operation, name="canonicalize"):
+ input: Operand[IntegerType[32]]
+ output: Result[IntegerType[32]] = infer_result()
+
+ @staticmethod
+ def get_canonicalization_patterns(patterns):
+ def replace_with_input(op, rewriter):
+ rewriter.replace_op(op, [op.input])
+
+ patterns.add(CanonicalizeOp, replace_with_input)
+
+ class SinkOp(TestCanonicalization.Operation, name="sink"):
+ input: Operand[IntegerType[32]]
+
+ with Context(), Location.unknown():
+ TestCanonicalization.load()
+
+ i32 = IntegerType.get_signless(32)
+ module = Module.create()
+ with InsertionPoint(module.body):
+ one = arith.constant(i32, 1)
+ canonicalized = CanonicalizeOp(one)
+ SinkOp(canonicalized)
+
+ assert module.operation.verify()
+ PassManager.parse("builtin.module(canonicalize)").run(module.operation)
+
+ # CHECK: module {
+ # CHECK: %[[ONE:.*]] = arith.constant 1 : i32
+ # CHECK-NOT: "ext_canonicalization.canonicalize"
+ # CHECK: "ext_canonicalization.sink"(%[[ONE]]) : (i32) -> ()
+ # CHECK: }
+ print(module)
+
+
# CHECK: TEST: testExtDialectWithPure
@run
def testExtDialectWithPure():
>From c5f2f4027114009f926306ad7164ab31553ee85c Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Sat, 19 Sep 2026 13:56:17 +0800
Subject: [PATCH 2/2] [MLIR][Python] Inline rewrite pattern set construction
and destruction
IRCore.cpp is built into MLIRPythonSupport, while Rewrite.cpp is built into
the _mlir extension that depends on it. Constructing PyRewritePatternSet in
the canonicalization callback therefore introduced unresolved references
in the support library and failed to link on Windows.
Define the constructors and destructor inline in Rewrite.h so the support
library can use the wrapper without depending on extension-module symbols.
Validated by rebuilding MLIRPythonModules, checking the support library for
unresolved PyRewritePatternSet symbols, and running the ext, rewrite, and
Transform pattern descriptor Python tests.
---
mlir/lib/Bindings/Python/Rewrite.cpp | 11 -----------
mlir/lib/Bindings/Python/Rewrite.h | 14 +++++++++++---
2 files changed, 11 insertions(+), 14 deletions(-)
diff --git a/mlir/lib/Bindings/Python/Rewrite.cpp b/mlir/lib/Bindings/Python/Rewrite.cpp
index f799dd33aa71e..30d5f290bc489 100644
--- a/mlir/lib/Bindings/Python/Rewrite.cpp
+++ b/mlir/lib/Bindings/Python/Rewrite.cpp
@@ -69,17 +69,6 @@ class PyPatternRewriter : public PyRewriterBase<PyPatternRewriter> {
// PyRewritePatternSet
//===----------------------------------------------------------------------===//
-PyRewritePatternSet::PyRewritePatternSet(MlirContext ctx)
- : patterns(mlirRewritePatternSetCreate(ctx)), owned(true) {}
-
-PyRewritePatternSet::PyRewritePatternSet(MlirRewritePatternSet patterns)
- : patterns(patterns), owned(false) {}
-
-PyRewritePatternSet::~PyRewritePatternSet() {
- if (owned && patterns.ptr)
- mlirRewritePatternSetDestroy(patterns);
-}
-
MlirRewritePatternSet PyRewritePatternSet::get() const { return patterns; }
bool PyRewritePatternSet::isOwned() const { return owned; }
diff --git a/mlir/lib/Bindings/Python/Rewrite.h b/mlir/lib/Bindings/Python/Rewrite.h
index fc48e0a254832..f60a75a8841b5 100644
--- a/mlir/lib/Bindings/Python/Rewrite.h
+++ b/mlir/lib/Bindings/Python/Rewrite.h
@@ -82,13 +82,21 @@ class MLIR_PYTHON_API_EXPORTED PyRewriterBase {
class PyTypeConverter;
class MLIR_PYTHON_API_EXPORTED PyRewritePatternSet {
public:
+ // Keep construction and destruction inline so the support library can use
+ // this wrapper without depending on extension-module symbols.
+
/// Create an owned pattern set.
- PyRewritePatternSet(MlirContext ctx);
+ PyRewritePatternSet(MlirContext ctx)
+ : patterns(mlirRewritePatternSetCreate(ctx)), owned(true) {}
/// Create a non-owning reference to an existing pattern set.
- PyRewritePatternSet(MlirRewritePatternSet patterns);
+ PyRewritePatternSet(MlirRewritePatternSet patterns)
+ : patterns(patterns), owned(false) {}
- ~PyRewritePatternSet();
+ ~PyRewritePatternSet() {
+ if (owned && patterns.ptr)
+ mlirRewritePatternSetDestroy(patterns);
+ }
MlirRewritePatternSet get() const;
More information about the Mlir-commits
mailing list