[mlir][ods] OpFormat: fix type inference issues

[Mogball]

Summary

This patch fixes issues with generating assembly format parsers for
operations that use the operands directive or which have unnamed
arguments or results.

This patch also fixes a function in OpAsmParser that always produced
an error when trying to resolve variadic operands with the same type.

Fixes #51841

[Mogball]

rebase

mlir/include/mlir/IR/OpImplementation.h

public:
   /// Resolve an operand to an SSA value, emitting an error on failure.
   virtual ParseResult resolveOperand(const UnresolvedOperand &operand,
                                      Type type,
                                      SmallVectorImpl<Value> &result) = 0;
 
   /// Resolve a list of operands to SSA values, emitting an error on failure, or
   /// appending the results to the list on success. This method should be used
   /// when all operands have the same type.
-  ParseResult resolveOperands(ArrayRef<UnresolvedOperand> operands, Type type,
+  template <typename Operands = ArrayRef<UnresolvedOperand>>
+  ParseResult resolveOperands(Operands &&operands, Type type,
                               SmallVectorImpl<Value> &result) {
-    for (auto elt : operands)
-      if (resolveOperand(elt, type, result))
+    for (const UnresolvedOperand &operand : operands)
+      if (resolveOperand(operand, type, result))
         return failure();
     return success();
   }
+  template <typename Operands = ArrayRef<UnresolvedOperand>>
+  ParseResult resolveOperands(Operands &&operands, Type type, SMLoc loc,
+                              SmallVectorImpl<Value> &result) {
+    return resolveOperands(std::forward<Operands>(operands), type, result);
+  }
 
   /// Resolve a list of operands and a list of operand types to SSA values,
   /// emitting an error and returning failure, or appending the results
   /// to the list on success.
-  ParseResult resolveOperands(ArrayRef<UnresolvedOperand> operands,
-                              ArrayRef<Type> types, SMLoc loc,
-                              SmallVectorImpl<Value> &result) {
-    if (operands.size() != types.size())
-      return emitError(loc)
-             << operands.size() << " operands present, but expected "
-             << types.size();
-
-    for (unsigned i = 0, e = operands.size(); i != e; ++i)
-      if (resolveOperand(operands[i], types[i], result))
-        return failure();
-    return success();
-  }
-  template <typename Operands>
-  ParseResult resolveOperands(Operands &&operands, Type type, SMLoc loc,
-                              SmallVectorImpl<Value> &result) {
-    return resolveOperands(std::forward<Operands>(operands),
-                           ArrayRef<Type>(type), loc, result);
-  }
-  template <typename Operands, typename Types>
+  template <typename Operands = ArrayRef<UnresolvedOperand>,
+            typename Types = ArrayRef<Type>>
   std::enable_if_t<!std::is_convertible<Types, Type>::value, ParseResult>
   resolveOperands(Operands &&operands, Types &&types, SMLoc loc,
                   SmallVectorImpl<Value> &result) {
     size_t operandSize = std::distance(operands.begin(), operands.end());
     size_t typeSize = std::distance(types.begin(), types.end());
     if (operandSize != typeSize)
       return emitError(loc)
              << operandSize << " operands present, but expected " << typeSize;
 
-    for (auto it : llvm::zip(operands, types))
-      if (resolveOperand(std::get<0>(it), std::get<1>(it), result))
+    for (auto [operand, type] : llvm::zip(operands, types))
+      if (resolveOperand(operand, type, result))
         return failure();
     return success();
   }
 
   /// Parses an affine map attribute where dims and symbols are SSA operands.
   /// Operand values must come from single-result sources, and be valid
   /// dimensions/symbol identifiers according to mlir::isValidDim/Symbol.
   virtual ParseResult

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

def FormatTwoVariadicOperandsNoBuildableTypeOp
   }];
 }
 
 def FormatInferVariadicTypeFromNonVariadic
     : TEST_Op<"format_infer_variadic_type_from_non_variadic",
               [SameOperandsAndResultType]> {
   let arguments = (ins Variadic<AnyType>:$args);
   let results = (outs AnyType:$result);
-  let assemblyFormat = "$args attr-dict `:` type($result)";
+  let assemblyFormat = "operands attr-dict `:` type($result)";
 }
 
 def FormatOptionalUnitAttr : TEST_Op<"format_optional_unit_attribute"> {
   let arguments = (ins UnitAttr:$is_optional);
   let assemblyFormat = "(`is_optional` $is_optional^)? attr-dict";
 }
 
 def FormatOptionalUnitAttrNoElide

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

 }]>, Arguments<(ins AnyMemRef:$operand)>, Results<(outs AnyMemRef:$result)>;
 def ZCoverageValidC : TestFormat_Op<[{
   operands functional-type(operands, results) attr-dict
 }]>, Arguments<(ins AnyMemRef:$operand)>, Results<(outs AnyMemRef:$result)>;
 
 // Check that we can infer type equalities from certain traits.
 def ZCoverageValidD : TestFormat_Op<[{
   operands type($result) attr-dict
-}], [SameOperandsAndResultType]>, Arguments<(ins AnyMemRef:$operand)>,
+}], [SameOperandsAndResultType]>, Arguments<(ins AnyMemRef)>,
      Results<(outs AnyMemRef:$result)>;
 def ZCoverageValidE : TestFormat_Op<[{
   $operand type($operand) attr-dict
 }], [SameOperandsAndResultType]>, Arguments<(ins AnyMemRef:$operand)>,
