[mlir][python] Infer result types in generated constructors whenever possible

[ftynse]

Summary

In several cases, operation result types can be unambiguously inferred from
operands and attributes at operation construction time. Stop requiring the user
to provide these types as arguments in the ODS-generated constructors in Python
bindings. In particular, handle the SameOperandAndResultTypes and
FirstAttrDerivedResultType traits as well as InferTypeOpInterface using the
recently added interface support. This is a significant usability improvement
for IR construction, similar to what C++ ODS provides.

Depends On D111656

[stellaraccident]

Would you mind sending a PSA on discourse: There is a lot of code out there that may need result type elided with this change.

[ftynse]

https://llvm.discourse.group/t/psa-some-python-op-constructors-will-no-longer-need-explicit-return-types/4482

[gysit]

LGTM!

[ftynse]

Rebase.

[ftynse]

Fix mismerge.

mlir/include/mlir/TableGen/Operator.h

public:
   attribute_iterator attribute_end() const;
   llvm::iterator_range<attribute_iterator> getAttributes() const;
 
   int getNumAttributes() const { return attributes.size(); }
   int getNumNativeAttributes() const { return numNativeAttributes; }
 
   // Op attribute accessors.
   NamedAttribute &getAttribute(int index) { return attributes[index]; }
+  const NamedAttribute &getAttribute(int index) const {
+    return attributes[index];
+  }
 
   // Op operand iterators.
   value_iterator operand_begin();
   value_iterator operand_end();
   value_range getOperands();
 
   int getNumOperands() const { return operands.size(); }
   NamedTypeConstraint &getOperand(int index) { return operands[index]; }

mlir/lib/Bindings/Python/IRModule.h

private:
   // Trampoline constructor that avoids null initializing members while
   // looking up parents.
   PyInsertionPoint(PyBlock block, llvm::Optional<PyOperationRef> refOperation)
       : refOperation(std::move(refOperation)), block(std::move(block)) {}
 
   llvm::Optional<PyOperationRef> refOperation;
   PyBlock block;
 };
