[Mlir-commits] [mlir] [MLIR][Python] Ensure traits are attached when `load(register=False)` (PR #181174)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Twice (PragmaTwice)

<details>
<summary>Changes</summary>

Currently, when calling `.load(register=False)`, `op._attach_traits()` isn’t executed. This PR ensures traits are attached regardless of whether `register` is `True` or `False`.


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


1 Files Affected:

- (modified) mlir/python/mlir/dialects/ext.py (+6-3) 


``````````diff
diff --git a/mlir/python/mlir/dialects/ext.py b/mlir/python/mlir/dialects/ext.py
index d80e9eed6a483..ac1b3065336e8 100644
--- a/mlir/python/mlir/dialects/ext.py
+++ b/mlir/python/mlir/dialects/ext.py
@@ -27,17 +27,18 @@
 
 __all__ = [
     "Dialect",
+    "Operation",
     "Operand",
     "Result",
+    "Region",
     "register_dialect",
     "register_operation",
-    "Region",
-    "Operation",
 ]
 
 Operand = ir.Value
 Result = ir.OpResult
 Region = ir.Region
+
 register_dialect = _cext.register_dialect
 register_operation = _cext.register_operation
 
@@ -528,10 +529,12 @@ def load(cls, register=True, reload=False) -> None:
 
         irdl.load_dialects(cls._mlir_module)
 
+        for op in cls.operations:
+            op._attach_traits()
+
         if register:
             register_dialect(cls)
 
             register_dialect_operation = register_operation(cls)
             for op in cls.operations:
-                op._attach_traits()
                 register_dialect_operation(op)

``````````

</details>


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