[mlir] Expand operand adapter to take attributes

[jpienaar]

Summary

• Enables using with more variadic sized operands;
• Generate convenience accessors for attributes;
  • The accessor are named the same as their name in ODS and returns attribute type (not convenience type) and no derived attributes.

This would also allow verifying argument constraints before the op is even
created (and dedupe some checking between shapes and created ops). This does
not change the name of adaptor nor does it require it except for ops with
variadic operands to keep this change smaller. The named attributes created
such could potentially also replace the xAttr variants being generated.

I considered creating separate adapter for attributes but decided against that
given operands queries require attributes in general (and definitely for
verification of operands and attributes).

[antiagainst]

Nice, thanks Jacques!

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

73–77

This is not really code for calculation. What about naming it as "adapterSegmentSizeAttrInitCode" or something?

74

Nit: "missing segment size attribute for op" ?

77

Similarly here. What about "opSegmentSizeAttrInitCode"?

1860

Nit: we can do std::string sizeAttrInit = formatv() to avoid the .str() at the end and be more explicit about the type.

1872

Do we want to construct the default value here like https://github.com/llvm/llvm-project/blob/master/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp#L368-L378?

[jpienaar]

Handling default attributes in adapter as in op

[jpienaar]

Updated, thanks

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

1872

Good point, this might also provide a different path to fixing that todo :)

mlir/include/mlir/TableGen/OpClass.h

protected:
   SmallVector<std::string, 4> fields;
 };
 
 // Class for holding an op for C++ code emission
 class OpClass : public Class {
 public:
   explicit OpClass(StringRef name, StringRef extraClassDeclaration = "");
 
-  // Sets whether this OpClass should generate the using directive for its
-  // associate operand adaptor class.
-  void setHasOperandAdaptorClass(bool has);
-
   // Adds an op trait.
   void addTrait(Twine trait);
 
   // Writes this op's class as a declaration to the given `os`.  Redefines
   // Class::writeDeclTo to also emit traits and extra class declarations.
   void writeDeclTo(raw_ostream &os) const;
 
 private:
   StringRef extraClassDeclaration;
   SmallVector<std::string, 4> traitsVec;
   StringSet<> traitsSet;
-  bool hasOperandAdaptor;
 };
 
 } // namespace tblgen
 } // namespace mlir
 
 #endif // MLIR_TABLEGEN_OPCLASS_H_

mlir/lib/TableGen/OpClass.cpp

void tblgen::Class::writeDefTo(raw_ostream &os) const {
   }
 }
 
 //===----------------------------------------------------------------------===//
 // OpClass definitions
 //===----------------------------------------------------------------------===//
 
 tblgen::OpClass::OpClass(StringRef name, StringRef extraClassDeclaration)
