[mlir] Add support for parsing optional Attribute values.

[rriddle]

Summary

This adds a parseOptionalAttribute method to the OpAsmParser that allows for parsing optional attributes, in a similar fashion to how optional types are parsed. This also enables the use of attribute values as the first element of an assembly format optional group.

[jpienaar]

Nice, thanks (just small nits)

mlir/lib/Parser/AttributeParser.cpp

194

Is there a reason for the ordering or could we keep it sorted?

mlir/tools/mlir-tblgen/OpFormatGen.cpp

635–639

Nit: firstElement is not that much longer, lets just spell it out (I was going to say firstElem so that we use element's actual abbreviation but then it was even closer to full).

646

Could we do this one first and so keep these sorted?

1954

Same re sorted

[rriddle]

Resolve comments

[rriddle]

mlir/lib/Parser/AttributeParser.cpp

194

No particular reason, this was just the order from parseAttribute.

mlir/docs/OpDefinitions.md

 #### Optional Groups
 
 In certain situations operations may have "optional" information, e.g.
 attributes or an empty set of variadic operands. In these situations a section
 of the assembly format can be marked as `optional` based on the presence of this
 information. An optional group is defined by wrapping a set of elements within
 `()` followed by a `?` and has the following requirements:
 
-*   The first element of the group must either be a literal or an operand.
+*   The first element of the group must either be a literal, attribute, or an
+    operand.
     -   This is because the first element must be optionally parsable.
 *   Exactly one argument variable within the group must be marked as the anchor
     of the group.
     -   The anchor is the element whose presence controls whether the group
         should be printed/parsed.
     -   An element is marked as the anchor by adding a trailing `^`.
     -   The first element is *not* required to be the anchor of the group.
 *   Literals, variables, and type directives are the only valid elements within

mlir/include/mlir/IR/OpImplementation.h

public:
 
   /// Parse an arbitrary attribute of a given type and return it in result. This
   /// also adds the attribute to the specified attribute list with the specified
   /// name.
   virtual ParseResult parseAttribute(Attribute &result, Type type,
                                      StringRef attrName,
                                      NamedAttrList &attrs) = 0;
 
+  /// Parse an optional attribute.
+  virtual OptionalParseResult parseOptionalAttribute(Attribute &result,
+                                                     Type type,
+                                                     StringRef attrName,
+                                                     NamedAttrList &attrs) = 0;
+  OptionalParseResult parseOptionalAttribute(Attribute &result,
+                                             StringRef attrName,
+                                             NamedAttrList &attrs) {
+    return parseOptionalAttribute(result, Type(), attrName, attrs);
+  }
+
   /// Parse an attribute of a specific kind and type.
   template <typename AttrType>
   ParseResult parseAttribute(AttrType &result, Type type, StringRef attrName,
                              NamedAttrList &attrs) {
     llvm::SMLoc loc = getCurrentLocation();
 
     // Parse any kind of attribute.
     Attribute attr;

mlir/lib/Parser/AttributeParser.cpp

Attribute Parser::parseAttribute(Type type) {
   default:
     // Parse a type attribute.
     if (Type type = parseType())
       return TypeAttr::get(type);
     return nullptr;
   }
 }
 
+/// Parse an optional attribute with the provided type.
+OptionalParseResult Parser::parseOptionalAttribute(Attribute &attribute,
+                                                   Type type) {
+  switch (getToken().getKind()) {
+  case Token::at_identifier:
+  case Token::floatliteral:
+  case Token::integer:
+  case Token::hash_identifier:
+  case Token::kw_affine_map:
+  case Token::kw_affine_set:
+  case Token::kw_dense:
+  case Token::kw_false:
+  case Token::kw_loc:
+  case Token::kw_opaque:
+  case Token::kw_sparse:
+  case Token::kw_true:
+  case Token::kw_unit:
+  case Token::l_brace:
+  case Token::l_square:
+  case Token::minus:
+  case Token::string:
+    attribute = parseAttribute(type);
+    return success(attribute != nullptr);
+
+  default:
+    // Parse an optional type attribute.
+    Type type;
+    OptionalParseResult result = parseOptionalType(type);
+    if (result.hasValue() && succeeded(*result))
+      attribute = TypeAttr::get(type);
+    return result;
+  }
+}
+
 /// Attribute dictionary.
 ///
 ///   attribute-dict ::= `{` `}`
 ///                    | `{` attribute-entry (`,` attribute-entry)* `}`
 ///   attribute-entry ::= (bare-id | string-literal) `=` attribute-value
 ///
 ParseResult Parser::parseAttributeDict(NamedAttrList &attributes) {
   if (parseToken(Token::l_brace, "expected '{' in attribute dictionary"))

mlir/lib/Parser/Parser.h

public:
 
   //===--------------------------------------------------------------------===//
   // Attribute Parsing
   //===--------------------------------------------------------------------===//
 
   /// Parse an arbitrary attribute with an optional type.
   Attribute parseAttribute(Type type = {});
 
+  /// Parse an optional attribute with the provided type.
+  OptionalParseResult parseOptionalAttribute(Attribute &attribute,
+                                             Type type = {});
+
   /// Parse an attribute dictionary.
   ParseResult parseAttributeDict(NamedAttrList &attributes);
 
   /// Parse an extended attribute.
   Attribute parseExtendedAttr(Type type);
 
   /// Parse a float attribute.
   Attribute parseFloatAttr(Type type, bool isNegative);

mlir/lib/Parser/Parser.cpp

ParseResult parseAttribute(Attribute &result, Type type, StringRef attrName,
     result = parser.parseAttribute(type);
     if (!result)
       return failure();
 
     attrs.push_back(parser.builder.getNamedAttr(attrName, result));
     return success();
   }
 
+  /// Parse an optional attribute.
+  OptionalParseResult parseOptionalAttribute(Attribute &result, Type type,
+                                             StringRef attrName,
+                                             NamedAttrList &attrs) override {
+    OptionalParseResult parseResult =
+        parser.parseOptionalAttribute(result, type);
+    if (parseResult.hasValue() && succeeded(*parseResult))
+      attrs.push_back(parser.builder.getNamedAttr(attrName, result));
+    return parseResult;
+  }
+
   /// Parse a named dictionary into 'result' if it is present.
   ParseResult parseOptionalAttrDict(NamedAttrList &result) override {
     if (parser.getToken().isNot(Token::l_brace))
       return success();
     return parser.parseAttributeDict(result);
   }
 
   /// Parse a named dictionary into 'result' if the `attributes` keyword is

mlir/test/lib/Dialect/Test/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 optional attributes that are within the syntax.
-def FormatOptAttrOp : TEST_Op<"format_opt_attr_op"> {
+def FormatOptAttrAOp : TEST_Op<"format_opt_attr_op_a"> {
   let arguments = (ins OptionalAttr<I64Attr>:$opt_attr);
-  let assemblyFormat = "(`(`$opt_attr^`)`)? attr-dict";
+  let assemblyFormat = "(`(` $opt_attr^ `)` )? attr-dict";
 }
+def FormatOptAttrBOp : TEST_Op<"format_opt_attr_op_b"> {
+  let arguments = (ins OptionalAttr<I64Attr>:$opt_attr);
+  let assemblyFormat = "($opt_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, OptionalAttr<I64Attr>:$opt_attr);
   let assemblyFormat = "attr-dict-with-keyword";
 }
 
 // Test that we don't need to provide types in the format if they are buildable.

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

 // CHECK: error: expected directive, literal, variable, or optional group
 def OptionalInvalidB : TestFormat_Op<"optional_invalid_b", [{
   () attr-dict
 }]>, Arguments<(ins OptionalAttr<I64Attr>:$attr)>;
 // CHECK: error: optional group specified no anchor element
 def OptionalInvalidC : TestFormat_Op<"optional_invalid_c", [{
   ($attr)? attr-dict
 }]>, Arguments<(ins OptionalAttr<I64Attr>:$attr)>;
-// CHECK: error: first element of an operand group must be a literal or operand
+// CHECK: error: first element of an operand group must be an attribute, literal, or operand
 def OptionalInvalidD : TestFormat_Op<"optional_invalid_d", [{
-  ($attr^)? attr-dict
-}]>, Arguments<(ins OptionalAttr<I64Attr>:$attr)>;
+  (type($operand) $operand^)? attr-dict
+}]>, Arguments<(ins Optional<I64>:$operand)>;
 // CHECK: error: type directive can only refer to variables within the optional group
 def OptionalInvalidE : TestFormat_Op<"optional_invalid_e", [{
   (`,` $attr^ type(operands))? attr-dict
 }]>, Arguments<(ins OptionalAttr<I64Attr>:$attr)>;
 // CHECK: error: only one element can be marked as the anchor of an optional group
 def OptionalInvalidF : TestFormat_Op<"optional_invalid_f", [{
   ($attr^ $attr2^) attr-dict
 }]>, Arguments<(ins OptionalAttr<I64Attr>:$attr, OptionalAttr<I64Attr>:$attr2)>;

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

 // RUN: mlir-opt -allow-unregistered-dialect %s | mlir-opt -allow-unregistered-dialect -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_opt_attr_op(10)
+// CHECK: test.format_opt_attr_op_a(10)
 // CHECK-NOT: {opt_attr
-test.format_opt_attr_op(10)
+test.format_opt_attr_op_a(10)
+test.format_opt_attr_op_a
+
+// CHECK: test.format_opt_attr_op_b 10
+// CHECK-NOT: {opt_attr
+test.format_opt_attr_op_b 10
+test.format_opt_attr_op_b
 
 // CHECK: test.format_attr_dict_w_keyword attributes {attr = 10 : i64}
 test.format_attr_dict_w_keyword attributes {attr = 10 : i64}
 
 // CHECK: test.format_attr_dict_w_keyword attributes {attr = 10 : i64, opt_attr = 10 : i64}
 test.format_attr_dict_w_keyword attributes {attr = 10 : i64, opt_attr = 10 : i64}
 
 // CHECK: test.format_buildable_type_op %[[I64]]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

 /// {0}: The storage type of the attribute.
 /// {1}: The name of the attribute.
 /// {2}: The type for the attribute.
 const char *const attrParserCode = R"(
   {0} {1}Attr;
   if (parser.parseAttribute({1}Attr{2}, "{1}", result.attributes))
     return failure();
 )";
