Skip to content

[MLIR][Python] Add optional results parameter for building op with inferable result types #156818

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Sep 5, 2025
Merged

[MLIR][Python] Add optional results parameter for building op with inferable result types #156818

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
689a092
[MLIR][Python] Add optional results parameter for building op with in…
PragmaTwice Sep 4, 2025
8ed02d8
Merge branch 'main' into mlir-python-results
PragmaTwice Sep 4, 2025
7d434c9
add test case for results parameter
PragmaTwice Sep 4, 2025
4e72e8f
fix test
PragmaTwice Sep 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
67 changes: 33 additions & 34 deletions mlir/test/mlir-tblgen/op-python-bindings.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,12 +23,12 @@ def AttrSizedOperandsOp : TestOp<"attr_sized_operands",
[AttrSizedOperandSegments]> {
// CHECK: def __init__(self, variadic1, non_variadic, *, variadic2=None, loc=None, ip=None):
// CHECK: operands = []
// CHECK: results = []
// CHECK: attributes = {}
// CHECK: regions = None
// CHECK: operands.append(_get_op_results_or_values(variadic1))
// CHECK: operands.append(non_variadic)
// CHECK: operands.append(variadic2)
// CHECK: results = []
// CHECK: _ods_successors = None
// CHECK: super().__init__(
// CHECK: self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS,
Expand Down Expand Up @@ -71,9 +71,9 @@ def AttrSizedResultsOp : TestOp<"attr_sized_results",
[AttrSizedResultSegments]> {
// CHECK: def __init__(self, variadic1, non_variadic, variadic2, *, loc=None, ip=None):
// CHECK: operands = []
// CHECK: results = []
// CHECK: attributes = {}
// CHECK: regions = None
// CHECK: results = []
// CHECK: if variadic1 is not None: results.append(variadic1)
// CHECK: results.append(non_variadic)
// CHECK: results.append(variadic2)
Expand Down Expand Up @@ -120,7 +120,6 @@ def AttrSizedResultsOp : TestOp<"attr_sized_results",
def AttributedOp : TestOp<"attributed_op"> {
// CHECK: def __init__(self, i32attr, in_, *, optionalF32Attr=None, unitAttr=None, loc=None, ip=None):
// CHECK: operands = []
// CHECK: results = []
// CHECK: attributes = {}
// CHECK: regions = None
// CHECK: attributes["i32attr"] = (i32attr if (
Expand All @@ -131,6 +130,7 @@ def AttributedOp : TestOp<"attributed_op"> {
// CHECK: if bool(unitAttr): attributes["unitAttr"] = _ods_ir.UnitAttr.get(
// CHECK: _ods_get_default_loc_context(loc))
// CHECK: attributes["in"] = (in_
// CHECK: results = []
// CHECK: _ods_successors = None
// CHECK: super().__init__(
// CHECK: self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS,
Expand Down Expand Up @@ -170,14 +170,14 @@ def AttributedOp : TestOp<"attributed_op"> {
def AttributedOpWithOperands : TestOp<"attributed_op_with_operands"> {
// CHECK: def __init__(self, _gen_arg_0, _gen_arg_2, *, in_=None, is_=None, loc=None, ip=None):
// CHECK: operands = []
// CHECK: results = []
// CHECK: attributes = {}
// CHECK: regions = None
// CHECK: operands.append(_gen_arg_0)
// CHECK: operands.append(_gen_arg_2)
// CHECK: if bool(in_): attributes["in"] = _ods_ir.UnitAttr.get(
// CHECK: _ods_get_default_loc_context(loc))
// CHECK: if is_ is not None: attributes["is"] = (is_
// CHECK: results = []
// CHECK: _ods_successors = None
// CHECK: super().__init__(
// CHECK: self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS,
Expand Down Expand Up @@ -205,11 +205,11 @@ def AttributedOpWithOperands : TestOp<"attributed_op_with_operands"> {
def DefaultValuedAttrsOp : TestOp<"default_valued_attrs"> {
// CHECK: def __init__(self, *, arr=None, unsupported=None, loc=None, ip=None):
// CHECK: operands = []
// CHECK: results = []
// CHECK: attributes = {}
// CHECK: regions = None
// CHECK: if arr is not None: attributes["arr"] = (arr
// CHECK: if unsupported is not None: attributes["unsupported"] = (unsupported
// CHECK: results = []
// CHECK: _ods_successors = None
// CHECK: super().__init__(
// CHECK: self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS,
Expand All @@ -226,21 +226,21 @@ def DefaultValuedAttrsOp : TestOp<"default_valued_attrs"> {

// CHECK-LABEL: OPERATION_NAME = "test.derive_result_types_op"
def DeriveResultTypesOp : TestOp<"derive_result_types_op", [FirstAttrDerivedResultType]> {
// CHECK: def __init__(self, type_, *, loc=None, ip=None):
// CHECK: def __init__(self, type_, *, results=None, loc=None, ip=None):
// CHECK: operands = []
// CHECK: results = []
// CHECK: _ods_result_type_source_attr = attributes["type"]
// CHECK: _ods_derived_result_type = (
// CHECK: if results is None:
// CHECK: _ods_result_type_source_attr = attributes["type"]
// CHECK: _ods_derived_result_type = (
// CHECK: _ods_ir.TypeAttr(_ods_result_type_source_attr).value
// CHECK: if _ods_ir.TypeAttr.isinstance(_ods_result_type_source_attr) else
// CHECK: _ods_result_type_source_attr.type)
// CHECK: results.extend([_ods_derived_result_type] * 2)
// CHECK: results = [_ods_derived_result_type] * 2
let arguments = (ins TypeAttr:$type);
let results = (outs AnyType:$res, AnyType);
}

// CHECK: def derive_result_types_op(type_, *, loc=None, ip=None)
// CHECK: return DeriveResultTypesOp(type_=type_, loc=loc, ip=ip).results
// CHECK: def derive_result_types_op(type_, *, results=None, loc=None, ip=None)
// CHECK: return DeriveResultTypesOp(type_=type_, results=results, loc=loc, ip=ip).results

// CHECK-LABEL: OPERATION_NAME = "test.derive_result_types_variadic_op"
def DeriveResultTypesVariadicOp : TestOp<"derive_result_types_variadic_op", [FirstAttrDerivedResultType]> {
Expand All @@ -258,9 +258,9 @@ def DeriveResultTypesVariadicOp : TestOp<"derive_result_types_variadic_op", [Fir
def EmptyOp : TestOp<"empty">;
// CHECK: def __init__(self, *, loc=None, ip=None):
// CHECK: operands = []
// CHECK: results = []
// CHECK: attributes = {}
// CHECK: regions = None
// CHECK: results = []
// CHECK: _ods_successors = None
// CHECK: super().__init__(
// CHECK: self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS,
Expand All @@ -272,44 +272,44 @@ def EmptyOp : TestOp<"empty">;

// CHECK-LABEL: OPERATION_NAME = "test.infer_result_types_implied_op"
def InferResultTypesImpliedOp : TestOp<"infer_result_types_implied_op"> {
// CHECK: def __init__(self, *, loc=None, ip=None):
// CHECK: def __init__(self, *, results=None, loc=None, ip=None):
// CHECK: _ods_context = _ods_get_default_loc_context(loc)
// CHECK: super().__init__(
// CHECK: self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS,
// CHECK: attributes=attributes, operands=operands,
// CHECK: attributes=attributes, results=results, operands=operands,
// CHECK: successors=_ods_successors, regions=regions, loc=loc, ip=ip)
let results = (outs I32:$i32, F32:$f32);
}

// CHECK: def infer_result_types_implied_op(*, loc=None, ip=None)
// CHECK: return InferResultTypesImpliedOp(loc=loc, ip=ip).results
// CHECK: def infer_result_types_implied_op(*, results=None, loc=None, ip=None)
// CHECK: return InferResultTypesImpliedOp(results=results, loc=loc, ip=ip).results

// CHECK-LABEL: OPERATION_NAME = "test.infer_result_types_op"
def InferResultTypesOp : TestOp<"infer_result_types_op", [InferTypeOpInterface]> {
// CHECK: def __init__(self, *, loc=None, ip=None):
// CHECK: def __init__(self, *, results=None, loc=None, ip=None):
// CHECK: operands = []
// CHECK: super().__init__(
// CHECK: self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS,
// CHECK: attributes=attributes, operands=operands,
// CHECK: attributes=attributes, results=results, operands=operands,
// CHECK: successors=_ods_successors, regions=regions, loc=loc, ip=ip)
let results = (outs AnyType, AnyType, AnyType);
}

// CHECK: def infer_result_types_op(*, loc=None, ip=None)
// CHECK: return InferResultTypesOp(loc=loc, ip=ip).results
// CHECK: def infer_result_types_op(*, results=None, loc=None, ip=None)
// CHECK: return InferResultTypesOp(results=results, loc=loc, ip=ip).results

// CHECK: @_ods_cext.register_operation(_Dialect)
// CHECK: class MissingNamesOp(_ods_ir.OpView):
// CHECK-LABEL: OPERATION_NAME = "test.missing_names"
def MissingNamesOp : TestOp<"missing_names"> {
// CHECK: def __init__(self, i32, _gen_res_1, i64, _gen_arg_0, f32, _gen_arg_2, *, loc=None, ip=None):
// CHECK: operands = []
// CHECK: results = []
// CHECK: attributes = {}
// CHECK: regions = None
// CHECK: operands.append(_gen_arg_0)
// CHECK: operands.append(f32)
// CHECK: operands.append(_gen_arg_2)
// CHECK: results = []
// CHECK: results.append(i32)
// CHECK: results.append(_gen_res_1)
// CHECK: results.append(i64)
Expand Down Expand Up @@ -346,11 +346,11 @@ def OneOptionalOperandOp : TestOp<"one_optional_operand"> {
let arguments = (ins AnyType:$non_optional, Optional<AnyType>:$optional);
// CHECK: def __init__(self, non_optional, *, optional=None, loc=None, ip=None):
// CHECK: operands = []
// CHECK: results = []
// CHECK: attributes = {}
// CHECK: regions = None
// CHECK: operands.append(non_optional)
// CHECK: if optional is not None: operands.append(optional)
// CHECK: results = []
// CHECK: _ods_successors = None
// CHECK: super().__init__(
// CHECK: self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS
Expand All @@ -377,11 +377,11 @@ def OneOptionalOperandOp : TestOp<"one_optional_operand"> {
def OneVariadicOperandOp : TestOp<"one_variadic_operand"> {
// CHECK: def __init__(self, non_variadic, variadic, *, loc=None, ip=None):
// CHECK: operands = []
// CHECK: results = []
// CHECK: attributes = {}
// CHECK: regions = None
// CHECK: operands.append(non_variadic)
// CHECK: operands.extend(_get_op_results_or_values(variadic))
// CHECK: results = []
// CHECK: _ods_successors = None
// CHECK: super().__init__(
// CHECK: self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS
Expand Down Expand Up @@ -410,9 +410,9 @@ def OneVariadicOperandOp : TestOp<"one_variadic_operand"> {
def OneVariadicResultOp : TestOp<"one_variadic_result"> {
// CHECK: def __init__(self, variadic, non_variadic, *, loc=None, ip=None):
// CHECK: operands = []
// CHECK: results = []
// CHECK: attributes = {}
// CHECK: regions = None
// CHECK: results = []
// CHECK: results.extend(variadic)
// CHECK: results.append(non_variadic)
// CHECK: _ods_successors = None
Expand Down Expand Up @@ -442,10 +442,10 @@ def OneVariadicResultOp : TestOp<"one_variadic_result"> {
def PythonKeywordOp : TestOp<"python_keyword"> {
// CHECK: def __init__(self, in_, *, loc=None, ip=None):
// CHECK: operands = []
// CHECK: results = []
// CHECK: attributes = {}
// CHECK: regions = None
// CHECK: operands.append(in_)
// CHECK: results = []
// CHECK: _ods_successors = None
// CHECK: super().__init__(
// CHECK: self.OPERATION_NAME, self._ODS_REGIONS, self._ODS_OPERAND_SEGMENTS, self._ODS_RESULT_SEGMENTS
Expand All @@ -463,17 +463,16 @@ def PythonKeywordOp : TestOp<"python_keyword"> {

// CHECK-LABEL: OPERATION_NAME = "test.same_results"
def SameResultsOp : TestOp<"same_results", [SameOperandsAndResultType]> {
// CHECK: def __init__(self, in1, in2, *, loc=None, ip=None):
// CHECK: def __init__(self, in1, in2, *, results=None, loc=None, ip=None):
// CHECK: operands = []
// CHECK: results = []
// CHECK: operands.append
// CHECK: results.extend([operands[0].type] * 1)
// CHECK: if results is None: results = [operands[0].type] * 1
let arguments = (ins AnyType:$in1, AnyType:$in2);
let results = (outs AnyType:$res);
}

// CHECK: def same_results(in1, in2, *, loc=None, ip=None)
// CHECK: return SameResultsOp(in1=in1, in2=in2, loc=loc, ip=ip)
// CHECK: def same_results(in1, in2, *, results=None, loc=None, ip=None)
// CHECK: return SameResultsOp(in1=in1, in2=in2, results=results, loc=loc, ip=ip)

// CHECK-LABEL: OPERATION_NAME = "test.same_results_variadic"
def SameResultsVariadicOp : TestOp<"same_results_variadic", [SameOperandsAndResultType]> {
Expand Down Expand Up @@ -544,11 +543,11 @@ def SameVariadicResultSizeOp : TestOp<"same_variadic_result",
def SimpleOp : TestOp<"simple"> {
// CHECK: def __init__(self, i64, f64, i32, f32, *, loc=None, ip=None):
// CHECK: operands = []
// CHECK: results = []
// CHECK: attributes = {}
// CHECK: regions = None
// CHECK: operands.append(i32)
// CHECK: operands.append(f32)
// CHECK: results = []
// CHECK: results.append(i64)
// CHECK: results.append(f64)
// CHECK: _ods_successors = None
Expand Down Expand Up @@ -584,9 +583,9 @@ def SimpleOp : TestOp<"simple"> {
def VariadicAndNormalRegionOp : TestOp<"variadic_and_normal_region"> {
// CHECK: def __init__(self, num_variadic, *, loc=None, ip=None):
// CHECK: operands = []
// CHECK: results = []
// CHECK: attributes = {}
// CHECK: regions = None
// CHECK: results = []
// CHECK: _ods_successors = None
// CHECK: regions = 2 + num_variadic
// CHECK: super().__init__(
Expand All @@ -612,9 +611,9 @@ def VariadicAndNormalRegionOp : TestOp<"variadic_and_normal_region"> {
def VariadicRegionOp : TestOp<"variadic_region"> {
// CHECK: def __init__(self, num_variadic, *, loc=None, ip=None):
// CHECK: operands = []
// CHECK: results = []
// CHECK: attributes = {}
// CHECK: regions = None
// CHECK: results = []
// CHECK: _ods_successors = None
// CHECK: regions = 0 + num_variadic
// CHECK: super().__init__(
Expand Down
13 changes: 13 additions & 0 deletions mlir/test/python/dialects/python_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -283,6 +283,10 @@ def resultTypesDefinedByTraits():
module = Module.create()
with InsertionPoint(module.body):
inferred = test.InferResultsOp()

# CHECK: i32 i64
print(inferred.single.type, inferred.doubled.type)

same = test.SameOperandAndResultTypeOp([inferred.results[0]])
# CHECK-COUNT-2: i32
print(same.one.type)
Expand All @@ -309,6 +313,15 @@ def resultTypesDefinedByTraits():
# CHECK: index
print(implied.index.type)

# provide the result types to avoid inferring them
f64 = F64Type.get()
no_imply = test.InferResultsImpliedOp(results=[f64, f64, f64])
# CHECK-COUNT-3: f64
print(no_imply.integer.type, no_imply.flt.type, no_imply.index.type)

no_infer = test.InferResultsOp(results=[F32Type.get(), IndexType.get()])
# CHECK: f32 index
print(no_infer.single.type, no_infer.doubled.type)

# CHECK-LABEL: TEST: testOptionalOperandOp
@run
Expand Down
6 changes: 3 additions & 3 deletions mlir/test/python/ir/auto_location.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ def testInferLocations():
_cext.globals.register_traceback_file_inclusion(_arith_ops_gen.__file__)
three = arith.constant(IndexType.get(), 3)
# fmt: off
# CHECK: loc(callsite("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]_arith_ops_gen.py":405:4 to :218) at callsite("testInferLocations"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":52:16 to :50) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":13:4 to :7) at "<module>"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":26:1 to :4)))))
# CHECK: loc(callsite("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]_arith_ops_gen.py":397:4 to :235) at callsite("testInferLocations"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":52:16 to :50) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":13:4 to :7) at "<module>"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":26:1 to :4)))))
# fmt: on
print(three.location)

Expand All @@ -60,14 +60,14 @@ def foo():
print(four.location)

# fmt: off
# CHECK: loc(callsite("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]_arith_ops_gen.py":405:4 to :218) at callsite("testInferLocations.<locals>.foo"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":59:19 to :53) at callsite("testInferLocations"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":65:8 to :13) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":13:4 to :7) at "<module>"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":26:1 to :4))))))
# CHECK: loc(callsite("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]_arith_ops_gen.py":397:4 to :235) at callsite("testInferLocations.<locals>.foo"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":59:19 to :53) at callsite("testInferLocations"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":65:8 to :13) at callsite("run"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":13:4 to :7) at "<module>"("{{.*}}[[SEP]]test[[SEP]]python[[SEP]]ir[[SEP]]auto_location.py":26:1 to :4))))))
# fmt: on
foo()

_cext.globals.register_traceback_file_exclusion(__file__)

# fmt: off
# CHECK: loc("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]_arith_ops_gen.py":405:4 to :218))
# CHECK: loc("ConstantOp.__init__"("{{.*}}[[SEP]]mlir[[SEP]]dialects[[SEP]]_arith_ops_gen.py":397:4 to :235))
# fmt: on
foo()

Expand Down
Loading