[mlir][OpFormatGen] Add support for anchoring optional groups with types

[rriddle]

Summary

This revision adds support for using either operand or result types to anchor an optional group. It also removes the arbitrary restriction that type directives must refer to variables in the same group, which is overly limiting for a declarative format syntax.

Fixes PR#48784

[jpienaar]

Thanks!

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

234

Could we add a negative test to clearly show differnce between what a and b allows?

mlir/tools/mlir-tblgen/OpFormatGen.cpp

1761

Could this have been a type switch?

1770

Could you expand the request here?

[lattner]

Nice, thank you River!

mlir/docs/OpDefinitions.md

 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 attribute, literal, operand,
     or region.
     -   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.
+*   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.
     -   An element is marked as the anchor by adding a trailing `^`.
     -   The first element is *not* required to be the anchor of the group.
     -   When a non-variadic region anchors a group, the detector for printing
         the group is if the region is empty.
 *   Literals, variables, custom directives, and type directives are the only
     valid elements within the group.
     -   Any attribute variable may be used, but only optional attributes can be
         marked as the anchor.
-    -   Only variadic or optional operand arguments can be used.
+    -   Only variadic or optional results and operand arguments and can be used.
     -   All region variables can be used. When a non-variable length region is
         used, if the group is not present the region is empty.
