[mlir][ods] Fix OpFormatGen sometimes not calling inferReturnTypes

[jpienaar]

Nice, good catch

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

"MultiResultOpEnum", "Multi-result op kinds", [
     MultiResultOpKind4, MultiResultOpKind5, MultiResultOpKind6
   ]>;
 
 def ThreeResultOp : TEST_Op<"three_result"> {
   let arguments = (ins MultiResultOpEnum:$kind);
   let results = (outs I32:$result1, F32:$result2, F32:$result3);
 }
 
-def AnotherThreeResultOp : TEST_Op<"another_three_result", [DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
+def AnotherThreeResultOp
+    : TEST_Op<"another_three_result",
+              [DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
   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);
 }
static ::mlir::LogicalResult inferReturnTypes(::mlir::MLIRContext *context,
           ::mlir::DictionaryAttr attributes, ::mlir::RegionRange regions,
           ::llvm::SmallVectorImpl<::mlir::Type> &inferredReturnTypes) {
       inferredReturnTypes.assign({::mlir::IntegerType::get(context, 16)});
       return ::mlir::success();
     }
    }];
 }
 
+// Base class for testing mixing allOperandTypes, allOperands, and
+// inferResultTypes.
+class FormatInferAllTypesBaseOp<string mnemonic, list<OpTrait> traits = []>
+    : TEST_Op<mnemonic, [InferTypeOpInterface] # traits> {
+  let arguments = (ins Variadic<AnyType>:$args);
+  let results = (outs Variadic<AnyType>:$outs);
+  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) {
+      ::mlir::TypeRange operandTypes = operands.getTypes();
+      inferredReturnTypes.assign(operandTypes.begin(), operandTypes.end());
+      return ::mlir::success();
+    }
+   }];
+}
+
+// Test inferReturnTypes is called when allOperandTypes and allOperands is true.
+def FormatInferTypeAllOperandsAndTypesOp
+    : FormatInferAllTypesBaseOp<"format_infer_type_all_operands_and_types"> {
+  let assemblyFormat = "`(` operands `)` attr-dict `:` type(operands)";
+}
+
+// Test inferReturnTypes is called when allOperandTypes is true and there is one
+// ODS operand.
+def FormatInferTypeAllOperandsAndTypesOneOperandOp
+    : FormatInferAllTypesBaseOp<"format_infer_type_all_types_one_operand"> {
+  let assemblyFormat = "`(` $args `)` attr-dict `:` type(operands)";
+}
+
+// Test inferReturnTypes is called when allOperandTypes is true and there are
+// more than one ODS operands.
+def FormatInferTypeAllOperandsAndTypesTwoOperandsOp
+    : FormatInferAllTypesBaseOp<"format_infer_type_all_types_two_operands",
+                                [SameVariadicOperandSize]> {
+  let arguments = (ins Variadic<AnyType>:$args0, Variadic<AnyType>:$args1);
+  let assemblyFormat = "`(` $args0 `)` `(` $args1 `)` attr-dict `:` type(operands)";
+}
+
+// Test inferReturnTypes is called when allOperands is true and operand types
+// are separately specified.
+def FormatInferTypeAllTypesOp
+    : FormatInferAllTypesBaseOp<"format_infer_type_all_types"> {
+  let assemblyFormat = "`(` operands `)` attr-dict `:` type($args)";
+}
+
 //===----------------------------------------------------------------------===//
 // Test SideEffects
 //===----------------------------------------------------------------------===//
 
 def SideEffectOp : TEST_Op<"side_effect_op",
     [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
      DeclareOpInterfaceMethods<TestEffectOpInterface>]> {
   let results = (outs AnyType:$result);

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

 
 //===----------------------------------------------------------------------===//
 // InferTypeOpInterface type inference
 //===----------------------------------------------------------------------===//
 
 // CHECK: test.format_infer_type
 %ignored_res7 = test.format_infer_type
 
+// CHECK: test.format_infer_type_all_operands_and_types(%[[I64]], %[[I32]]) : i64, i32
+%ignored_res8:2 = test.format_infer_type_all_operands_and_types(%i64, %i32) : i64, i32
+
+// CHECK: test.format_infer_type_all_types_one_operand(%[[I64]], %[[I32]]) : i64, i32
+%ignored_res9:2 = test.format_infer_type_all_types_one_operand(%i64, %i32) : i64, i32
+
+// CHECK: test.format_infer_type_all_types_two_operands(%[[I64]], %[[I32]]) (%[[I64]], %[[I32]]) : i64, i32, i64, i32
+%ignored_res10:4 = test.format_infer_type_all_types_two_operands(%i64, %i32) (%i64, %i32) : i64, i32, i64, i32
+
+// CHECK: test.format_infer_type_all_types(%[[I64]], %[[I32]]) : i64, i32
+%ignored_res11:2 = test.format_infer_type_all_types(%i64, %i32) : i64, i32
+
 //===----------------------------------------------------------------------===//
 // Check DefaultValuedStrAttr
 //===----------------------------------------------------------------------===//
 
 // CHECK: test.has_str_value
 test.has_str_value {}

mlir/tools/mlir-tblgen/OpFormatGen.cpp

hasSingleBlockTrait =
         hasImplicitTermTrait || op.getTrait("::mlir::OpTrait::SingleBlock");
   }
 
   /// Generate the operation parser from this format.
   void genParser(Operator &op, OpClass &opClass);
   /// Generate the parser code for a specific format element.
   void genElementParser(Element *element, MethodBody &body,
                         FmtContext &attrTypeCtx);
-  /// Generate the c++ to resolve the types of operands and results during
+  /// Generate the C++ to resolve the types of operands and results during
   /// parsing.
   void genParserTypeResolution(Operator &op, MethodBody &body);
