[mlir][ODS] Support result type inference in custom assembly format

[trilorez]

Summary

Operations that have the InferTypeOpInterface trait can now omit the return
types in their custom assembly formats.

[trilorez]

Ran clang-format, thought I had ran it before

[Mogball]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

1560

Can you move hoist this to the top as a const char *const and add a small snippet of documentation like the other code snippets?

[trilorez]

Moved infer return type parser code snippet to a constant

[rriddle]

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

2028–2041

Can you simplify this to just infer a constant type? That would remove the need to add inputs.

mlir/tools/mlir-tblgen/OpFormatGen.cpp

693

1560

nit: Drop the trivial braces here.

2728

2737

What about situations where only some types can be resolved this way? Should we just default to using inferResultTypes if the user doesn't specify any results in the format?

[lattner]

This looks great @trilorez ! Thank you for working on this

[trilorez]

Simplified FormatInferTypeOp to remove inputs, which led to a fix for handling
zero operand operations.

Addressed other pr comments.

[trilorez]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

2737

The resolvers for result types can only be set by AllTypesMatch, SameOperandsAndResultType, and TypesMatchWith, which look like they can't have partial type inference.

For other cases, like one of several result types being manually specified, I'm not sure how partial type inference would work. I suppose already inferred types could be passed to inferReturnTypes along with the operands, but then they'd need to either be in order or pass along their index.

[trilorez]

Added InferTypeOpInterface section to OpDefintions.md

[rriddle]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

2737

I was more saying that if the user didn't explicitly add a type to the format, we could just always use inferResultTypes(for all of the types, ignoring the other resolvers). I think it would be problematic if the user had one result that was a constant(e.g. Index or something), but had others that required being referred. Right now we would never give them a chance to infer the other result types.

[trilorez]

Prefer inferring result types over using buildable type builders

[trilorez]

Infer all result types if none are specified

[trilorez]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

2737

Ah, I think I understand now. I believe the latest update addresses this.

[rriddle]

LG, thanks.

mlir/docs/OpDefinitions.md

 
 There are many operations that have known type equality constraints registered
 as traits on the operation; for example the true, false, and result values of a
 `select` operation often have the same type. The assembly format may inspect
 these equal constraints to discern the types of missing variables. The currently
 supported traits are: `AllTypesMatch`, `TypesMatchWith`, `SameTypeOperands`, and
 `SameOperandsAndResultType`.
 
+*   InferTypeOpInterface
+
+Operations that implement `InferTypeOpInterface` can omit their result types in
+their assembly format since the result types can be inferred from the operands.
+
 ### `hasCanonicalizer`
 
 This boolean field indicate whether canonicalization patterns have been defined
 for this operation. If it is `1`, then `::getCanonicalizationPatterns()` should
 be defined.
 
 ### `hasCanonicalizeMethod`
 

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

TypesMatchWith<"tuple result type matches operand type", "value", "result",
         "::mlir::TupleType::get($_ctxt, $_self)">
   ]> {
   let arguments = (ins AnyType:$value);
   let results = (outs AnyType:$result);
   let assemblyFormat = "attr-dict $value `:` type($value)";
 }
 
 //===----------------------------------------------------------------------===//
+// InferTypeOpInterface type inference in assembly format
+
+def FormatInferTypeOp : TEST_Op<"format_infer_type", [InferTypeOpInterface]> {
+  let results = (outs AnyType);
+  let assemblyFormat = "attr-dict";
+
+  let extraClassDeclaration = [{
+    static ::mlir::LogicalResult inferReturnTypes(::mlir::MLIRContext *context,
+          ::llvm::Optional<::mlir::Location> location, ::mlir::ValueRange operands,
+          ::mlir::DictionaryAttr attributes, ::mlir::RegionRange regions,
+          ::llvm::SmallVectorImpl<::mlir::Type> &inferredReturnTypes) {
+      inferredReturnTypes.assign({::mlir::IntegerType::get(context, 16)});
+      return ::mlir::success();
+    }
+   }];
+}
+
+//===----------------------------------------------------------------------===//
 // Test SideEffects
 //===----------------------------------------------------------------------===//
 
 def SideEffectOp : TEST_Op<"side_effect_op",
     [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
      DeclareOpInterfaceMethods<TestEffectOpInterface>]> {
   let results = (outs AnyType:$result);
 }

