[Mlir-commits] [mlir] 0d18c2d - [mlir][ods] Allow custom directives to anchor optional groups in op assembly formats

[llvmlistbot at llvm.org]

Author: Mogball
Date: 2023-08-31T03:20:53Z
New Revision: 0d18c2d90f8f69e6517a8851b94c43be7b87575f

URL: https://github.com/llvm/llvm-project/commit/0d18c2d90f8f69e6517a8851b94c43be7b87575f
DIFF: https://github.com/llvm/llvm-project/commit/0d18c2d90f8f69e6517a8851b94c43be7b87575f.diff

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

`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.

Reviewed By: rriddle

Differential Revision: https://reviews.llvm.org/D159243

Added: 
    

Modified: 
    mlir/test/mlir-tblgen/op-format-invalid.td
    mlir/test/mlir-tblgen/op-format-spec.td
    mlir/test/mlir-tblgen/op-format.td
    mlir/tools/mlir-tblgen/OpFormatGen.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/test/mlir-tblgen/op-format-invalid.td b/mlir/test/mlir-tblgen/op-format-invalid.td
index b038ddd5717b97..210241caef3163 100644
--- a/mlir/test/mlir-tblgen/op-format-invalid.td
+++ b/mlir/test/mlir-tblgen/op-format-invalid.td
@@ -389,7 +389,7 @@ def OptionalInvalidJ : TestFormat_Op<[{
 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)>;

diff  --git a/mlir/test/mlir-tblgen/op-format-spec.td b/mlir/test/mlir-tblgen/op-format-spec.td
index d546cb97658f24..5d786cce5f2801 100644
--- a/mlir/test/mlir-tblgen/op-format-spec.td
+++ b/mlir/test/mlir-tblgen/op-format-spec.td
@@ -41,6 +41,9 @@ def DirectiveCustomValidB : TestFormat_Op<[{
 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

diff  --git a/mlir/test/mlir-tblgen/op-format.td b/mlir/test/mlir-tblgen/op-format.td
index e5452e02240416..c098a52f15b3b0 100644
--- a/mlir/test/mlir-tblgen/op-format.td
+++ b/mlir/test/mlir-tblgen/op-format.td
@@ -83,3 +83,10 @@ def OptionalGroupB : TestFormat_Op<[{
 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)>;

diff  --git a/mlir/tools/mlir-tblgen/OpFormatGen.cpp b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
index 3f65aba21fa628..3d045044103c14 100644
--- a/mlir/tools/mlir-tblgen/OpFormatGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
@@ -2008,19 +2008,19 @@ static void genOptionalGroupPrinterAnchor(FormatElement *anchor,
         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())
@@ -2037,6 +2037,18 @@ static void genOptionalGroupPrinterAnchor(FormatElement *anchor,
           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 << ')';
       });
 }
 
@@ -3434,15 +3446,28 @@ LogicalResult OpFormatParser::verifyOptionalGroupElement(SMLoc loc,
         return verifyOptionalGroupElement(loc, ele->getResults(),
                                           /*isAnchor=*/false);
       })
-      // Literals, whitespace, and custom directives may be used, but they can't
-      // anchor the group.
-      .Case<LiteralElement, WhitespaceElement, CustomDirective,
-            FunctionalTypeDirective, OptionalElement>([&](FormatElement *) {
-        if (isAnchor)
-          return emitError(loc, "only variables and types can be used "
-                                "to anchor an optional group");
+      .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, 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");