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

[llvmlistbot at llvm.org]

Author: Maksim Levental
Date: 2026-09-16T15:50:05Z
New Revision: f84aee6257bd37def74ce179252ce38c19323fc8

URL: https://github.com/llvm/llvm-project/commit/f84aee6257bd37def74ce179252ce38c19323fc8
DIFF: https://github.com/llvm/llvm-project/commit/f84aee6257bd37def74ce179252ce38c19323fc8.diff

LOG: [mlir][python] Register linalg.ElementwiseOp wrapper as op view (#223806)

Register the wrapper for `ElementwiseOp` with
`@_ods_cext.register_operation(_Dialect, replace=True)` so the exported
class is the one the registry instantiates.

Fixes #223673

Assisted by: Claude

Added: 
    

Modified: 
    mlir/python/mlir/dialects/linalg/__init__.py
    mlir/test/python/dialects/linalg/ops.py

Removed: 
    


################################################################################
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..127a812dd83e3 100644
--- a/mlir/test/python/dialects/linalg/ops.py
+++ b/mlir/test/python/dialects/linalg/ops.py
@@ -884,6 +884,25 @@ 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 <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]
+        assert isinstance(op.opview, linalg.ElementwiseOp)
+
+
 @run
 def testReduceOp():
     with Context(), Location.unknown():