[mlir][OpAsmFormat] Add support for an "else" group on optional elements

[rriddle]

Summary

The "else" group of an optional element is a collection of elements that get parsed/printed when the anchor of the main element group is *not* present. This is useful when there is a special syntax when an element is not present. The new syntax for an optional element is shown below:

optional-group: `(` elements `)` (`:` `(` else-elements `)`)? `?`

An example of how this might be used is shown below:

tablegen
def FooOp : ... {
  let arguments = (ins UnitAttr:$foo);

  let assemblyFormat = "attr-dict (`foo_is_present` $foo^):(`foo_is_absent`)?";
}

would be formatted as such:

mlir
// When the `foo` attribute is present:
foo.op foo_is_present

// When the `foo` attribute is not present:
foo.op foo_is_absent

[mehdi_amini]

Sweet! Thanks :)

mlir/docs/OpDefinitions.md

 
 When a variable is optional, the provided value may be null.
 
 #### 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:
+information. An optional group is defined as follows:
+
+```
+optional-group: `(` elements `)` (`:` `(` else-elements `)`)? `?`
+```
+
+The `elements` of an optional group have the following requirements:
 
 *   The first element of the group must either be a attribute, literal, operand,
     or region.
     -   This is because the first element must be optionally parsable.
 *   Exactly one argument variable or type directive 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.

 ```mlir
 // When the unit attribute is present:
 foo.op is_read_only
 
 // When the unit attribute is not present:
 foo.op
 ```
 
+##### Optional "else" Group
+
+Optional groups also have support for an "else" group of elements. These are
+elements that are parsed/printed if the `anchor` element of the optional group
+is *not* present. Unlike the main element group, the "else" group has no
+restriction on the first element and none of the elements may act as the
+`anchor` for the optional. An example is shown below:
+
+```tablegen
+def FooOp : ... {
+  let arguments = (ins UnitAttr:$foo);
+
+  let assemblyFormat = "attr-dict (`foo_is_present` $foo^):(`foo_is_absent`)?";
+}
+```
+
+would be formatted as such:
+
+```mlir
+// When the `foo` attribute is present:
+foo.op foo_is_present
+
+// When the `foo` attribute is not present:
+foo.op foo_is_absent
+```
+
 #### Requirements
 
 The format specification has a certain set of requirements that must be adhered
 to:
 
 1.  The output and operation name are never shown as they are fixed and cannot
     be altered.
 1.  All operands within the operation must appear within the format, either

mlir/test/lib/Dialect/Test/TestOps.td

def FormatOptionalUnitAttrNoElide
   let assemblyFormat = "($is_optional^)? attr-dict";
 }
 
 def FormatOptionalEnumAttr : TEST_Op<"format_optional_enum_attr"> {
   let arguments = (ins OptionalAttr<SomeI64Enum>:$attr);
   let assemblyFormat = "($attr^)? attr-dict";
 }
 
