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

[Maksim Levental]

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

>From dc5d67c1c1984b2950939f3a53e4e2bda38d9de0 Mon Sep 17 00:00:00 2001
From: makslevental <maksim.levental at gmail.com>
Date: Tue, 15 Sep 2026 13:08:01 -0700
Subject: [PATCH 1/2] [mlir][python] Register linalg.ElementwiseOp wrapper as
 op view

The hand-written ElementwiseOp wrapper subclassed the generated OpView but
was not registered in the op-view registry, so ops reconstructed from the
registry (e.g. Operation.opview) were instances of the generated base class
rather than the publicly exported subclass. This made
isinstance(op.opview, linalg.ElementwiseOp) and match/case dispatch return
False for real linalg.elementwise ops.

Register the wrapper with replace=True (as pdl.py does) so the exported
class is the one the registry instantiates.

Fixes #223673
---
 mlir/python/mlir/dialects/linalg/__init__.py |  7 +++----
 mlir/test/python/dialects/linalg/ops.py      | 21 ++++++++++++++++++++
 2 files changed, 24 insertions(+), 4 deletions(-)

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():

>From 60fc5b53085c3304a9e76e890abe8441da32c92a Mon Sep 17 00:00:00 2001
From: makslevental <maksim.levental at gmail.com>
Date: Tue, 15 Sep 2026 13:10:40 -0700
Subject: [PATCH 2/2] Remove comment per review

---
 mlir/test/python/dialects/linalg/ops.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/mlir/test/python/dialects/linalg/ops.py b/mlir/test/python/dialects/linalg/ops.py
index 592492995b84e..6c4cfb8e25809 100644
--- a/mlir/test/python/dialects/linalg/ops.py
+++ b/mlir/test/python/dialects/linalg/ops.py
@@ -900,8 +900,6 @@ def testElementwiseOpView():
         )
         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)