[Mlir-commits] [mlir] c59aec0 - [mlir][OpFormatGen] Add support for resolving variadic types from non-variadic

[River Riddle]

Author: River Riddle
Date: 2020-07-01T22:27:08-07:00
New Revision: c59aec0ca1edac409d8789956049ae13af24e370

URL: https://github.com/llvm/llvm-project/commit/c59aec0ca1edac409d8789956049ae13af24e370
DIFF: https://github.com/llvm/llvm-project/commit/c59aec0ca1edac409d8789956049ae13af24e370.diff

LOG: [mlir][OpFormatGen] Add support for resolving variadic types from non-variadic

This enables better support for traits such as SameOperandsAndResultType, and other situations in which a variadic operand may be resolved from a non-variadic.

Differential Revision: https://reviews.llvm.org/D83011

Added: 
    

Modified: 
    mlir/test/lib/Dialect/Test/TestOps.td
    mlir/test/mlir-tblgen/op-format.mlir
    mlir/tools/mlir-tblgen/OpFormatGen.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 906060c18305..14b2a7851be1 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -1344,6 +1344,14 @@ def FormatTwoVariadicOperandsNoBuildableTypeOp
   }];
 }
 
+def FormatInferVariadicTypeFromNonVariadic
+    : TEST_Op<"format_infer_variadic_type_from_non_variadic",
+              [SameOperandsAndResultType]> {
+  let arguments = (ins Variadic<AnyType>:$operands);
+  let results = (outs AnyType:$result);
+  let assemblyFormat = "$operands attr-dict `:` type($result)";
+}
+
 //===----------------------------------------------------------------------===//
 // Test SideEffects
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/test/mlir-tblgen/op-format.mlir b/mlir/test/mlir-tblgen/op-format.mlir
index 066e548e1708..157f6cbf4159 100644
--- a/mlir/test/mlir-tblgen/op-format.mlir
+++ b/mlir/test/mlir-tblgen/op-format.mlir
@@ -101,3 +101,10 @@ test.format_optional_operand_result_b_op( : ) : i64
 
 // CHECK: test.format_optional_operand_result_b_op : i64
 test.format_optional_operand_result_b_op : i64
+
+//===----------------------------------------------------------------------===//
+// Format trait type inference
+//===----------------------------------------------------------------------===//
+
+// CHECK: test.format_infer_variadic_type_from_non_variadic %[[I64]], %[[I64]] : i64
+test.format_infer_variadic_type_from_non_variadic %i64, %i64 : i64

diff  --git a/mlir/tools/mlir-tblgen/OpFormatGen.cpp b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
index f8618061da6e..8628cf252712 100644
--- a/mlir/tools/mlir-tblgen/OpFormatGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
@@ -840,10 +840,28 @@ void OperationFormat::genParserTypeResolution(Operator &op,
   for (unsigned i = 0, e = op.getNumOperands(); i != e; ++i) {
     NamedTypeConstraint &operand = op.getOperand(i);
     body << "  if (parser.resolveOperands(" << operand.name << "Operands, ";
-    emitTypeResolver(operandTypes[i], operand.name);
 
-    // If this isn't a buildable type, verify the sizes match by adding the loc.
-    if (!operandTypes[i].getBuilderIdx())
+    // 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 failure();\n";
   }