+def FormatOptionalWithElse : TEST_Op<"format_optional_else"> {
+  let arguments = (ins UnitAttr:$isFirstBranchPresent);
+  let assemblyFormat = "(`then` $isFirstBranchPresent^):(`else`)? attr-dict";
+}
+
 //===----------------------------------------------------------------------===//
 // Custom Directives
 
 def FormatCustomDirectiveOperands
     : TEST_Op<"format_custom_directive_operands", [AttrSizedOperandSegments]> {
   let arguments = (ins I64:$operand, Optional<I64>:$optOperand,
                        Variadic<I64>:$varOperands);
   let assemblyFormat = [{

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

 // CHECK: error: only variables and types can be used to anchor an optional group
 def OptionalInvalidL : TestFormat_Op<[{
   (custom<MyDirective>($arg)^)?
 }]>, Arguments<(ins I64:$arg)>;
 // CHECK: error: only variables and types can be used to anchor an optional group
 def OptionalInvalidM : TestFormat_Op<[{
   (` `^)?
 }]>, Arguments<(ins)>;
+// CHECK: error: expected '(' to start else branch of optional group
+def OptionalInvalidN : TestFormat_Op<[{
+  ($arg^):
+}]>, Arguments<(ins Variadic<I64>:$arg)>;
+// CHECK: error: expected directive, literal, variable, or optional group
+def OptionalInvalidO : TestFormat_Op<[{
+  ($arg^):(`test`
+}]>, Arguments<(ins Variadic<I64>:$arg)>;
+// CHECK: error: expected '?' after optional group
+def OptionalInvalidP : TestFormat_Op<[{
+  ($arg^):(`test`)
+}]>, Arguments<(ins Variadic<I64>:$arg)>;
 
 // CHECK-NOT: error
 def OptionalValidA : TestFormat_Op<[{
   (` ` `` $arg^)?
 }]>;
 
 //===----------------------------------------------------------------------===//
 // Variables

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

 
 // CHECK: test.format_optional_result_b_op : i64 -> i64, i64
 test.format_optional_result_b_op : i64 -> i64, i64
 
 // CHECK: test.format_optional_result_c_op : (i64) -> (i64, i64)
 test.format_optional_result_c_op : (i64) -> (i64, i64)
 
 //===----------------------------------------------------------------------===//
+// Format optional with else
+//===----------------------------------------------------------------------===//
+
+// CHECK: test.format_optional_else then
+test.format_optional_else then
+
+// CHECK: test.format_optional_else else
+test.format_optional_else else
+
+//===----------------------------------------------------------------------===//
 // Format custom directives
 //===----------------------------------------------------------------------===//
 
 // CHECK: test.format_custom_directive_operands %[[I64]], %[[I64]] -> (%[[I64]])
 test.format_custom_directive_operands %i64, %i64 -> (%i64)
 
 // CHECK: test.format_custom_directive_operands %[[I64]] -> (%[[I64]])
 test.format_custom_directive_operands %i64 -> (%i64)

mlir/tools/mlir-tblgen/OpFormatGen.cpp

 };
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // OptionalElement
 
 namespace {
 /// This class represents a group of elements that are optionally emitted based
-/// upon an optional variable of the operation.
+/// upon an optional variable of the operation, and a group of elements that are
+/// emotted when the anchor element is not present.
 class OptionalElement : public Element {
 public:
-  OptionalElement(std::vector<std::unique_ptr<Element>> &&elements,
+  OptionalElement(std::vector<std::unique_ptr<Element>> &&thenElements,
+                  std::vector<std::unique_ptr<Element>> &&elseElements,
                   unsigned anchor, unsigned parseStart)
-      : Element{Kind::Optional}, elements(std::move(elements)), anchor(anchor),
+      : Element{Kind::Optional}, thenElements(std::move(thenElements)),
+        elseElements(std::move(elseElements)), anchor(anchor),
         parseStart(parseStart) {}
   static bool classof(const Element *element) {
     return element->getKind() == Kind::Optional;
   }
 
-  /// Return the nested elements of this grouping.
-  auto getElements() const { return llvm::make_pointee_range(elements); }
+  /// Return the `then` elements of this grouping.
+  auto getThenElements() const {
+    return llvm::make_pointee_range(thenElements);
+  }
+
+  /// Return the `else` elements of this grouping.
+  auto getElseElements() const {
+    return llvm::make_pointee_range(elseElements);
+  }
 
   /// Return the anchor of this optional group.
-  Element *getAnchor() const { return elements[anchor].get(); }
+  Element *getAnchor() const { return thenElements[anchor].get(); }
 
   /// Return the index of the first element that needs to be parsed.
   unsigned getParseStart() const { return parseStart; }
 
 private:
-  /// The child elements of this optional.
-  std::vector<std::unique_ptr<Element>> elements;
+  /// The child elements of `then` branch of this optional.
+  std::vector<std::unique_ptr<Element>> thenElements;
+  /// The child elements of `else` branch of this optional.
+  std::vector<std::unique_ptr<Element>> elseElements;
   /// The index of the element that acts as the anchor for the optional group.
   unsigned anchor;
   /// The index of the first element that is parsed (is not a
   /// WhitespaceElement).
   unsigned parseStart;
 };
 } // end anonymous namespace
 
body << (StringRef)StringSwitch<StringRef>(value)
               .Case("?", "Question()")
               .Case("+", "Plus()")
               .Case("*", "Star()");
 }
 
 /// Generate the storage code required for parsing the given element.
 static void genElementParserStorage(Element *element, OpMethodBody &body) {
   if (auto *optional = dyn_cast<OptionalElement>(element)) {
-    auto elements = optional->getElements();
+    auto elements = optional->getThenElements();
 
     // If the anchor is a unit attribute, it won't be parsed directly so elide
     // it.
     auto *anchor = dyn_cast<AttributeVariable>(optional->getAnchor());
     Element *elidedAnchorElement = nullptr;
     if (anchor && anchor != &*elements.begin() && anchor->isUnitAttr())
       elidedAnchorElement = anchor;
     for (auto &childElement : elements)
       if (&childElement != elidedAnchorElement)
         genElementParserStorage(&childElement, body);
+    for (auto &childElement : optional->getElseElements())
+      genElementParserStorage(&childElement, body);
 
   } else if (auto *custom = dyn_cast<CustomDirective>(element)) {
     for (auto &paramElement : custom->getArguments())
       genElementParserStorage(&paramElement, body);
 
   } else if (isa<OperandsDirective>(element)) {
     body << "  ::mlir::SmallVector<::mlir::OpAsmParser::OperandType, 4> "
             "allOperands;\n";
void OperationFormat::genParser(Operator &op, OpClass &opClass) {
 
   body << "  return ::mlir::success();\n";
 }
 
 void OperationFormat::genElementParser(Element *element, OpMethodBody &body,
                                        FmtContext &attrTypeCtx) {
   /// Optional Group.
   if (auto *optional = dyn_cast<OptionalElement>(element)) {
-    auto elements =
-        llvm::drop_begin(optional->getElements(), optional->getParseStart());
+    auto elements = llvm::drop_begin(optional->getThenElements(),
+                                     optional->getParseStart());
 
     // Generate a special optional parser for the first element to gate the
     // parsing of the rest of the elements.
     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)) {
if (anchorAttr && anchorAttr != firstElement && anchorAttr->isUnitAttr()) {
            << "\", parser.getBuilder().getUnitAttr());\n";
     }
 
     // Generate the rest of the elements normally.
     for (Element &childElement : llvm::drop_begin(elements, 1)) {
       if (&childElement != elidedAnchorElement)
         genElementParser(&childElement, body, attrTypeCtx);
     }
-    body << "  }\n";
+    body << "  }";
+
+    // Generate the else elements.
+    auto elseElements = optional->getElseElements();
+    if (!elseElements.empty()) {
+      body << " else {\n";
+      for (Element &childElement : elseElements)
+        genElementParser(&childElement, body, attrTypeCtx);
+      body << "  }";
+    }
+    body << "\n";
 
     /// Literals.
   } else if (LiteralElement *literal = dyn_cast<LiteralElement>(element)) {
     body << "  if (parser.parse";
     genLiteralParser(literal->getLiteral(), body);
     body << ")\n    return ::mlir::failure();\n";
 
     /// Whitespaces.
void OperationFormat::genElementPrinter(Element *element, OpMethodBody &body,
   // Emit an optional group.
   if (OptionalElement *optional = dyn_cast<OptionalElement>(element)) {
     // Emit the check for the presence of the anchor element.
     Element *anchor = optional->getAnchor();
     genOptionalGroupPrinterAnchor(anchor, body);
 
     // If the anchor is a unit attribute, we don't need to print it. When
     // parsing, we will add this attribute if this group is present.
-    auto elements = optional->getElements();
+    auto elements = optional->getThenElements();
     Element *elidedAnchorElement = nullptr;
     auto *anchorAttr = dyn_cast<AttributeVariable>(anchor);
     if (anchorAttr && anchorAttr != &*elements.begin() &&
         anchorAttr->isUnitAttr()) {
       elidedAnchorElement = anchorAttr;
     }
 
     // Emit each of the elements.
     for (Element &childElement : elements) {
       if (&childElement != elidedAnchorElement) {
         genElementPrinter(&childElement, body, op, shouldEmitSpace,
                           lastWasPunctuation);
       }
     }
-    body << "  }\n";
+    body << "  }";
+
+    // Emit each of the else elements.
+    auto elseElements = optional->getElseElements();
+    if (!elseElements.empty()) {
+      body << " else {\n";
+      for (Element &childElement : elseElements) {
+        genElementPrinter(&childElement, body, op, shouldEmitSpace,
+                          lastWasPunctuation);
+      }
+      body << "  }";
+    }
+
+    body << "\n";
     return;
   }
 
   // Emit the attribute dictionary.
   if (auto *attrDict = dyn_cast<AttrDictDirective>(element)) {
     genAttrDictPrinter(*this, op, body, attrDict->isWithKeyword());
     lastWasPunctuation = false;
     return;
enum Kind {
     // Markers.
     eof,
     error,
 
     // Tokens with no info.
     l_paren,
     r_paren,
     caret,
+    colon,
     comma,
     equal,
     less,
     greater,
     question,
 
     // Keywords.
     keyword_start,
default:
     return emitError(tokStart, "unexpected character");
   case EOF:
     // Return EOF denoting the end of lexing.
     return formToken(Token::eof, tokStart);
 
   // Lex punctuation.
   case '^':
     return formToken(Token::caret, tokStart);
+  case ':':
+    return formToken(Token::colon, tokStart);
   case ',':
     return formToken(Token::comma, tokStart);
   case '=':
     return formToken(Token::equal, tokStart);
   case '<':
     return formToken(Token::less, tokStart);
   case '>':
     return formToken(Token::greater, tokStart);
LogicalResult FormatParser::verifyAttributes(
     SmallVectorImpl<std::pair<ElementsIterT, ElementsIterT>> &iteratorStack) {
   auto &stackIt = iteratorStack.back();
   ElementsIterT &it = stackIt.first, e = stackIt.second;
   while (it != e) {
     Element *element = &*(it++);
 
     // Traverse into optional groups.
     if (auto *optional = dyn_cast<OptionalElement>(element)) {
-      auto elements = optional->getElements();
-      iteratorStack.emplace_back(elements.begin(), elements.end());
+      auto thenElements = optional->getThenElements();
+      iteratorStack.emplace_back(thenElements.begin(), thenElements.end());
+
+      auto elseElements = optional->getElseElements();
+      iteratorStack.emplace_back(elseElements.begin(), elseElements.end());
       return ::mlir::success();
     }
 
     // We are checking for an attribute element followed by a `:`, so there is
     // no need to check the end.
     if (it == e && iteratorStack.size() == 1)
       break;
 
LogicalResult FormatParser::parseOptional(std::unique_ptr<Element> &element,
                                           ParserContext context) {
   llvm::SMLoc curLoc = curToken.getLoc();
   if (context != TopLevelContext)
     return emitError(curLoc, "optional groups can only be used as top-level "
                              "elements");
   consumeToken();
 
   // Parse the child elements for this optional group.
-  std::vector<std::unique_ptr<Element>> elements;
+  std::vector<std::unique_ptr<Element>> thenElements, elseElements;
   Optional<unsigned> anchorIdx;
   do {
-    if (failed(parseOptionalChildElement(elements, anchorIdx)))
+    if (failed(parseOptionalChildElement(thenElements, anchorIdx)))
       return ::mlir::failure();
   } while (curToken.getKind() != Token::r_paren);
   consumeToken();
+
+  // Parse the `else` elements of this optional group.
+  if (curToken.getKind() == Token::colon) {
+    consumeToken();
+    if (failed(parseToken(Token::l_paren, "expected '(' to start else branch "
+                                          "of optional group")))
+      return failure();
+    do {
+      llvm::SMLoc childLoc = curToken.getLoc();
+      elseElements.push_back({});
+      if (failed(parseElement(elseElements.back(), TopLevelContext)) ||
+          failed(verifyOptionalChildElement(elseElements.back().get(), childLoc,
+                                            /*isAnchor=*/false)))
+        return failure();
+    } while (curToken.getKind() != Token::r_paren);
+    consumeToken();
+  }
+
   if (failed(parseToken(Token::question, "expected '?' after optional group")))
     return ::mlir::failure();
 
   // The optional group is required to have an anchor.
   if (!anchorIdx)
     return emitError(curLoc, "optional group specified no anchor element");
 
   // The first parsable element of the group must be able to be parsed in an
   // optional fashion.
-  auto parseBegin = llvm::find_if_not(elements, [](auto &element) {
+  auto parseBegin = llvm::find_if_not(thenElements, [](auto &element) {
     return isa<WhitespaceElement>(element.get());
   });
   Element *firstElement = parseBegin->get();
   if (!isa<AttributeVariable>(firstElement) &&
       !isa<LiteralElement>(firstElement) &&
       !isa<OperandVariable>(firstElement) && !isa<RegionVariable>(firstElement))
     return emitError(curLoc,
                      "first parsable element of an operand group must be "
                      "an attribute, literal, operand, or region");
 
-  auto parseStart = parseBegin - elements.begin();
-  element = std::make_unique<OptionalElement>(std::move(elements), *anchorIdx,
-                                              parseStart);
+  auto parseStart = parseBegin - thenElements.begin();
+  element = std::make_unique<OptionalElement>(
+      std::move(thenElements), std::move(elseElements), *anchorIdx, parseStart);
   return ::mlir::success();
 }
 
 LogicalResult FormatParser::parseOptionalChildElement(
     std::vector<std::unique_ptr<Element>> &childElements,
     Optional<unsigned> &anchorIdx) {
   llvm::SMLoc childLoc = curToken.getLoc();
   childElements.push_back({});