[mlir][ods] Allow custom directives to anchor optional groups in op assembly formats

[Mogball]

Summary

custom directives were allowed to anchor optional groups in attribute
and type assembly formats already, but the feature was never extended to
operation assembly formats. This patch adds the feature to op assembly
formats too.

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

 // CHECK: error: only literals, types, and variables can be used within an optional group
 def OptionalInvalidJ : TestFormat_Op<[{
   (attr-dict^)?
 }]>;
 // CHECK: error: expected '?' after optional group
 def OptionalInvalidK : TestFormat_Op<[{
   ($arg^)
 }]>, Arguments<(ins Variadic<I64>:$arg)>;
-// CHECK: error: only variables and types can be used to anchor an optional group
+// CHECK: error: only variable length operands can be used within 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

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

def DirectiveCustomValidA : TestFormat_Op<[{
   custom<MyDirective>($operand) attr-dict
 }]>, Arguments<(ins Optional<I64>:$operand)>;
 def DirectiveCustomValidB : TestFormat_Op<[{
   custom<MyDirective>($operand, type($operand), type($result)) attr-dict
 }]>, Arguments<(ins I64:$operand)>, Results<(outs I64:$result)>;
 def DirectiveCustomValidC : TestFormat_Op<[{
   custom<MyDirective>($attr) attr-dict
 }]>, Arguments<(ins I64Attr:$attr)>;
+def DirectiveCustomValidD : TestFormat_Op<[{
+  (`(` custom<MyDirective>($operand)^ `)`)? attr-dict
+}]>, Arguments<(ins Optional<I64>:$operand)>;
 
 //===----------------------------------------------------------------------===//
 // functional-type
 
 // CHECK-NOT: error
 def DirectiveFunctionalTypeValid : TestFormat_Op<[{
   functional-type(operands, results) attr-dict
 }]>;

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

 
 //       CHECK: if (getAAttr() && getAAttr() != ::mlir::OpBuilder((*this)->getContext()).getStringAttr("default")) {
 //  CHECK-NEXT:   odsPrinter << ' ';
 //  CHECK-NEXT:   odsPrinter.printAttributeWithoutType(getAAttr());
 //  CHECK-NEXT: }
 def OptionalGroupC : TestFormat_Op<[{
   ($a^)? attr-dict
 }]>, Arguments<(ins DefaultValuedStrAttr<StrAttr, "default">:$a)>;
+
+// CHECK-LABEL: OptionalGroupD::print
+// CHECK-NEXT: if (((getA()) || (getB()))) {
+// CHECK-NEXT:   odsPrinter << "("
+def OptionalGroupD : TestFormat_Op<[{
+  (`(` custom<Custom>($a, $b)^ `)`)? attr-dict
+}], [AttrSizedOperandSegments]>, Arguments<(ins Optional<I64>:$a, Optional<I64>:$b)>;

mlir/tools/mlir-tblgen/OpFormatGen.cpp

TypeSwitch<FormatElement *>(anchor)
       .Case<OperandVariable, ResultVariable>([&](auto *element) {
         const NamedTypeConstraint *var = element->getVar();
         std::string name = op.getGetterName(var->name);
         if (var->isOptional())
           body << name << "()";
         else if (var->isVariadic())
           body << "!" << name << "().empty()";
       })
-      .Case<RegionVariable>([&](RegionVariable *element) {
+      .Case([&](RegionVariable *element) {
         const NamedRegion *var = element->getVar();
         std::string name = op.getGetterName(var->name);
         // TODO: Add a check for optional regions here when ODS supports it.
         body << "!" << name << "().empty()";
       })
-      .Case<TypeDirective>([&](TypeDirective *element) {
+      .Case([&](TypeDirective *element) {
         genOptionalGroupPrinterAnchor(element->getArg(), op, body);
       })
-      .Case<FunctionalTypeDirective>([&](FunctionalTypeDirective *element) {
+      .Case([&](FunctionalTypeDirective *element) {
         genOptionalGroupPrinterAnchor(element->getInputs(), op, body);
       })
-      .Case<AttributeVariable>([&](AttributeVariable *element) {
+      .Case([&](AttributeVariable *element) {
         Attribute attr = element->getVar()->attr;
         body << op.getGetterName(element->getVar()->name) << "Attr()";
         if (attr.isOptional())
           return; // done
         if (attr.hasDefaultValue()) {
           // Consider a default-valued attribute as present if it's not the
           // default value.
           FmtContext fctx;
           fctx.withBuilder("::mlir::OpBuilder((*this)->getContext())");
           body << " && " << op.getGetterName(element->getVar()->name)
                << "Attr() != "
                << tgfmt(attr.getConstBuilderTemplate(), &fctx,
                         attr.getDefaultValue());
           return;
         }
         llvm_unreachable("attribute must be optional or default-valued");
+      })
+      .Case([&](CustomDirective *ele) {
+        body << '(';
+        llvm::interleave(
+            ele->getArguments(), body,
+            [&](FormatElement *child) {
+              body << '(';
+              genOptionalGroupPrinterAnchor(child, op, body);
+              body << ')';
+            },
+            " || ");
+        body << ')';
       });
 }
 
 void collect(FormatElement *element,
              SmallVectorImpl<VariableElement *> &variables) {
   TypeSwitch<FormatElement *>(element)
       .Case([&](VariableElement *var) { variables.emplace_back(var); })
       .Case([&](CustomDirective *ele) {
return TypeSwitch<FormatElement *, LogicalResult>(element)
       })
       .Case([&](FunctionalTypeDirective *ele) {
         if (failed(verifyOptionalGroupElement(loc, ele->getInputs(),
                                               /*isAnchor=*/false)))
           return failure();
         return verifyOptionalGroupElement(loc, ele->getResults(),
                                           /*isAnchor=*/false);
       })
+      .Case([&](CustomDirective *ele) {
+        if (!isAnchor)
+          return success();
+        // Verify each child as being valid in an optional group. They are all
+        // potential anchors if the custom directive was marked as one.
+        for (FormatElement *child : ele->getArguments()) {
+          if (isa<RefDirective>(child))
+            continue;
+          if (failed(verifyOptionalGroupElement(loc, child, /*isAnchor=*/true)))
+            return failure();
+        }
+        return success();
+      })
       // Literals, whitespace, and custom directives may be used, but they can't
       // anchor the group.
-      .Case<LiteralElement, WhitespaceElement, CustomDirective,
-            FunctionalTypeDirective, OptionalElement>([&](FormatElement *) {
+      .Case<LiteralElement, WhitespaceElement, OptionalElement>(
+          [&](FormatElement *) {
             if (isAnchor)
               return emitError(loc, "only variables and types can be used "
                                     "to anchor an optional group");
             return success();
           })
       .Default([&](FormatElement *) {
         return emitError(loc, "only literals, types, and variables can be "
                               "used within an optional group");
       });
 }
 
 //===----------------------------------------------------------------------===//
 // Interface