-  /// Generate the c++ to resolve regions during parsing.
+  /// Generate the C++ to resolve the types of the operands during parsing.
+  void genParserOperandTypeResolution(
+      Operator &op, MethodBody &body,
+      function_ref<void(TypeResolution &, StringRef)> emitTypeResolver);
+  /// Generate the C++ to resolve regions during parsing.
   void genParserRegionResolution(Operator &op, MethodBody &body);
-  /// Generate the c++ to resolve successors during parsing.
+  /// Generate the C++ to resolve successors during parsing.
   void genParserSuccessorResolution(Operator &op, MethodBody &body);
-  /// Generate the c++ to handling variadic segment size traits.
+  /// Generate the C++ to handling variadic segment size traits.
   void genParserVariadicSegmentResolution(Operator &op, MethodBody &body);
 
   /// Generate the operation printer from this format.
   void genPrinter(Operator &op, OpClass &opClass);
 
   /// Generate the printer code for a specific format element.
   void genElementPrinter(Element *element, MethodBody &body, Operator &op,
                          bool &shouldEmitSpace, bool &lastWasPunctuation);
if (allResultTypes) {
       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) {
-    // Handle return type inference here if there are no operands
-    if (infersResultTypes)
-      body << formatv(inferReturnTypesParserCode, op.getCppClassName());
-    return;
-  }
+  // Emit the operand type resolutions.
+  genParserOperandTypeResolution(op, body, emitTypeResolver);
+
+  // Handle return type inference once all operands have been resolved
 
-  // Handle the case where all operand types are in one group.
+void OperationFormat::genParserOperandTypeResolution(
+    Operator &op, MethodBody &body,
+    function_ref<void(TypeResolution &, StringRef)> emitTypeResolver) {
+  // Early exit if there are no operands.
+  if (op.getNumOperands() == 0)
+    return;
+
+  // Handle the case where all operand types are grouped together with
+  // "types(operands)".
   if (allOperandTypes) {
-    // If we have all operands together, use the full operand list directly.
+    // If `operands` was specified, use the full operand list directly.
     if (allOperands) {
       body << "  if (parser.resolveOperands(allOperands, allOperandTypes, "
               "allOperandLoc, result.operands))\n"
               "    return ::mlir::failure();\n";
       return;
     }
 
     // Otherwise, use llvm::concat to merge the disjoint operand lists together.
     // llvm::concat does not allow the case of a single range, so guard it here.
     body << "  if (parser.resolveOperands(";
     if (op.getNumOperands() > 1) {
       body << "::llvm::concat<const ::mlir::OpAsmParser::OperandType>(";
       llvm::interleaveComma(op.getOperands(), body, [&](auto &operand) {
         body << operand.name << "Operands";
       });
       body << ")";
     } else {
       body << op.operand_begin()->name << "Operands";
     }
     body << ", allOperandTypes, parser.getNameLoc(), result.operands))\n"
          << "    return ::mlir::failure();\n";
     return;
   }
-  // Handle the case where all of the operands were grouped together.
+
+  // Handle the case where all operands are grouped together with "operands".
   if (allOperands) {
     body << "  if (parser.resolveOperands(allOperands, ";
 
     // Group all of the operand types together to perform the resolution all at
     // once. Use llvm::concat to perform the merge. llvm::concat does not allow
     // the case of a single range, so guard it here.
     if (op.getNumOperands() > 1) {
       body << "::llvm::concat<const Type>(";
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,
                                                 MethodBody &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) {
body << llvm::formatv(enumAttrBeginPrinterCode,
                         (var->attr.isOptional() ? "*" : "") +
                             op.getGetterName(var->name),
                         enumAttr.getSymbolToStringFnName());
 
   // Get a string containing all of the cases that can't be represented with a
   // keyword.
   llvm::BitVector nonKeywordCases(cases.size());
   bool hasStrCase = false;
-  for (auto it : llvm::enumerate(cases)) {
+  for (auto &it : llvm::enumerate(cases)) {
     hasStrCase = it.value().isStrCase();
     if (!canFormatStringAsKeyword(it.value().getStr()))
       nonKeywordCases.set(it.index());
   }
 
   // If this is a string enum, use the case string to determine which cases
   // need to use the string form.
   if (hasStrCase) {
body << "    _odsPrinter << caseValueStr;\n"
             "  }\n";
     return;
   }
 
   // Otherwise if this is a bit enum attribute, don't allow cases that may
   // overlap with other cases. For simplicity sake, only allow cases with a
   // single bit value.
   if (enumAttr.isBitEnum()) {
-    for (auto it : llvm::enumerate(cases)) {
+    for (auto &it : llvm::enumerate(cases)) {
       int64_t value = it.value().getValue();
       if (value < 0 || !llvm::isPowerOf2_64(value))
         nonKeywordCases.set(it.index());
     }
   }
 
   // If there are any cases that can't be used with a keyword, switch on the
   // case value to determine when to print in the string form.
   if (nonKeywordCases.any()) {
     body << "    switch (caseValue) {\n";
     StringRef cppNamespace = enumAttr.getCppNamespace();
     StringRef enumName = enumAttr.getEnumClassName();
-    for (auto it : llvm::enumerate(cases)) {
+    for (auto &it : llvm::enumerate(cases)) {
       if (nonKeywordCases.test(it.index()))
         continue;
       StringRef symbol = it.value().getSymbol();
       body << llvm::formatv("    case {0}::{1}::{2}:\n", cppNamespace, enumName,
                             llvm::isDigit(symbol.front()) ? ("_" + symbol)
                                                           : symbol);
     }
     body << "      _odsPrinter << caseValueStr;\n"