[mlir] Generate parser/printers for enums

[rriddle]

Summary

This greatly simplifies composing enums in attribute/type printers,
which currently reimplement these functions as needed.

[Mogball]

mlir/test/mlir-tblgen/enums-gen.td

47

Can you add a test for a non keyword case?

mlir/tools/mlir-tblgen/EnumsGen.cpp

79

You can use structured bindings on enumerate now, I'm pretty sure

[jpienaar]

Nice!

[rriddle]

mlir/tools/mlir-tblgen/EnumsGen.cpp

79

nice!!

mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td

 class LLVM_Attr<string name> : AttrDef<LLVM_Dialect, name>;
 
 // Attribute definition for the LLVM Linkage enum.
 def LinkageAttr : LLVM_Attr<"Linkage"> {
   let mnemonic = "linkage";
   let parameters = (ins
     "linkage::Linkage":$linkage
   );
-  let hasCustomAssemblyFormat = 1;
+  let assemblyFormat = "`<` $linkage `>`";
 }
 
 // Attribute definition for the LLVM Linkage enum.
 def CConvAttr : LLVM_Attr<"CConv"> {
   let mnemonic = "cconv";
   let parameters = (ins
     "CConv":$CallingConv
   );
-  let hasCustomAssemblyFormat = 1;
+  let assemblyFormat = "`<` $CallingConv `>`";
 }
 
 def LoopOptionsAttr : LLVM_Attr<"LoopOptions"> {
   let mnemonic = "loopopts";
 
   let description = [{
     This attributes encapsulates "loop options". It is means to decorate
     branches that are "latches" (loop backedges) and maps to the `!llvm.loop`

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

return builder.create<LLVM::GEPOp>(
       ArrayRef<GEPArg>{0, 0});
 }
 
 bool mlir::LLVM::satisfiesLLVMModule(Operation *op) {
   return op->hasTrait<OpTrait::SymbolTable>() &&
          op->hasTrait<OpTrait::IsIsolatedFromAbove>();
 }
 
-void LinkageAttr::print(AsmPrinter &printer) const {
-  printer << "<";
-  if (static_cast<uint64_t>(getLinkage()) <= getMaxEnumValForLinkage())
-    printer << stringifyEnum(getLinkage());
-  else
-    printer << static_cast<uint64_t>(getLinkage());
-  printer << ">";
-}
-
-Attribute LinkageAttr::parse(AsmParser &parser, Type type) {
-  StringRef elemName;
-  if (parser.parseLess() || parser.parseKeyword(&elemName) ||
-      parser.parseGreater())
-    return {};
-  auto elem = linkage::symbolizeLinkage(elemName);
-  if (!elem) {
-    parser.emitError(parser.getNameLoc(), "Unknown linkage: ") << elemName;
-    return {};
-  }
-  Linkage linkage = *elem;
-  return LinkageAttr::get(parser.getContext(), linkage);
-}
-
-void CConvAttr::print(AsmPrinter &printer) const {
-  printer << "<";
-  if (static_cast<uint64_t>(getCallingConv()) <= cconv::getMaxEnumValForCConv())
-    printer << stringifyEnum(getCallingConv());
-  else
-    printer << "INVALID_cc_" << static_cast<uint64_t>(getCallingConv());
-  printer << ">";
-}
-
-Attribute CConvAttr::parse(AsmParser &parser, Type type) {
-  StringRef convName;
-
-  if (parser.parseLess() || parser.parseKeyword(&convName) ||
-      parser.parseGreater())
-    return {};
-  auto cconv = cconv::symbolizeCConv(convName);
-  if (!cconv) {
-    parser.emitError(parser.getNameLoc(), "unknown calling convention: ")
-        << convName;
-    return {};
-  }
-  CConv cconvVal = *cconv;
-  return CConvAttr::get(parser.getContext(), cconvVal);
-}
-
 LoopOptionsAttrBuilder::LoopOptionsAttrBuilder(LoopOptionsAttr attr)
     : options(attr.getOptions().begin(), attr.getOptions().end()) {}
 
 template <typename T>
 LoopOptionsAttrBuilder &LoopOptionsAttrBuilder::setOption(LoopOptionCase tag,
                                                           Optional<T> value) {
   auto option = llvm::find_if(
       options, [tag](auto option) { return option.first == tag; });

mlir/test/Dialect/LLVMIR/func.mlir

 module {
   // expected-error@+1 {{custom op 'llvm.func' expected valid '@'-identifier for symbol name}}
   llvm.func cc_12 @unknown_calling_convention()
 }
 
 // -----
 
 module {
-  // expected-error@+2 {{unknown calling convention: cc_12}}
   "llvm.func"() ({
+  // expected-error @below {{invalid Calling Conventions specification: cc_12}}
+  // expected-error @below {{failed to parse CConvAttr parameter 'CallingConv' which is to be a `CConv`}}
   }) {sym_name = "generic_unknown_calling_convention", CConv = #llvm.cconv<cc_12>, function_type = !llvm.func<i64 (i64, i64)>} : () -> ()
 }
 
 // -----
 
 module {
   // expected-error@+3 {{'llvm.readnone' is permitted only on FunctionOpInterface operations}}
   "llvm.func"() ({

mlir/test/mlir-tblgen/enums-gen.td

 // DECL: Bit2 = 4,
 // DECL: Bit3 = 8,
 // DECL: }
 
 // DECL: ::llvm::Optional<MyBitEnum> symbolizeMyBitEnum(uint32_t);
 // DECL: std::string stringifyMyBitEnum(MyBitEnum);
 // DECL: ::llvm::Optional<MyBitEnum> symbolizeMyBitEnum(::llvm::StringRef);
 
+// DECL: struct FieldParser<::MyBitEnum, ::MyBitEnum> {
+// DECL:   template <typename ParserT>
+// DECL:   static FailureOr<::MyBitEnum> parse(ParserT &parser) {
+// DECL:     // Parse the keyword/string containing the enum.
+// DECL:     std::string enumKeyword;
+// DECL:     auto loc = parser.getCurrentLocation();
+// DECL:     if (failed(parser.parseOptionalKeywordOrString(&enumKeyword)))
+// DECL:       return parser.emitError(loc, "expected keyword for An example bit enum");
+// DECL:     // Symbolize the keyword.
+// DECL:     if (::llvm::Optional<::MyBitEnum> attr = ::symbolizeEnum<::MyBitEnum>(enumKeyword))
+// DECL:       return *attr;
+// DECL:     return parser.emitError(loc, "invalid An example bit enum specification: ") << enumKeyword;
+// DECL:   }
+
+// DECL: inline ::llvm::raw_ostream &operator<<(::llvm::raw_ostream &p, ::MyBitEnum value) {
+// DECL:   auto valueStr = stringifyEnum(value);
+// DECL:   return p << valueStr;
+
 // DEF-LABEL: std::string stringifyMyBitEnum
 // DEF: auto val = static_cast<uint32_t>
 // DEF: if (val == 0) return "None";
 // DEF: if (1u == (1u & val))
 // DEF-NEXT: push_back("tagged")
 // DEF: if (2u == (2u & val))
 // DEF-NEXT: push_back("Bit1")
 
 // DEF-LABEL: ::llvm::Optional<MyBitEnum> symbolizeMyBitEnum(::llvm::StringRef str)
 // DEF: if (str == "None") return MyBitEnum::None;
 // DEF: .Case("tagged", 1)
 // DEF: .Case("Bit1", 2)
+
+// Test enum printer generation for non non-keyword enums.
+
+def NonKeywordBit: I32BitEnumAttrCaseBit<"Bit0", 0, "tag-ged">;
+def MyMixedNonKeywordBitEnum: I32BitEnumAttr<"MyMixedNonKeywordBitEnum", "An example bit enum", [
+    NonKeywordBit,
+    Bit1
+  ]> {
+  let genSpecializedAttr = 0;
+}
+
+def MyNonKeywordBitEnum: I32BitEnumAttr<"MyNonKeywordBitEnum", "An example bit enum", [
+    NonKeywordBit
+  ]> {
+  let genSpecializedAttr = 0;
+}
+
+// DECL: inline ::llvm::raw_ostream &operator<<(::llvm::raw_ostream &p, ::MyMixedNonKeywordBitEnum value) {
+// DECL: auto valueStr = stringifyEnum(value);
+// DECL:   switch (value) {
+// DECL:   case ::MyMixedNonKeywordBitEnum::Bit1:
+// DECL:     break;
+// DECL:   default:
+// DECL:     return p << '"' << valueStr << '"';
+// DECL:   }
+// DECL:   return p << valueStr;
+// DECL: }
+
+// DECL: inline ::llvm::raw_ostream &operator<<(::llvm::raw_ostream &p, ::MyNonKeywordBitEnum value) {
+// DECL:   auto valueStr = stringifyEnum(value);
+// DECL:   return p << '"' << valueStr << '"';

mlir/tools/mlir-tblgen/EnumsGen.cpp

 //===- EnumsGen.cpp - MLIR enum utility generator -------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // EnumsGen generates common utility functions for enums.
 //
 //===----------------------------------------------------------------------===//
 
+#include "FormatGen.h"
 #include "mlir/TableGen/Attribute.h"
 #include "mlir/TableGen/Format.h"
 #include "mlir/TableGen/GenInfo.h"
+#include "llvm/ADT/BitVector.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/TableGen/Error.h"
 #include "llvm/TableGen/Record.h"
 #include "llvm/TableGen/TableGenBackend.h"
 
if (value >= 0) {
       os << formatv("  {0} = {1},\n", symbol, value);
     } else {
       os << formatv("  {0},\n", symbol);
     }
   }
   os << "};\n\n";
 }
 
-static void emitDenseMapInfo(StringRef enumName, std::string underlyingType,
+static void emitParserPrinter(const EnumAttr &enumAttr, StringRef qualName,
+                              StringRef cppNamespace, raw_ostream &os) {
+  if (enumAttr.getUnderlyingType().empty() ||
+      enumAttr.getConstBuilderTemplate().empty())
+    return;
+  auto cases = enumAttr.getAllCases();
+
+  // Check which cases shouldn't be printed using a keyword.
+  llvm::BitVector nonKeywordCases(cases.size());
+  for (auto [index, caseVal] : llvm::enumerate(cases))
+    if (!mlir::tblgen::canFormatStringAsKeyword(caseVal.getStr()))
+      nonKeywordCases.set(index);
+
+  // If this is a bit enum attribute, don't allow cases that may overlap with
+  // other cases. For simplicity sake, only allow cases with a single bit value.
+  if (enumAttr.isBitEnum()) {
+    for (auto [index, caseVal] : llvm::enumerate(cases)) {
+      int64_t value = caseVal.getValue();
+      if (value < 0 || (value != 0 && !llvm::isPowerOf2_64(value)))
+        nonKeywordCases.set(index);
+    }
+  }
+
+  // Generate the parser and the start of the printer for the enum.
+  const char *parsedAndPrinterStart = R"(
+namespace mlir {
+template <typename T, typename>
+struct FieldParser;
+
+template<>
+struct FieldParser<{0}, {0}> {{
+  template <typename ParserT>
+  static FailureOr<{0}> parse(ParserT &parser) {{
+    // Parse the keyword/string containing the enum.
+    std::string enumKeyword;
+    auto loc = parser.getCurrentLocation();
+    if (failed(parser.parseOptionalKeywordOrString(&enumKeyword)))
+      return parser.emitError(loc, "expected keyword for {2}");
+
+    // Symbolize the keyword.
+    if (::llvm::Optional<{0}> attr = {1}::symbolizeEnum<{0}>(enumKeyword))
+      return *attr;
+    return parser.emitError(loc, "invalid {2} specification: ") << enumKeyword;
+  }
+};
+} // namespace mlir
+
+namespace llvm {
+inline ::llvm::raw_ostream &operator<<(::llvm::raw_ostream &p, {0} value) {{
+  auto valueStr = stringifyEnum(value);
+)";
+  os << formatv(parsedAndPrinterStart, qualName, cppNamespace,
+                enumAttr.getSummary());
+
+  // If all cases require a string, always wrap.
+  if (nonKeywordCases.all()) {
+    os << "  return p << '\"' << valueStr << '\"';\n"
+          "}\n"
+          "} // namespace llvm\n";
+    return;
+  }
+
+  // If there are any cases that can't be used with a keyword, switch on the
+  // case value to determine when to print in the string form.
+  if (nonKeywordCases.any()) {
+    os << "  switch (value) {\n";
+    for (auto &it : llvm::enumerate(cases)) {
+      if (nonKeywordCases.test(it.index()))
+        continue;
+      StringRef symbol = it.value().getSymbol();
+      os << llvm::formatv("    case {0}::{1}:\n", qualName,
+                          llvm::isDigit(symbol.front()) ? ("_" + symbol)
+                                                        : symbol);
+    }
+    os << "    break;\n"
+          "  default:\n"
+          "    return p << '\"' << valueStr << '\"';\n"
+          "  }\n";
+  }
+  os << "  return p << valueStr;\n"
+        "}\n"
+        "} // namespace llvm\n";
+}
+
+static void emitDenseMapInfo(StringRef qualName, std::string underlyingType,
                              StringRef cppNamespace, raw_ostream &os) {
-  std::string qualName =
-      std::string(formatv("{0}::{1}", cppNamespace, enumName));
   if (underlyingType.empty())
     underlyingType =
         std::string(formatv("std::underlying_type_t<{0}>", qualName));
 
   const char *const mapInfo = R"(
 namespace llvm {
 template<> struct DenseMapInfo<{0}> {{
   using StorageInfo = ::llvm::DenseMapInfo<{1}>;
if (enumAttr.genSpecializedAttr()) {
     StringRef attrClassName = enumAttr.getSpecializedAttrClassName();
     StringRef baseAttrClassName = "IntegerAttr";
     os << formatv(attrClassDecl, enumName, attrClassName, baseAttrClassName);
   }
 
   for (auto ns : llvm::reverse(namespaces))
     os << "} // namespace " << ns << "\n";
 
+  // Generate a generic parser and printer for the enum.
+  std::string qualName =
+      std::string(formatv("{0}::{1}", cppNamespace, enumName));
+  emitParserPrinter(enumAttr, qualName, cppNamespace, os);
+
   // Emit DenseMapInfo for this enum class
-  emitDenseMapInfo(enumName, underlyingType, cppNamespace, os);
+  emitDenseMapInfo(qualName, underlyingType, cppNamespace, os);
 }
 
 static bool emitEnumDecls(const RecordKeeper &recordKeeper, raw_ostream &os) {
   llvm::emitSourceFileHeader("Enum Utility Declarations", os);
 
   auto defs = recordKeeper.getAllDerivedDefinitionsIfDefined("EnumAttrInfo");
   for (const auto *def : defs)
     emitEnumDecl(*def, os);

mlir/tools/mlir-tblgen/FormatGen.cpp

if (value.size() != 1 && value != "->")
     return true;
   if (lastWasPunctuation)
     return !StringRef(">)}],").contains(value.front());
   return !StringRef("<>(){}[],").contains(value.front());
 }
 
 bool mlir::tblgen::canFormatStringAsKeyword(
     StringRef value, function_ref<void(Twine)> emitError) {
+  if (value.empty()) {
+    if (emitError)
+      emitError("keywords cannot be empty");
+    return false;
+  }
   if (!isalpha(value.front()) && value.front() != '_') {
     if (emitError)
       emitError("valid keyword starts with a letter or '_'");
     return false;
   }
   if (!llvm::all_of(value.drop_front(), [](char c) {
         return isalnum(c) || c == '_' || c == '$' || c == '.';
       })) {