mlir/test/mlir-tblgen/op-format-spec.td

 // RUN: mlir-tblgen -gen-op-decls -asmformat-error-is-fatal=false -I %S/../../include %s -o=%t 2>&1 | FileCheck %s
 
 // This file contains tests for the specification of the declarative op format.
 
 include "mlir/IR/OpBase.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
 
 def TestDialect : Dialect {
   let name = "test";
 }
 class TestFormat_Op<string fmt, list<OpTrait> traits = []>
     : Op<TestDialect, "format_op", traits> {
   let assemblyFormat = fmt;
 }

 def ZCoverageValidG : TestFormat_Op<[{
   operands type($other) attr-dict
 }], [AllTypesMatch<["operand", "other"]>]>,
      Arguments<(ins AnyMemRef:$operand, AnyMemRef:$other)>;
 def ZCoverageValidH : TestFormat_Op<[{
   operands type($result) attr-dict
 }], [AllTypesMatch<["operand", "result"]>]>,
      Arguments<(ins AnyMemRef:$operand)>, Results<(outs AnyMemRef:$result)>;
-
+def ZCoverageValidI : TestFormat_Op<[{
+  operands type(operands) attr-dict
+}], [InferTypeOpInterface]>, Arguments<(ins Variadic<I64>:$inputs)>, Results<(outs I64:$result)>;

mlir/test/mlir-tblgen/op-format.mlir

 // CHECK: test.format_types_match_variadic %[[I64]], %[[I64]], %[[I64]] : i64, i64, i64
 %ignored_res4:3 = test.format_types_match_variadic %i64, %i64, %i64 : i64, i64, i64
 
 // CHECK: test.format_types_match_attr 1 : i64
 %ignored_res5 = test.format_types_match_attr 1 : i64
 
 // CHECK: test.format_types_match_context %[[I64]] : i64
 %ignored_res6 = test.format_types_match_context %i64 : i64
+
+//===----------------------------------------------------------------------===//
+// InferTypeOpInterface type inference
+//===----------------------------------------------------------------------===//
+
+// CHECK: test.format_infer_type
+%ignored_res7 = test.format_infer_type

mlir/tools/mlir-tblgen/OpFormatGen.cpp

private:
     /// the variable or the attribute that this type is resolved to.
     ConstArgument resolver;
     /// If the type is resolved based upon another operand or result, this is
     /// a transformer to apply to the variable when resolving.
     Optional<StringRef> variableTransformer;
   };
 
   OperationFormat(const Operator &op)
