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

[llvmlistbot at llvm.org]

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

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`.


>From c04cbc9079d4e20b87f923e0206910d0ef0ec17c Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice at apache.org>
Date: Fri, 13 Feb 2026 00:08:49 +0800
Subject: [PATCH] [MLIR][Python] Ensure traits are attached when
 load(register=False)

---
 mlir/python/mlir/dialects/ext.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

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)