[Mlir-commits] [mlir] [MLIR][Python] Refine trait support in MLIR Python (PR #180550)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Feb 9 18:05:21 PST 2026
https://github.com/PragmaTwice updated https://github.com/llvm/llvm-project/pull/180550
>From c9c545a8467ee4d535fe893fa7f97c27dfb02974 Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Mon, 9 Feb 2026 23:45:32 +0800
Subject: [PATCH 1/3] [MLIR][Python] Refine trait support in MLIR Python
---
mlir/include/mlir-c/ExtensibleDialect.h | 8 ++++++--
mlir/lib/Bindings/Python/IRCore.cpp | 9 +++++----
mlir/lib/CAPI/IR/ExtensibleDialect.cpp | 4 ++--
3 files changed, 13 insertions(+), 8 deletions(-)
diff --git a/mlir/include/mlir-c/ExtensibleDialect.h b/mlir/include/mlir-c/ExtensibleDialect.h
index 98457805f57c0..b0e03371da963 100644
--- a/mlir/include/mlir-c/ExtensibleDialect.h
+++ b/mlir/include/mlir-c/ExtensibleDialect.h
@@ -33,6 +33,8 @@ extern "C" {
DEFINE_C_API_STRUCT(MlirDynamicOpTrait, void);
+#undef DEFINE_C_API_STRUCT
+
/// Attach a dynamic op trait to the given operation name.
/// Note that the operation name must be modeled by dynamic dialect and must be
/// registered.
@@ -43,10 +45,12 @@ mlirDynamicOpTraitAttach(MlirDynamicOpTrait dynamicOpTrait,
MlirStringRef opName, MlirContext context);
/// Get the dynamic op trait that indicates the operation is a terminator.
-MLIR_CAPI_EXPORTED MlirDynamicOpTrait mlirDynamicOpTraitGetIsTerminator(void);
+MLIR_CAPI_EXPORTED MlirDynamicOpTrait
+mlirDynamicOpTraitCreateIsTerminator(void);
/// Get the dynamic op trait that indicates regions have no terminator.
-MLIR_CAPI_EXPORTED MlirDynamicOpTrait mlirDynamicOpTraitGetNoTerminator(void);
+MLIR_CAPI_EXPORTED MlirDynamicOpTrait
+mlirDynamicOpTraitCreateNoTerminator(void);
/// Destroy the dynamic op trait.
MLIR_CAPI_EXPORTED void
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 6f03f334e34b4..3c0e46ff36007 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2517,9 +2517,8 @@ void PyOpAdaptor::bind(nb::module_ &m) {
static MlirLogicalResult verifyTraitByMethod(MlirOperation op, void *userData,
const char *methodName) {
nb::handle targetObj(static_cast<PyObject *>(userData));
- if (!nb::hasattr(targetObj, methodName)) {
+ if (!nb::hasattr(targetObj, methodName))
return mlirLogicalResultSuccess();
- }
PyMlirContextRef ctx = PyMlirContext::forContext(mlirOperationGetContext(op));
nb::object opView = PyOperation::forOperation(ctx, op)->createOpView();
bool success = nb::cast<bool>(targetObj.attr(methodName)(opView));
@@ -2566,6 +2565,8 @@ bool PyDynamicOpTrait::attach(const nb::object &opName,
return verifyTraitByMethod(op, userData, "verify_region");
};
+ // To ensure that the same dynamic trait gets the same TypeID despite how many
+ // times `attach` is called, we store it as an attribute on the target class.
constexpr const char *typeIDAttr = "_TYPE_ID";
if (!nb::hasattr(target, typeIDAttr)) {
nb::setattr(target, typeIDAttr,
@@ -2593,7 +2594,7 @@ void PyDynamicOpTrait::bind(nb::module_ &m) {
bool PyDynamicOpTraits::IsTerminator::attach(const nb::object &opName,
PyMlirContext &context) {
- MlirDynamicOpTrait trait = mlirDynamicOpTraitGetIsTerminator();
+ MlirDynamicOpTrait trait = mlirDynamicOpTraitCreateIsTerminator();
return attachOpTrait(opName, trait, context);
}
@@ -2611,7 +2612,7 @@ void PyDynamicOpTraits::IsTerminator::bind(nb::module_ &m) {
bool PyDynamicOpTraits::NoTerminator::attach(const nb::object &opName,
PyMlirContext &context) {
- MlirDynamicOpTrait trait = mlirDynamicOpTraitGetNoTerminator();
+ MlirDynamicOpTrait trait = mlirDynamicOpTraitCreateNoTerminator();
return attachOpTrait(opName, trait, context);
}
diff --git a/mlir/lib/CAPI/IR/ExtensibleDialect.cpp b/mlir/lib/CAPI/IR/ExtensibleDialect.cpp
index f3239d996a0e6..983c34d68e6b9 100644
--- a/mlir/lib/CAPI/IR/ExtensibleDialect.cpp
+++ b/mlir/lib/CAPI/IR/ExtensibleDialect.cpp
@@ -36,11 +36,11 @@ bool mlirDynamicOpTraitAttach(MlirDynamicOpTrait dynamicOpTrait,
return static_cast<DynamicOpDefinition *>(impl)->addTrait(std::move(trait));
}
-MlirDynamicOpTrait mlirDynamicOpTraitGetIsTerminator() {
+MlirDynamicOpTrait mlirDynamicOpTraitCreateIsTerminator() {
return wrap(new DynamicOpTraits::IsTerminator());
}
-MlirDynamicOpTrait mlirDynamicOpTraitGetNoTerminator() {
+MlirDynamicOpTrait mlirDynamicOpTraitCreateNoTerminator() {
return wrap(new DynamicOpTraits::NoTerminator());
}
>From 8276736fb1444a4163247c62a6475687c00543d2 Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Tue, 10 Feb 2026 00:19:16 +0800
Subject: [PATCH 2/3] fix
---
mlir/include/mlir-c/ExtensibleDialect.h | 4 ++--
mlir/lib/Bindings/Python/IRCore.cpp | 4 ++--
mlir/lib/CAPI/IR/ExtensibleDialect.cpp | 4 ++--
3 files changed, 6 insertions(+), 6 deletions(-)
diff --git a/mlir/include/mlir-c/ExtensibleDialect.h b/mlir/include/mlir-c/ExtensibleDialect.h
index b0e03371da963..6901b2ec15d37 100644
--- a/mlir/include/mlir-c/ExtensibleDialect.h
+++ b/mlir/include/mlir-c/ExtensibleDialect.h
@@ -46,11 +46,11 @@ mlirDynamicOpTraitAttach(MlirDynamicOpTrait dynamicOpTrait,
/// Get the dynamic op trait that indicates the operation is a terminator.
MLIR_CAPI_EXPORTED MlirDynamicOpTrait
-mlirDynamicOpTraitCreateIsTerminator(void);
+mlirDynamicOpTraitIsTerminatorCreate(void);
/// Get the dynamic op trait that indicates regions have no terminator.
MLIR_CAPI_EXPORTED MlirDynamicOpTrait
-mlirDynamicOpTraitCreateNoTerminator(void);
+mlirDynamicOpTraitNoTerminatorCreate(void);
/// Destroy the dynamic op trait.
MLIR_CAPI_EXPORTED void
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 3c0e46ff36007..9edc981220f8b 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2594,7 +2594,7 @@ void PyDynamicOpTrait::bind(nb::module_ &m) {
bool PyDynamicOpTraits::IsTerminator::attach(const nb::object &opName,
PyMlirContext &context) {
- MlirDynamicOpTrait trait = mlirDynamicOpTraitCreateIsTerminator();
+ MlirDynamicOpTrait trait = mlirDynamicOpTraitIsTerminatorCreate();
return attachOpTrait(opName, trait, context);
}
@@ -2612,7 +2612,7 @@ void PyDynamicOpTraits::IsTerminator::bind(nb::module_ &m) {
bool PyDynamicOpTraits::NoTerminator::attach(const nb::object &opName,
PyMlirContext &context) {
- MlirDynamicOpTrait trait = mlirDynamicOpTraitCreateNoTerminator();
+ MlirDynamicOpTrait trait = mlirDynamicOpTraitNoTerminatorCreate();
return attachOpTrait(opName, trait, context);
}
diff --git a/mlir/lib/CAPI/IR/ExtensibleDialect.cpp b/mlir/lib/CAPI/IR/ExtensibleDialect.cpp
index 983c34d68e6b9..e5fb7a9fc3563 100644
--- a/mlir/lib/CAPI/IR/ExtensibleDialect.cpp
+++ b/mlir/lib/CAPI/IR/ExtensibleDialect.cpp
@@ -36,11 +36,11 @@ bool mlirDynamicOpTraitAttach(MlirDynamicOpTrait dynamicOpTrait,
return static_cast<DynamicOpDefinition *>(impl)->addTrait(std::move(trait));
}
-MlirDynamicOpTrait mlirDynamicOpTraitCreateIsTerminator() {
+MlirDynamicOpTrait mlirDynamicOpTraitIsTerminatorCreate() {
return wrap(new DynamicOpTraits::IsTerminator());
}
-MlirDynamicOpTrait mlirDynamicOpTraitCreateNoTerminator() {
+MlirDynamicOpTrait mlirDynamicOpTraitNoTerminatorCreate() {
return wrap(new DynamicOpTraits::NoTerminator());
}
>From 9972482a32ba09ff54f74a15653eb011ce6b5c5c Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Tue, 10 Feb 2026 10:05:07 +0800
Subject: [PATCH 3/3] reword comment
---
mlir/lib/CAPI/IR/ExtensibleDialect.cpp | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/CAPI/IR/ExtensibleDialect.cpp b/mlir/lib/CAPI/IR/ExtensibleDialect.cpp
index e5fb7a9fc3563..2d6608ff1d1cd 100644
--- a/mlir/lib/CAPI/IR/ExtensibleDialect.cpp
+++ b/mlir/lib/CAPI/IR/ExtensibleDialect.cpp
@@ -31,8 +31,8 @@ bool mlirDynamicOpTraitAttach(MlirDynamicOpTrait dynamicOpTrait,
static_cast<RegisteredOperationNameWithImpl &>(*opNameFound).getImpl();
std::unique_ptr<DynamicOpTrait> trait(unwrap(dynamicOpTrait));
- // TODO: we should check whether the `impl` is a DynamicOpDefinition here
- // via llvm-style RTTI.
+ // 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));
}
More information about the Mlir-commits
mailing list