[mlir] ODS: support operations with resizable operand lists

[ftynse]

Summary

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.

[jpienaar]

How do we refer to resizable operands? E.g., OpE has resizeable operands, how do you refer to them in a pattern?

[jpienaar]

Or is the expectation that the last operand is always variadic and that is the resizable one?

[ftynse]

In D78292#1987137, @jpienaar wrote:

How do we refer to resizable operands? E.g., OpE has resizeable operands, how do you refer to them in a pattern?

The same way you usually do. It's up to your op how its operands are split into groups and how those groups are referred to. Being able to resize the list doesn't necessarily mean append operands. In most cases, I'd expect there is a variadic group. But I can also see an op with one optional operand. The use case I have has 3 variadic groups and 2 of them are modified.

[mehdi_amini]

Can we handle this with a traits on the op itself instead?

This implementation does not seem to be robust with respect to custom builders, and I suspect that with a trait we could just change the generic create method to query this traits and set the right flag.

mlir/include/mlir/IR/OpBase.td

class HasParent<string op>
     : ParamNativeOpTrait<"HasParent", op>;
 
 // Op result type is derived from the first attribute. If the attribute is an
 // subclass of `TypeAttrBase`, its value is used, otherwise, the type of the
 // attribute content is used.
 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.
 
 // All variadic operands of the op have the same number of values.
 // A variadic operand contains an array of values whose array size is only
 // known at runtime. This trait requires all variadic operands of an op
 // to have the same array size.
 def SameVariadicOperandSize : GenInternalOpTrait<"SameVariadicOperandSize">;

mlir/include/mlir/TableGen/Operator.h

public:
   StringRef getArgName(int index) const;
   var_decorator_range getArgDecorators(int index) const;
 
   // Returns the trait wrapper for the given MLIR C++ `trait`.
   // TODO: We should add a C++ wrapper class for TableGen OpTrait instead of
   // 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;
   const_region_iterator region_end() const;
   llvm::iterator_range<const_region_iterator> getRegions() const;
 
   // Returns the number of regions.
   unsigned getNumRegions() const;

mlir/lib/TableGen/Operator.cpp

for (const auto &t : traits) {
     } else if (auto opTrait = dyn_cast<tblgen::InterfaceOpTrait>(&t)) {
       if (opTrait->getTrait() == trait)
         return opTrait;
     }
   }
   return nullptr;
 }
 
+bool tblgen::Operator::hasResizableOperandList() const {
+  return getTrait("OpTrait::ResizableOperandList") != nullptr;
+}
+
 auto tblgen::Operator::region_begin() const -> const_region_iterator {
   return regions.begin();
 }
 auto tblgen::Operator::region_end() const -> const_region_iterator {
   return regions.end();
 }
 auto tblgen::Operator::getRegions() const
     -> llvm::iterator_range<const_region_iterator> {

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

 
 // CHECK-LABEL: Value OpD::input2
 // CHECK-NEXT: return *getODSOperands(1).begin();
 
 // CHECK-LABEL: OpD::build
 // 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

auto emit = [&](AttrParamKind attrType, TypeParamKind paramKind,
                   bool inferType) {
     std::string paramList;
     llvm::SmallVector<std::string, 4> resultNames;
     buildParamList(paramList, resultNames, paramKind, attrType);
 
     auto &m =
         opClass.newMethod("void", "build", paramList, OpMethod::MP_Static);
     auto &body = m.body();
+
     genCodeForAddingArgAndRegionForBuilder(
         body, /*isRawValueAttr=*/attrType == AttrParamKind::UnwrappedValue);
 
     // Push all result types to the operation state
 
     if (inferType) {
       // Generate builder that infers type too.
       // TODO(jpienaar): Subsume this with general checking if type can be
std::string params =
       std::string("Builder *odsBuilder, OperationState &") + builderOpState +
       ", ValueRange operands, ArrayRef<NamedAttribute> attributes";
   if (op.getNumVariadicRegions())
     params += ", unsigned numRegions";
   auto &m = opClass.newMethod("void", "build", params, OpMethod::MP_Static);
   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";
 
   // Create the correct number of regions
   if (int numRegions = op.getNumRegions()) {
     body << llvm::formatv(
         "  for (unsigned i = 0; i != {0}; ++i)\n",
void OpEmitter::genUseAttrAsResultTypeBuilder() {
   body << "    if (attr.first != \"" << namedAttr.name << "\") continue;\n";
   if (namedAttr.attr.isTypeAttr()) {
     resultType = "attr.second.cast<TypeAttr>().getValue()";
   } else {
     resultType = "attr.second.getType()";
   }
 
   // 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";
 
   // Result types
   SmallVector<std::string, 2> resultTypes(op.getNumResults(), resultType);
   body << "    " << builderOpState << ".addTypes({"
        << llvm::join(resultTypes, ", ") << "});\n";
   body << "  }\n";
void OpEmitter::genCollectiveParamBuilder() {
   auto &body = m.body();
 
   // Operands
   if (numVariadicOperands == 0 || numNonVariadicOperands != 0)
     body << "  assert(operands.size()"
          << (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";
 
   // Create the correct number of regions
   if (int numRegions = op.getNumRegions()) {
     body << llvm::formatv(
         "  for (unsigned i = 0; i != {0}; ++i)\n",
void OpEmitter::genCodeForAddingArgAndRegionForBuilder(OpMethodBody &body,
                                                        bool isRawValueAttr) {
   // Push all operands to the result.
   for (int i = 0, e = op.getNumOperands(); i < e; ++i) {
     std::string argName = getArgumentName(op, i);
     if (op.getOperand(i).isOptional())
       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")) {
     body << "  " << builderOpState
          << ".addAttribute(\"operand_segment_sizes\", "
             "odsBuilder->getI32VectorAttr({";
     llvm::interleaveComma(
         llvm::seq<int>(0, op.getNumOperands()), body, [&](int i) {

mlir/tools/mlir-tblgen/OpFormatGen.cpp

for (auto &element : elements)
     genElementParser(element.get(), body, attrTypeCtx);
 
   // Generate the code to resolve the operand/result types and successors now
   // that they have been parsed.
   genParserTypeResolution(op, body);
   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";
 }
 
 void OperationFormat::genParserTypeResolution(Operator &op,
                                               OpMethodBody &body) {
   // If any of type resolutions use transformed variables, make sure that the
   // types of those variables are resolved.
   SmallPtrSet<const NamedTypeConstraint *, 8> verifiedVariables;