-    : Class(name), extraClassDeclaration(extraClassDeclaration),
-      hasOperandAdaptor(true) {}
-
-void tblgen::OpClass::setHasOperandAdaptorClass(bool has) {
-  hasOperandAdaptor = has;
-}
+    : Class(name), extraClassDeclaration(extraClassDeclaration) {}
 
 void tblgen::OpClass::addTrait(Twine trait) {
   auto traitStr = trait.str();
   if (traitsSet.insert(traitStr).second)
     traitsVec.push_back(std::move(traitStr));
 }
 
 void tblgen::OpClass::writeDeclTo(raw_ostream &os) const {
   os << "class " << className << " : public Op<" << className;
   for (const auto &trait : traitsVec)
     os << ", " << trait;
   os << "> {\npublic:\n";
   os << "  using Op::Op;\n";
-  if (hasOperandAdaptor)
-    os << "  using OperandAdaptor = " << className << "OperandAdaptor;\n";
 
   bool hasPrivateMethod = false;
   for (const auto &method : methods) {
     if (!method.isPrivate()) {
       method.writeDeclTo(os);
       os << "\n";
     } else {
       hasPrivateMethod = true;

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

 }
 
 // CHECK: class AOp;
 
 // CHECK-LABEL: NS::AOp declarations
 
 // CHECK: class AOpOperandAdaptor {
 // CHECK: public:
-// CHECK:   AOpOperandAdaptor(ArrayRef<Value> values);
+// CHECK:   AOpOperandAdaptor(ArrayRef<Value> values
 // CHECK:   ArrayRef<Value> getODSOperands(unsigned index);
 // CHECK:   Value a();
 // CHECK:   ArrayRef<Value> b();
+// CHECK:   IntegerAttr attr1();
+// CHECL:   FloatAttr attr2();
 // CHECK: private:
-// CHECK:   ArrayRef<Value> tblgen_operands;
+// CHECK:   ArrayRef<Value> odsOperands;
 // CHECK: };
 
 // CHECK: class AOp : public Op<AOp, OpTrait::AtLeastNRegions<1>::Impl, OpTrait::AtLeastNResults<1>::Impl, OpTrait::ZeroSuccessor, OpTrait::AtLeastNOperands<1>::Impl, OpTrait::IsIsolatedFromAbove
 // CHECK-NOT: OpTrait::IsIsolatedFromAbove
 // CHECK: public:
 // CHECK:   using Op::Op;
 // CHECK:   using OperandAdaptor = AOpOperandAdaptor;
 // CHECK:   static StringRef getOperationName();

 // CHECK:   void print(OpAsmPrinter &p);
 // CHECK:   LogicalResult verify();
 // CHECK:   static void getCanonicalizationPatterns(OwningRewritePatternList &results, MLIRContext *context);
 // CHECK:   LogicalResult fold(ArrayRef<Attribute> operands, SmallVectorImpl<OpFoldResult> &results);
 // CHECK:   // Display a graph for debugging purposes.
 // CHECK:   void displayGraph();
 // CHECK: };
 
+// Check AttrSizedOperandSegments
+// ---
+
+def NS_AttrSizedOperandOp : NS_Op<"attr_sized_operands",
+                                 [AttrSizedOperandSegments]> {
+  let arguments = (ins
+    Variadic<I32>:$a,
+    Variadic<I32>:$b,
+    I32:$c,
+    Variadic<I32>:$d,
+    I32ElementsAttr:$operand_segment_sizes
+  );
+}
+
+// CHECK-LABEL: AttrSizedOperandOpOperandAdaptor(
+// CHECK-SAME:    ArrayRef<Value> values
+// CHECK-SAME:    DictionaryAttr attrs
+// CHECK:  ArrayRef<Value> a();
+// CHECK:  ArrayRef<Value> b();
+// CHECK:  Value c();
+// CHECK:  ArrayRef<Value> d();
+// CHECK:  DenseIntElementsAttr operand_segment_sizes();
+
 // Check op trait for different number of operands
 // ---
 
 def NS_BOp : NS_Op<"op_with_no_operand", []> {
   let arguments = (ins);
 }
 
 // CHECK-LABEL: NS::BOp declarations

 
 // CHECK-LABEL: NS::_AOp declarations
 // CHECK: class _AOp : public Op<_AOp
 
 def _BOp : NS_Op<"_op_with_leading_underscore_and_no_namespace", []>;
 
 // CHECK-LABEL: _BOp declarations
 // CHECK: class _BOp : public Op<_BOp
+

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

 
 def OpA : NS_Op<"one_normal_operand_op", []> {
   let arguments = (ins I32:$input);
 }
 
 // CHECK-LABEL: OpA definitions
 
 // CHECK:      OpAOperandAdaptor::OpAOperandAdaptor
-// CHECK-NEXT: tblgen_operands = values
+// CHECK-NEXT: odsOperands = values
 
 // CHECK:      void OpA::build
 // CHECK:        Value input
 // CHECK:        odsState.addOperands(input);
 
 // CHECK:      void OpA::build
 // CHECK:        ValueRange operands
 // CHECK:        assert(operands.size() == 1u && "mismatched number of parameters");

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

 )";
 
 // The logic to calculate the actual value range for a declared operand/result
 // of an op with variadic operands/results. Note that this logic is assumes
 // the op has an attribute specifying the size of each operand/result segment
 // (variadic or not).
 //
 // {0}: The name of the attribute specifying the segment sizes.
-const char *attrSizedSegmentValueRangeCalcCode = R"(
+const char *adapterSegmentSizeAttrInitCode = R"(
+  assert(odsAttrs && "missing segment size attribute for op");
+  auto sizeAttr = odsAttrs.get("{0}").cast<DenseIntElementsAttr>();
+)";
+const char *opSegmentSizeAttrInitCode = R"(
   auto sizeAttr = getAttrOfType<DenseIntElementsAttr>("{0}");
