[mlir][ods] Support using custom directives as first optional group element

[Mogball]

Summary

This adds support for using a custom directive as the first optional
group element. The first optional group element guards the parsing of
the rest of the optional group. This can be done for custom directives
by expecting the parse function to return an OptionalParseResult
instead of a ParseResult.

Depends on D159243

mlir/test/mlir-tblgen/attr-or-type-format.td

 
 def TypeN : TestType<"TestP"> {
   let parameters = (ins OptionalParameter<"int">:$a,
                         OptionalParameter<"int">:$b);
   let mnemonic = "type_n";
   let assemblyFormat = "`<` (`?`) : (struct($a, $b)^)? `>`";
 }
 
+// TYPE-LABEL: TestQType::parse
+// TYPE: if (auto result = [&]() -> ::mlir::OptionalParseResult {
+// TYPE:     auto odsCustomResult = parseAB(odsParser
+// TYPE:     if (!odsCustomResult) return {};
+// TYPE:     if (::mlir::failed(*odsCustomResult)) return ::mlir::failure();
+// TYPE:   return ::mlir::success();
+// TYPE: }(); result.has_value() && ::mlir::failed(*result)) {
+// TYPE:   return {};
+// TYPE: } else if (result.has_value()) {
+// TYPE:   // Parse literal 'y'
+// TYPE: } else {
+// TYPE:   // Parse literal 'x'
+def TypeO : TestType<"TestQ"> {
+  let parameters = (ins OptionalParameter<"int">:$a);
+  let mnemonic = "type_o";
+  let assemblyFormat = "(custom<AB>($a)^ `x`) : (`y`)?";
+}
+
 // DEFAULT_TYPE_PARSER: TestDialect::parseType(::mlir::DialectAsmParser &parser)
 // DEFAULT_TYPE_PARSER: auto parseResult = parseOptionalDynamicType(mnemonic, parser, genType);
 // DEFAULT_TYPE_PARSER: if (parseResult.has_value()) {
 // DEFAULT_TYPE_PARSER:   if (::mlir::succeeded(parseResult.value()))
-// DEFAULT_TYPE_PARSER:     return genType;
+// DEFAULT_TYPE_PARSER:     return genType;
 

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

 // CHECK: error: expected literal, variable, directive, or optional group
 def OptionalInvalidB : TestFormat_Op<[{
   () attr-dict
 }]>, Arguments<(ins OptionalAttr<I64Attr>:$attr)>;
 // CHECK: error: optional group has no anchor element
 def OptionalInvalidC : TestFormat_Op<[{
   ($attr)? attr-dict
 }]>, Arguments<(ins OptionalAttr<I64Attr>:$attr)>;
