[mlir][ods] Optional Attribute or Type Parameters

[Mogball]

Summary

Implements optional attribute or type parameters, including support for such parameters in the assembly format struct directive. Also implements optional groups.

Depends on D117971

[mehdi_amini]

mlir/docs/Tutorials/DefiningAttributesAndTypes.md

1 ↗

(On Diff #403104)

Seeing this file path again, that reminds me: wasn't it supposed to be converted to a doc like the other instead of being under "tutorial"?

[rriddle]

mlir/docs/Tutorials/DefiningAttributesAndTypes.md

1 ↗

(On Diff #403104)

Oooh, never pushed that for review. Let me send that out tomorrow.

[rriddle]

mlir/docs/Tutorials/DefiningAttributesAndTypes.md

504 ↗

(On Diff #403104)

?

mlir/include/mlir/IR/OpBase.td

3119 ↗

(On Diff #403104)

mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp

263–267

333–336

Avoid floating statements.

378

This seems quadratic, why not cache this computation off somewhere.

407

nit: Why not use ods as prefix instead of _?

527

Is the *_result_{0} here presenting a parse failure?

547–548

742–743

[Mogball]

mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp

527

This condition guards on "no value". E.g. FailureOr<Attribute> = failure() is parse failure but FailureOr<Attribute> = nullptr is no error but no value either.

[Mogball]

review comments

[Mogball]

review comments

[Mogball]

@rriddle ping .-.

[rriddle]

The windows build is failing BTW, might want to look into that.

[Mogball]

fix windows build

[Mogball]

another one

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

 
 def TestTypeOptionalGroupStruct : Test_Type<"TestTypeOptionalGroupStruct"> {
   let parameters = (ins OptionalParameter<"mlir::Optional<int>">:$a,
                         OptionalParameter<"mlir::Optional<int>">:$b);
   let mnemonic = "optional_group_struct";
   let assemblyFormat = "`<` (`(` struct(params)^ `)`) : (`x`)? `>`";
 }
 
+def TestTypeSpaces : Test_Type<"TestTypeSpaceS"> {
+  let parameters = (ins "int":$a, "int":$b);
+  let mnemonic = "spaces";
+  let assemblyFormat = "`<` ` ` $a `\\n` `(` `)` `` `(` `)` $b `>`";
+}
+
 #endif // TEST_TYPEDEFS

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

 // CHECK: !test.optional_group<(5) 6>
 // CHECK: !test.optional_group<x 6>
 // CHECK: !test.optional_group_params<x>
 // CHECK: !test.optional_group_params<(5)>
 // CHECK: !test.optional_group_params<(5, 6)>
 // CHECK: !test.optional_group_struct<x>
 // CHECK: !test.optional_group_struct<(b = 5)>
 // CHECK: !test.optional_group_struct<(a = 10, b = 5)>
+// CHECK: !test.spaces< 5
+// CHECK-NEXT: ()() 6>
 func private @test_roundtrip_default_parsers_struct(
   !test.no_parser<255, [1, 2, 3, 4, 5], "foobar", 4>
 ) -> (
   !test.struct_capture_all<v3 = 3, v1 = 1, v2 = 2, v0 = 0>,
   !test.optional_param<, 6>,
   !test.optional_param<5, 6>,
   !test.optional_params<"a">,
   !test.optional_params<5, "a">,
) -> (
   !test.all_optional_struct<b = 10, a = 5>,
   !test.optional_group<(5) 6>,
   !test.optional_group<x 6>,
   !test.optional_group_params<x>,
   !test.optional_group_params<(5)>,
   !test.optional_group_params<(5, 6)>,
   !test.optional_group_struct<x>,
   !test.optional_group_struct<(b = 5)>,
-  !test.optional_group_struct<(b = 5, a = 10)>
+  !test.optional_group_struct<(b = 5, a = 10)>,
+  !test.spaces<5 ()() 6>
 )

mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp

void genCommaSeparatedPrinter(ArrayRef<ParameterElement *> params,
                                 function_ref<void(ParameterElement *)> extra);
   /// Generate the printer code for a `params` directive.
   void genParamsPrinter(ParamsDirective *el, FmtContext &ctx, MethodBody &os);
   /// Generate the printer code for a `struct` directive.
   void genStructPrinter(StructDirective *el, FmtContext &ctx, MethodBody &os);
   /// Generate the printer code for an optional group.
   void genOptionalGroupPrinter(OptionalElement *el, FmtContext &ctx,
                                MethodBody &os);
+  /// Generate a printer (or space eraser) for a whitespace element.
+  void genWhitespacePrinter(WhitespaceElement *el, FmtContext &ctx,
+                            MethodBody &os);
 
   /// The ODS definition of the attribute or type whose format is being used to
   /// generate a parser and printer.
   const AttrOrTypeDef &def;
   /// The list of top-level format elements returned by the assembly format
   /// parser.
   std::vector<FormatElement *> elements;
 
os << tgfmt("::llvm::SMLoc $_loc = $_parser.getCurrentLocation();\n"
               &ctx);
 
   // Generate call to each parameter parser.
   for (FormatElement *el : elements)
     genElementParser(el, ctx, os);
 
   // Emit an assert for each mandatory parameter. Triggering an assert means
   // the generated parser is incorrect (i.e. there is a bug in this code).
   for (const AttrOrTypeParameter &param : params) {
     if (param.isOptional())
       continue;
     os << formatv("assert(::mlir::succeeded(_result_{0}));\n", param.getName());
   }
 
   // Generate call to the attribute or type builder. Use the checked getter
   // if one was generated.
   if (def.genVerifyDecl()) {
     os << tgfmt("return $_parser.getChecked<$0>($_loc, $_parser.getContext()",
                 &ctx, def.getCppClassName());
   } else {
     os << tgfmt("return $0::get($_parser.getContext()", &ctx,
void DefFormat::genElementParser(FormatElement *el, FmtContext &ctx,
   if (auto *var = dyn_cast<ParameterElement>(el))
     return genVariableParser(var, ctx, os);
   if (auto *params = dyn_cast<ParamsDirective>(el))
     return genParamsParser(params, ctx, os);
   if (auto *strct = dyn_cast<StructDirective>(el))
     return genStructParser(strct, ctx, os);
   if (auto *optional = dyn_cast<OptionalElement>(el))
     return genOptionalGroupParser(optional, ctx, os);
+  if (isa<WhitespaceElement>(el))
+    return;
 
   llvm_unreachable("unknown format element");
 }
 
 void DefFormat::genLiteralParser(StringRef value, FmtContext &ctx,
                                  MethodBody &os, bool isOptional) {
   os << "// Parse literal '" << value << "'\n";
   os << tgfmt("if ($_parser.parse", &ctx);
os << StringSwitch<StringRef>(value)
               .Case(")", "RParen")
               .Case("[", "LSquare")
               .Case("]", "RSquare")
               .Case("?", "Question")
               .Case("+", "Plus")
               .Case("*", "Star")
        << "()";
   }
   if (isOptional)
     // Leave the `if` unclosed to guard optional groups.
     return;
   // Parser will emit an error
   os << ") return {};\n";
 }
 
 void DefFormat::genVariableParser(ParameterElement *el, FmtContext &ctx,
                                   MethodBody &os) {
   // Check for a custom parser. Use the default attribute parser otherwise.
   const AttrOrTypeParameter &param = el->getParam();
   auto customParser = param.getParser();
void DefFormat::genParamsParser(ParamsDirective *el, FmtContext &ctx,
   auto betweenFn = [&](IteratorT it, IteratorT end) {
     ParameterElement *el = *std::prev(it);
     // Parse a comma if the last optional parameter had a value.
     if (el->isOptional()) {
       os << formatv("if (::mlir::succeeded(_result_{0}) && *_result_{0}) {{\n",
                     el->getName());
       os.indent();
     }
     if (std::any_of(it, end, paramNotOptional)) {
       genLiteralParser(",", ctx, os);
     } else {
       genLiteralParser(",", ctx, os, /*isOptional=*/true);
       os << ") break;\n";
     }
     if (el->isOptional())
       os.unindent() << "}\n";
   };

 }
 
 void DefFormat::genStructParser(StructDirective *el, FmtContext &ctx,
                                 MethodBody &os) {
   // Loop declaration for struct parser with only required parameters.
   //
   // $0: Number of expected parameters.
   const char *const loopHeader = R"(
   for (unsigned _index = 0; _index < $0; ++_index) {
 )";
 
   // Loop body start for struct parser.
   const char *const loopStart = R"(
     ::llvm::StringRef _paramKey;
     if ($_parser.parseKeyword(&_paramKey)) {
       $_parser.emitError($_parser.getCurrentLocation(),
                          "expected a parameter name in struct");
)";
   if (numOptional) {
     if (numOptional == el->getNumParams()) {
       os << "}\n";
       os.unindent() << "}\n";
     } else {
       os << tgfmt("} while(!$_parser.parseOptionalComma());\n", &ctx);
       for (ParameterElement *param : el->getParams()) {
         if (param->isOptional())
           continue;
         os.getStream().printReindented(
             strfmt(checkParam, param->getName(), tgfmt(parserErrorStr, &ctx)));
       }
     }
   } else {
     // Because the loop loops N times and each non-failing iteration sets 1 of
     // N flags, successfully exiting the loop means that all parameters have
     // 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::genOptionalGroupParser(OptionalElement *el, FmtContext &ctx,
                                        MethodBody &os) {
   ArrayRef<FormatElement *> elements =
       el->getThenElements().drop_front(el->getParseStart());
 
   FormatElement *first = elements.front();
   const auto guardOn = [&](auto params) {
     os << "if (!(";
     llvm::interleave(
         params, os,
         [&](ParameterElement *el) {
           os << formatv("(::mlir::succeeded(_result_{0}) && *_result_{0})",
                         el->getName());
         },
         " || ");
void DefFormat::genElementPrinter(FormatElement *el, FmtContext &ctx,
   if (auto *params = dyn_cast<ParamsDirective>(el))
     return genParamsPrinter(params, ctx, os);
   if (auto *strct = dyn_cast<StructDirective>(el))
     return genStructPrinter(strct, ctx, os);
   if (auto *var = dyn_cast<ParameterElement>(el))
     return genVariablePrinter(var, ctx, os);
   if (auto *optional = dyn_cast<OptionalElement>(el))
     return genOptionalGroupPrinter(optional, ctx, os);
+  if (auto *whitespace = dyn_cast<WhitespaceElement>(el))
+    return genWhitespacePrinter(whitespace, ctx, os);
 
   llvm::PrintFatalError("unsupported format element");
 }
 
 void DefFormat::genLiteralPrinter(StringRef value, FmtContext &ctx,
                                   MethodBody &os) {
   // Don't insert a space before certain punctuation.
   bool needSpace =
void DefFormat::genOptionalGroupPrinter(OptionalElement *el, FmtContext &ctx,
   } else if (auto *params = dyn_cast<ParamsDirective>(anchor)) {
     guardOn(params->getParams());
   } else {
     auto *strct = dyn_cast<StructDirective>(anchor);
     guardOn(strct->getParams());
   }
   // Generate the printer for the contained elements.
   {
     llvm::SaveAndRestore<bool> shouldEmitSpaceFlag(shouldEmitSpace),
         lastWasPunctuationFlag(lastWasPunctuation);
     for (FormatElement *element : el->getThenElements())
       genElementPrinter(element, ctx, os);
   }
   os.unindent() << "} else {\n";
   os.indent();
   for (FormatElement *element : el->getElseElements())
     genElementPrinter(element, ctx, os);
   os.unindent() << "}\n";
 }
 
+void DefFormat::genWhitespacePrinter(WhitespaceElement *el, FmtContext &ctx,
+                                     MethodBody &os) {
+  if (el->getValue() == "\\n") {
+    os << tgfmt("$_printer << '\\n';\n", &ctx);
+  } else if (!el->getValue().empty()) {
+    os << tgfmt("$_printer << \"$0\";\n", &ctx, el->getValue());
+  } else {
+    lastWasPunctuation = true;
+  }
+  shouldEmitSpace = false;
+}
+
 //===----------------------------------------------------------------------===//
 // DefFormatParser
 //===----------------------------------------------------------------------===//
 
 namespace {
 class DefFormatParser : public FormatParser {
 public:
   DefFormatParser(llvm::SourceMgr &mgr, const AttrOrTypeDef &def)