+const char *const optionalAttrParserCode = R"(
+  {0} {1}Attr;
+  {
+    ::mlir::OptionalParseResult parseResult =
+      parser.parseOptionalAttribute({1}Attr{2}, "{1}", result.attributes);
+    if (parseResult.hasValue() && failed(*parseResult))
+      return failure();
+  }
+)";
 
 /// The code snippet used to generate a parser call for an enum attribute.
 ///
 /// {0}: The name of the attribute.
 /// {1}: The c++ namespace for the enum symbolize functions.
 /// {2}: The function to symbolize a string of the enum.
 /// {3}: The constant builder call to create an attribute of the enum type.
 const char *const enumAttrParserCode = R"(
   {
     ::mlir::StringAttr attrVal;
     ::mlir::NamedAttrList attrStorage;
     auto loc = parser.getCurrentLocation();
     if (parser.parseAttribute(attrVal, parser.getBuilder().getNoneType(),
                               "{0}", attrStorage))
       return failure();
 
     auto attrOptional = {1}::{2}(attrVal.getValue());
     if (!attrOptional)
       return parser.emitError(loc, "invalid ")
              << "{0} attribute specification: " << attrVal;
 
     result.addAttribute("{0}", {3});
   }
 )";
+const char *const optionalEnumAttrParserCode = R"(
+  Attribute {0}Attr;
+  {
+    ::mlir::StringAttr attrVal;
+    ::mlir::NamedAttrList attrStorage;
+    auto loc = parser.getCurrentLocation();
+
+    ::mlir::OptionalParseResult parseResult =
+      parser.parseOptionalAttribute(attrVal, parser.getBuilder().getNoneType(),
+                                    "{0}", attrStorage);
+    if (parseResult.hasValue()) {
+      if (failed(*parseResult))
+        return failure();
+
+      auto attrOptional = {1}::{2}(attrVal.getValue());
+      if (!attrOptional)
+        return parser.emitError(loc, "invalid ")
+               << "{0} attribute specification: " << attrVal;
+
+      {0}Attr = {3};
+      result.addAttribute("{0}", {0}Attr);
+    }
+  }
+)";
 
 /// The code snippet used to generate a parser call for an operand.
 ///
 /// {0}: The name of the operand.
 const char *const variadicOperandParserCode = R"(
   if (parser.parseOperandList({0}Operands))
     return failure();
 )";

 static void genElementParser(Element *element, OpMethodBody &body,
                              FmtContext &attrTypeCtx) {
   /// Optional Group.
   if (auto *optional = dyn_cast<OptionalElement>(element)) {
     auto elements = optional->getElements();
 
     // Generate a special optional parser for the first element to gate the
     // parsing of the rest of the elements.
-    if (auto *literal = dyn_cast<LiteralElement>(&*elements.begin())) {
+    Element *firstElement = &*elements.begin();
+    if (auto *attrVar = dyn_cast<AttributeVariable>(firstElement)) {
+      genElementParser(attrVar, body, attrTypeCtx);
+      body << "  if (" << attrVar->getVar()->name << "Attr) {\n";
+    } else if (auto *literal = dyn_cast<LiteralElement>(firstElement)) {
       body << "  if (succeeded(parser.parseOptional";
       genLiteralParser(literal->getLiteral(), body);
       body << ")) {\n";
-    } else if (auto *opVar = dyn_cast<OperandVariable>(&*elements.begin())) {
+    } else if (auto *opVar = dyn_cast<OperandVariable>(firstElement)) {
       genElementParser(opVar, body, attrTypeCtx);
       body << "  if (!" << opVar->getVar()->name << "Operands.empty()) {\n";
     }
 
     // Generate the rest of the elements normally.
     for (auto &childElement : llvm::drop_begin(elements, 1))
       genElementParser(&childElement, body, attrTypeCtx);
     body << "  }\n";
 
     /// Literals.
   } else if (LiteralElement *literal = dyn_cast<LiteralElement>(element)) {
if (canFormatEnumAttr(var)) {
       // Generate the code for building an attribute for this enum.
       std::string attrBuilderStr;
       {
         llvm::raw_string_ostream os(attrBuilderStr);
         os << tgfmt(enumAttr.getConstBuilderTemplate(), &attrTypeCtx,
                     "attrOptional.getValue()");
       }
 
-      body << formatv(enumAttrParserCode, var->name, enumAttr.getCppNamespace(),
+      body << formatv(var->attr.isOptional() ? optionalEnumAttrParserCode
+                                             : enumAttrParserCode,
+                      var->name, enumAttr.getCppNamespace(),
                       enumAttr.getStringToSymbolFnName(), attrBuilderStr);
       return;
     }
 
     // If this attribute has a buildable type, use that when parsing the
     // attribute.
     std::string attrTypeStr;
     if (Optional<StringRef> typeBuilder = attr->getTypeBuilder()) {
       llvm::raw_string_ostream os(attrTypeStr);
       os << ", " << tgfmt(*typeBuilder, &attrTypeCtx);
     }
 
-    body << formatv(attrParserCode, var->attr.getStorageType(), var->name,
-                    attrTypeStr);
+    body << formatv(var->attr.isOptional() ? optionalAttrParserCode
+                                           : attrParserCode,
+                    var->attr.getStorageType(), var->name, attrTypeStr);
   } else if (auto *operand = dyn_cast<OperandVariable>(element)) {
     ArgumentLengthKind lengthKind = getArgumentLengthKind(operand->getVar());
     StringRef name = operand->getVar()->name;
     if (lengthKind == ArgumentLengthKind::Variadic)
       body << llvm::formatv(variadicOperandParserCode, name);
     else if (lengthKind == ArgumentLengthKind::Optional)
       body << llvm::formatv(optionalOperandParserCode, name);
     else
if (failed(parseToken(Token::question, "expected '?' after optional group")))
     return failure();
 
   // The optional group is required to have an anchor.
   if (!anchorIdx)
     return emitError(curLoc, "optional group specified no anchor element");
 
   // The first element of the group must be one that can be parsed/printed in an
   // optional fashion.
-  if (!isa<LiteralElement>(&*elements.front()) &&
-      !isa<OperandVariable>(&*elements.front()))
-    return emitError(curLoc, "first element of an operand group must be a "
-                             "literal or operand");
+  Element *firstElement = &*elements.front();
+  if (!isa<AttributeVariable>(firstElement) &&
+      !isa<LiteralElement>(firstElement) && !isa<OperandVariable>(firstElement))
+    return emitError(curLoc, "first element of an operand group must be an "
+                             "attribute, literal, or operand");
 
   // After parsing all of the elements, ensure that all type directives refer
   // only to elements within the group.
   auto checkTypeOperand = [&](Element *typeEle) {
     auto *opVar = dyn_cast<OperandVariable>(typeEle);
     const NamedTypeConstraint *var = opVar ? opVar->getVar() : nullptr;
     if (!seenVariables.count(var))
       return emitError(curLoc, "type directive can only refer to variables "