-// CHECK: error: first parsable element of an optional group must be a literal or variable
+// CHECK: error: first parsable element of an optional group must be a literal, variable, or custom directive
 def OptionalInvalidD : TestFormat_Op<[{
   (type($operand) $operand^)? attr-dict
 }]>, Arguments<(ins Optional<I64>:$operand)>;
 // CHECK: error: only literals, types, and variables can be used within an optional group
 def OptionalInvalidE : TestFormat_Op<[{
   (`,` $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

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::parse
+// CHECK: if (auto result = [&]() -> ::mlir::OptionalParseResult {
+// CHECK:   auto odsResult = parseCustom(parser, aOperand, bOperand);
+// CHECK:   if (!odsResult) return {};
+// CHECK:   if (::mlir::failed(*odsResult)) return ::mlir::failure();
+// CHECK:   return ::mlir::success();
+// CHECK: }(); result.has_value() && ::mlir::failed(*result)) {
+// CHECK:   return ::mlir::failure();
+// CHECK: } else if (result.has_value()) {
+
 // CHECK-LABEL: OptionalGroupD::print
 // CHECK-NEXT: if (((getA()) || (getB()))) {
-// CHECK-NEXT:   odsPrinter << "("
+// CHECK-NEXT:   odsPrinter << ' ';
+// CHECK-NEXT:   printCustom
 def OptionalGroupD : TestFormat_Op<[{
-  (`(` custom<Custom>($a, $b)^ `)`)? attr-dict
+  (custom<Custom>($a, $b)^)? attr-dict
 }], [AttrSizedOperandSegments]>, Arguments<(ins Optional<I64>:$a, Optional<I64>:$b)>;

mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp

void genLiteralParser(StringRef value, FmtContext &ctx, MethodBody &os,
                         bool isOptional = false);
   /// Generate the parser code for a variable.
   void genVariableParser(ParameterElement *el, FmtContext &ctx, MethodBody &os);
   /// Generate the parser code for a `params` directive.
   void genParamsParser(ParamsDirective *el, FmtContext &ctx, MethodBody &os);
   /// Generate the parser code for a `struct` directive.
   void genStructParser(StructDirective *el, FmtContext &ctx, MethodBody &os);
   /// Generate the parser code for a `custom` directive.
-  void genCustomParser(CustomDirective *el, FmtContext &ctx, MethodBody &os);
+  void genCustomParser(CustomDirective *el, FmtContext &ctx, MethodBody &os,
+                       bool isOptional = false);
   /// Generate the parser code for an optional group.
   void genOptionalGroupParser(OptionalElement *el, FmtContext &ctx,
                               MethodBody &os);
 
   /// Generate the printer code for a specific format element.
   void genElementPrinter(FormatElement *el, FmtContext &ctx, MethodBody &os);
   /// Generate the printer code for a literal.
   void genLiteralPrinter(StringRef value, FmtContext &ctx, MethodBody &os);
if (numOptional) {
     // been seen. `parseOptionalComma` would cause issues with any formats that
     // use "struct(...) `,`" beacuse structs aren't sounded by braces.
     os.getStream().printReindented(strfmt(loopTerminator, el->getNumParams()));
   }
   os.unindent() << "}\n";
 }
 
 void DefFormat::genCustomParser(CustomDirective *el, FmtContext &ctx,
-                                MethodBody &os) {
+                                MethodBody &os, bool isOptional) {
   os << "{\n";
   os.indent();
 
   // Bound variables are passed directly to the parser as `FailureOr<T> &`.
   // Referenced variables are passed as `T`. The custom parser fails if it
   // returns failure or if any of the required parameters failed.
   os << tgfmt("auto odsCustomLoc = $_parser.getCurrentLocation();\n", &ctx);
   os << "(void)odsCustomLoc;\n";
   os << tgfmt("auto odsCustomResult = parse$0($_parser", &ctx, el->getName());
   os.indent();
   for (FormatElement *arg : el->getArguments()) {
     os << ",\n";
     if (auto *param = dyn_cast<ParameterElement>(arg))
       os << "::mlir::detail::unwrapForCustomParse(_result_" << param->getName()
          << ")";
     else if (auto *ref = dyn_cast<RefDirective>(arg))
       os << "*_result_" << cast<ParameterElement>(ref->getArg())->getName();
     else
       os << tgfmt(cast<StringElement>(arg)->getValue(), &ctx);
   }
   os.unindent() << ");\n";
+  if (isOptional) {
+    os << "if (!odsCustomResult) return {};\n";
+    os << "if (::mlir::failed(*odsCustomResult)) return ::mlir::failure();\n";
+  } else {
     os << "if (::mlir::failed(odsCustomResult)) return {};\n";
+  }
   for (FormatElement *arg : el->getArguments()) {
     if (auto *param = dyn_cast<ParameterElement>(arg)) {
       if (param->isOptional())
         continue;
       os << formatv("if (::mlir::failed(_result_{0})) {{\n", param->getName());
       os.indent() << tgfmt("$_parser.emitError(odsCustomLoc, ", &ctx)
                   << "\"custom parser failed to parse parameter '"
                   << param->getName() << "'\");\n";
-      os << "return {};\n";
+      os << "return " << (isOptional ? "::mlir::failure()" : "{}") << ";\n";
       os.unindent() << "}\n";
     }
   }
 
   os.unindent() << "}\n";
 }
 
 void DefFormat::genOptionalGroupParser(OptionalElement *el, FmtContext &ctx,
if (auto *literal = dyn_cast<LiteralElement>(first)) {
     genLiteralParser(literal->getSpelling(), ctx, os, /*isOptional=*/true);
     os << ") {\n";
   } else if (auto *param = dyn_cast<ParameterElement>(first)) {
     genVariableParser(param, ctx, os);
     guardOn(llvm::ArrayRef(param));
   } else if (auto *params = dyn_cast<ParamsDirective>(first)) {
     genParamsParser(params, ctx, os);
     guardOn(params->getParams());
+  } else if (auto *custom = dyn_cast<CustomDirective>(first)) {
+    os << "if (auto result = [&]() -> ::mlir::OptionalParseResult {\n";
+    os.indent();
+    genCustomParser(custom, ctx, os, /*isOptional=*/true);
+    os << "return ::mlir::success();\n";
+    os.unindent();
+    os << "}(); result.has_value() && ::mlir::failed(*result)) {\n";
+    os.indent();
+    os << "return {};\n";
+    os.unindent();
+    os << "} else if (result.has_value()) {\n";
   } else {
     auto *strct = cast<StructDirective>(first);
     genStructParser(strct, ctx, os);
     guardOn(params->getParams());
   }
   os.indent();
 
   // Generate the parsers for the rest of the thenElements.

mlir/tools/mlir-tblgen/FormatGen.cpp

FailureOr<FormatElement *> FormatParser::parseOptionalGroup(Context ctx) {
   auto isWhitespace = [](FormatElement *element) {
     return isa<WhitespaceElement>(element);
   };
   auto thenParseBegin = llvm::find_if_not(thenElements, isWhitespace);
   auto elseParseBegin = llvm::find_if_not(elseElements, isWhitespace);
   unsigned thenParseStart = std::distance(thenElements.begin(), thenParseBegin);
   unsigned elseParseStart = std::distance(elseElements.begin(), elseParseBegin);
 
-  if (!isa<LiteralElement, VariableElement>(*thenParseBegin)) {
+  if (!isa<LiteralElement, VariableElement, CustomDirective>(*thenParseBegin)) {
     return emitError(loc, "first parsable element of an optional group must be "
-                          "a literal or variable");
+                          "a literal, variable, or custom directive");
   }
   return create<OptionalElement>(std::move(thenElements),
                                  std::move(elseElements), thenParseStart,
                                  elseParseStart, anchor, inverted);
 }
 
 FailureOr<FormatElement *> FormatParser::parseCustomDirective(SMLoc loc,
                                                               Context ctx) {

mlir/tools/mlir-tblgen/OpFormatGen.cpp

static void genCustomParameterParser(FormatElement *param, MethodBody &body) {
   } else {
     llvm_unreachable("unknown custom directive parameter");
   }
 }
 
 /// Generate the parser for a custom directive.
 static void genCustomDirectiveParser(CustomDirective *dir, MethodBody &body,
                                      bool useProperties,
-                                     StringRef opCppClassName) {
+                                     StringRef opCppClassName,
+                                     bool isOptional = false) {
   body << "  {\n";
 
   // Preprocess the directive variables.
   // * Add a local variable for optional operands and types. This provides a
   //   better API to the user defined parser methods.
   // * Set the location of operand variables.
   for (FormatElement *param : dir->getArguments()) {
     if (auto *operand = dyn_cast<OperandVariable>(param)) {
if (auto *operand = dyn_cast<OperandVariable>(param)) {
           body << llvm::formatv("    ::mlir::Type {0}Type = {0}Types.empty() ? "
                                 "::mlir::Type() : {0}Types[0];\n",
                                 listName);
         }
       }
     }
   }
 
-  body << "    if (parse" << dir->getName() << "(parser";
+  body << "    auto odsResult = parse" << dir->getName() << "(parser";
   for (FormatElement *param : dir->getArguments()) {
     body << ", ";
     genCustomParameterParser(param, body);
   }
+  body << ");\n";
 
-  body << "))\n"
-       << "      return ::mlir::failure();\n";
+  if (isOptional) {
+    body << "    if (!odsResult) return {};\n"
+         << "    if (::mlir::failed(*odsResult)) return ::mlir::failure();\n";
+  } else {
+    body << "    if (odsResult) return ::mlir::failure();\n";
+  }
 
   // After parsing, add handling for any of the optional constructs.
   for (FormatElement *param : dir->getArguments()) {
     if (auto *attr = dyn_cast<AttributeVariable>(param)) {
       const NamedAttribute *var = attr->getVar();
       if (var->attr.isOptional() || var->attr.hasDefaultValue())
         body << llvm::formatv("    if ({0}Attr)\n  ", var->name);
       if (useProperties) {
if (auto *attrVar = dyn_cast<AttributeVariable>(firstElement)) {
         body << llvm::formatv(optionalRegionParserCode, region->name);
         body << "  if (!" << region->name << "Region->empty()) {\n  ";
         if (hasImplicitTermTrait)
           body << llvm::formatv(regionEnsureTerminatorParserCode, region->name);
         else if (hasSingleBlockTrait)
           body << llvm::formatv(regionEnsureSingleBlockParserCode,
                                 region->name);
       }
+    } else if (auto *custom = dyn_cast<CustomDirective>(firstElement)) {
+      body << "  if (auto result = [&]() -> ::mlir::OptionalParseResult {\n";
+      genCustomDirectiveParser(custom, body, useProperties, opCppClassName,
+                               /*isOptional=*/true);
+      body << "    return ::mlir::success();\n"
+           << "  }(); result.has_value() && ::mlir::failed(*result)) {\n"
+           << "    return ::mlir::failure();\n"
+           << "  } else if (result.has_value()) {\n";
     }
 
     genElementParsers(firstElement, thenElements.drop_front(),
                       /*thenGroup=*/true);
     body << "  }";
 
     // Generate the else elements.
     auto elseElements = optional->getElseElements();