[Mlir-commits] [mlir] [mlir][python] Register linalg.ElementwiseOp wrapper as op view (PR #223806)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-linalg

Author: Maksim Levental (makslevental)

<details>
<summary>Changes</summary>

The hand-written `ElementwiseOp` wrapper in `mlir/python/mlir/dialects/linalg/__init__.py` subclasses the generated OpView but was not registered in the op-view registry. As a result, ops reconstructed from the registry (e.g. `Operation.opview`, `OpOperand.owner`, IR walks) were instances of the generated base class `_linalg_ops_gen.ElementwiseOp`, whereas the public symbol `linalg.ElementwiseOp` referred to the subclass.

Because the exported name was a *subclass* of the instantiated class, `isinstance(op.opview, linalg.ElementwiseOp)` and `match`/`case linalg.ElementwiseOp()` always evaluated to `False` for real `linalg.elementwise` ops, silently breaking type-based dispatch.

This registers the wrapper with `@<!-- -->_ods_cext.register_operation(_Dialect, replace=True)` (the same pattern used in `pdl.py`) so the exported class is the one the registry instantiates.

Fixes #<!-- -->223673

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


2 Files Affected:

- (modified) mlir/python/mlir/dialects/linalg/__init__.py (+3-4) 
- (modified) mlir/test/python/dialects/linalg/ops.py (+21) 


``````````diff
diff --git a/mlir/python/mlir/dialects/linalg/__init__.py b/mlir/python/mlir/dialects/linalg/__init__.py
index 1f6d1a68fbbb8..e15f4532cc3c8 100644
--- a/mlir/python/mlir/dialects/linalg/__init__.py
+++ b/mlir/python/mlir/dialects/linalg/__init__.py
@@ -63,6 +63,7 @@
     get_op_result_or_value as _get_op_result_or_value,
     get_op_result_or_op_results as _get_op_result_or_op_results,
     _dispatch_mixed_values,
+    _cext as _ods_cext,
 )
 from ...extras.meta import region_op
 
@@ -233,7 +234,8 @@ def contract(
 # Extend and shadow the TableGen-derived version to make sure correct default
 # indexing_maps are derived (as there is no mechanism for doing so given the
 # Python API bypasses the C++-builders).
-class ElementwiseOp_(ElementwiseOp):
+ at _ods_cext.register_operation(_Dialect, replace=True)
+class ElementwiseOp(ElementwiseOp):
     def __init__(
         self,
         result_tensors,
@@ -265,9 +267,6 @@ def __init__(
         )
 
 
-ElementwiseOp = ElementwiseOp_
-
-
 def elementwise(
     *ins: Union[Operation, OpView, Value],
     outs: Sequence[Union[Operation, OpView, Value]],
diff --git a/mlir/test/python/dialects/linalg/ops.py b/mlir/test/python/dialects/linalg/ops.py
index 1368e6ac495ab..592492995b84e 100644
--- a/mlir/test/python/dialects/linalg/ops.py
+++ b/mlir/test/python/dialects/linalg/ops.py
@@ -884,6 +884,27 @@ def elementwise_op(
         print(module)
 
 
+# CHECK-LABEL: TEST: testElementwiseOpView
+ at run
+def testElementwiseOpView():
+    with Context(), Location.unknown():
+        module = Module.parse(
+            """
+            func.func @f(%c: tensor<4x4xf32>) {
+              %0 = linalg.elementwise kind=#linalg.elementwise_kind<add>
+                   ins(%c, %c : tensor<4x4xf32>, tensor<4x4xf32>)
+                   outs(%c : tensor<4x4xf32>) -> tensor<4x4xf32>
+              return
+            }
+            """
+        )
+        func_op = module.body.operations[0]
+        op = func_op.regions[0].blocks[0].operations[0]
+        # The op view reconstructed from the registry must be an instance of the
+        # publicly exported class.
+        assert isinstance(op.opview, linalg.ElementwiseOp)
+
+
 @run
 def testReduceOp():
     with Context(), Location.unknown():

``````````

</details>


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