[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:32:30 PDT 2026
https://github.com/PragmaTwice created https://github.com/llvm/llvm-project/pull/224806
This PR allows Python-defined operations to provide canonicalization patterns through a `get_canonicalization_patterns` hook:
```python
class CanonicalizeOp(MyDialect.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)
```
The hook is attached when the dialect is loaded, allowing the `canonicalize` pass to collect and apply these patterns.
This also introduces dynamic operation definition lookup and hook registration in the C API and Python bindings. The registered hook owns the callback reference and releases it when replaced or destroyed.
Assisted-by: Codex
>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] [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():
More information about the Mlir-commits
mailing list