[mlir][ods] Add simple generator for return types

[jpienaar]

Summary

Take advantage of equality constrains to generate the type inference interface.
This is used for equality and trivially built types. The type inference method
is only generated when no type inference trait is specified already. Avoids the
need for custom build methods for simple cases that have been increasing. More
work needed to consolidate more, but simplifies the specification and can
iterate implementation side post.

This reorders verification that changes some test error messages (two verifications that overlap and the original verification isn't reached to be triggered).

[silvas]

I feel like there must be some documentation somewhere that should be updated to describe this behavior?

[silvas]

Awesome! Thanks so much for doing this :)

mlir/lib/TableGen/Operator.cpp

258

This comment seems a bit out of date?

263

any_of(concat(operands, results), [](op) { return op.IsVariadic()}) would be a bit more succicnt.

297

maybe a small comment explaining how we intend to use EquivalenceClasses.

[jpienaar]

Added documentation and address review comments

mlir/docs/OpDefinitions.md

-# Table-driven Operation Definition Specification (ODS)
+# Operation Definition Specification (ODS)
 
 In addition to specializing the `mlir::Op` C++ template, MLIR also supports
 defining operations in a table-driven manner. This is achieved via
 [TableGen][TableGen], which is both a generic language and its tooling to
 maintain records of domain-specific information. Facts regarding an operation
 are specified concisely into a TableGen record, which will be expanded into an
 equivalent `mlir::Op` C++ template specialization at compiler build time.
 

 
 // Each operand/attribute has a separate parameter but result type is aggregate.
 static void build(OpBuilder &odsBuilder, OperationState &odsState,
                   ArrayRef<Type> resultTypes,
                   Value i32_operand, Value f32_operand, ...,
                   IntegerAttr i32_attr, FloatAttr f32_attr, ...);
 
 // All operands/attributes have aggregate parameters.
-// Generated if InferTypeOpInterface interface is specified.
+// Generated if return type can be inferred.
 static void build(OpBuilder &odsBuilder, OperationState &odsState,
-                  ValueRange operands,
-                  ArrayRef<NamedAttribute> attributes);
+                  ValueRange operands, ArrayRef<NamedAttribute> attributes);
 
 // (And manually specified builders depending on the specific op.)
 ```
 
 The first form provides basic uniformity so that we can create ops using the
 same form regardless of the exact op. This is particularly useful for
 implementing declarative pattern rewrites.
 
 The second and third forms are good for use in manually written code given that
 they provide better guarantee via signatures.
 
 The third form will be generated if any of the op's attribute has different
 `Attr.returnType` from `Attr.storageType` and we know how to build an attribute
 from an unwrapped value (i.e., `Attr.constBuilderCall` is defined.)
 Additionally, for the third form, if an attribute appearing later in the
 `arguments` list has a default value, the default value will be supplied in the
 declaration. This works for `BoolAttr`, `StrAttr`, `EnumAttr` for now and the
 list can grow in the future. So if possible, default valued attribute should be
 placed at the end of the `arguments` list to leverage this feature. (This
 behavior is essentially due to C++ function parameter default value placement
 restrictions.) Otherwise, the builder of the third form will still be generated
 but default values for the attributes not at the end of the `arguments` list
 will not be supplied in the builder's signature.
 
+ODS will generate a builder that doesn't require return type specified if
+
+*   Op implements InferTypeOpInterface interface;
+*   All return types are either buildable types or are the same as a given
+    operand (e.g., `AllTypesMatch` constraint between operand and result);
+
 And there may potentially exist other builders depending on the specific op;
 please refer to the
 [generated C++ file](#run-mlir-tblgen-to-see-the-generated-content) for the
 complete list.
 
 #### Custom builder methods
 
 However, if the above cases cannot satisfy all needs, you can define additional

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

def Shape_BroadcastOp : Shape_Op<"broadcast",
 
   let arguments = (ins Shape_ShapeType:$lhs, Shape_ShapeType:$rhs,
                    OptionalAttr<StrAttr>:$error);
   let results = (outs Shape_ShapeType:$result);
 
   let hasFolder = 1;
 }
 
-def Shape_ConstShapeOp : Shape_Op<"const_shape",
-    [ConstantLike,
-     NoSideEffect,
-     DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
+def Shape_ConstShapeOp : Shape_Op<"const_shape", [ConstantLike, NoSideEffect]> {
   let summary = "Creates a constant of !shape.shape type.";
   let description = [{
     Creates a !shape.shape with rank given by the length of `shape` and with
     dimension sizes given by the values of `shape`.
 
     ```mlir
     %0 = shape.const_shape []
     %1 = shape.const_shape [1, 2, 3]

mlir/include/mlir/TableGen/Attribute.h