-    -   The operands to a type directive must be defined within the optional
-        group.
 
 An example of an operation with an optional group is `std.return`, which has a
 variadic number of operands.
 
 ```tablegen
 def ReturnOp : ... {
   let arguments = (ins Variadic<AnyType>:$operands);
 

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

def FormatOptionalOperandResultAOp : FormatOptionalOperandResultOpBase<"a", [{
   (`[` $variadic^ `]`)? attr-dict
 }]>;
 
 def FormatOptionalOperandResultBOp : FormatOptionalOperandResultOpBase<"b", [{
   (`(` $optional^ `:` type($optional) `)`)? `:` type($optional_res)
   (`[` $variadic^ `]`)? attr-dict
 }]>;
 
+// Test optional result type formatting.
+class FormatOptionalResultOpBase<string suffix, string fmt>
+    : TEST_Op<"format_optional_result_" # suffix # "_op",
+              [AttrSizedResultSegments]> {
+  let results = (outs Optional<I64>:$optional, Variadic<I64>:$variadic);
+  let assemblyFormat = fmt;
+}
+def FormatOptionalResultAOp : FormatOptionalResultOpBase<"a", [{
+  (`:` type($optional)^ `->` type($variadic))? attr-dict
+}]>;
+
+def FormatOptionalResultBOp : FormatOptionalResultOpBase<"b", [{
+  (`:` type($optional) `->` type($variadic)^)? attr-dict
+}]>;
+
+def FormatOptionalResultCOp : FormatOptionalResultOpBase<"c", [{
+  (`:` functional-type($optional, $variadic)^)? attr-dict
+}]>;
+
 def FormatTwoVariadicOperandsNoBuildableTypeOp
     : TEST_Op<"format_two_variadic_operands_no_buildable_type_op",
               [AttrSizedOperandSegments]> {
   let arguments = (ins Variadic<AnyType>:$a,
                        Variadic<AnyType>:$b);
   let assemblyFormat = [{
     `(` $a `:` type($a) `)` `->` `(` $b `:` type($b) `)`  attr-dict
   }];

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

 // 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 parsable element of an operand group must be an attribute, literal, operand, or region
 def OptionalInvalidD : TestFormat_Op<"optional_invalid_d", [{
   (type($operand) $operand^)? attr-dict
 }]>, Arguments<(ins Optional<I64>:$operand)>;
-// CHECK: error: type directive can only refer to variables within the optional group
+// CHECK: error: only literals, types, and variables can be used within an 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)>;
 // CHECK: error: only optional attributes can be used to anchor an optional group
 def OptionalInvalidG : TestFormat_Op<"optional_invalid_g", [{
   ($attr^) attr-dict
 }]>, Arguments<(ins I64Attr:$attr)>;
 // CHECK: error: only variable length operands can be used within an optional group
 def OptionalInvalidH : TestFormat_Op<"optional_invalid_h", [{
   ($arg^) attr-dict
 }]>, Arguments<(ins I64:$arg)>;
-// CHECK: error: only variables can be used to anchor an optional group
+// CHECK: error: only literals, types, and variables can be used within an optional group
 def OptionalInvalidI : TestFormat_Op<"optional_invalid_i", [{
-  ($arg type($arg)^) attr-dict
+  (functional-type($arg, results)^)? attr-dict
 }]>, Arguments<(ins Variadic<I64>:$arg)>;
 // CHECK: error: only literals, types, and variables can be used within an optional group
 def OptionalInvalidJ : TestFormat_Op<"optional_invalid_j", [{
   (attr-dict)
 }]>;
 // CHECK: error: expected '?' after optional group
 def OptionalInvalidK : TestFormat_Op<"optional_invalid_k", [{
   ($arg^)
 }]>, Arguments<(ins Variadic<I64>:$arg)>;
-// CHECK: error: only variables can be used to anchor an optional group
+// CHECK: error: only variables and types can be used to anchor an optional group
 def OptionalInvalidL : TestFormat_Op<"optional_invalid_l", [{
   (custom<MyDirective>($arg)^)?
 }]>, Arguments<(ins I64:$arg)>;
-// CHECK: error: only variables can be used to anchor an optional group
+// CHECK: error: only variables and types can be used to anchor an optional group
 def OptionalInvalidM : TestFormat_Op<"optional_invalid_m", [{
   (` `^)?
 }]>, Arguments<(ins)>;
 
 // CHECK-NOT: error
 def OptionalValidA : TestFormat_Op<"optional_valid_a", [{
   (` ` `` $arg^)?
 }]>;

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

 
 // CHECK: test.format_optional_operand_result_b_op : i64
 test.format_optional_operand_result_b_op( : ) : i64
 
 // CHECK: test.format_optional_operand_result_b_op : i64
 test.format_optional_operand_result_b_op : i64
 
 //===----------------------------------------------------------------------===//
+// Format optional results
+//===----------------------------------------------------------------------===//
+
+// CHECK: test.format_optional_result_a_op
+test.format_optional_result_a_op
+
+// CHECK: test.format_optional_result_a_op : i64 -> i64, i64
+test.format_optional_result_a_op : i64 -> i64, i64
+
+// CHECK: test.format_optional_result_b_op
+test.format_optional_result_b_op
+
+// 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 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

body << "      p << caseValueStr;\n"
             "  }\n";
     return;
   }
 
   body << "    p << caseValueStr;\n"
           "  }\n";
 }
 
+/// Generate the check for the anchor of an optional group.
+static void genOptionalGroupPrinterAnchor(Element *anchor, OpMethodBody &body) {
+  TypeSwitch<Element *>(anchor)
+      .Case<OperandVariable, ResultVariable>([&](auto *element) {
+        const NamedTypeConstraint *var = element->getVar();
+        if (var->isOptional())
+          body << "  if (" << var->name << "()) {\n";
+        else if (var->isVariadic())
+          body << "  if (!" << var->name << "().empty()) {\n";
+      })
+      .Case<RegionVariable>([&](RegionVariable *element) {
+        const NamedRegion *var = element->getVar();
+        // TODO: Add a check for optional regions here when ODS supports it.
+        body << "  if (!" << var->name << "().empty()) {\n";
+      })
+      .Case<TypeDirective>([&](TypeDirective *element) {
+        genOptionalGroupPrinterAnchor(element->getOperand(), body);
+      })
+      .Case<FunctionalTypeDirective>([&](FunctionalTypeDirective *element) {
+        genOptionalGroupPrinterAnchor(element->getInputs(), body);
+      })
+      .Case<AttributeVariable>([&](AttributeVariable *attr) {
+        body << "  if ((*this)->getAttr(\"" << attr->getVar()->name
+             << "\")) {\n";
+      });
+}
+
 void OperationFormat::genElementPrinter(Element *element, OpMethodBody &body,
                                         Operator &op, bool &shouldEmitSpace,
                                         bool &lastWasPunctuation) {
   if (LiteralElement *literal = dyn_cast<LiteralElement>(element))
     return genLiteralPrinter(literal->getLiteral(), body, shouldEmitSpace,
                              lastWasPunctuation);
 
   // Emit a whitespace element.
   if (isa<NewlineElement>(element)) {
     body << "  p.printNewline();\n";
     return;
   }
   if (SpaceElement *space = dyn_cast<SpaceElement>(element))
     return genSpacePrinter(space->getValue(), body, shouldEmitSpace,
                            lastWasPunctuation);
 
   // 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();
-    if (auto *operand = dyn_cast<OperandVariable>(anchor)) {
-      const NamedTypeConstraint *var = operand->getVar();
-      if (var->isOptional())
-        body << "  if (" << var->name << "()) {\n";
-      else if (var->isVariadic())
-        body << "  if (!" << var->name << "().empty()) {\n";
-    } else if (auto *region = dyn_cast<RegionVariable>(anchor)) {
-      const NamedRegion *var = region->getVar();
-      // TODO: Add a check for optional here when ODS supports it.
-      body << "  if (!" << var->name << "().empty()) {\n";
-
-    } else {
-      body << "  if ((*this)->getAttr(\""
-           << cast<AttributeVariable>(anchor)->getVar()->name << "\")) {\n";
-    }
+    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();
     Element *elidedAnchorElement = nullptr;
     auto *anchorAttr = dyn_cast<AttributeVariable>(anchor);
     if (anchorAttr && anchorAttr != &*elements.begin() &&
         anchorAttr->isUnitAttr()) {
LogicalResult parseVariable(std::unique_ptr<Element> &element,
                               bool isTopLevel);
   LogicalResult parseDirective(std::unique_ptr<Element> &element,
                                bool isTopLevel);
   LogicalResult parseLiteral(std::unique_ptr<Element> &element);
   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);
+  LogicalResult verifyOptionalChildElement(Element *element,
+                                           llvm::SMLoc childLoc, bool isAnchor);
 
   /// Parse the various different directives.
   LogicalResult parseAttrDictDirective(std::unique_ptr<Element> &element,
                                        llvm::SMLoc loc, bool isTopLevel,
                                        bool withKeyword);
   LogicalResult parseCustomDirective(std::unique_ptr<Element> &element,
                                      llvm::SMLoc loc, bool isTopLevel);
   LogicalResult parseCustomDirectiveParameter(
private:
   // during parsing.
   bool hasAttrDict = false;
   bool hasAllRegions = false, hasAllSuccessors = false;
   llvm::SmallBitVector seenOperandTypes, seenResultTypes;
   llvm::SmallSetVector<const NamedAttribute *, 8> seenAttrs;
   llvm::DenseSet<const NamedTypeConstraint *> seenOperands;
   llvm::DenseSet<const NamedRegion *> seenRegions;
   llvm::DenseSet<const NamedSuccessor *> seenSuccessors;
-  llvm::DenseSet<const NamedTypeConstraint *> optionalVariables;
 };
 } // end anonymous namespace
 
 LogicalResult FormatParser::parse() {
   llvm::SMLoc loc = curToken.getLoc();
 
   // Parse each of the format elements into the main format.
   while (curToken.getKind() != Token::eof) {
LogicalResult FormatParser::parseOptional(std::unique_ptr<Element> &element,
   llvm::SMLoc curLoc = curToken.getLoc();
   if (!isTopLevel)
     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;
-  SmallPtrSet<const NamedTypeConstraint *, 8> seenVariables;
   Optional<unsigned> anchorIdx;
   do {
-    if (failed(parseOptionalChildElement(elements, seenVariables, anchorIdx)))
+    if (failed(parseOptionalChildElement(elements, anchorIdx)))
       return ::mlir::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) {
     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");
 
-  // 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 "
-                               "within the optional group");
-    return ::mlir::success();
-  };
-  for (auto &ele : elements) {
-    if (auto *typeEle = dyn_cast<TypeRefDirective>(ele.get())) {
-      if (failed(checkTypeOperand(typeEle->getOperand())))
-        return failure();
-    } else if (auto *typeEle = dyn_cast<TypeDirective>(ele.get())) {
-      if (failed(checkTypeOperand(typeEle->getOperand())))
-        return ::mlir::failure();
-    } else if (auto *typeEle = dyn_cast<FunctionalTypeDirective>(ele.get())) {
-      if (failed(checkTypeOperand(typeEle->getInputs())) ||
-          failed(checkTypeOperand(typeEle->getResults())))
-        return ::mlir::failure();
-    }
-  }
-
-  optionalVariables.insert(seenVariables.begin(), seenVariables.end());
   auto parseStart = parseBegin - elements.begin();
   element = std::make_unique<OptionalElement>(std::move(elements), *anchorIdx,
                                               parseStart);
   return ::mlir::success();
 }
 
 LogicalResult FormatParser::parseOptionalChildElement(
     std::vector<std::unique_ptr<Element>> &childElements,
-    SmallPtrSetImpl<const NamedTypeConstraint *> &seenVariables,
     Optional<unsigned> &anchorIdx) {
   llvm::SMLoc childLoc = curToken.getLoc();
   childElements.push_back({});
   if (failed(parseElement(childElements.back(), /*isTopLevel=*/true)))
     return ::mlir::failure();
 
   // Check to see if this element is the anchor of the optional group.
   bool isAnchor = curToken.getKind() == Token::caret;
   if (isAnchor) {
     if (anchorIdx)
       return emitError(childLoc, "only one element can be marked as the anchor "
                                  "of an optional group");
     anchorIdx = childElements.size() - 1;
     consumeToken();
   }
 
-  return TypeSwitch<Element *, LogicalResult>(childElements.back().get())
+  return verifyOptionalChildElement(childElements.back().get(), childLoc,
+                                    isAnchor);
+}
+
+LogicalResult FormatParser::verifyOptionalChildElement(Element *element,
+                                                       llvm::SMLoc childLoc,
+                                                       bool isAnchor) {
+  return TypeSwitch<Element *, LogicalResult>(element)
       // All attributes can be within the optional group, but only optional
       // attributes can be the anchor.
       .Case([&](AttributeVariable *attrEle) {
         if (isAnchor && !attrEle->getVar()->attr.isOptional())
           return emitError(childLoc, "only optional attributes can be used to "
                                      "anchor an optional group");
         return ::mlir::success();
       })
       // Only optional-like(i.e. variadic) operands can be within an optional
       // group.
       .Case<OperandVariable>([&](OperandVariable *ele) {
         if (!ele->getVar()->isVariableLength())
           return emitError(childLoc, "only variable length operands can be "
                                      "used within an optional group");
-        seenVariables.insert(ele->getVar());
+        return ::mlir::success();
+      })
+      // Only optional-like(i.e. variadic) results can be within an optional
+      // group.
+      .Case<ResultVariable>([&](ResultVariable *ele) {
+        if (!ele->getVar()->isVariableLength())
+          return emitError(childLoc, "only variable length results can be "
+                                     "used within an optional group");
         return ::mlir::success();
       })
       .Case<RegionVariable>([&](RegionVariable *) {
         // TODO: When ODS has proper support for marking "optional" regions, add
         // a check here.
         return ::mlir::success();
       })
-      // Literals, whitespace, custom directives, and type directives may be
-      // used, but they can't anchor the group.
+      .Case<TypeDirective>([&](TypeDirective *ele) {
+        return verifyOptionalChildElement(ele->getOperand(), childLoc,
+                                          /*isAnchor=*/false);
+      })
+      .Case<FunctionalTypeDirective>([&](FunctionalTypeDirective *ele) {
+        if (failed(verifyOptionalChildElement(ele->getInputs(), childLoc,
+                                              /*isAnchor=*/false)))
+          return failure();
+        return verifyOptionalChildElement(ele->getResults(), childLoc,
+                                          /*isAnchor=*/false);
+      })
+      // Literals, whitespace, and custom directives may be used, but they can't
+      // anchor the group.
       .Case<LiteralElement, WhitespaceElement, CustomDirective,
-            FunctionalTypeDirective, OptionalElement, TypeRefDirective,
-            TypeDirective>([&](Element *) {
+            FunctionalTypeDirective, OptionalElement, TypeRefDirective>(
+          [&](Element *) {
             if (isAnchor)
-          return emitError(childLoc, "only variables can be used to anchor "
-                                     "an optional group");
+              return emitError(childLoc, "only variables and types can be used "
+                                         "to anchor an optional group");
             return ::mlir::success();
           })
       .Default([&](Element *) {
         return emitError(childLoc, "only literals, types, and variables can be "
                                    "used within an optional group");
       });
 }
 
 LogicalResult
 FormatParser::parseAttrDictDirective(std::unique_ptr<Element> &element,