[Mlir-commits] [mlir] [MLIR][Python] Refine trait support in MLIR Python (PR #180550)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Twice (PragmaTwice)

<details>
<summary>Changes</summary>

This PR is mainly to address review suggestions in #<!-- -->179705.

---
Full diff: https://github.com/llvm/llvm-project/pull/180550.diff


3 Files Affected:

- (modified) mlir/include/mlir-c/ExtensibleDialect.h (+6-2) 
- (modified) mlir/lib/Bindings/Python/IRCore.cpp (+5-4) 
- (modified) mlir/lib/CAPI/IR/ExtensibleDialect.cpp (+2-2) 


``````````diff
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());
 }
 

``````````

</details>


https://github.com/llvm/llvm-project/pull/180550