+/// Wrapper around the generic MlirType.
+/// The lifetime of a type is bound by the PyContext that created it.
+class PyType : public BaseContextObject {
+public:
+  PyType(PyMlirContextRef contextRef, MlirType type)
+      : BaseContextObject(std::move(contextRef)), type(type) {}
+  bool operator==(const PyType &other);
+  operator MlirType() const { return type; }
+  MlirType get() const { return type; }
+
+  /// Gets a capsule wrapping the void* within the MlirType.
+  pybind11::object getCapsule();
+
+  /// Creates a PyType from the MlirType wrapped by a capsule.
+  /// Note that PyType instances are uniqued, so the returned object
+  /// may be a pre-existing object. Ownership of the underlying MlirType
+  /// is taken by calling this function.
+  static PyType createFromCapsule(pybind11::object capsule);
+
+private:
+  MlirType type;
+};
+
+/// CRTP base classes for Python types that subclass Type and should be
+/// castable from it (i.e. via something like IntegerType(t)).
+/// By default, type class hierarchies are one level deep (i.e. a
+/// concrete type class extends PyType); however, intermediate python-visible
+/// base classes can be modeled by specifying a BaseTy.
+template <typename DerivedTy, typename BaseTy = PyType>
+class PyConcreteType : public BaseTy {
+public:
+  // Derived classes must define statics for:
+  //   IsAFunctionTy isaFunction
+  //   const char *pyClassName
+  using ClassTy = pybind11::class_<DerivedTy, BaseTy>;
+  using IsAFunctionTy = bool (*)(MlirType);
+
+  PyConcreteType() = default;
+  PyConcreteType(PyMlirContextRef contextRef, MlirType t)
+      : BaseTy(std::move(contextRef), t) {}
+  PyConcreteType(PyType &orig)
+      : PyConcreteType(orig.getContext(), castFrom(orig)) {}
+
+  static MlirType castFrom(PyType &orig) {
+    if (!DerivedTy::isaFunction(orig)) {
+      auto origRepr = pybind11::repr(pybind11::cast(orig)).cast<std::string>();
+      throw SetPyError(PyExc_ValueError, llvm::Twine("Cannot cast type to ") +
+                                             DerivedTy::pyClassName +
+                                             " (from " + origRepr + ")");
+    }
+    return orig;
+  }
+
+  static void bind(pybind11::module &m) {
+    auto cls = ClassTy(m, DerivedTy::pyClassName, pybind11::module_local());
+    cls.def(pybind11::init<PyType &>(), pybind11::keep_alive<0, 1>());
+    cls.def_static("isinstance", [](PyType &otherType) -> bool {
+      return DerivedTy::isaFunction(otherType);
+    });
+    DerivedTy::bindDerived(cls);
+  }
+
+  /// Implemented by derived classes to add methods to the Python subclass.
+  static void bindDerived(ClassTy &m) {}
+};
 
 /// Wrapper around the generic MlirAttribute.
 /// The lifetime of a type is bound by the PyContext that created it.
 class PyAttribute : public BaseContextObject {
 public:
   PyAttribute(PyMlirContextRef contextRef, MlirAttribute attr)
       : BaseContextObject(std::move(contextRef)), attr(attr) {}
   bool operator==(const PyAttribute &other);
public:
 
   static void bind(pybind11::module &m) {
     auto cls = ClassTy(m, DerivedTy::pyClassName, pybind11::buffer_protocol(),
                        pybind11::module_local());
     cls.def(pybind11::init<PyAttribute &>(), pybind11::keep_alive<0, 1>());
     cls.def_static("isinstance", [](PyAttribute &otherAttr) -> bool {
       return DerivedTy::isaFunction(otherAttr);
     });
-    DerivedTy::bindDerived(cls);
-  }
-
-  /// Implemented by derived classes to add methods to the Python subclass.
-  static void bindDerived(ClassTy &m) {}
-};
-
-/// Wrapper around the generic MlirType.
-/// The lifetime of a type is bound by the PyContext that created it.
-class PyType : public BaseContextObject {
-public:
-  PyType(PyMlirContextRef contextRef, MlirType type)
-      : BaseContextObject(std::move(contextRef)), type(type) {}
-  bool operator==(const PyType &other);
-  operator MlirType() const { return type; }
-  MlirType get() const { return type; }
-
-  /// Gets a capsule wrapping the void* within the MlirType.
-  pybind11::object getCapsule();
-
-  /// Creates a PyType from the MlirType wrapped by a capsule.
-  /// Note that PyType instances are uniqued, so the returned object
-  /// may be a pre-existing object. Ownership of the underlying MlirType
-  /// is taken by calling this function.
-  static PyType createFromCapsule(pybind11::object capsule);
-
-private:
-  MlirType type;
-};
-
-/// CRTP base classes for Python types that subclass Type and should be
-/// castable from it (i.e. via something like IntegerType(t)).
-/// By default, type class hierarchies are one level deep (i.e. a
-/// concrete type class extends PyType); however, intermediate python-visible
-/// base classes can be modeled by specifying a BaseTy.
-template <typename DerivedTy, typename BaseTy = PyType>
-class PyConcreteType : public BaseTy {
-public:
-  // Derived classes must define statics for:
-  //   IsAFunctionTy isaFunction
-  //   const char *pyClassName
-  using ClassTy = pybind11::class_<DerivedTy, BaseTy>;
-  using IsAFunctionTy = bool (*)(MlirType);
-
-  PyConcreteType() = default;
-  PyConcreteType(PyMlirContextRef contextRef, MlirType t)
-      : BaseTy(std::move(contextRef), t) {}
-  PyConcreteType(PyType &orig)
-      : PyConcreteType(orig.getContext(), castFrom(orig)) {}
-
-  static MlirType castFrom(PyType &orig) {
-    if (!DerivedTy::isaFunction(orig)) {
-      auto origRepr = pybind11::repr(pybind11::cast(orig)).cast<std::string>();
-      throw SetPyError(PyExc_ValueError, llvm::Twine("Cannot cast type to ") +
-                                             DerivedTy::pyClassName +
-                                             " (from " + origRepr + ")");
-    }
-    return orig;
-  }
-
-  static void bind(pybind11::module &m) {
-    auto cls = ClassTy(m, DerivedTy::pyClassName, pybind11::module_local());
-    cls.def(pybind11::init<PyType &>(), pybind11::keep_alive<0, 1>());
-    cls.def_static("isinstance", [](PyType &otherType) -> bool {
-      return DerivedTy::isaFunction(otherType);
+    cls.def_property_readonly("type", [](PyAttribute &attr) {
+      return PyType(attr.getContext(), mlirAttributeGetType(attr));
     });
     DerivedTy::bindDerived(cls);
   }
 
   /// Implemented by derived classes to add methods to the Python subclass.
   static void bindDerived(ClassTy &m) {}
 };
 

mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py

if expr.scalar_arg:
         raise ValueError(f"Argument {expr.scalar_arg.arg} is not bound for "
                          f"this structured op.")
     elif expr.scalar_const:
       value_attr = Attribute.parse(expr.scalar_const.value)
       return arith.ConstantOp(value_attr.type, value_attr).result
     elif expr.scalar_index:
       dim_attr = IntegerAttr.get(
           IntegerType.get_signless(64), expr.scalar_index.dim)
-      return linalg.IndexOp(IndexType.get(), dim_attr).result
+      return linalg.IndexOp(dim_attr).result
     elif expr.scalar_apply:
       try:
         fn = getattr(self, f"_eval_{expr.scalar_apply.fn_name}")
       except AttributeError:
         raise ValueError(
             f"Function '{expr.scalar_apply.fn_name}' is not a known "
             "scalar body function")
       operand_values = [
for n in output_names:
         output_values.append(self.yield_mapping[n])
       except KeyError:
         raise ValueError(f"Body assignments do not assign all outputs: "
                          f"missing '{n}'")
     linalg.YieldOp(output_values)
 
   def _eval_add(self, lhs: Value, rhs: Value) -> Value:
     if _is_floating_point_type(lhs.type):
-      return arith.AddFOp(lhs.type, lhs, rhs).result
+      return arith.AddFOp(lhs, rhs).result
     if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
-      return arith.AddIOp(lhs.type, lhs, rhs).result
+      return arith.AddIOp(lhs, rhs).result
     raise NotImplementedError("Unsupported 'add' operand: {lhs}")
 
   def _eval_exp(self, x: Value) -> Value:
     if _is_floating_point_type(x.type):
-      return math.ExpOp(x.type, x).result
+      return math.ExpOp(x).result
     raise NotImplementedError("Unsupported 'exp' operand: {x}")
 
   def _eval_log(self, x: Value) -> Value:
     if _is_floating_point_type(x.type):
-      return math.LogOp(x.type, x).result
+      return math.LogOp(x).result
     raise NotImplementedError("Unsupported 'log' operand: {x}")
 
   def _eval_sub(self, lhs: Value, rhs: Value) -> Value:
     if _is_floating_point_type(lhs.type):
-      return arith.SubFOp(lhs.type, lhs, rhs).result
+      return arith.SubFOp(lhs, rhs).result
     if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
-      return arith.SubIOp(lhs.type, lhs, rhs).result
+      return arith.SubIOp(lhs, rhs).result
     raise NotImplementedError("Unsupported 'sub' operand: {lhs}")
 
   def _eval_mul(self, lhs: Value, rhs: Value) -> Value:
     if _is_floating_point_type(lhs.type):
-      return arith.MulFOp(lhs.type, lhs, rhs).result
+      return arith.MulFOp(lhs, rhs).result
     if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
-      return arith.MulIOp(lhs.type, lhs, rhs).result
+      return arith.MulIOp(lhs, rhs).result
     raise NotImplementedError("Unsupported 'mul' operand: {lhs}")
 
   def _eval_max(self, lhs: Value, rhs: Value) -> Value:
     if _is_floating_point_type(lhs.type):
-      return std.MaxFOp(lhs.type, lhs, rhs).result
+      return std.MaxFOp(lhs, rhs).result
     if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
-      return std.MaxSIOp(lhs.type, lhs, rhs).result
+      return std.MaxSIOp(lhs, rhs).result
     raise NotImplementedError("Unsupported 'max' operand: {lhs}")
 
   def _eval_max_unsigned(self, lhs: Value, rhs: Value) -> Value:
     if _is_floating_point_type(lhs.type):
-      return std.MaxFOp(lhs.type, lhs, rhs).result
+      return std.MaxFOp(lhs, rhs).result
     if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
-      return std.MaxUIOp(lhs.type, lhs, rhs).result
+      return std.MaxUIOp(lhs, rhs).result
     raise NotImplementedError("Unsupported 'max_unsigned' operand: {lhs}")
 
   def _eval_min(self, lhs: Value, rhs: Value) -> Value:
     if _is_floating_point_type(lhs.type):
-      return std.MinFOp(lhs.type, lhs, rhs).result
+      return std.MinFOp(lhs, rhs).result
     if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
-      return std.MinSIOp(lhs.type, lhs, rhs).result
+      return std.MinSIOp(lhs, rhs).result
     raise NotImplementedError("Unsupported 'min' operand: {lhs}")
 
   def _eval_min_unsigned(self, lhs: Value, rhs: Value) -> Value:
     if _is_floating_point_type(lhs.type):
-      return std.MinFOp(lhs.type, lhs, rhs).result
+      return std.MinFOp(lhs, rhs).result
     if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
-      return std.MinUIOp(lhs.type, lhs, rhs).result
+      return std.MinUIOp(lhs, rhs).result
     raise NotImplementedError("Unsupported 'min_unsigned' operand: {lhs}")
 
 
 def _infer_structured_outs(
     op_config: LinalgStructuredOpConfig,
     in_arg_defs: Sequence[OperandDefConfig], ins: Sequence[Value],
     out_arg_defs: Sequence[OperandDefConfig],
     outs: Union[Sequence[Value], OpResultList]) -> Tuple[ValueList, List[Type]]:

mlir/test/mlir-tblgen/op-python-bindings.td

 // RUN: mlir-tblgen -gen-python-op-bindings -bind-dialect=test -I %S/../../include %s | FileCheck %s
 
 include "mlir/IR/OpBase.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Bindings/Python/Attributes.td"
 
 // CHECK: @_ods_cext.register_dialect
 // CHECK: class _Dialect(_ods_ir.Dialect):
   // CHECK: DIALECT_NAMESPACE = "test"
   // CHECK: pass
 def Test_Dialect : Dialect {
   let name = "test";
def AttributedOpWithOperands : TestOp<"attributed_op_with_operands"> {
   // CHECK: @builtins.property
   // CHECK: def is_(self):
   // CHECK:   if "is" not in self.operation.attributes:
   // CHECK:     return None
   // CHECK:   return _ods_ir.FloatAttr(self.operation.attributes["is"])
   let arguments = (ins I32, UnitAttr:$in, F32, OptionalAttr<F32Attr>:$is);
 }
 
+// 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:   operands = []
+  // CHECK:   results = []
+  // 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)
+  let arguments = (ins TypeAttr:$type);
+  let results = (outs AnyType:$res, AnyType);
+}
+
+// CHECK-LABEL: OPERATION_NAME = "test.derive_result_types_variadic_op"
+def DeriveResultTypesVariadicOp : TestOp<"derive_result_types_variadic_op", [FirstAttrDerivedResultType]> {
+  // CHECK: def __init__(self, res, _gen_res_1, type, *, loc=None, ip=None):
+  let arguments = (ins TypeAttr:$type);
+  let results = (outs AnyType:$res, Variadic<AnyType>);
+}
 
 // CHECK: @_ods_cext.register_operation(_Dialect)
 // CHECK: class EmptyOp(_ods_ir.OpView):
 // CHECK-LABEL: OPERATION_NAME = "test.empty"
 def EmptyOp : TestOp<"empty">;
   // CHECK: def __init__(self, *, loc=None, ip=None):
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
   // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
   // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
 
+
+// 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:    operands = []
+  // CHECK:    results = []
+  // CHECK:    _ods_context = _ods_get_default_loc_context(loc)
+  // CHECK:    results = _ods_ir.InferTypeOpInterface(InferResultTypesImpliedOp).inferReturnTypes(
+  // CHECK:        operands=operands,
+  // CHECK:        attributes=_ods_ir.DictAttr.get(attributes, context=_ods_context),
+  // CHECK:        context=_ods_context,
+  // CHECK:        loc=loc)
+  let results = (outs I32:$i32, F32:$f32);
+}
+
+// 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:    operands = []
+  // CHECK:    results = []
+  // CHECK:    _ods_context = _ods_get_default_loc_context(loc)
+  // CHECK:    results = _ods_ir.InferTypeOpInterface(InferResultTypesOp).inferReturnTypes(
+  // CHECK:        operands=operands,
+  // CHECK:        attributes=_ods_ir.DictAttr.get(attributes, context=_ods_context),
+  // CHECK:        context=_ods_context,
+  // CHECK:        loc=loc)
+  let results = (outs AnyType, AnyType, AnyType);
+}
+
 // 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:   results.append(i32)
-  // CHECK:   results.append(_gen_res_1)
-  // CHECK:   results.append(i64)
   // CHECK:   operands.append(_get_op_result_or_value(_gen_arg_0))
   // CHECK:   operands.append(_get_op_result_or_value(f32))
   // CHECK:   operands.append(_get_op_result_or_value(_gen_arg_2))
+  // CHECK:   results.append(i32)
+  // CHECK:   results.append(_gen_res_1)
+  // CHECK:   results.append(i64)
   // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
   // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
 
   // CHECK: @builtins.property
   // CHECK: def f32(self):
   // CHECK:   return self.operation.operands[1]
   let arguments = (ins I32, F32:$f32, I64);
 
   // CHECK: @builtins.property
   // CHECK: def i32(self):
   // CHECK:   return self.operation.results[0]
   //
   // CHECK: @builtins.property
   // CHECK: def i64(self):
   // CHECK:   return self.operation.results[2]
-  let results = (outs I32:$i32, F32, I64:$i64);
+  let results = (outs I32:$i32, AnyFloat, I64:$i64);
 }
 
 // CHECK: @_ods_cext.register_operation(_Dialect)
 // CHECK: class OneVariadicOperandOp(_ods_ir.OpView):
 // CHECK-LABEL: OPERATION_NAME = "test.one_variadic_operand"
 // CHECK-NOT: _ODS_OPERAND_SEGMENTS
 // CHECK-NOT: _ODS_RESULT_SEGMENTS
 def OneVariadicOperandOp : TestOp<"one_variadic_operand"> {
def PythonKeywordOp : TestOp<"python_keyword"> {
   // CHECK:     attributes=attributes, results=results, operands=operands,
   // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
 
   // CHECK: @builtins.property
   // CHECK: def in_(self):
   // CHECK:   return self.operation.operands[0]
   let arguments = (ins AnyType:$in);
 }
+// CHECK-LABEL: OPERATION_NAME = "test.same_results"
+def SameResultsOp : TestOp<"same_results", [SameOperandsAndResultType]> {
+  // CHECK: def __init__(self, in1, in2, *, loc=None, ip=None):
+  // CHECK: operands = []
+  // CHECK: results = []
+  // CHECK: operands.append
+  // CHECK: results.extend([operands[0].type] * 1)
+  let arguments = (ins AnyType:$in1, AnyType:$in2);
+  let results = (outs AnyType:$res);
+}
+
+// CHECK-LABEL: OPERATION_NAME = "test.same_results_variadic"
+def SameResultsVariadicOp : TestOp<"same_results_variadic", [SameOperandsAndResultType]> {
+  // CHECK: def __init__(self, res, in1, in2, *, loc=None, ip=None):
+  let arguments = (ins AnyType:$in1, AnyType:$in2);
+  let results = (outs Variadic<AnyType>:$res);
+}
+
 
 // CHECK: @_ods_cext.register_operation(_Dialect)
 // CHECK: class SameVariadicOperandSizeOp(_ods_ir.OpView):
 // CHECK-LABEL: OPERATION_NAME = "test.same_variadic_operand"
 def SameVariadicOperandSizeOp : TestOp<"same_variadic_operand",
                                        [SameVariadicOperandSize]> {
   // CHECK: @builtins.property
   // CHECK: def variadic1(self):

 // CHECK: @_ods_cext.register_operation(_Dialect)
 // CHECK: class SimpleOp(_ods_ir.OpView):
 // CHECK-LABEL: OPERATION_NAME = "test.simple"
 def SimpleOp : TestOp<"simple"> {
   // CHECK: def __init__(self, i64, f64, i32, f32, *, loc=None, ip=None):
   // CHECK:   operands = []
   // CHECK:   results = []
   // CHECK:   attributes = {}
-  // CHECK:   results.append(i64)
-  // CHECK:   results.append(f64)
   // CHECK:   operands.append(_get_op_result_or_value(i32))
   // CHECK:   operands.append(_get_op_result_or_value(f32))
+  // CHECK:   results.append(i64)
+  // CHECK:   results.append(f64)
   // CHECK:   _ods_successors = None
   // CHECK:   super().__init__(self.build_generic(
   // CHECK:     attributes=attributes, results=results, operands=operands,
   // CHECK:     successors=_ods_successors, loc=loc, ip=ip))
 
   // CHECK: @builtins.property
   // CHECK: def i32(self):
   // CHECK:   return self.operation.operands[0]
   //
   // CHECK: @builtins.property
   // CHECK: def f32(self):
   // CHECK:   return self.operation.operands[1]
   let arguments = (ins I32:$i32, F32:$f32);
 
   // CHECK: @builtins.property
   // CHECK: def i64(self):
   // CHECK:   return self.operation.results[0]
   //
   // CHECK: @builtins.property
   // CHECK: def f64(self):
   // CHECK:   return self.operation.results[1]
-  let results = (outs I64:$i64, F64:$f64);
+  let results = (outs I64:$i64, AnyFloat:$f64);
 }
 
 // CHECK: @_ods_cext.register_operation(_Dialect)
 // CHECK: class WithSuccessorsOp(_ods_ir.OpView):
 // CHECK-LABEL: OPERATION_NAME = "test.with_successors"
 def WithSuccessorsOp : TestOp<"with_successors"> {
   // CHECK-NOT:  _ods_successors = None
   // CHECK:      _ods_successors = []
   // CHECK-NEXT: _ods_successors.append(successor)
   // CHECK-NEXT: _ods_successors.extend(successors)
   let successors = (successor AnySuccessor:$successor,
                               VariadicSuccessor<AnySuccessor>:$successors);
 }
+

mlir/test/python/dialects/math.py

 # CHECK-LABEL: TEST: testMathOps
 @run
 def testMathOps():
   with Context() as ctx, Location.unknown():
     module = Module.create()
     with InsertionPoint(module.body):
       @builtin.FuncOp.from_py_func(F32Type.get())
       def emit_sqrt(arg):
-        return mlir_math.SqrtOp(F32Type.get(), arg)
+        return mlir_math.SqrtOp(arg)
 
     # CHECK-LABEL: func @emit_sqrt(
     # CHECK-SAME:                  %[[ARG:.*]]: f32) {
     # CHECK:         math.sqrt %[[ARG]] : f32
     # CHECK:         return
     # CHECK:       }
     print(module)

mlir/test/python/dialects/python_test.py

 
 # CHECK-LABEL: TEST: inferReturnTypes
 @run
 def inferReturnTypes():
   with Context() as ctx, Location.unknown(ctx):
     test.register_python_test_dialect(ctx)
     module = Module.create()
     with InsertionPoint(module.body):
-      op = test.InferResultsOp(
-          IntegerType.get_signless(32), IntegerType.get_signless(64))
+      op = test.InferResultsOp()
       dummy = test.DummyOp()
 
     # CHECK: [Type(i32), Type(i64)]
     iface = InferTypeOpInterface(op)
     print(iface.inferReturnTypes())
 
     # CHECK: [Type(i32), Type(i64)]
     iface_static = InferTypeOpInterface(test.InferResultsOp)
else:
       assert False, "not expected dummy op to implement the interface"
 
     try:
       InferTypeOpInterface(test.DummyOp)
     except ValueError:
       pass
     else:
       assert False, "not expected dummy op class to implement the interface"
+
+
+# CHECK-LABEL: TEST: resultTypesDefinedByTraits
+@run
+def resultTypesDefinedByTraits():
+  with Context() as ctx, Location.unknown(ctx):
+    test.register_python_test_dialect(ctx)
+    module = Module.create()
+    with InsertionPoint(module.body):
+      inferred = test.InferResultsOp()
+      same = test.SameOperandAndResultTypeOp([inferred.results[0]])
+      # CHECK-COUNT-2: i32
+      print(same.one.type)
+      print(same.two.type)
+
+      first_type_attr = test.FirstAttrDeriveTypeAttrOp(
+          inferred.results[1], TypeAttr.get(IndexType.get()))
+      # CHECK-COUNT-2: index
+      print(first_type_attr.one.type)
+      print(first_type_attr.two.type)
+
+      first_attr = test.FirstAttrDeriveAttrOp(
+          FloatAttr.get(F32Type.get(), 3.14))
+      # CHECK-COUNT-3: f32
+      print(first_attr.one.type)
+      print(first_attr.two.type)
+      print(first_attr.three.type)
+
+      implied = test.InferResultsImpliedOp()
+      # CHECK: i32
+      print(implied.integer.type)
+      # CHECK: f64
+      print(implied.flt.type)
+      # CHECK: index
+      print(implied.index.type)

mlir/test/python/dialects/shape.py

 
 
 # CHECK-LABEL: TEST: testConstShape
 @run
 def testConstShape():
   with Context() as ctx, Location.unknown():
     module = Module.create()
     f32 = F32Type.get()
-    indexT = IndexType.get()
     with InsertionPoint(module.body):
       @builtin.FuncOp.from_py_func(
           RankedTensorType.get((12, -1), f32))
       def const_shape_tensor(arg):
-        return shape.ConstShapeOp(RankedTensorType.get((2,), indexT),
-                                  DenseElementsAttr.get(np.array([10, 20])))
+        return shape.ConstShapeOp(DenseElementsAttr.get(np.array([10, 20])))
 
     # CHECK-LABEL: func @const_shape_tensor(%arg0: tensor<12x?xf32>)
     # CHECK: shape.const_shape [10, 20] : tensor<2xindex>
     print(module)
-

mlir/test/python/ir/dialects.py

with Context() as ctx, Location.unknown():
     ctx.allow_unregistered_dialects = True
     m = Module.create()
 
     with InsertionPoint(m.body):
       f32 = F32Type.get()
       # Create via dialects context collection.
       input1 = createInput()
       input2 = createInput()
-      op1 = ctx.dialects.arith.AddFOp(input1.type, input1, input2)
+      op1 = ctx.dialects.arith.AddFOp(input1, input2)
 
       # Create via an import
       from mlir.dialects.arith import AddFOp
-      AddFOp(input1.type, input1, op1.result)
+      AddFOp(input1, op1.result)
 
   # CHECK: %[[INPUT0:.*]] = "pytest_dummy.intinput"
   # CHECK: %[[INPUT1:.*]] = "pytest_dummy.intinput"
   # CHECK: %[[R0:.*]] = arith.addf %[[INPUT0]], %[[INPUT1]] : f32
   # CHECK: %[[R1:.*]] = arith.addf %[[INPUT0]], %[[R0]] : f32
   m.operation.print()
 
 

mlir/test/python/python_test_ops.td

static ::mlir::LogicalResult inferReturnTypes(
       ::mlir::Builder b(context);
       inferredReturnTypes.push_back(b.getI32Type());
       inferredReturnTypes.push_back(b.getI64Type());
       return ::mlir::success();
     }
   }];
 }
 