-      : allOperands(false), allOperandTypes(false), allResultTypes(false) {
+      : allOperands(false), allOperandTypes(false), allResultTypes(false),
+        infersResultTypes(false) {
     operandTypes.resize(op.getNumOperands(), TypeResolution());
     resultTypes.resize(op.getNumResults(), TypeResolution());
 
     hasImplicitTermTrait = llvm::any_of(op.getTraits(), [](const Trait &trait) {
       return trait.getDef().isSubClassOf("SingleBlockImplicitTerminator");
     });
 
     hasSingleBlockTrait =
struct OperationFormat {
 
   /// The various elements in this format.
   std::vector<std::unique_ptr<Element>> elements;
 
   /// A flag indicating if all operand/result types were seen. If the format
   /// contains these, it can not contain individual type resolvers.
   bool allOperands, allOperandTypes, allResultTypes;
 
+  /// A flag indicating if this operation infers its result types
+  bool infersResultTypes;
+
   /// A flag indicating if this operation has the SingleBlockImplicitTerminator
   /// trait.
   bool hasImplicitTermTrait;
 
   /// A flag indicating if this operation has the SingleBlock trait.
   bool hasSingleBlockTrait;
 
   /// A map of buildable types to indices.

 const char *const functionalTypeParserCode = R"(
   ::mlir::FunctionType {0}__{1}_functionType;
   if (parser.parseType({0}__{1}_functionType))
     return ::mlir::failure();
   {0}Types = {0}__{1}_functionType.getInputs();
   {1}Types = {0}__{1}_functionType.getResults();
 )";
 
+/// The code snippet used to generate a parser call to infer return types.
+///
+/// {0}: The operation class name
+const char *const inferReturnTypesParserCode = R"(
+  ::llvm::SmallVector<::mlir::Type> inferredReturnTypes;
+  if (::mlir::failed({0}::inferReturnTypes(parser.getContext(),
+      result.location, result.operands,
+      result.attributes.getDictionary(parser.getContext()),
+      result.regions, inferredReturnTypes)))
+    return ::mlir::failure();
+  result.addTypes(inferredReturnTypes);
+)";
+
 /// The code snippet used to generate a parser call for a region list.
 ///
 /// {0}: The name for the region list.
 const char *regionListParserCode = R"(
   {
     std::unique_ptr<::mlir::Region> region;
     auto firstRegionResult = parser.parseOptionalRegion(region);
     if (firstRegionResult.hasValue()) {
if (Optional<int> val = resolver.getBuilderIdx()) {
       else
         body << attr->name << "Attr.getType()";
     } else {
       body << curVar << "Types";
     }
   };
 
   // Resolve each of the result types.
+  if (!infersResultTypes) {
     if (allResultTypes) {
       body << "  result.addTypes(allResultTypes);\n";
     } else {
       for (unsigned i = 0, e = op.getNumResults(); i != e; ++i) {
         body << "  result.addTypes(";
         emitTypeResolver(resultTypes[i], op.getResultName(i));
         body << ");\n";
       }
     }
+  }
 
   // Early exit if there are no operands.
-  if (op.getNumOperands() == 0)
+  if (op.getNumOperands() == 0) {
+    // Handle return type inference here if there are no operands
+    if (infersResultTypes)
+      body << formatv(inferReturnTypesParserCode, op.getCppClassName());
     return;
+  }
 
   // Handle the case where all operand types are in one group.
   if (allOperandTypes) {
     // If we have all operands together, use the full operand list directly.
     if (allOperands) {
       body << "  if (parser.resolveOperands(allOperands, allOperandTypes, "
               "allOperandLoc, result.operands))\n"
               "    return ::mlir::failure();\n";
for (unsigned i = 0, e = op.getNumOperands(); i != e; ++i) {
 
     // Check to see if the sizes between the types and operands must match. If
     // they do, provide the operand location to select the proper resolution
     // overload.
     if (verifyOperandAndTypeSize)
       body << ", " << operand.name << "OperandsLoc";
     body << ", result.operands))\n    return ::mlir::failure();\n";
   }
+
+  // Handle return type inference once all operands have been resolved
+  if (infersResultTypes)
+    body << formatv(inferReturnTypesParserCode, op.getCppClassName());
 }
 
 void OperationFormat::genParserRegionResolution(Operator &op,
                                                 OpMethodBody &body) {
   // Check for the case where all regions were parsed.
   bool hasAllRegions = llvm::any_of(
       elements, [](auto &elt) { return isa<RegionsDirective>(elt.get()); });
   if (hasAllRegions) {
private:
   Token curToken;
   OperationFormat &fmt;
   Operator &op;
 
   // The following are various bits of format state used for verification
   // during parsing.
   bool hasAttrDict = false;
   bool hasAllRegions = false, hasAllSuccessors = false;
+  bool canInferResultTypes = false;
   llvm::SmallBitVector seenOperandTypes, seenResultTypes;
   llvm::SmallSetVector<const NamedAttribute *, 8> seenAttrs;
   llvm::DenseSet<const NamedTypeConstraint *> seenOperands;
   llvm::DenseSet<const NamedRegion *> seenRegions;
   llvm::DenseSet<const NamedSuccessor *> seenSuccessors;
 };
 } // end anonymous namespace
 
if (def.isSubClassOf("AllTypesMatch")) {
       handleAllTypesMatchConstraint(def.getValueAsListOfStrings("values"),
                                     variableTyResolver);
     } else if (def.getName() == "SameTypeOperands") {
       handleSameTypesConstraint(variableTyResolver, /*includeResults=*/false);
     } else if (def.getName() == "SameOperandsAndResultType") {
       handleSameTypesConstraint(variableTyResolver, /*includeResults=*/true);
     } else if (def.isSubClassOf("TypesMatchWith")) {
       handleTypesMatchConstraint(variableTyResolver, def);
+    } else if (def.getName() == "InferTypeOpInterface" &&
+               !op.allResultTypesKnown()) {
+      canInferResultTypes = true;
     }
   }
 
   // Verify the state of the various operation components.
   if (failed(verifyAttributes(loc)) ||
       failed(verifyResults(loc, variableTyResolver)) ||
       failed(verifyOperands(loc, variableTyResolver)) ||
       failed(verifyRegions(loc)) || failed(verifySuccessors(loc)))

 
 LogicalResult FormatParser::verifyResults(
     llvm::SMLoc loc,
     llvm::StringMap<TypeResolutionInstance> &variableTyResolver) {
   // If we format all of the types together, there is nothing to check.
   if (fmt.allResultTypes)
     return ::mlir::success();
 
+  // If no result types are specified and we can infer them, infer all result
+  // types
+  if (op.getNumResults() > 0 && seenResultTypes.count() == 0 &&
+      canInferResultTypes) {
+    fmt.infersResultTypes = true;
+    return ::mlir::success();
+  }
+
   // Check that all of the result types can be inferred.
   auto &buildableTypes = fmt.buildableTypes;
   for (unsigned i = 0, e = op.getNumResults(); i != e; ++i) {
     if (seenResultTypes.test(i))
       continue;
 
     // Check to see if we can infer this type from another variable.
     auto varResolverIt = variableTyResolver.find(op.getResultName(i));
     if (varResolverIt != variableTyResolver.end()) {
       TypeResolutionInstance resolver = varResolverIt->second;
       fmt.resultTypes[i].setResolver(resolver.resolver, resolver.transformer);
       continue;
     }
 
     // If the result is not variable length, allow for the case where the type
     // has a builder that we can use.
     NamedTypeConstraint &result = op.getResult(i);
     Optional<StringRef> builder = result.constraint.getBuilderCall();
     if (!builder || result.isVariableLength()) {
       return emitErrorAndNote(