[mlir] Better handling for bit groups in enum parser/printer

[rriddle]

Summary

We currently wrap all multi-bit cases with a string, but this is
overly restrictive. This commit refactors to use keywords when
we know they are valid, and only degrade to string when the validity
of the bitgroup is unknown.

[Mogball]

mlir/tools/mlir-tblgen/EnumsGen.cpp

143

Why not SmallVector?

146

nit: I feel like !llvm::has_single_bit should come before test

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

 // RUN: mlir-tblgen -gen-enum-decls -I %S/../../include %s | FileCheck %s --check-prefix=DECL
 // RUN: mlir-tblgen -gen-enum-defs -I %S/../../include %s | FileCheck %s --check-prefix=DEF
 
 include "mlir/IR/EnumAttr.td"
 include "mlir/IR/OpBase.td"
 
 // Test bit enums
 def None: I32BitEnumAttrCaseNone<"None">;
 def Bit0: I32BitEnumAttrCaseBit<"Bit0", 0, "tagged">;
 def Bit1: I32BitEnumAttrCaseBit<"Bit1", 1>;
 def Bit2: I32BitEnumAttrCaseBit<"Bit2", 2>;
 def Bit3: I32BitEnumAttrCaseBit<"Bit3", 3>;
+def BitGroup: I32BitEnumAttrCaseGroup<"BitGroup", [
+  Bit0, Bit1
+]>;
 
 def MyBitEnum: I32BitEnumAttr<"MyBitEnum", "An example bit enum",
-                           [None, Bit0, Bit1, Bit2, Bit3]> {
+                           [None, Bit0, Bit1, Bit2, Bit3, BitGroup]> {
   let genSpecializedAttr = 0;
 }
 
 // DECL-LABEL: enum class MyBitEnum : uint32_t
 // DECL: None = 0,
 // DECL: Bit0 = 1,
 // DECL: Bit1 = 2,
 // DECL: Bit2 = 4,

 // 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:   switch (value) {
+// DECL:   case ::MyBitEnum::BitGroup:
+// DECL:     return p << valueStr;
+// DECL:   default:
+// DECL:     break;
+// DECL:   }
+// DECL:   auto underlyingValue = static_cast<std::make_unsigned_t<::MyBitEnum>>(value);
+// DECL:   if (underlyingValue && !llvm::has_single_bit(underlyingValue))
+// DECL:     return p << '"' << valueStr << '"';
 // 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))

mlir/tools/mlir-tblgen/EnumsGen.cpp

static void emitParserPrinter(const EnumAttr &enumAttr, StringRef qualName,
   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}> {{
)";
   // 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,
+      os << llvm::formatv("  case {0}::{1}:\n", qualName,
                           llvm::isDigit(symbol.front()) ? ("_" + symbol)
                                                         : symbol);
     }
     os << "    break;\n"
           "  default:\n"
           "    return p << '\"' << valueStr << '\"';\n"
           "  }\n";
+
+    // If this is a bit enum, conservatively print the string form if the value
+    // is not a power of two (i.e. not a single bit case) and not a known case.
+  } else if (enumAttr.isBitEnum()) {
+    // Process the known multi-bit cases that use valid keywords.
+    llvm::SmallVector<EnumAttrCase *> validMultiBitCases;
+    for (auto [index, caseVal] : llvm::enumerate(cases)) {
+      uint64_t value = caseVal.getValue();
+      if (value && !llvm::has_single_bit(value) && !nonKeywordCases.test(index))
+        validMultiBitCases.push_back(&caseVal);
+    }
+    if (!validMultiBitCases.empty()) {
+      os << "  switch (value) {\n";
+      for (EnumAttrCase *caseVal : validMultiBitCases) {
+        StringRef symbol = caseVal->getSymbol();
+        os << llvm::formatv("  case {0}::{1}:\n", qualName,
+                            llvm::isDigit(symbol.front()) ? ("_" + symbol)
+                                                          : symbol);
+      }
+      os << "    return p << valueStr;\n"
+            "  default:\n"
+            "    break;\n"
+            "  }\n";
+    }
+
+    // All other multi-bit cases should be printed as strings.
+    os << formatv("  auto underlyingValue = "
+                  "static_cast<std::make_unsigned_t<{0}>>(value);\n",
+                  qualName);
+    os << "  if (underlyingValue && !llvm::has_single_bit(underlyingValue))\n"
+          "    return p << '\"' << valueStr << '\"';\n";
   }
   os << "  return p << valueStr;\n"
         "}\n"
         "} // namespace llvm\n";
 }
 
 static void emitDenseMapInfo(StringRef qualName, std::string underlyingType,
                              StringRef cppNamespace, raw_ostream &os) {