+// If all result types are buildable, the InferTypeOpInterface is implied and is
+// autogenerated by C++ ODS.
+def InferResultsImpliedOp : TestOp<"infer_results_implied_op"> {
+  let results = (outs I32:$integer, F64:$flt, Index:$index);
+}
+
+def SameOperandAndResultTypeOp : TestOp<"same_operand_and_result_type_op",
+                                        [SameOperandsAndResultType]> {
+  let arguments = (ins Variadic<AnyType>);
+  let results = (outs AnyType:$one, AnyType:$two);
+}
+
+def FirstAttrDeriveTypeAttrOp : TestOp<"first_attr_derive_type_attr_op",
+                               [FirstAttrDerivedResultType]> {
+  let arguments = (ins AnyType:$input, TypeAttr:$type);
+  let results = (outs AnyType:$one, AnyType:$two);
+}
+
+def FirstAttrDeriveAttrOp : TestOp<"first_attr_derive_attr_op",
+                               [FirstAttrDerivedResultType]> {
+  let arguments = (ins AnyAttr:$iattr);
+  let results = (outs AnyType:$one, AnyType:$two, AnyType:$three);
+}
+
 #endif // PYTHON_TEST_OPS

mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp

 ///   {0} is the value to initialize the successors list to.
 constexpr const char *initSuccessorsTemplate = R"Py(_ods_successors = {0})Py";
 
 /// Template to append or extend the list of successors in the builder.
 ///   {0} is the list method ('append' or 'extend');
 ///   {1} is the value to add.
 constexpr const char *addSuccessorTemplate = R"Py(_ods_successors.{0}({1}))Py";
 
