llvm / llvm/llvm-project

[MLIR][Python bindings] linalg.ElementwiseOp class shadowing breaks op isinstance matching

Open
#223,673 3 comments 0 reactions 1 assignee Claimed by @makslevental View on GitHub
mlir:python
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.