[MLIR][Python bindings] linalg.ElementwiseOp class shadowing breaks op isinstance matching
- Dominant language
- LLVM
- Stars
- 40.5k
- Forks
- 18.7k
- PR merge metrics
- PR metrics pending
Description
## Summary
In the MLIR Python bindings, `mlir.dialects.linalg.ElementwiseOp` is **not** the class that op views are actually instantiated as. When an operation is reconstructed into an op view (e.g. via `OpOperand.owner`, `Operation.opview`, or during an IR walk), the returned object is an instance of the *generated base* class `mlir.dialects._linalg_ops_gen.ElementwiseOp`, whereas `linalg.ElementwiseOp` refers to a hand-written **subclass** (`ElementwiseOp_`).
Because the exported name is a subclass of the instantiated class, `isinstance(op.opview, linalg.ElementwiseOp)` and `match`/`case linalg.ElementwiseOp()` always evaluate to **False** for real `linalg.elementwise` ops.
## Root cause
[`mlir/python/mlir/dialects/linalg/__init__.py`](https://github.com/llvm/llvm-project/blob/184ef35f1d4fd0433b50c953f0495f4199d22223/mlir/python/mlir/dialects/linalg/__init__.py#L236) defines a convenience wrapper:
```python
class ElementwiseOp_(ElementwiseOp): # subclasses generated _linalg_ops_gen.ElementwiseOp
def __init__(self, ...): ...
# and re-exports it as the public name:
ElementwiseOp = ElementwiseOp_
```
The op-view registry, however, maps the operation name `linalg.elementwise` to the **generated base** class. So:
- Instances from the registry -> `_linalg_ops_gen.ElementwiseOp`
- Public symbol `linalg.ElementwiseOp` -> `ElementwiseOp_` (a subclass)
- `isinstance(base_instance, subclass)` -> `False`
MRO confirms the direction of the relationship:
```
mlir.dialects.linalg.ElementwiseOp_
-> mlir.dialects._linalg_ops_gen.ElementwiseOp
-> mlir._mlir_libs._mlir.ir.OpView
```
Note this only affects ops with hand-written wrapper subclasses. Ops without a wrapper (e.g. `FillOp`, `GenericOp`) satisfy `linalg.FillOp is _linalg_ops_gen.FillOp`, so `isinstance` works for them.
## Reproduction
```python
from mlir import ir
from mlir.dialects import linalg
from mlir.dialects import _linalg_ops_gen as gen
src = '''
func.func @f(%a: tensor<4x4xf16>, %b: tensor<4x4xf16>, %c: tensor<4x4xf32>, %bias: tensor<4x4xf32>) {
%0 = linalg.elementwise ins(%bias, %c : tensor<4x4xf32>, tensor<4x4xf32>)
outs(%c : tensor<4x4xf32>) -> tensor<4x4xf32>
return
}
'''
with ir.Context(), ir.Location.unknown():
m = ir.Module.parse(src)
op = m.body.operations[0].regions[0].blocks[0].operations[0]
print(type(op.opview)) # _linalg_ops_gen.ElementwiseOp
print(isinstance(op.opview, linalg.ElementwiseOp)) # False <-- surprising
print(isinstance(op.opview, gen.ElementwiseOp)) # True
print(linalg.ElementwiseOp is gen.ElementwiseOp) # False
print(linalg.FillOp is gen.FillOp) # True (unaffected)
```
## Expected
`isinstance(op.opview, linalg.ElementwiseOp)` should be `True` for a `linalg.elementwise` op -- the publicly exported op class should match the class produced by the op-view registry.
## Impact
Any downstream code doing type-based dispatch on op views breaks silently (no error, just a missed match):
- `isinstance(op.opview, linalg.ElementwiseOp)`
- `match op: case linalg.ElementwiseOp(): ...`
This regressed for us after named elementwise ops (e.g. `linalg.add`) were removed in favor of `linalg.elementwise `; the replacement op is one that carries a wrapper subclass, exposing the mismatch.
## Suggested fixes (either one)
1. Register the wrapper subclass in the op-view registry so reconstructed views use `ElementwiseOp_` (i.e. make the exported class the one the registry instantiates).
2. Have the wrapper only add a custom builder (`@classmethod`/factory) instead of subclassing, keeping `linalg.ElementwiseOp is _linalg_ops_gen.ElementwiseOp`.
## Workaround
Match on the operation name instead of the class:
```python
op.operation.name == "linalg.elementwise"
```
Assisted-by: Claude
Contributor guide
Assessment
This issue has not been assessed yet.