diff --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td --- a/mlir/include/mlir/IR/OpBase.td +++ b/mlir/include/mlir/IR/OpBase.td @@ -1653,6 +1653,10 @@ 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. diff --git a/mlir/include/mlir/TableGen/Operator.h b/mlir/include/mlir/TableGen/Operator.h --- a/mlir/include/mlir/TableGen/Operator.h +++ b/mlir/include/mlir/TableGen/Operator.h @@ -165,6 +165,9 @@ // 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; diff --git a/mlir/lib/TableGen/Operator.cpp b/mlir/lib/TableGen/Operator.cpp --- a/mlir/lib/TableGen/Operator.cpp +++ b/mlir/lib/TableGen/Operator.cpp @@ -169,6 +169,10 @@ return nullptr; } +bool tblgen::Operator::hasResizableOperandList() const { + return getTrait("OpTrait::ResizableOperandList") != nullptr; +} + auto tblgen::Operator::region_begin() const -> const_region_iterator { return regions.begin(); } diff --git a/mlir/test/mlir-tblgen/op-operand.td b/mlir/test/mlir-tblgen/op-operand.td --- a/mlir/test/mlir-tblgen/op-operand.td +++ b/mlir/test/mlir-tblgen/op-operand.td @@ -58,3 +58,23 @@ // 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:$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 +// CHECK: odsState.setOperandListToResizable() + +// CHECK-LABEL: OpE::build(Builder *, OperationState +// CHECK: odsState.setOperandListToResizable() + +// CHECK-LABEL: OpE::parse +// CHECK: result.setOperandListToResizable() diff --git a/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp b/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp --- a/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp +++ b/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp @@ -685,6 +685,7 @@ auto &m = opClass.newMethod("void", "build", paramList, OpMethod::MP_Static); auto &body = m.body(); + genCodeForAddingArgAndRegionForBuilder( body, /*isRawValueAttr=*/attrType == AttrParamKind::UnwrappedValue); @@ -762,7 +763,9 @@ 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 @@ } // 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 @@ << (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 @@ 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")) { diff --git a/mlir/tools/mlir-tblgen/OpFormatGen.cpp b/mlir/tools/mlir-tblgen/OpFormatGen.cpp --- a/mlir/tools/mlir-tblgen/OpFormatGen.cpp +++ b/mlir/tools/mlir-tblgen/OpFormatGen.cpp @@ -706,6 +706,10 @@ 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"; }