[Mlir-commits] [mlir] [MLIR][Python] Support bare Operand and Result annotations (PR #220216)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Sep 1 04:19:40 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Twice (PragmaTwice)
<details>
<summary>Changes</summary>
Operation definitions currently need to spell unconstrained operands and results as `Operand[Any]` and `Result[Any]`. Since `Any` is the natural default when no type constraint is intended, this is unnecessarily verbose and may require importing `Any` solely for these annotations.
This change allows bare `Operand` and `Result` annotations as a simpler equivalent spelling.
Previously:
```python
from typing import Any
class MyOp(MyDialect.Operation, name="my_op"):
lhs: Operand[Any]
rhs: Operand[Any]
res: Result[Any]
```
This can now be written as:
```python
class MyOp(MyDialect.Operation, name="my_op"):
lhs: Operand
rhs: Operand
res: Result
```
The existing `Operand[Any]` and `Result[Any]` spellings remain supported.
Assited-by: Codex/GPT 5.6 Sol
---
Full diff: https://github.com/llvm/llvm-project/pull/220216.diff
2 Files Affected:
- (modified) mlir/python/mlir/dialects/ext.py (+10-2)
- (modified) mlir/test/python/dialects/ext.py (+95)
``````````diff
diff --git a/mlir/python/mlir/dialects/ext.py b/mlir/python/mlir/dialects/ext.py
index ba321ffc32c32..491008f86023f 100644
--- a/mlir/python/mlir/dialects/ext.py
+++ b/mlir/python/mlir/dialects/ext.py
@@ -246,12 +246,20 @@ def from_type_hint(name, type_, specifier) -> "FieldDef":
type_ = get_args(type_)[0]
origin = get_origin(type_)
+ args = get_args(type_)
+
+ # An unsubscripted operand or result is shorthand for one constrained
+ # by `Any`.
+ if type_ is Operand or type_ is Result:
+ origin = type_
+ args = (Any,)
+
if origin is ir.OpResult:
if specifier.type_ and specifier.type_ is not Result:
raise TypeError(
f"only `result` field specifier can be used for result fields"
)
- constraint = get_args(type_)[0]
+ constraint = args[0]
return ResultDef(
name,
variadicity,
@@ -271,7 +279,7 @@ def from_type_hint(name, type_, specifier) -> "FieldDef":
return OperandDef(
name,
variadicity,
- get_args(type_)[0],
+ args[0],
param_kind=specifier.param_kind,
default_is_none=specifier.default_is_none,
)
diff --git a/mlir/test/python/dialects/ext.py b/mlir/test/python/dialects/ext.py
index d46060c654a7e..bbaca38dd130d 100644
--- a/mlir/test/python/dialects/ext.py
+++ b/mlir/test/python/dialects/ext.py
@@ -105,6 +105,101 @@ class AddOp(Operation, dialect=MyInt, name="add"):
print(adaptor1.rhs)
+# CHECK: TEST: testBareOperandAndResult
+ at run
+def testBareOperandAndResult():
+ class TestBare(Dialect, name="ext_bare"):
+ pass
+
+ class BareOp(TestBare.Operation, name="bare"):
+ lhs: Operand
+ rhs: Operand
+ res: Result
+
+ class BareSpecifierOp(TestBare.Operation, name="bare_specifier"):
+ arg: Operand = operand(kw_only=True)
+ res: Result = result(kw_only=True)
+
+ class BareVariadicOp(TestBare.Operation, name="bare_variadic"):
+ ins: Sequence[Operand]
+ outs: Sequence[Result]
+ maybe_in: Optional[Operand] = None
+ maybe_out: Optional[Result] = None
+
+ with Context(), Location.unknown():
+ TestBare.load()
+
+ # CHECK: irdl.dialect @ext_bare {
+ # CHECK: irdl.operation @bare {
+ # CHECK: %0 = irdl.any
+ # CHECK: %1 = irdl.any
+ # CHECK: irdl.operands(lhs: %0, rhs: %1)
+ # CHECK: %2 = irdl.any
+ # CHECK: irdl.results(res: %2)
+ # CHECK: }
+ # CHECK: irdl.operation @bare_specifier {
+ # CHECK: %0 = irdl.any
+ # CHECK: irdl.operands(arg: %0)
+ # CHECK: %1 = irdl.any
+ # CHECK: irdl.results(res: %1)
+ # CHECK: }
+ # CHECK: irdl.operation @bare_variadic {
+ # CHECK: %0 = irdl.any
+ # CHECK: %1 = irdl.any
+ # CHECK: irdl.operands(ins: variadic %0, maybe_in: optional %1)
+ # CHECK: %2 = irdl.any
+ # CHECK: %3 = irdl.any
+ # CHECK: irdl.results(outs: variadic %2, maybe_out: optional %3)
+ # CHECK: }
+ # CHECK: }
+ print(TestBare._mlir_module)
+
+ # CHECK: (self, /, lhs, rhs, res, *, loc=None, ip=None)
+ print(BareOp.__init__.__signature__)
+ # CHECK: (self, /, *, arg, res, loc=None, ip=None)
+ print(BareSpecifierOp.__init__.__signature__)
+ # CHECK: (self, /, ins, outs, *, maybe_in=None, maybe_out=None, loc=None, ip=None)
+ print(BareVariadicOp.__init__.__signature__)
+
+ i32 = IntegerType.get_signless(32)
+ f32 = F32Type.get()
+ module = Module.create()
+ with InsertionPoint(module.body):
+ ione = arith.constant(i32, 1)
+ fone = arith.constant(f32, 1.0)
+ bare = BareOp(ione, fone, f32)
+ bare_specifier = BareSpecifierOp(arg=fone, res=i32)
+ variadic = BareVariadicOp(
+ [ione, ione],
+ [i32, i32],
+ maybe_in=fone,
+ maybe_out=f32,
+ )
+ empty = BareVariadicOp([], [])
+
+ assert module.operation.verify()
+ # CHECK: "ext_bare.bare"(%c1_i32, %cst) : (i32, f32) -> f32
+ # CHECK: "ext_bare.bare_specifier"(%cst) : (f32) -> i32
+ # CHECK: "ext_bare.bare_variadic"(%c1_i32, %c1_i32, %cst) {operandSegmentSizes = array<i32: 2, 1>, resultSegmentSizes = array<i32: 2, 1>} : (i32, i32, f32) -> (i32, i32, f32)
+ # CHECK: "ext_bare.bare_variadic"() {operandSegmentSizes = array<i32: 0, 0>, resultSegmentSizes = array<i32: 0, 0>} : () -> ()
+ print(module)
+
+ # Different bare operands are constrained independently.
+ # CHECK: i32 f32 f32
+ print(bare.lhs.type, bare.rhs.type, bare.res.type)
+ # CHECK: f32 i32
+ print(bare_specifier.arg.type, bare_specifier.res.type)
+ # CHECK: 2 f32 2 f32
+ print(
+ len(variadic.ins),
+ variadic.maybe_in.type,
+ len(variadic.outs),
+ variadic.maybe_out.type,
+ )
+ # CHECK: 0 None 0 None
+ print(len(empty.ins), empty.maybe_in, len(empty.outs), empty.maybe_out)
+
+
# CHECK: TEST: testDialectLoadInMultipleContexts
@run
def testDialectLoadInMultipleContexts():
``````````
</details>
https://github.com/llvm/llvm-project/pull/220216
More information about the Mlir-commits
mailing list