public:
   StringRef getStructClassName() const;
 
   // Returns the C++ namespaces this struct class should be placed in.
   StringRef getCppNamespace() const;
 
   std::vector<StructFieldAttr> getAllFields() const;
 };
 
+// Name of infer type op interface.
+extern const char *inferTypeOpInterface;
+
 } // end namespace tblgen
 } // end namespace mlir
 
 #endif // MLIR_TABLEGEN_ATTRIBUTE_H_

mlir/include/mlir/TableGen/Operator.h

 #include "mlir/TableGen/Attribute.h"
 #include "mlir/TableGen/Dialect.h"
 #include "mlir/TableGen/OpTrait.h"
 #include "mlir/TableGen/Region.h"
 #include "mlir/TableGen/Successor.h"
 #include "mlir/TableGen/Type.h"
 #include "llvm/ADT/PointerUnion.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/SMLoc.h"
 
 namespace llvm {
 class CodeInit;
 class DefInit;
 class Record;
 class StringInit;
public:
 
   // Returns the dialect of the op.
   const Dialect &getDialect() const { return dialect; }
 
   // Prints the contents in this operator to the given `os`. This is used for
   // debugging purposes.
   void print(llvm::raw_ostream &os) const;
 
+  // Return whether all the result types are known.
+  bool allResultTypesKnown() const { return allResultsHaveKnownTypes; };
+
+  // Pair representing either a index to an argument or a type constraint. Only
+  // one of these entries should have the non-default value.
+  struct ArgOrType {
+    explicit ArgOrType(int index) : index(index), constraint(None) {}
+    explicit ArgOrType(TypeConstraint constraint)
+        : index(None), constraint(constraint) {}
+    bool isArg() const {
+      assert(constraint.hasValue() ^ index.hasValue());
+      return index.hasValue();
+    }
+    bool isType() const {
+      assert(constraint.hasValue() ^ index.hasValue());
+      return constraint.hasValue();
+    }
+
+    int getArg() const { return *index; }
+    TypeConstraint getType() const { return *constraint; }
+
+  private:
+    Optional<int> index;
+    Optional<TypeConstraint> constraint;
+  };
+
+  // Return all arguments or type constraints with same type as result[index].
+  // Requires: all result types are known.
+  ArrayRef<ArgOrType> getSameTypeAsResult(int index) const;
+
 private:
   // Populates the vectors containing operands, attributes, results and traits.
   void populateOpStructure();
 
+  // Populates type inference info (mostly equality) with input a mapping from
+  // names to indices for arguments and results.
+  void populateTypeInferenceInfo(
+      const llvm::StringMap<int> &argumentsAndResultsIndex);
+
   // The dialect of this op.
   Dialect dialect;
 
   // The unqualified C++ class name of the op.
   StringRef cppClassName;
 
   // The operands of the op.
   SmallVector<NamedTypeConstraint, 4> operands;
private:
   SmallVector<NamedSuccessor, 0> successors;
 
   // The traits of the op.
   SmallVector<OpTrait, 4> traits;
 
   // The regions of this op.
   SmallVector<NamedRegion, 1> regions;
 
+  // The argument with the same type as the result.
+  SmallVector<SmallVector<ArgOrType, 2>, 4> resultTypeMapping;
+
   // The number of native attributes stored in the leading positions of
   // `attributes`.
   int numNativeAttributes;
 
   // The TableGen definition of this op.
   const llvm::Record &def;
+
+  // Whether the type of all results are known.
+  bool allResultsHaveKnownTypes;
 };
 
 } // end namespace tblgen
 } // end namespace mlir
 
 #endif // MLIR_TABLEGEN_OPERATOR_H_

mlir/lib/Dialect/Shape/IR/Shape.cpp

static ParseResult parseConstShapeOp(OpAsmParser &parser,
   result.addAttribute("shape", builder.getI64TensorAttr(ints));
 
   result.types.push_back(ShapeType::get(builder.getContext()));
   return success();
 }
 
 OpFoldResult ConstShapeOp::fold(ArrayRef<Attribute>) { return shape(); }
 
-LogicalResult
-ConstShapeOp::inferReturnTypes(MLIRContext *context,
-                               Optional<Location> location, ValueRange operands,
-                               DictionaryAttr attributes, RegionRange regions,
-                               SmallVectorImpl<Type> &inferredReturnTypes) {
-  inferredReturnTypes.push_back(ShapeType::get(context));
-  return success();
-}
-
 //===----------------------------------------------------------------------===//
 // ConstSizeOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 ConstSizeOp::inferReturnTypes(MLIRContext *context, Optional<Location> location,
                               ValueRange operands, DictionaryAttr attributes,
                               RegionRange regions,

mlir/lib/TableGen/Attribute.cpp