-/// Populates `builderArgs` with the Python-compatible names of builder function
-/// arguments, first the results, then the intermixed attributes and operands in
-/// the same order as they appear in the `arguments` field of the op definition.
-/// Additionally, `operandNames` is populated with names of operands in their
-/// order of appearance.
+/// Returns true if the SameArgumentAndResultTypes trait can be used to infer
+/// result types of the given operation.
+static bool hasSameArgumentAndResultTypes(const Operator &op) {
+  return op.getTrait("::mlir::OpTrait::SameOperandsAndResultType") &&
+         op.getNumVariableLengthResults() == 0;
+}
+
+/// Returns true if the FirstAttrDerivedResultType trait can be used to infer
+/// result types of the given operation.
+static bool hasFirstAttrDerivedResultTypes(const Operator &op) {
+  return op.getTrait("::mlir::OpTrait::FirstAttrDerivedResultType") &&
+         op.getNumVariableLengthResults() == 0;
+}
+
+/// Returns true if the InferTypeOpInterface can be used to infer result types
+/// of the given operation.
+static bool hasInferTypeInterface(const Operator &op) {
+  return op.getTrait("::mlir::InferTypeOpInterface::Trait") &&
+         op.getNumRegions() == 0;
+}
+
+/// Returns true if there is a trait or interface that can be used to infer
+/// result types of the given operation.
+static bool canInferType(const Operator &op) {
+  return hasSameArgumentAndResultTypes(op) ||
+         hasFirstAttrDerivedResultTypes(op) || hasInferTypeInterface(op);
+}
+
+/// Populates `builderArgs` with result names if the builder is expected to
+/// accept them as arguments.
 static void
