[mlir] ODS: support operations with resizable operand lists

MLIR supports operations with resizable operand lists, but this property must
be indicated during the construction of such operations. It can be done
programmatically by calling a function on OperationState. Introduce an
ODS-internal trait `ResizableOperandList` to indicate such operations are use
it when generating the bodies of various `build` functions as well as the
`parse` function when the declarative assembly format is used.

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

mlir/include/mlir/IR/OpBase.td

@@ -1653,6 +1653,10 @@ class HasParent<string op>
 def FirstAttrDerivedResultType :
   GenInternalOpTrait<"FirstAttrDerivedResultType">;
 
+// Op has a resizable operand list. Auto-generated build and parse functions
+// should construct it as such.
+def ResizableOperandList : GenInternalOpTrait<"ResizableOperandList">;
+
 // TODO(antiagainst): Turn the following into normal traits and generate
 // verification for them.
 

mlir/include/mlir/TableGen/Operator.h

@@ -165,6 +165,9 @@ public:
   // requiring the raw MLIR trait here.
   const OpTrait *getTrait(llvm::StringRef trait) const;
 
+  // Returns "true" if Op has a ResizableOperandList trait.
+  bool hasResizableOperandList() const;
+
   // Regions.
   using const_region_iterator = const NamedRegion *;
   const_region_iterator region_begin() const;

mlir/lib/TableGen/Operator.cpp

@@ -169,6 +169,10 @@ const tblgen::OpTrait *tblgen::Operator::getTrait(StringRef trait) const {
   return nullptr;
 }
 
+bool tblgen::Operator::hasResizableOperandList() const {
+  return getTrait("OpTrait::ResizableOperandList") != nullptr;
+}
+
 auto tblgen::Operator::region_begin() const -> const_region_iterator {
   return regions.begin();
 }

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

@@ -58,3 +58,23 @@ def OpD : NS_Op<"mix_variadic_and_normal_inputs_op", [SameVariadicOperandSize]>
 // CHECK-NEXT: odsState.addOperands(input1);
 // CHECK-NEXT: odsState.addOperands(input2);
 // CHECK-NEXT: odsState.addOperands(input3);
+// CHECK-NOT: odsState.setOperandListToResizable
+
+// Check that resizable operand list flag is set up correctly in all generated
+// builders and in the parser.
+def OpE : NS_Op<"resizable_operand_list", [ResizableOperandList]> {
+  let arguments = (ins Variadic<AnyType>:$input);
+  let assemblyFormat = "$input attr-dict `:` type($input)";
+}
+
+// CHECK-LABEL: OpE::build(Builder *odsBuilder, OperationState &odsState, ValueRange
+// CHECK: odsState.setOperandListToResizable()
+
+// CHECK-LABEL: OpE::build(Builder *odsBuilder, OperationState &odsState, ArrayRef<Type>
+// CHECK: odsState.setOperandListToResizable()
+
+// CHECK-LABEL: OpE::build(Builder *, OperationState
+// CHECK: odsState.setOperandListToResizable()
+
+// CHECK-LABEL: OpE::parse
+// CHECK: result.setOperandListToResizable()

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

@@ -685,6 +685,7 @@ void OpEmitter::genSeparateArgParamBuilder() {
     auto &m =
         opClass.newMethod("void", "build", paramList, OpMethod::MP_Static);
     auto &body = m.body();
+
     genCodeForAddingArgAndRegionForBuilder(
         body, /*isRawValueAttr=*/attrType == AttrParamKind::UnwrappedValue);
 
@@ -762,7 +763,9 @@ void OpEmitter::genUseOperandAsResultTypeCollectiveParamBuilder() {
   auto &body = m.body();
 
   // Operands
-  body << "  " << builderOpState << ".addOperands(operands);\n\n";
+  body << "  " << builderOpState << ".addOperands(operands);\n";
+  if (op.hasResizableOperandList())
+    body << formatv("  {0}.setOperandListToResizable();\n\n", builderOpState);
 
   // Attributes
   body << "  " << builderOpState << ".addAttributes(attributes);\n";
@@ -843,7 +846,10 @@ void OpEmitter::genUseAttrAsResultTypeBuilder() {
   }
 
   // Operands
-  body << "  " << builderOpState << ".addOperands(operands);\n\n";
+  body << "  " << builderOpState << ".addOperands(operands);\n";
+  if (op.hasResizableOperandList())
+    body << formatv("  {0}.setOperandListToResizable();\n\n", builderOpState);
+
   // Attributes
   body << "  " << builderOpState << ".addAttributes(attributes);\n";
 
@@ -929,7 +935,9 @@ void OpEmitter::genCollectiveParamBuilder() {
          << (numVariadicOperands != 0 ? " >= " : " == ")
          << numNonVariadicOperands
          << "u && \"mismatched number of parameters\");\n";
-  body << "  " << builderOpState << ".addOperands(operands);\n\n";
+  body << "  " << builderOpState << ".addOperands(operands);\n";
+  if (op.hasResizableOperandList())
+    body << formatv("  {0}.setOperandListToResizable();\n\n", builderOpState);
 
   // Attributes
   body << "  " << builderOpState << ".addAttributes(attributes);\n";
@@ -1099,6 +1107,8 @@ void OpEmitter::genCodeForAddingArgAndRegionForBuilder(OpMethodBody &body,
       body << "  if (" << argName << ")\n  ";
     body << "  " << builderOpState << ".addOperands(" << argName << ");\n";
   }
+  if (op.hasResizableOperandList())
+    body << formatv("  {0}.setOperandListToResizable();\n", builderOpState);
 
   // If the operation has the operand segment size attribute, add it here.
   if (op.getTrait("OpTrait::AttrSizedOperandSegments")) {

mlir/tools/mlir-tblgen/OpFormatGen.cpp

@@ -706,6 +706,10 @@ void OperationFormat::genParser(Operator &op, OpClass &opClass) {
   genParserSuccessorResolution(op, body);
   genParserVariadicSegmentResolution(op, body);
 
+  // Mark the operation as having resizable operand list if required.
+  if (op.hasResizableOperandList())
+    body << "  result.setOperandListToResizable();\n";
+
   body << "  return success();\n";
 }