tblgen::StructAttr::getAllFields() const {
   attributes.reserve(inits->size());
 
   for (const llvm::Init *init : *inits) {
     attributes.emplace_back(cast<llvm::DefInit>(init));
   }
 
   return attributes;
 }
+
+const char *mlir::tblgen::inferTypeOpInterface = "InferTypeOpInterface";

mlir/lib/TableGen/Operator.cpp

 //===- Operator.cpp - Operator class --------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Operator wrapper to simplify using TableGen Record defining a MLIR Op.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/TableGen/Operator.h"
 #include "mlir/TableGen/OpTrait.h"
 #include "mlir/TableGen/Predicate.h"
 #include "mlir/TableGen/Type.h"
+#include "llvm/ADT/EquivalenceClasses.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/TableGen/Error.h"
 #include "llvm/TableGen/Record.h"
 
 #define DEBUG_TYPE "mlir-tblgen-operator"
Record *arg =
       cast<DefInit>(def.getValueAsDag("arguments")->getArg(index))->getDef();
   if (!arg->isSubClassOf("OpVariable"))
     return var_decorator_range(nullptr, nullptr);
   return *arg->getValueAsListInit("decorators");
 }
 
 const tblgen::OpTrait *tblgen::Operator::getTrait(StringRef trait) const {
   for (const auto &t : traits) {
-    if (auto opTrait = dyn_cast<tblgen::NativeOpTrait>(&t)) {
+    if (const auto *opTrait = dyn_cast<tblgen::NativeOpTrait>(&t)) {
       if (opTrait->getTrait() == trait)
         return opTrait;
-    } else if (auto opTrait = dyn_cast<tblgen::InternalOpTrait>(&t)) {
+    } else if (const auto *opTrait = dyn_cast<tblgen::InternalOpTrait>(&t)) {
       if (opTrait->getTrait() == trait)
         return opTrait;
-    } else if (auto opTrait = dyn_cast<tblgen::InterfaceOpTrait>(&t)) {
+    } else if (const auto *opTrait = dyn_cast<tblgen::InterfaceOpTrait>(&t)) {
       if (opTrait->getTrait() == trait)
         return opTrait;
     }
   }
   return nullptr;
 }
 
 auto tblgen::Operator::region_begin() const -> const_region_iterator {

 auto tblgen::Operator::getOperands() -> value_range {
   return {operand_begin(), operand_end()};
 }
 
 auto tblgen::Operator::getArg(int index) const -> Argument {
   return arguments[index];
 }
 
+// Mapping from result index to combined argument and result index. Arguments
+// are indexed to match getArg index, while the result indexes are mapped to
+// avoid overlap.
+static int resultIndex(int i) { return -1 - i; }
+
+bool tblgen::Operator::isVariadic() const {
+  return any_of(llvm::concat<const NamedTypeConstraint>(operands, results),
+                [](const NamedTypeConstraint &op) { return op.isVariadic(); });
+}
+
+void tblgen::Operator::populateTypeInferenceInfo(
+    const llvm::StringMap<int> &argumentsAndResultsIndex) {
+  // If the type inference op interface is not registered, then do not attempt
+  // to determine if the result types an be inferred.
+  auto &recordKeeper = def.getRecords();
+  auto *inferTrait = recordKeeper.getDef(inferTypeOpInterface);
+  allResultsHaveKnownTypes = false;
+  if (!inferTrait)
+    return;
+
+  // If there are no results, the skip this else the build method generated
+  // overlaps with another autogenerated builder.
+  if (getNumResults() == 0)
+    return;
+
+  // Skip for ops with variadic operands/results.
+  // TODO: This can be relaxed.
+  if (isVariadic())
+    return;
+
+  // Skip cases currently being custom generated.
+  // TODO: Remove special cases.
+  if (getTrait("OpTrait::SameOperandsAndResultType"))
+    return;
+
+  // We create equivalence classes of argument/result types where arguments
+  // and results are mapped into the same index space and indices corresponding
+  // to the same type are in the same equivalence class.
+  llvm::EquivalenceClasses<int> ecs;
+  resultTypeMapping.resize(getNumResults());
+  // Captures the argument whose type matches a given result type. Preference
+  // towards capturing operands first before attributes.
+  auto captureMapping = [&](int i) {
+    bool found = false;
+    ecs.insert(resultIndex(i));
+    auto mi = ecs.findLeader(resultIndex(i));
+    for (auto me = ecs.member_end(); mi != me; ++mi) {
+      if (*mi < 0) {
+        auto tc = getResultTypeConstraint(i);
+        if (tc.getBuilderCall().hasValue()) {
+          resultTypeMapping[i].emplace_back(tc);
+          found = true;
+        }
+        continue;
+      }
+
+      if (auto *attr = getArg(*mi).dyn_cast<NamedAttribute *>()) {
+        // TODO: Handle attributes.
+        continue;
+      } else {
+        resultTypeMapping[i].emplace_back(*mi);
+        found = true;
+      }
+    }
+    return found;
+  };
+
+  for (const OpTrait &trait : traits) {
+    const llvm::Record &def = trait.getDef();
+    // If the infer type op interface was manually added, then treat it as
+    // intention that the op needs special handling.
+    // TODO: Reconsider whether to always generate, this is more conservative
+    // and keeps existing behavior so starting that way for now.
+    if (def.isSubClassOf(
+            llvm::formatv("{0}::Trait", inferTypeOpInterface).str()))
+      return;
+    if (const auto *opTrait = dyn_cast<tblgen::InterfaceOpTrait>(&trait))
+      if (opTrait->getTrait().startswith(inferTypeOpInterface))
+        return;
+
+    if (!def.isSubClassOf("AllTypesMatch"))
+      continue;
+
+    auto values = def.getValueAsListOfStrings("values");
+    auto root = argumentsAndResultsIndex.lookup(values.front());
+    for (StringRef str : values)
+      ecs.unionSets(argumentsAndResultsIndex.lookup(str), root);
+  }
+
+  // Verifies that all output types have a corresponding known input type
+  // and chooses matching operand or attribute (in that order) that
+  // matches it.
+  allResultsHaveKnownTypes =
+      all_of(llvm::seq<int>(0, getNumResults()), captureMapping);
+
+  // If the types could be computed, then add type inference trait.
+  if (allResultsHaveKnownTypes)
+    traits.push_back(OpTrait::create(inferTrait->getDefInit()));
+}
+
 void tblgen::Operator::populateOpStructure() {
   auto &recordKeeper = def.getRecords();
-  auto typeConstraintClass = recordKeeper.getClass("TypeConstraint");
-  auto attrClass = recordKeeper.getClass("Attr");
-  auto derivedAttrClass = recordKeeper.getClass("DerivedAttr");
-  auto opVarClass = recordKeeper.getClass("OpVariable");
+  auto *typeConstraintClass = recordKeeper.getClass("TypeConstraint");
+  auto *attrClass = recordKeeper.getClass("Attr");
+  auto *derivedAttrClass = recordKeeper.getClass("DerivedAttr");
+  auto *opVarClass = recordKeeper.getClass("OpVariable");
   numNativeAttributes = 0;
 
   DagInit *argumentValues = def.getValueAsDag("arguments");
   unsigned numArgs = argumentValues->getNumArgs();
 
+  // Mapping from name of to argument or result index. Arguments are indexed
+  // to match getArg index, while the results are negatively indexed.
+  llvm::StringMap<int> argumentsAndResultsIndex;
+
   // Handle operands and native attributes.
   for (unsigned i = 0; i != numArgs; ++i) {
-    auto arg = argumentValues->getArg(i);
+    auto *arg = argumentValues->getArg(i);
     auto givenName = argumentValues->getArgNameStr(i);
-    auto argDefInit = dyn_cast<DefInit>(arg);
+    auto *argDefInit = dyn_cast<DefInit>(arg);
     if (!argDefInit)
       PrintFatalError(def.getLoc(),
                       Twine("undefined type for argument #") + Twine(i));
     Record *argDef = argDefInit->getDef();
     if (argDef->isSubClassOf(opVarClass))
       argDef = argDef->getValueAsDef("constraint");
 
     if (argDef->isSubClassOf(typeConstraintClass)) {
       operands.push_back(
           NamedTypeConstraint{givenName, TypeConstraint(argDef)});
     } else if (argDef->isSubClassOf(attrClass)) {
       if (givenName.empty())
         PrintFatalError(argDef->getLoc(), "attributes must be named");
       if (argDef->isSubClassOf(derivedAttrClass))
         PrintFatalError(argDef->getLoc(),
                         "derived attributes not allowed in argument list");
       attributes.push_back({givenName, Attribute(argDef)});
       ++numNativeAttributes;
     } else {
       PrintFatalError(def.getLoc(), "unexpected def type; only defs deriving "
                                     "from TypeConstraint or Attr are allowed");
     }
+    if (!givenName.empty())
+      argumentsAndResultsIndex[givenName] = i;
   }
 
   // Handle derived attributes.
   for (const auto &val : def.getValues()) {
     if (auto *record = dyn_cast<llvm::RecordRecTy>(val.getType())) {
       if (!record->isSubClassOf(attrClass))
         continue;
       if (!record->isSubClassOf(derivedAttrClass))
for (unsigned i = 0, e = resultsDag->getNumArgs(); i < e; ++i) {
     if (!resultInit) {
       PrintFatalError(def.getLoc(),
                       Twine("undefined type for result #") + Twine(i));
     }
     auto *resultDef = resultInit->getDef();
     if (resultDef->isSubClassOf(opVarClass))
       resultDef = resultDef->getValueAsDef("constraint");
     results.push_back({name, TypeConstraint(resultDef)});
+    if (!name.empty())
+      argumentsAndResultsIndex[name] = resultIndex(i);
   }
 
   // Handle successors
   auto *successorsDag = def.getValueAsDag("successors");
   auto *successorsOp = dyn_cast<DefInit>(successorsDag->getOperator());
   if (!successorsOp || successorsOp->getDef()->getName() != "successor") {
     PrintFatalError(def.getLoc(),
                     "'successors' must have 'successor' directive");
for (unsigned i = 0, e = successorsDag->getNumArgs(); i < e; ++i) {
     // Only support variadic successors if it is the last one for now.
     if (i != e - 1 && successor.isVariadic())
       PrintFatalError(def.getLoc(), "only the last successor can be variadic");
     successors.push_back({name, successor});
   }
 
   // Create list of traits, skipping over duplicates: appending to lists in
   // tablegen is easy, making them unique less so, so dedupe here.
-  if (auto traitList = def.getValueAsListInit("traits")) {
+  if (auto *traitList = def.getValueAsListInit("traits")) {
     // This is uniquing based on pointers of the trait.
     SmallPtrSet<const llvm::Init *, 32> traitSet;
     traits.reserve(traitSet.size());
-    for (auto traitInit : *traitList) {
+    for (auto *traitInit : *traitList) {
       // Keep traits in the same order while skipping over duplicates.
       if (traitSet.insert(traitInit).second)
         traits.push_back(OpTrait::create(traitInit));
     }
   }
 
+  populateTypeInferenceInfo(argumentsAndResultsIndex);
+
   // Handle regions
   auto *regionsDag = def.getValueAsDag("regions");
   auto *regionsOp = dyn_cast<DefInit>(regionsDag->getOperator());
   if (!regionsOp || regionsOp->getDef()->getName() != "region") {
     PrintFatalError(def.getLoc(), "'regions' must have 'region' directive");
   }
 
   for (unsigned i = 0, e = regionsDag->getNumArgs(); i < e; ++i) {
for (unsigned i = 0, e = regionsDag->getNumArgs(); i < e; ++i) {
     }
 
     regions.push_back({name, region});
   }
 
   LLVM_DEBUG(print(llvm::dbgs()));
 }
 
+auto tblgen::Operator::getSameTypeAsResult(int index) const
+    -> ArrayRef<ArgOrType> {
+  assert(allResultTypesKnown());
+  return resultTypeMapping[index];
+}
+
 ArrayRef<llvm::SMLoc> tblgen::Operator::getLoc() const { return def.getLoc(); }
 
 bool tblgen::Operator::hasDescription() const {
   return def.getValue("description") != nullptr;
 }
 
 StringRef tblgen::Operator::getDescription() const {
   return def.getValueAsString("description");

mlir/test/lib/Dialect/Test/TestOps.td

 // Test symbol binding.
 def OpSymbolBindingA : TEST_Op<"symbol_binding_a", []> {
   let arguments = (ins I32:$operand, I64Attr:$attr);
   let results = (outs I32);
 }
 def OpSymbolBindingB : TEST_Op<"symbol_binding_b", []> {
   let arguments = (ins I32:$operand);
   let results = (outs I32);
-
-  let builders = [
-    OpBuilder<
-      "OpBuilder &builder, OperationState &state, Value operand",
-      [{
-        state.types.assign({builder.getIntegerType(32)});
-        state.addOperands({operand});
-      }]>
-  ];
 }
 def OpSymbolBindingC : TEST_Op<"symbol_binding_c", []> {
   let arguments = (ins I32:$operand);
   let results = (outs I32);
   let builders = OpSymbolBindingB.builders;
 }
 def OpSymbolBindingD : TEST_Op<"symbol_binding_d", []> {
   let arguments = (ins I32:$input1, I32:$input2, I64Attr:$attr);

 def AnotherThreeResultOp : TEST_Op<"another_three_result"> {
   let arguments = (ins MultiResultOpEnum:$kind);
   let results = (outs I32:$result1, F32:$result2, F32:$result3);
 }
 
 def TwoResultOp : TEST_Op<"two_result"> {
   let arguments = (ins MultiResultOpEnum:$kind);
   let results = (outs I32:$result1, F32:$result2);
-
-  let builders = [
-    OpBuilder<
-      "OpBuilder &builder, OperationState &state, IntegerAttr kind",
-      [{
-        auto i32 = builder.getIntegerType(32);
-        auto f32 = builder.getF32Type();
-        state.types.assign({i32, f32});
-        state.addAttribute("kind", kind);
-      }]>
-  ];
 }
 
 def AnotherTwoResultOp : TEST_Op<"another_two_result"> {
   let arguments = (ins MultiResultOpEnum:$kind);
   let results = (outs F32:$result1, F32:$result2);
 }
 
 def OneResultOp1 : TEST_Op<"one_result1"> {

mlir/test/mlir-tblgen/op-decl.td

 // RUN: mlir-tblgen -gen-op-decls -I %S/../../include %s | FileCheck --dump-input-on-failure %s
 
 include "mlir/IR/OpBase.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 def Test_Dialect : Dialect {
   let name = "test";
   let cppNamespace = "NS";
 }
 class NS_Op<string mnemonic, list<OpTrait> traits> :
     Op<Test_Dialect, mnemonic, traits>;
def NS_AOp : NS_Op<"a_op", [IsolatedFromAbove, IsolatedFromAbove]> {
   let hasFolder = 1;
 
   let extraClassDeclaration = [{
     // Display a graph for debugging purposes.
     void displayGraph();
   }];
 }
 
-// CHECK: class AOp;
-
 // CHECK-LABEL: NS::AOp declarations
 
 // CHECK: class AOpOperandAdaptor {
 // CHECK: public:
 // CHECK:   AOpOperandAdaptor(ArrayRef<Value> values
 // CHECK:   ArrayRef<Value> getODSOperands(unsigned index);
 // CHECK:   Value a();
 // CHECK:   ArrayRef<Value> b();

 }
 
 // CHECK-LABEL: NS::EOp declarations
 // CHECK:   Value a();
 // CHECK:   ::mlir::MutableOperandRange aMutable();
 // CHECK:   Value b();
 // CHECK:   static void build(OpBuilder &odsBuilder, OperationState &odsState, /*optional*/Type b, /*optional*/Value a)
 
+// Check that all types match constraint results in generating builder.
+// ---
+
+def NS_FOp : NS_Op<"op_with_all_types_constraint",
+    [AllTypesMatch<["a", "b"]>]> {
+  let arguments = (ins AnyType:$a);
+  let results = (outs AnyType:$b);
+}
+
+// CHECK-LABEL: class FOp :
+// CHECK: static LogicalResult inferReturnTypes
+
+def NS_GOp : NS_Op<"op_with_fixed_return_type", []> {
+  let arguments = (ins AnyType:$a);
+  let results = (outs I32:$b);
+}
+
+// CHECK-LABEL: class GOp :
+// CHECK: static LogicalResult inferReturnTypes
+
 // Check that default builders can be suppressed.
 // ---
 
 def NS_SkipDefaultBuildersOp : NS_Op<"skip_default_builders", []> {
   let skipDefaultBuilders = 1;
   let builders = [OpBuilder<"Value val">];
 }
 

mlir/test/mlir-tblgen/types.mlir

func @operand_rank_equals_result_size_failure(%arg : tensor<1x2x3x4xi32>) {
   // expected-error@+1 {{failed to verify that operand rank equals result size}}
   "test.operand_rank_equals_result_size"(%arg) : (tensor<1x2x3x4xi32>) -> tensor<2xi32>
   return
 }
 
 // -----
 
 func @same_types_element_mismatch(%arg0: tensor<* x i32>, %arg1: tensor<* x f32>) {
-  // expected-error@+1 {{all of {x, res} have same type}}
+  // expected-error@+1 {{type incompatible with return type of operation}}
   "test.operand0_and_result_have_same_type"(%arg0, %arg1) : (tensor<* x i32>, tensor<* x f32>) -> tensor<* x f32>
   return
 }
 
 // -----
 
 func @same_types_shape_mismatch(%arg0: tensor<1x2xi32>, %arg1: tensor<2x1xi32>) {
-  // expected-error@+1 {{all of {x, res} have same type}}
+  // expected-error@+1 {{type incompatible with return type of operation}}
   "test.operand0_and_result_have_same_type"(%arg0, %arg1) : (tensor<1x2xi32>, tensor<2x1xi32>) -> tensor<2x1xi32>
   return
 }
 
 // -----
 
 func @does_not_have_i32(%arg0: tensor<1x2xi32>, %arg1: none) {
   // expected-error@+1 {{either both none type operands or first is not none}}

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

private:
   void genSuccessorVerifier(OpMethodBody &body);
 
   // Generates the traits used by the object.
   void genTraits();
 
   // Generate the OpInterface methods.
   void genOpInterfaceMethods();
 
+  // Generate op interface method.
+  void genOpInterfaceMethod(const tblgen::InterfaceOpTrait *trait);
+
   // Generate the side effect interface methods.
   void genSideEffectInterfaceMethods();
 
+  // Generate the type inference interface methods.
+  void genTypeInterfaceMethods();
+
 private:
   // The TableGen record for this op.
   // TODO(antiagainst,zinenko): OpEmitter should not have a Record directly,
   // it should rather go through the Operator for better abstraction.
   const Record &def;
 
   // The wrapper operator class for querying information from this op.
   Operator op;
 
   // The C++ code builder for this op
   OpClass opClass;
 
   // The format context for verification code generation.
   FmtContext verifyCtx;
 };
 } // end anonymous namespace
 
 OpEmitter::OpEmitter(const Operator &op)
     : def(op.getDef()), op(op),
       opClass(op.getCppClassName(), op.getExtraClassDeclaration()) {
   verifyCtx.withOp("(*this->getOperation())");
 
   genTraits();
+
   // Generate C++ code for various op methods. The order here determines the
   // methods in the generated file.
   genOpAsmInterface();
   genOpNameGetter();
   genNamedOperandGetters();
   genNamedOperandSetters();
   genNamedResultGetters();
   genNamedRegionGetters();
   genNamedSuccessorGetters();
   genAttrGetters();
   genAttrSetters();
   genBuilder();
   genParser();
   genPrinter();
   genVerifier();
   genCanonicalizerDecls();
   genFolderDecls();
   genOpInterfaceMethods();
   generateOpFormat(op, opClass);
   genSideEffectInterfaceMethods();
+  genTypeInterfaceMethods();
 }
 
 void OpEmitter::emitDecl(const Operator &op, raw_ostream &os) {
   OpEmitter(op).emitDecl(os);
 }
 
 void OpEmitter::emitDef(const Operator &op, raw_ostream &os) {
   OpEmitter(op).emitDef(os);
for (int i = 0, e = op.getNumNativeAttributes(); i < e; ++i) {
     if (canUseUnwrappedRawValue(namedAttr.attr)) {
       canGenerate = true;
       break;
     }
   }
   return canGenerate;
 }
 
+static bool canInferType(Operator &op) {
+  return op.getTrait("InferTypeOpInterface::Trait") && op.getNumRegions() == 0;
+}
+
 void OpEmitter::genSeparateArgParamBuilder() {
   SmallVector<AttrParamKind, 2> attrBuilderType;
   attrBuilderType.push_back(AttrParamKind::WrappedAttr);
   if (canGenerateUnwrappedBuilder(op))
     attrBuilderType.push_back(AttrParamKind::UnwrappedValue);
 
   // Emit with separate builders with or without unwrapped attributes and/or
   // inferring result type.
case TypeParamKind::Collective:
            << (op.getNumResults() - op.getNumVariableLengthResults())
            << "u && \"mismatched number of results\");\n";
       body << "  " << builderOpState << ".addTypes(resultTypes);\n";
       return;
     }
     llvm_unreachable("unhandled TypeParamKind");
   };
 
-  bool canInferType =
-      op.getTrait("InferTypeOpInterface::Trait") && op.getNumRegions() == 0;
   for (auto attrType : attrBuilderType) {
     emit(attrType, TypeParamKind::Separate, /*inferType=*/false);
-    if (canInferType)
+    if (canInferType(op))
       emit(attrType, TypeParamKind::None, /*inferType=*/true);
     // Emit separate arg build with collective type, unless there is only one
     // variadic result, in which case the above would have already generated
     // the same build method.
     if (!(op.getNumResults() == 1 && op.getResult(0).isVariableLength()))
       emit(attrType, TypeParamKind::Collective, /*inferType=*/false);
   }
 }
void OpEmitter::genCollectiveParamBuilder() {
   // Result types
   if (numVariadicResults == 0 || numNonVariadicResults != 0)
     body << "  assert(resultTypes.size()"
          << (numVariadicResults != 0 ? " >= " : " == ") << numNonVariadicResults
          << "u && \"mismatched number of return types\");\n";
   body << "  " << builderOpState << ".addTypes(resultTypes);\n";
 
   // Generate builder that infers type too.
-  // TODO(jpienaar): Subsume this with general checking if type can be inferred
-  // automatically.
   // TODO(jpienaar): Expand to handle regions and successors.
-  if (op.getTrait("InferTypeOpInterface::Trait") && op.getNumRegions() == 0 &&
-      op.getNumSuccessors() == 0)
+  if (canInferType(op) && op.getNumSuccessors() == 0)
     genInferredTypeCollectiveParamBuilder();
 }
 
 void OpEmitter::buildParamList(std::string &paramList,
                                SmallVectorImpl<std::string> &resultTypeNames,
                                TypeParamKind typeParamKind,
                                AttrParamKind attrParamKind) {
   resultTypeNames.clear();
if (hasSingleResult) {
       const char *const params = "ArrayRef<Attribute> operands, "
                                  "SmallVectorImpl<OpFoldResult> &results";
       opClass.newMethod("LogicalResult", "fold", params, OpMethod::MP_None,
                         /*declOnly=*/true);
     }
   }
 }
 
-void OpEmitter::genOpInterfaceMethods() {
-  for (const auto &trait : op.getTraits()) {
-    auto opTrait = dyn_cast<tblgen::InterfaceOpTrait>(&trait);
-    if (!opTrait || !opTrait->shouldDeclareMethods())
-      continue;
-    auto interface = opTrait->getOpInterface();
+void OpEmitter::genOpInterfaceMethod(const tblgen::InterfaceOpTrait *opTrait) {
 
-    // Get the set of methods that should always be declared.
-    auto alwaysDeclaredMethodsVec = opTrait->getAlwaysDeclaredMethods();
-    llvm::StringSet<> alwaysDeclaredMethods;
-    alwaysDeclaredMethods.insert(alwaysDeclaredMethodsVec.begin(),
-                                 alwaysDeclaredMethodsVec.end());
 
-    for (const OpInterfaceMethod &method : interface.getMethods()) {
-      // Don't declare if the method has a body.
-      if (method.getBody())
-        continue;
-      // Don't declare if the method has a default implementation and the op
-      // didn't request that it always be declared.
-      if (method.getDefaultImplementation() &&
-          !alwaysDeclaredMethods.count(method.getName()))
-        continue;
 
-      std::string args;
-      llvm::raw_string_ostream os(args);
-      interleaveComma(method.getArguments(), os,
-                      [&](const OpInterfaceMethod::Argument &arg) {
-                        os << arg.type << " " << arg.name;
-                      });
-      opClass.newMethod(method.getReturnType(), method.getName(), os.str(),
-                        method.isStatic() ? OpMethod::MP_Static
-                                          : OpMethod::MP_None,
-                        /*declOnly=*/true);
-    }
-  }
+
+void OpEmitter::genOpInterfaceMethods() {
+  for (const auto &trait : op.getTraits()) {
+    if (const auto *opTrait = dyn_cast<tblgen::InterfaceOpTrait>(&trait))
+      if (opTrait->shouldDeclareMethods())
+        genOpInterfaceMethod(opTrait);
+  }
 }
 
 void OpEmitter::genSideEffectInterfaceMethods() {
   enum EffectKind { Operand, Result, Static };
   struct EffectLocation {
     /// The effect applied.
     SideEffect effect;
 
for (auto &location : it.second) {
       // If the effect isn't static, it has a specific value attached to it.
       if (location.kind != EffectKind::Static)
         body << ", value";
       body << ", " << location.effect.getResource() << "::get());\n";
     }
   }
 }
 
+void OpEmitter::genTypeInterfaceMethods() {
+  if (!op.allResultTypesKnown())
+    return;
+
+  auto &method = opClass.newMethod(
+      "LogicalResult", "inferReturnTypes",
+      "MLIRContext* context, Optional<Location> location, "
+      "ValueRange operands, DictionaryAttr attributes, RegionRange regions, "
+      "SmallVectorImpl<Type>& inferredReturnTypes",
+      OpMethod::MP_Static,
+      /*declOnly=*/false);
+  auto &os = method.body();
+  os << "  inferredReturnTypes.resize(" << op.getNumResults() << ");\n";
+
+  FmtContext fctx;
+  fctx.withBuilder("odsBuilder");
+  os << "  Builder odsBuilder(context);\n";
+
+  auto emitType =
+      [&](const tblgen::Operator::ArgOrType &type) -> OpMethodBody & {
+    if (type.isArg()) {
+      auto argIndex = type.getArg();
+      assert(!op.getArg(argIndex).is<NamedAttribute *>());
+      return os << "operands[" << argIndex << "].getType()";
+    } else {
+      return os << tgfmt(*type.getType().getBuilderCall(), &fctx);
+    }
+  };
+
+  for (int i = 0, e = op.getNumResults(); i != e; ++i) {
+    os << "  inferredReturnTypes[" << i << "] = ";
+    auto types = op.getSameTypeAsResult(i);
+    emitType(types[0]) << ";\n";
+    if (types.size() == 1)
+      continue;
+    // TODO: We could verify equality here, but skipping that for verification.
+  }
+  os << "  return success();";
+}
+
 void OpEmitter::genParser() {
   if (!hasStringAttribute(def, "parser") ||
       hasStringAttribute(def, "assemblyFormat"))
     return;
 
   auto &method = opClass.newMethod(
       "ParseResult", "parse", "OpAsmParser &parser, OperationState &result",
       OpMethod::MP_Static);