-populateBuilderArgs(const Operator &op,
-                    llvm::SmallVectorImpl<std::string> &builderArgs,
-                    llvm::SmallVectorImpl<std::string> &operandNames,
-                    llvm::SmallVectorImpl<std::string> &successorArgNames) {
+populateBuilderArgsResults(const Operator &op,
+                           llvm::SmallVectorImpl<std::string> &builderArgs) {
+  if (canInferType(op))
+    return;
+
   for (int i = 0, e = op.getNumResults(); i < e; ++i) {
     std::string name = op.getResultName(i).str();
     if (name.empty()) {
       if (op.getNumResults() == 1) {
         // Special case for one result, make the default name be 'result'
         // to properly match the built-in result accessor.
         name = "result";
       } else {
         name = llvm::formatv("_gen_res_{0}", i);
       }
     }
     name = sanitizeName(name);
     builderArgs.push_back(name);
   }
+}
+
+/// Populates `builderArgs` with the Python-compatible names of builder function
+/// arguments using intermixed attributes and operands in the same order as they
+/// appear in the `arguments` field of the op definition. Additionally,
+/// `operandNames` is populated with names of operands in their order of
+/// appearance.
+static void
+populateBuilderArgs(const Operator &op,
+                    llvm::SmallVectorImpl<std::string> &builderArgs,
+                    llvm::SmallVectorImpl<std::string> &operandNames,
+                    llvm::SmallVectorImpl<std::string> &successorArgNames) {
+
   for (int i = 0, e = op.getNumArgs(); i < e; ++i) {
     std::string name = op.getArgName(i).str();
     if (name.empty())
       name = llvm::formatv("_gen_arg_{0}", i);
     name = sanitizeName(name);
     builderArgs.push_back(name);
     if (!op.getArg(i).is<NamedAttribute *>())
       operandNames.push_back(name);
if (!element.isVariableLength()) {
         formatString = multiOperandAppendTemplate;
       }
     }
 
     builderLines.push_back(llvm::formatv(formatString.data(), name));
   }
 }
 
