[Mlir-commits] [mlir] [MLIR][Python] Allow passing dialect as a class keyword argument (PR #182465)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Twice (PragmaTwice)

<details>
<summary>Changes</summary>

Previously, we constructed new ops using the pattern `class MyOp(MyInt.Operation)`.

Now we’ve added a new pattern: `class MyOp(Operation, dialect=MyInt)`, which allows more flexible composition. For example:
```python
class BinOpBase(Operation): # it can be used in any dialect!
  res: Result[Any]
  lhs: Operand[Any]
  rhs: Operand[Any]
  
class AddOp(BinOpBase, dialect=MyInt):
  ...
```


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


2 Files Affected:

- (modified) mlir/python/mlir/dialects/ext.py (+17-3) 
- (modified) mlir/test/python/dialects/ext.py (+1-1) 


``````````diff
diff --git a/mlir/python/mlir/dialects/ext.py b/mlir/python/mlir/dialects/ext.py
index ac1b3065336e8..b7dd1af8afb96 100644
--- a/mlir/python/mlir/dialects/ext.py
+++ b/mlir/python/mlir/dialects/ext.py
@@ -216,7 +216,12 @@ def __init__(*args, **kwargs):
 
     @classmethod
     def __init_subclass__(
-        cls, *, name: str | None = None, traits: list[type] | None = None, **kwargs
+        cls,
+        *,
+        name: str | None = None,
+        traits: list[type] | None = None,
+        dialect: type | None = None,
+        **kwargs,
     ):
         """
         This method is to perform all magic to make a `Operation` subclass works like a dataclass, like:
@@ -251,9 +256,18 @@ def __init_subclass__(
         if not name:
             return
 
+        if dialect:
+            if hasattr(cls, "_dialect_name") or hasattr(cls, "_dialect_obj"):
+                raise RuntimeError(
+                    f"This operation has already been attached to dialect '{cls._dialect_name}'."
+                )
+            cls._dialect_obj = dialect
+            cls._dialect_name = dialect.DIALECT_NAMESPACE
+
         if not hasattr(cls, "_dialect_name") or not hasattr(cls, "_dialect_obj"):
             raise RuntimeError(
-                "Operation subclasses must inherit from a Dialect's Operation subclass"
+                "Operation subclasses must either inherit from a Dialect's Operation subclass "
+                "or provide the dialect as a class keyword argument."
             )
 
         op_name = name
@@ -278,7 +292,7 @@ def _variadicity_to_segment(variadicity: Variadicity) -> int:
     @staticmethod
     def _generate_segments(
         operands_or_results: List[Union[OperandDef, ResultDef]],
-    ) -> List[int]:
+    ) -> List[int] | None:
         if any(i.variadicity != Variadicity.single for i in operands_or_results):
             return [
                 Operation._variadicity_to_segment(i.variadicity)
diff --git a/mlir/test/python/dialects/ext.py b/mlir/test/python/dialects/ext.py
index d300f0b0442ae..196af91e511ec 100644
--- a/mlir/test/python/dialects/ext.py
+++ b/mlir/test/python/dialects/ext.py
@@ -24,7 +24,7 @@ class ConstantOp(MyInt.Operation, name="constant"):
         value: IntegerAttr
         cst: Result[i32]
 
-    class AddOp(MyInt.Operation, name="add"):
+    class AddOp(Operation, dialect=MyInt, name="add"):
         lhs: Operand[i32]
         rhs: Operand[i32]
         res: Result[i32]

``````````

</details>


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