[mlir][DeclarativeParser] Add an 'attr-dict-with-keyword' directive

[rriddle]

Summary

This matches the '(print|parse)OptionalAttrDictWithKeyword' functionality provided by the assembly parser/printer.

Depends On D74681

[jpienaar]

mlir/docs/OpDefinitions.md

593

I'm not sure I follow this, so attributes is just some string? Why do we need to treat attributes special vs just have "attributes attr-dic" ?

[rriddle]

mlir/docs/OpDefinitions.md

593

We do this because the attributes keyword is optionally printed in the same context as the dictionary. The OpAsmPrinter gained support a while ago to do this to allow differentiation between region and attribute dictionaries.

[jpienaar]

LGTM, Seems a bit too specialized but OK matches current behavior and not too invasive.

[antiagainst]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

134

Nit: explicit?

[rriddle]

Resolve comments.

mlir/docs/OpDefinitions.md

 
 The format is comprised of three components:
 
 #### Directives
 
 A directive is a type of builtin function, with an optional set of arguments.
 The available directives are as follows:
 
-* `attr-dict`
-  -  Represents the attribute dictionary of the operation.
+*   `attr-dict`
+
+    -   Represents the attribute dictionary of the operation.
 
-* `functional-type` ( inputs , results )
-  -  Formats the `inputs` and `results` arguments as a
+*   `attr-dict-with-keyword`
+
+    -   Represents the attribute dictionary of the operation, but prefixes the
+        dictionary with an `attributes` keyword.
+
+*   `functional-type` ( inputs , results )
+
+    -   Formats the `inputs` and `results` arguments as a
         [function type](LangRef.md#function-type).
-  -  The constraints on `inputs` and `results` are the same as the `input` of
+    -   The constraints on `inputs` and `results` are the same as the `input` of
         the `type` directive.
 
-* `operands`
-  -  Represents all of the operands of an operation.
+*   `operands`
+
+    -   Represents all of the operands of an operation.
 
-* `results`
-  -  Represents all of the results of an operation.
+*   `results`
+
+    -   Represents all of the results of an operation.
 
-* `type` ( input )
-  - Represents the type of the given input.
-  - `input` must be either an operand or result [variable](#variables), the
+*   `type` ( input )
+
+    -   Represents the type of the given input.
+    -   `input` must be either an operand or result [variable](#variables), the
         `operands` directive, or the `results` directive.
 
 #### Literals
 
 A literal is either a keyword or punctuation surrounded by \`\`.
 
 The following are the set of valid punctuation:
   `:`, `,`, `=`, `<`, `>`, `(`, `)`, `[`, `]`, `->`
 

mlir/test/lib/TestDialect/TestOps.td

 }
 
 // Test that we elide attributes that are within the syntax.
 def FormatAttrOp : TEST_Op<"format_attr_op"> {
   let arguments = (ins I64Attr:$attr);
   let assemblyFormat = "$attr attr-dict";
 }
 
+// Test that we elide attributes that are within the syntax.
+def FormatAttrDictWithKeywordOp : TEST_Op<"format_attr_dict_w_keyword"> {
+  let arguments = (ins I64Attr:$attr);
+  let assemblyFormat = "attr-dict-with-keyword";
+}
+
 // Test that we don't need to provide types in the format if they are buildable.
 def FormatBuildableTypeOp : TEST_Op<"format_buildable_type_op"> {
   let arguments = (ins I64:$buildable);
   let results = (outs I64:$buildable_res);
   let assemblyFormat = "$buildable attr-dict";
 }
 
 // Test various mixings of result type formatting.

mlir/test/mlir-tblgen/op-format-spec.td

 
 // CHECK: error: format missing 'attr-dict' directive
 def DirectiveAttrDictInvalidA : TestFormat_Op<"attrdict_invalid_a", [{
 }]>;
 // CHECK: error: 'attr-dict' directive has already been seen
 def DirectiveAttrDictInvalidB : TestFormat_Op<"attrdict_invalid_b", [{
   attr-dict attr-dict
 }]>;
-// CHECK: error: 'attr-dict' directive can only be used as a top-level directive
+// CHECK: error: 'attr-dict' directive has already been seen
 def DirectiveAttrDictInvalidC : TestFormat_Op<"attrdict_invalid_c", [{
+  attr-dict attr-dict-with-keyword
+}]>;
+// CHECK: error: 'attr-dict' directive can only be used as a top-level directive
+def DirectiveAttrDictInvalidD : TestFormat_Op<"attrdict_invalid_d", [{
   type(attr-dict)
 }]>;
 // CHECK-NOT: error
-def DirectiveAttrDictValid : TestFormat_Op<"attrdict_valid", [{
+def DirectiveAttrDictValidA : TestFormat_Op<"attrdict_valid_a", [{
   attr-dict
 }]>;
+def DirectiveAttrDictValidB : TestFormat_Op<"attrdict_valid_b", [{
+  attr-dict-with-keyword
+}]>;
 
 //===----------------------------------------------------------------------===//
 // functional-type
 
 // CHECK: error: 'functional-type' is only valid as a top-level directive
 def DirectiveFunctionalTypeInvalidA : TestFormat_Op<"functype_invalid_a", [{
   functional-type(functional-type)
 }]>;

mlir/test/mlir-tblgen/op-format.mlir

 // RUN: mlir-opt %s | mlir-opt -verify-diagnostics | FileCheck %s
 
 // CHECK: %[[I64:.*]] =
 %i64 = "foo.op"() : () -> (i64)
 // CHECK: %[[MEMREF:.*]] =
 %memref = "foo.op"() : () -> (memref<1xf64>)
 
 // CHECK: test.format_literal_op keyword_$. -> :, = <> () [] {foo.some_attr}
 test.format_literal_op keyword_$. -> :, = <> () [] {foo.some_attr}
 
 // CHECK: test.format_attr_op 10
 // CHECK-NOT: {attr
 test.format_attr_op 10
 
+// CHECK: test.format_attr_dict_w_keyword attributes {attr = 10 : i64}
+test.format_attr_dict_w_keyword attributes {attr = 10 : i64}
+
 // CHECK: test.format_buildable_type_op %[[I64]]
 %ignored = test.format_buildable_type_op %i64
 
 // CHECK: test.format_result_a_op memref<1xf64>
 %ignored_a:2 = test.format_result_a_op memref<1xf64>
 
 // CHECK: test.format_result_b_op i64, memref<1xf64>
 %ignored_b:2 = test.format_result_b_op i64, memref<1xf64>

mlir/tools/mlir-tblgen/OpFormatGen.cpp

 namespace {
 /// This class implements single kind directives.
 template <Element::Kind type>
 class DirectiveElement : public Element {
 public:
   DirectiveElement() : Element(type){};
   static bool classof(const Element *ele) { return ele->getKind() == type; }
 };
-/// This class represents the `attr-dict` directive. This directive represents
-/// the attribute dictionary of the operation.
-using AttrDictDirective = DirectiveElement<Element::Kind::AttrDictDirective>;
-
 /// This class represents the `operands` directive. This directive represents
 /// all of the operands of an operation.
 using OperandsDirective = DirectiveElement<Element::Kind::OperandsDirective>;
 
 /// This class represents the `results` directive. This directive represents
 /// all of the results of an operation.
 using ResultsDirective = DirectiveElement<Element::Kind::ResultsDirective>;
 
+/// This class represents the `attr-dict` directive. This directive represents
+/// the attribute dictionary of the operation.
+class AttrDictDirective
+    : public DirectiveElement<Element::Kind::AttrDictDirective> {
+public:
+  AttrDictDirective(bool withKeyword) : withKeyword(withKeyword) {}
+  bool isWithKeyword() const { return withKeyword; }
+
+private:
+  /// If the dictionary should be printed with the 'attributes' keyword.
+  bool withKeyword;
+};
+
 /// This class represents the `functional-type` directive. This directive takes
 /// two arguments and formats them, respectively, as the inputs and results of a
 /// FunctionType.
-struct FunctionalTypeDirective
+class FunctionalTypeDirective
     : public DirectiveElement<Element::Kind::FunctionalTypeDirective> {
 public:
   FunctionalTypeDirective(std::unique_ptr<Element> inputs,
                           std::unique_ptr<Element> results)
       : inputs(std::move(inputs)), results(std::move(results)) {}
   Element *getInputs() const { return inputs.get(); }
   Element *getResults() const { return results.get(); }
 
 private:
   /// The input and result arguments.
   std::unique_ptr<Element> inputs, results;
 };
 
 /// This class represents the `type` directive.
-struct TypeDirective : public DirectiveElement<Element::Kind::TypeDirective> {
+class TypeDirective : public DirectiveElement<Element::Kind::TypeDirective> {
 public:
   TypeDirective(std::unique_ptr<Element> arg) : operand(std::move(arg)) {}
   Element *getOperand() const { return operand.get(); }
 
 private:
   /// The operand that is used to format the directive.
   std::unique_ptr<Element> operand;
 };
if (auto *optional = dyn_cast<OptionalElement>(element)) {
     body << formatv(attrParserCode, var->attr.getStorageType(), var->name,
                     attrTypeStr);
   } else if (auto *operand = dyn_cast<OperandVariable>(element)) {
     bool isVariadic = operand->getVar()->isVariadic();
     body << formatv(isVariadic ? variadicOperandParserCode : operandParserCode,
                     operand->getVar()->name);
 
     /// Directives.
-  } else if (isa<AttrDictDirective>(element)) {
-    body << "  if (parser.parseOptionalAttrDict(result.attributes))\n"
+  } else if (auto *attrDict = dyn_cast<AttrDictDirective>(element)) {
+    body << "  if (parser.parseOptionalAttrDict"
+         << (attrDict->isWithKeyword() ? "WithKeyword" : "")
+         << "(result.attributes))\n"
          << "    return failure();\n";
   } else if (isa<OperandsDirective>(element)) {
     body << "  llvm::SMLoc allOperandLoc = parser.getCurrentLocation();\n"
          << "  SmallVector<OpAsmParser::OperandType, 4> allOperands;\n"
          << "  if (parser.parseOperandList(allOperands))\n"
          << "    return failure();\n";
   } else if (auto *dir = dyn_cast<TypeDirective>(element)) {
     bool isVariadic = false;
for (unsigned i = 0, e = op.getNumOperands(); i != e; ++i) {
     body << ", result.operands))\n    return failure();\n";
   }
 }
 
 //===----------------------------------------------------------------------===//
 // PrinterGen
 
 /// Generate the printer for the 'attr-dict' directive.
-static void genAttrDictPrinter(OperationFormat &fmt, OpMethodBody &body) {
+static void genAttrDictPrinter(OperationFormat &fmt, OpMethodBody &body,
+                               bool withKeyword) {
   // Collect all of the attributes used in the format, these will be elided.
   SmallVector<const NamedAttribute *, 1> usedAttributes;
   for (auto &it : fmt.elements)
     if (auto *attr = dyn_cast<AttributeVariable>(it.get()))
       usedAttributes.push_back(attr->getVar());
 
-  body << "  p.printOptionalAttrDict(getAttrs(), /*elidedAttrs=*/{";
+  body << "  p.printOptionalAttrDict" << (withKeyword ? "WithKeyword" : "")
+       << "(getAttrs(), /*elidedAttrs=*/{";
   interleaveComma(usedAttributes, body, [&](const NamedAttribute *attr) {
     body << "\"" << attr->name << "\"";
   });
   body << "});\n";
 }
 
 /// Generate the printer for a literal value. `shouldEmitSpace` is true if a
 /// space should be emitted before this element. `lastWasPunctuation` is true if
if (OptionalElement *optional = dyn_cast<OptionalElement>(element)) {
     for (Element &childElement : optional->getElements())
       genElementPrinter(&childElement, body, fmt, shouldEmitSpace,
                         lastWasPunctuation);
     body << "  }\n";
     return;
   }
 
   // Emit the attribute dictionary.
-  if (isa<AttrDictDirective>(element)) {
-    genAttrDictPrinter(fmt, body);
+  if (auto *attrDict = dyn_cast<AttrDictDirective>(element)) {
+    genAttrDictPrinter(fmt, body, attrDict->isWithKeyword());
     lastWasPunctuation = false;
     return;
   }
 
   // Optionally insert a space before the next element. The AttrDict printer
   // already adds a space as necessary.
   if (shouldEmitSpace || !lastWasPunctuation)
     body << "  p << \" \";\n";
enum Kind {
     r_paren,
     comma,
     equal,
     question,
 
     // Keywords.
     keyword_start,
     kw_attr_dict,
+    kw_attr_dict_w_keyword,
     kw_functional_type,
     kw_operands,
     kw_results,
     kw_type,
     keyword_end,
 
     // String valued tokens.
     identifier,

 
 Token FormatLexer::lexIdentifier(const char *tokStart) {
   // Match the rest of the identifier regex: [0-9a-zA-Z_\-]*
   while (isalnum(*curPtr) || *curPtr == '_' || *curPtr == '-')
     ++curPtr;
 
   // Check to see if this identifier is a keyword.
   StringRef str(tokStart, curPtr - tokStart);
-  Token::Kind kind = llvm::StringSwitch<Token::Kind>(str)
-                         .Case("attr-dict", Token::kw_attr_dict)
-                         .Case("functional-type", Token::kw_functional_type)
-                         .Case("operands", Token::kw_operands)
-                         .Case("results", Token::kw_results)
-                         .Case("type", Token::kw_type)
-                         .Default(Token::identifier);
+  Token::Kind kind =
+      llvm::StringSwitch<Token::Kind>(str)
+          .Case("attr-dict-with-keyword", Token::kw_attr_dict_w_keyword)
   return Token(kind, str);
 }
 
 //===----------------------------------------------------------------------===//
 // FormatParser
 //===----------------------------------------------------------------------===//
 
 /// Function to find an element within the given range that has the same name as
LogicalResult parseOptional(std::unique_ptr<Element> &element,
                               bool isTopLevel);
   LogicalResult parseOptionalChildElement(
       std::vector<std::unique_ptr<Element>> &childElements,
       SmallPtrSetImpl<const NamedTypeConstraint *> &seenVariables,
       Optional<unsigned> &anchorIdx);
 
   /// Parse the various different directives.
   LogicalResult parseAttrDictDirective(std::unique_ptr<Element> &element,
-                                       llvm::SMLoc loc, bool isTopLevel);
+                                       llvm::SMLoc loc, bool isTopLevel,
+                                       bool withKeyword);
   LogicalResult parseFunctionalTypeDirective(std::unique_ptr<Element> &element,
                                              Token tok, bool isTopLevel);
   LogicalResult parseOperandsDirective(std::unique_ptr<Element> &element,
                                        llvm::SMLoc loc, bool isTopLevel);
   LogicalResult parseResultsDirective(std::unique_ptr<Element> &element,
                                       llvm::SMLoc loc, bool isTopLevel);
   LogicalResult parseTypeDirective(std::unique_ptr<Element> &element, Token tok,
                                    bool isTopLevel);

 
 LogicalResult FormatParser::parseDirective(std::unique_ptr<Element> &element,
                                            bool isTopLevel) {
   Token dirTok = curToken;
   consumeToken();
 
   switch (dirTok.getKind()) {
   case Token::kw_attr_dict:
-    return parseAttrDictDirective(element, dirTok.getLoc(), isTopLevel);
+    return parseAttrDictDirective(element, dirTok.getLoc(), isTopLevel,
+                                  /*withKeyword=*/false);
+  case Token::kw_attr_dict_w_keyword:
+    return parseAttrDictDirective(element, dirTok.getLoc(), isTopLevel,
+                                  /*withKeyword=*/true);
   case Token::kw_functional_type:
     return parseFunctionalTypeDirective(element, dirTok, isTopLevel);
   case Token::kw_operands:
     return parseOperandsDirective(element, dirTok.getLoc(), isTopLevel);
   case Token::kw_results:
     return parseResultsDirective(element, dirTok.getLoc(), isTopLevel);
   case Token::kw_type:
     return parseTypeDirective(element, dirTok, isTopLevel);
return TypeSwitch<Element *, LogicalResult>(childElements.back().get())
       .Default([&](auto *) {
         return emitError(childLoc, "only literals, types, and variables can be "
                                    "used within an optional group");
       });
 }
 
 LogicalResult
 FormatParser::parseAttrDictDirective(std::unique_ptr<Element> &element,
-                                     llvm::SMLoc loc, bool isTopLevel) {
+                                     llvm::SMLoc loc, bool isTopLevel,
+                                     bool withKeyword) {
   if (!isTopLevel)
     return emitError(loc, "'attr-dict' directive can only be used as a "
                           "top-level directive");
   if (hasAttrDict)
     return emitError(loc, "'attr-dict' directive has already been seen");
 
   hasAttrDict = true;
-  element = std::make_unique<AttrDictDirective>();
+  element = std::make_unique<AttrDictDirective>(withKeyword);
   return success();
 }
 
 LogicalResult
 FormatParser::parseFunctionalTypeDirective(std::unique_ptr<Element> &element,
                                            Token tok, bool isTopLevel) {
   llvm::SMLoc loc = tok.getLoc();
   if (!isTopLevel)