+)";
+const char *attrSizedSegmentValueRangeCalcCode = R"(
   unsigned start = 0;
   for (unsigned i = 0; i < index; ++i)
     start += (*(sizeAttr.begin() + i)).getZExtValue();
   unsigned size = (*(sizeAttr.begin() + index)).getZExtValue();
-  return {{start, size};
+  return {start, size};
 )";
 
 // The logic to build a range of either operand or result values.
 //
 // {0}: The begin iterator of the actual values.
 // {1}: The call to generate the start and length of the value range.
 const char *valueRangeReturnCode = R"(
   auto valueRange = {1};

 
 // Generates the code to compute the start and end index of an operand or result
 // range.
 template <typename RangeT>
 static void
 generateValueRangeStartAndEnd(Class &opClass, StringRef methodName,
                               int numVariadic, int numNonVariadic,
                               StringRef rangeSizeCall, bool hasAttrSegmentSize,
-                              StringRef segmentSizeAttr, RangeT &&odsValues) {
+                              StringRef sizeAttrInit, RangeT &&odsValues) {
   auto &method = opClass.newMethod("std::pair<unsigned, unsigned>", methodName,
                                    "unsigned index");
 
   if (numVariadic == 0) {
     method.body() << "  return {index, 1};\n";
   } else if (hasAttrSegmentSize) {
-    method.body() << formatv(attrSizedSegmentValueRangeCalcCode,
-                             segmentSizeAttr);
+    method.body() << sizeAttrInit << attrSizedSegmentValueRangeCalcCode;
   } else {
     // Because the op can have arbitrarily interleaved variadic and non-variadic
     // operands, we need to embed a list in the "sink" getter method for
     // calculation at run-time.
     llvm::SmallVector<StringRef, 4> isVariadic;
     isVariadic.reserve(llvm::size(odsValues));
     for (auto &it : odsValues)
       isVariadic.push_back(it.isVariableLength() ? "true" : "false");

 // an iterator to the beginning of the operand range; use `rangeSizeCall` to
 // obtain the number of operands. `getOperandCallPattern` contains the code
 // necessary to obtain a single operand whose position will be substituted
 // instead of
 // "{0}" marker in the pattern.  Note that the pattern should work for any kind
 // of ops, in particular for one-operand ops that may not have the
 // `getOperand(unsigned)` method.
 static void generateNamedOperandGetters(const Operator &op, Class &opClass,
+                                        StringRef sizeAttrInit,
                                         StringRef rangeType,
                                         StringRef rangeBeginCall,
                                         StringRef rangeSizeCall,
                                         StringRef getOperandCallPattern) {
   const int numOperands = op.getNumOperands();
   const int numVariadicOperands = op.getNumVariableLengthOperands();
   const int numNormalOperands = numOperands - numVariadicOperands;
 
static void generateNamedOperandGetters(const Operator &op, Class &opClass,
   if (attrSizedOperands && sameVariadicSize) {
     PrintFatalError(op.getLoc(),
                     "op cannot have both 'AttrSizedOperandSegments' and "
                     "'SameVariadicOperandSize' traits");
   }
 
   // First emit a few "sink" getter methods upon which we layer all nicer named
   // getter methods.
-  generateValueRangeStartAndEnd(
-      opClass, "getODSOperandIndexAndLength", numVariadicOperands,
-      numNormalOperands, rangeSizeCall, attrSizedOperands,
-      "operand_segment_sizes", const_cast<Operator &>(op).getOperands());
+  generateValueRangeStartAndEnd(opClass, "getODSOperandIndexAndLength",
+                                numVariadicOperands, numNormalOperands,
+                                rangeSizeCall, attrSizedOperands, sizeAttrInit,
+                                const_cast<Operator &>(op).getOperands());
 
   auto &m = opClass.newMethod(rangeType, "getODSOperands", "unsigned index");
   m.body() << formatv(valueRangeReturnCode, rangeBeginCall,
                       "getODSOperandIndexAndLength(index)");
 
   // Then we emit nicer named getter methods by redirecting to the "sink" getter
   // method.
-
   for (int i = 0; i != numOperands; ++i) {
     const auto &operand = op.getOperand(i);
     if (operand.name.empty())
       continue;
 
     if (operand.isOptional()) {
       auto &m = opClass.newMethod("Value", operand.name);
       m.body() << "  auto operands = getODSOperands(" << i << ");\n"
                << "  return operands.empty() ? Value() : *operands.begin();";
     } else if (operand.isVariadic()) {
       auto &m = opClass.newMethod(rangeType, operand.name);
       m.body() << "  return getODSOperands(" << i << ");";
     } else {
       auto &m = opClass.newMethod("Value", operand.name);
       m.body() << "  return *getODSOperands(" << i << ").begin();";
     }
   }
 }
 
 void OpEmitter::genNamedOperandGetters() {
-  if (op.getTrait("OpTrait::AttrSizedOperandSegments"))
-    opClass.setHasOperandAdaptorClass(false);
-
   generateNamedOperandGetters(
-      op, opClass, /*rangeType=*/"Operation::operand_range",
+      op, opClass,
+      /*sizeAttrInit=*/
+      formatv(opSegmentSizeAttrInitCode, "operand_segment_sizes").str(),
+      /*rangeType=*/"Operation::operand_range",
       /*rangeBeginCall=*/"getOperation()->operand_begin()",
       /*rangeSizeCall=*/"getOperation()->getNumOperands()",
       /*getOperandCallPattern=*/"getOperation()->getOperand({0})");
 }
 
 void OpEmitter::genNamedOperandSetters() {
   auto *attrSizedOperands = op.getTrait("OpTrait::AttrSizedOperandSegments");
   for (int i = 0, e = op.getNumOperands(); i != e; ++i) {
if (attrSizedResults && sameVariadicSize) {
     PrintFatalError(op.getLoc(),
                     "op cannot have both 'AttrSizedResultSegments' and "
                     "'SameVariadicResultSize' traits");
   }
 
   generateValueRangeStartAndEnd(
       opClass, "getODSResultIndexAndLength", numVariadicResults,
       numNormalResults, "getOperation()->getNumResults()", attrSizedResults,
-      "result_segment_sizes", op.getResults());
+      formatv(opSegmentSizeAttrInitCode, "result_segment_sizes").str(),
+      op.getResults());
   auto &m = opClass.newMethod("Operation::result_range", "getODSResults",
                               "unsigned index");
   m.body() << formatv(valueRangeReturnCode, "getOperation()->result_begin()",
                       "getODSResultIndexAndLength(index)");
 
   for (int i = 0; i != numResults; ++i) {
     const auto &result = op.getResult(i);
     if (result.name.empty())
private:
   explicit OpOperandAdaptorEmitter(const Operator &op);
 
   Class adapterClass;
 };
 } // end namespace
 
 OpOperandAdaptorEmitter::OpOperandAdaptorEmitter(const Operator &op)
     : adapterClass(op.getCppClassName().str() + "OperandAdaptor") {
-  adapterClass.newField("ArrayRef<Value>", "tblgen_operands");
-  auto &constructor = adapterClass.newConstructor("ArrayRef<Value> values");
-  constructor.body() << "  tblgen_operands = values;\n";
-
-  generateNamedOperandGetters(op, adapterClass,
+  adapterClass.newField("ArrayRef<Value>", "odsOperands");
+  adapterClass.newField("DictionaryAttr", "odsAttrs");
+  const auto *attrSizedOperands =
+      op.getTrait("OpTrait::AttrSizedOperandSegments");
+  auto &constructor = adapterClass.newConstructor(
+      attrSizedOperands
+          ? "ArrayRef<Value> values, DictionaryAttr attrs"
+          : "ArrayRef<Value> values, DictionaryAttr attrs = nullptr");
+  constructor.body() << "  odsOperands = values;\n";
+  constructor.body() << "  odsAttrs = attrs;\n";
+
+  std::string sizeAttrInit =
+      formatv(adapterSegmentSizeAttrInitCode, "operand_segment_sizes");
+  generateNamedOperandGetters(op, adapterClass, sizeAttrInit,
                               /*rangeType=*/"ArrayRef<Value>",
-                              /*rangeBeginCall=*/"tblgen_operands.begin()",
-                              /*rangeSizeCall=*/"tblgen_operands.size()",
-                              /*getOperandCallPattern=*/"tblgen_operands[{0}]");
+                              /*rangeBeginCall=*/"odsOperands.begin()",
+                              /*rangeSizeCall=*/"odsOperands.size()",
+                              /*getOperandCallPattern=*/"odsOperands[{0}]");
+
+  FmtContext fctx;
+  fctx.withBuilder("mlir::Builder(odsAttrs.getContext())");
+
+  auto emitAttr = [&](StringRef name, Attribute attr) {
+    auto &body = adapterClass.newMethod(attr.getStorageType(), name).body();
+    body << "  assert(odsAttrs && \"no attributes when constructing adapter\");"
+         << "\n  " << attr.getStorageType() << " attr = "
+         << "odsAttrs.get(\"" << name << "\").";
+    if (attr.hasDefaultValue() || attr.isOptional())
+      body << "dyn_cast_or_null<";
+    else
+      body << "cast<";
+    body << attr.getStorageType() << ">();\n";
+
+    if (attr.hasDefaultValue()) {
+      // Use the default value if attribute is not set.
+      // TODO: this is inefficient, we are recreating the attribute for every
+      // call. This should be set instead.
+      std::string defaultValue = std::string(
+          tgfmt(attr.getConstBuilderTemplate(), &fctx, attr.getDefaultValue()));
+      body << "  if (!attr)\n    attr = " << defaultValue << ";\n";
+    }
+    body << "  return attr;\n";
+  };
+
+  for (auto &namedAttr : op.getAttributes()) {
+    const auto &name = namedAttr.name;
+    const auto &attr = namedAttr.attr;
+    if (!attr.isDerivedAttr())
+      emitAttr(name, attr);
+  }
 }
 
 void OpOperandAdaptorEmitter::emitDecl(const Operator &op, raw_ostream &os) {
   OpOperandAdaptorEmitter(op).adapterClass.writeDeclTo(os);
 }
 
 void OpOperandAdaptorEmitter::emitDef(const Operator &op, raw_ostream &os) {
   OpOperandAdaptorEmitter(op).adapterClass.writeDefTo(os);
 }
 
 // Emits the opcode enum and op classes.
 static void emitOpClasses(const std::vector<Record *> &defs, raw_ostream &os,
                           bool emitDecl) {
   IfDefScope scope("GET_OP_CLASSES", os);
   // First emit forward declaration for each class, this allows them to refer
   // to each others in traits for example.
   if (emitDecl) {
     for (auto *def : defs) {
       Operator op(*def);
       os << "class " << op.getCppClassName() << ";\n";
     }
   }
   for (auto *def : defs) {
     Operator op(*def);
-    const auto *attrSizedOperands =
-        op.getTrait("OpTrait::AttrSizedOperandSegments");
     if (emitDecl) {
       os << formatv(opCommentHeader, op.getQualCppClassName(), "declarations");
-      // We cannot generate the operand adaptor class if operand getters depend
-      // on an attribute.
-      if (!attrSizedOperands)
-        OpOperandAdaptorEmitter::emitDecl(op, os);
       OpEmitter::emitDecl(op, os);
     } else {
       os << formatv(opCommentHeader, op.getQualCppClassName(), "definitions");
-      if (!attrSizedOperands)
-        OpOperandAdaptorEmitter::emitDef(op, os);
       OpEmitter::emitDef(op, os);
     }
   }
 }
 
 // Emits a comma-separated list of the ops.
 static void emitOpList(const std::vector<Record *> &defs, raw_ostream &os) {
   IfDefScope scope("GET_OP_LIST", os);