+/// Python code template for deriving the operation result types from its
+/// attribute:
+///   - {0} is the name of the attribute from which to derive the types.
+constexpr const char *deriveTypeFromAttrTemplate =
+    R"PY(_ods_result_type_source_attr = attributes["{0}"]
+_ods_derived_result_type = (
+    _ods_ir.TypeAttr(_ods_result_type_source_attr).value
+    if _ods_ir.TypeAttr.isinstance(_ods_result_type_source_attr) else
+    _ods_result_type_source_attr.type))PY";
+
+/// Python code template appending {0} type {1} times to the results list.
+constexpr const char *appendSameResultsTemplate = "results.extend([{0}] * {1})";
+
+/// Python code template for inferring the operation results using the
+/// corresponding interface:
+///   - {0} is the name of the class for which the types are inferred.
+constexpr const char *inferTypeInterfaceTemplate =
+    R"PY(_ods_context = _ods_get_default_loc_context(loc)
+results = _ods_ir.InferTypeOpInterface({0}).inferReturnTypes(
+    operands=operands,
+    attributes=_ods_ir.DictAttr.get(attributes, context=_ods_context),
+    context=_ods_context,
+    loc=loc)
+)PY";
+
+/// Appends the given multiline string as individual strings into
+/// `builderLines`.
+static void appendLineByLine(StringRef string,
+                             llvm::SmallVectorImpl<std::string> &builderLines) {
+
+  std::pair<StringRef, StringRef> split = std::make_pair(string, string);
+  do {
+    split = split.second.split('\n');
+    builderLines.push_back(split.first.str());
+  } while (!split.second.empty());
+}
+
 /// Populates `builderLines` with additional lines that are required in the
 /// builder to set up op results.
 static void
 populateBuilderLinesResult(const Operator &op,
                            llvm::ArrayRef<std::string> names,
                            llvm::SmallVectorImpl<std::string> &builderLines) {
   bool sizedSegments = op.getTrait(attrSizedTraitForKind("result")) != nullptr;
 
+  if (hasSameArgumentAndResultTypes(op)) {
+    builderLines.push_back(llvm::formatv(
+        appendSameResultsTemplate, "operands[0].type", op.getNumResults()));
+    return;
+  }
+
+  if (hasFirstAttrDerivedResultTypes(op)) {
+    const NamedAttribute &firstAttr = op.getAttribute(0);
+    assert(!firstAttr.name.empty() && "unexpected empty name for the attribute "
+                                      "from which the type is derived");
+    appendLineByLine(
+        llvm::formatv(deriveTypeFromAttrTemplate, firstAttr.name).str(),
+        builderLines);
+    builderLines.push_back(llvm::formatv(appendSameResultsTemplate,
+                                         "_ods_derived_result_type",
+                                         op.getNumResults()));
+    return;
+  }
+
+  if (hasInferTypeInterface(op)) {
+    appendLineByLine(
+        llvm::formatv(inferTypeInterfaceTemplate, op.getCppClassName()).str(),
+        builderLines);
+    return;
+  }
+
   // For each element, find or generate a name.
   for (int i = 0, e = op.getNumResults(); i < e; ++i) {
     const NamedTypeConstraint &element = op.getResult(i);
     std::string name = names[i];
 
     // Choose the formatting string based on the element kind.
     llvm::StringRef formatString;
     if (!element.isVariableLength()) {
if (op.skipDefaultBuilders())
     return;
 
   llvm::SmallVector<std::string> builderArgs;
   llvm::SmallVector<std::string> builderLines;
   llvm::SmallVector<std::string> operandArgNames;
   llvm::SmallVector<std::string> successorArgNames;
   builderArgs.reserve(op.getNumOperands() + op.getNumResults() +
                       op.getNumNativeAttributes() + op.getNumSuccessors());
+  populateBuilderArgsResults(op, builderArgs);
+  size_t numResultArgs = builderArgs.size();
   populateBuilderArgs(op, builderArgs, operandArgNames, successorArgNames);
 
-  populateBuilderLinesResult(
-      op, llvm::makeArrayRef(builderArgs).take_front(op.getNumResults()),
-      builderLines);
   populateBuilderLinesOperand(op, operandArgNames, builderLines);
   populateBuilderLinesAttr(
-      op, llvm::makeArrayRef(builderArgs).drop_front(op.getNumResults()),
+      op, llvm::makeArrayRef(builderArgs).drop_front(numResultArgs),
+      builderLines);
+  populateBuilderLinesResult(
+      op, llvm::makeArrayRef(builderArgs).take_front(numResultArgs),
       builderLines);
   populateBuilderLinesSuccessors(op, successorArgNames, builderLines);
 
   builderArgs.push_back("*");
   builderArgs.push_back("loc=None");
   builderArgs.push_back("ip=None");
   os << llvm::formatv(initTemplate, llvm::join(builderArgs, ", "),
                       llvm::join(builderLines, "\n    "));