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

[Maksim Levental]

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

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

>From b7dc0f68134cb919c6bd0ab768b287248b39142a 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] [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 1f6d1a68fbbb8c..e15f4532cc3c8c 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 1368e6ac495ab6..592492995b84e5 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():