-     Results<(outs AnyMemRef:$result)>;
+     Results<(outs AnyMemRef)>;
 def ZCoverageValidF : TestFormat_Op<[{
   operands type($other) attr-dict
 }], [SameTypeOperands]>, Arguments<(ins AnyMemRef:$operand, AnyMemRef:$other)>;
 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/tools/mlir-tblgen/OpFormatGen.cpp

if (Optional<int> val = resolver.getBuilderIdx()) {
         fmtContext.addSubst("_ctxt", "parser.getContext()");
         if (var->isVariadic())
           fmtContext.withSelf(var->name + "Types");
         else
           fmtContext.withSelf(var->name + "Types[0]");
         body << tgfmt(*tform, &fmtContext);
       } else {
         body << var->name << "Types";
+        if (!var->isVariadic())
+          body << "[0]";
       }
     } else if (const NamedAttribute *attr = resolver.getAttribute()) {
       if (Optional<StringRef> tform = resolver.getVarTransformer())
         body << tgfmt(*tform,
                       &FmtContext().withSelf(attr->name + "Attr.getType()"));
       else
         body << attr->name << "Attr.getType()";
     } else {
if (op.getNumOperands() > 1) {
             emitTypeResolver(operandTypes[i], op.getOperand(i).name);
             body << ")";
           });
       body << ")";
     } else {
       emitTypeResolver(operandTypes.front(), op.getOperand(0).name);
     }
 
-    body << ", allOperandLoc, result.operands))\n"
-         << "    return ::mlir::failure();\n";
+    body << ", allOperandLoc, result.operands))\n    return "
+            "::mlir::failure();\n";
     return;
   }
 
   // The final case is the one where each of the operands types are resolved
   // separately.
   for (unsigned i = 0, e = op.getNumOperands(); i != e; ++i) {
     NamedTypeConstraint &operand = op.getOperand(i);
     body << "  if (parser.resolveOperands(" << operand.name << "Operands, ";
 
     // Resolve the type of this operand.
     TypeResolution &operandType = operandTypes[i];
     emitTypeResolver(operandType, operand.name);
 
-    // If the type is resolved by a non-variadic variable, index into the
-    // resolved type list. This allows for resolving the types of a variadic
-    // operand list from a non-variadic variable.
-    bool verifyOperandAndTypeSize = true;
-    if (auto *resolverVar = operandType.getVariable()) {
-      if (!resolverVar->isVariadic() && !operandType.getVarTransformer()) {
-        body << "[0]";
-        verifyOperandAndTypeSize = false;
-      }
-    } else {
-      verifyOperandAndTypeSize = !operandType.getBuilderIdx();
-    }
-
-    // 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";
+    body << ", " << operand.name
+         << "OperandsLoc, result.operands))\n    return ::mlir::failure();\n";
   }
 }
 
 void OperationFormat::genParserRegionResolution(Operator &op,
                                                 MethodBody &body) {
   // Check for the case where all regions were parsed.
   bool hasAllRegions = llvm::any_of(
       elements, [](FormatElement *elt) { return isa<RegionsDirective>(elt); });
if ((resolvedIt = seenOperandTypes.find_first()) != -1)
     resolver = &op.getOperand(resolvedIt);
   else if (includeResults && (resolvedIt = seenResultTypes.find_first()) != -1)
     resolver = &op.getResult(resolvedIt);
   else
     return;
 
   // Set the resolvers for each operand and result.
   for (unsigned i = 0, e = op.getNumOperands(); i != e; ++i)
-    if (!seenOperandTypes.test(i) && !op.getOperand(i).name.empty())
+    if (!seenOperandTypes.test(i))
       variableTyResolver[op.getOperand(i).name] = {resolver, llvm::None};
   if (includeResults) {
     for (unsigned i = 0, e = op.getNumResults(); i != e; ++i)
-      if (!seenResultTypes.test(i) && !op.getResultName(i).empty())
+      if (!seenResultTypes.test(i))
         variableTyResolver[op.getResultName(i)] = {resolver, llvm::None};
   }
 }
 
 void OpFormatParser::handleTypesMatchConstraint(
     llvm::StringMap<TypeResolutionInstance> &variableTyResolver,
     const llvm::Record &def) {
   StringRef lhsName = def.getValueAsString("lhs");