[mlir][ods] Add tablegen field for concise printing of BitEnum attributes

[jfurtek]

Summary

This diff introduces a tablegen field for bit enum attributes
(printBitEnumPrimaryGroups) to control printing when the enum uses "group"
cases. An example would be an implementation that uses a fastmath enum value
as an alias for individual fastmath flags. The proposed field would allow
printing of simply fast for the enum value, instead of the more verbose list
that would include fast as well as the individual flags (e.g. `reassoc,nnan,
ninf,nsz,arcp,contract,afn,fast`).

[rriddle]

mlir/tools/mlir-tblgen/EnumsGen.cpp

191

Can you document these?

192–193

Please format these as you would if you wrote it directly(e.g. proper newlines, formatting, etc.) We want the generated code to be in the same form as hand written code (when possible).

195

Can you drop the braces here?

218

Can you just use llvm::reverse here?

229

Drop the trivial braces here.

mlir/unittests/TableGen/enums.td

52–60

Can you reduce the number of lines here? This style feels different from the rest of the file.

[jfurtek]

• Correct formatting, add comments, use llvm::reverse()

[rriddle]

LG, thanks!

mlir/tools/mlir-tblgen/EnumsGen.cpp

216–217

Prefer to keep the code here readable as well.

223

Same here.

231

239

247

[jfurtek]

• Correct formatting, add comments, use llvm::reverse()
• Make formatting code more readable
• Properly punctuate comments.

[jfurtek]

• Rebase

mlir/include/mlir/IR/EnumAttr.td

class BitEnumAttr<I intType, string name, string summary,
   int validBits = !foldl(0, cases, value, bitcase, !or(value, bitcase.value));
 
   // We need to return a string because we may concatenate symbols for multiple
   // bits together.
   let symbolToStringFnRetType = "std::string";
 
   // The delimiter used to separate bit enum cases in strings.
   string separator = "|";
+
+  // Print the "primary group" only for bits that are members of case groups
+  // that have all bits present. When the value is 0, printing will display both
+  // both individual bit case names AND the names for all groups that the bit is
+  // contained in. When the value is 1, for each bit that is set AND is a member
+  // of a group with all bits set, only the "primary group" (i.e. the first
+  // group with all bits set in reverse declaration order) will be printed (for
+  // conciseness).
+  bit printBitEnumPrimaryGroups = 0;
 }
 
 class I32BitEnumAttr<string name, string summary,
                      list<BitEnumAttrCaseBase> cases>
     : BitEnumAttr<I32, name, summary, cases> {
   let underlyingType = "uint32_t";
 }
 

mlir/include/mlir/TableGen/Attribute.h

public:
   StringRef getMaxEnumValFnName() const;
 
   // Returns all allowed cases for this enum attribute.
   std::vector<EnumAttrCase> getAllCases() const;
 
   bool genSpecializedAttr() const;
   llvm::Record *getBaseAttrClass() const;
   StringRef getSpecializedAttrClassName() const;
+  bool printBitEnumPrimaryGroups() const;
 };
 
 class StructFieldAttr {
 public:
   explicit StructFieldAttr(const llvm::Record *record);
   explicit StructFieldAttr(const llvm::Record &record);
   explicit StructFieldAttr(const llvm::DefInit *init);
 

mlir/lib/TableGen/Attribute.cpp

 llvm::Record *EnumAttr::getBaseAttrClass() const {
   return def->getValueAsDef("baseAttrClass");
 }
 
 StringRef EnumAttr::getSpecializedAttrClassName() const {
   return def->getValueAsString("specializedAttrClassName");
 }
 
+bool EnumAttr::printBitEnumPrimaryGroups() const {
+  return def->getValueAsBit("printBitEnumPrimaryGroups");
+}
+
 StructFieldAttr::StructFieldAttr(const llvm::Record *record) : def(record) {
   assert(def->isSubClassOf("StructFieldAttr") &&
          "must be subclass of TableGen 'StructFieldAttr' class");
 }
 
 StructFieldAttr::StructFieldAttr(const llvm::Record &record)
     : StructFieldAttr(&record) {}
 

mlir/tools/mlir-tblgen/EnumsGen.cpp

static void emitSymToStrFnForBitEnum(const Record &enumDef, raw_ostream &os) {
   EnumAttr enumAttr(enumDef);
   StringRef enumName = enumAttr.getEnumClassName();
   StringRef symToStrFnName = enumAttr.getSymbolToStringFnName();
   StringRef symToStrFnRetType = enumAttr.getSymbolToStringFnRetType();
   StringRef separator = enumDef.getValueAsString("separator");
   auto enumerants = enumAttr.getAllCases();
   auto allBitsUnsetCase = getAllBitsUnsetCase(enumerants);
 
+  const char *const formatCompareRemove =
+      "  if ({0}u == ({0}u & val)) {{ strs.push_back(\"{1}\"); "
+      "val &= ~static_cast<{2}>({0}); }\n";
+  const char *const formatCompare =
+      "  if ({0}u == ({0}u & val)) {{ strs.push_back(\"{1}\"); }\n";
+
   os << formatv("{2} {1}({0} symbol) {{\n", enumName, symToStrFnName,
                 symToStrFnRetType);
 
   os << formatv("  auto val = static_cast<{0}>(symbol);\n",
                 enumAttr.getUnderlyingType());
   // If we have unknown bit set, return an empty string to signal errors.
   int64_t validBits = enumDef.getValueAsInt("validBits");
   os << formatv("  assert({0}u == ({0}u | val) && \"invalid bits set in bit "
                 "enum\");\n",
                 validBits);
   if (allBitsUnsetCase) {
     os << "  // Special case for all bits unset.\n";
     os << formatv("  if (val == 0) return \"{0}\";\n\n",
                   allBitsUnsetCase->getSymbol());
   }
   os << "  ::llvm::SmallVector<::llvm::StringRef, 2> strs;\n";
+  // Optionally elide bits that are members of groups that will be printed for
+  // more concise output.
+  if (enumAttr.printBitEnumPrimaryGroups()) {
+    os << "  // Print bit enum groups before individual bits\n";
+    // Check group bit patterns (in reverse)
+    for (auto it = enumerants.rbegin(); it != enumerants.rend(); ++it) {
+      auto enumerant = *it;
+      int64_t val = enumerant.getValue();
+      if (val && enumerant.getDef().isSubClassOf("BitEnumAttrCaseGroup")) {
+        os << formatv(formatCompareRemove, val, enumerant.getStr(),
+                      enumAttr.getUnderlyingType());
+      }
+    }
+    // Check non-group bit patterns
+    for (const auto &enumerant : enumerants) {
+      int64_t val = enumerant.getValue();
+      if (val && !enumerant.getDef().isSubClassOf("BitEnumAttrCaseGroup")) {
+        os << formatv(formatCompare, val, enumerant.getStr());
+      }
+    }
+  } else {
     for (const auto &enumerant : enumerants) {
       // Skip the special enumerant for None.
-    if (int64_t val = enumerant.getValue())
-      os << formatv(
-          "  if ({0}u == ({0}u & val)) {{ strs.push_back(\"{1}\"); }\n ", val,
-          enumerant.getStr());
+      if (int64_t val = enumerant.getValue()) {
+        // Emit code to check ALL cases (individual bits and groups)
+        os << formatv(formatCompare, val, enumerant.getStr());
+      }
+    }
   }
   os << formatv("  return ::llvm::join(strs, \"{0}\");\n", separator);
 
   os << "}\n\n";
 }
 
 static void emitStrToSymFnForIntEnum(const Record &enumDef, raw_ostream &os) {
   EnumAttr enumAttr(enumDef);
   StringRef enumName = enumAttr.getEnumClassName();
   StringRef strToSymFnName = enumAttr.getStringToSymbolFnName();
   auto enumerants = enumAttr.getAllCases();
 
   os << formatv("::llvm::Optional<{0}> {1}(::llvm::StringRef str) {{\n",
                 enumName, strToSymFnName);
   os << formatv("  return ::llvm::StringSwitch<::llvm::Optional<{0}>>(str)\n",

mlir/unittests/TableGen/EnumsGenTest.cpp

TEST(EnumsGenTest, GeneratedUnderlyingType) {
   bool v = std::is_same<uint32_t, std::underlying_type<I32Enum>::type>::value;
   EXPECT_TRUE(v);
 }
 
 TEST(EnumsGenTest, GeneratedBitEnumDefinition) {
   EXPECT_EQ(0u, static_cast<uint32_t>(BitEnumWithNone::None));
   EXPECT_EQ(1u, static_cast<uint32_t>(BitEnumWithNone::Bit0));
   EXPECT_EQ(8u, static_cast<uint32_t>(BitEnumWithNone::Bit3));
+
+  EXPECT_EQ(2u, static_cast<uint64_t>(BitEnum64_Test::Bit1));
+  EXPECT_EQ(144115188075855872u, static_cast<uint64_t>(BitEnum64_Test::Bit57));
 }
 
 TEST(EnumsGenTest, GeneratedSymbolToStringFnForBitEnum) {
   EXPECT_EQ(stringifyBitEnumWithNone(BitEnumWithNone::None), "None");
   EXPECT_EQ(stringifyBitEnumWithNone(BitEnumWithNone::Bit0), "Bit0");
   EXPECT_EQ(stringifyBitEnumWithNone(BitEnumWithNone::Bit3), "Bit3");
   EXPECT_EQ(
       stringifyBitEnumWithNone(BitEnumWithNone::Bit0 | BitEnumWithNone::Bit3),
       "Bit0|Bit3");
-  EXPECT_EQ(2u, static_cast<uint64_t>(BitEnum64_Test::Bit1));
-  EXPECT_EQ(144115188075855872u, static_cast<uint64_t>(BitEnum64_Test::Bit57));
+
+  EXPECT_EQ(stringifyBitEnum64_Test(BitEnum64_Test::Bit1), "Bit1");
+  EXPECT_EQ(
+      stringifyBitEnum64_Test(BitEnum64_Test::Bit1 | BitEnum64_Test::Bit57),
+      "Bit1|Bit57");
 }
 
 TEST(EnumsGenTest, GeneratedStringToSymbolForBitEnum) {
   EXPECT_EQ(symbolizeBitEnumWithNone("None"), BitEnumWithNone::None);
   EXPECT_EQ(symbolizeBitEnumWithNone("Bit0"), BitEnumWithNone::Bit0);
   EXPECT_EQ(symbolizeBitEnumWithNone("Bit3"), BitEnumWithNone::Bit3);
   EXPECT_EQ(symbolizeBitEnumWithNone("Bit3|Bit0"),
             BitEnumWithNone::Bit3 | BitEnumWithNone::Bit0);

 TEST(EnumsGenTest, GeneratedStringToSymbolForGroupedBitEnum) {
   EXPECT_EQ(symbolizeBitEnumWithGroup("Bit0"), BitEnumWithGroup::Bit0);
   EXPECT_EQ(symbolizeBitEnumWithGroup("Bit3"), BitEnumWithGroup::Bit3);
   EXPECT_EQ(symbolizeBitEnumWithGroup("Bit5"), llvm::None);
   EXPECT_EQ(symbolizeBitEnumWithGroup("Bit3|Bit0"),
             BitEnumWithGroup::Bit3 | BitEnumWithGroup::Bit0);
 }
 
+TEST(EnumsGenTest, GeneratedSymbolToStringFnForPrimaryGroupBitEnum) {
+  EXPECT_EQ(stringifyBitEnumPrimaryGroup(
+                BitEnumPrimaryGroup::Bit0 | BitEnumPrimaryGroup::Bit1 |
+                BitEnumPrimaryGroup::Bit2 | BitEnumPrimaryGroup::Bit3),
+            "Bits0To3");
+  EXPECT_EQ(stringifyBitEnumPrimaryGroup(BitEnumPrimaryGroup::Bit0 |
+                                         BitEnumPrimaryGroup::Bit2 |
+                                         BitEnumPrimaryGroup::Bit3),
+            "Bit0,Bit2,Bit3");
+  EXPECT_EQ(stringifyBitEnumPrimaryGroup(BitEnumPrimaryGroup::Bit0 |
+                                         BitEnumPrimaryGroup::Bit4 |
+                                         BitEnumPrimaryGroup::Bit5),
+            "Bits4And5,Bit0");
+  EXPECT_EQ(stringifyBitEnumPrimaryGroup(
+                BitEnumPrimaryGroup::Bit0 | BitEnumPrimaryGroup::Bit1 |
+                BitEnumPrimaryGroup::Bit2 | BitEnumPrimaryGroup::Bit3 |
+                BitEnumPrimaryGroup::Bit4 | BitEnumPrimaryGroup::Bit5),
+            "Bits0To5");
+}
+
 TEST(EnumsGenTest, GeneratedOperator) {
   EXPECT_TRUE(bitEnumContains(BitEnumWithNone::Bit0 | BitEnumWithNone::Bit3,
                               BitEnumWithNone::Bit0));
   EXPECT_FALSE(bitEnumContains(BitEnumWithNone::Bit0 & BitEnumWithNone::Bit3,
                                BitEnumWithNone::Bit0));
 }
 
 TEST(EnumsGenTest, GeneratedSymbolToCustomStringFn) {

mlir/unittests/TableGen/enums.td

 def BitEnumWithNone : I32BitEnumAttr<"BitEnumWithNone", "A test enum",
                                      [NoBits, Bit0, Bit3]>;
 
 def BitEnumWithoutNone : I32BitEnumAttr<"BitEnumWithoutNone", "A test enum",
                                         [Bit0, Bit3]>;
 
 def Bits0To3 : I32BitEnumAttrCaseGroup<"Bits0To3",
                                        [Bit0, Bit1, Bit2, Bit3]>;
+def Bits4And5 : I32BitEnumAttrCaseGroup<"Bits4And5",
+                                       [Bit4, Bit5]>;
+def Bits0To5 : I32BitEnumAttrCaseGroup<"Bits0To5",
+                                       [Bits0To3, Bits4And5]>;
 
 def BitEnumWithGroup : I32BitEnumAttr<"BitEnumWithGroup", "A test enum",
                                       [Bit0, Bit1, Bit2, Bit3, Bit4, Bits0To3]>;
 
+def BitEnumPrimaryGroup : I32BitEnumAttr<"BitEnumPrimaryGroup", "test enum",
+                                        [Bit0,
+                                         Bit1,
+                                         Bit2,
+                                         Bit3,
+                                         Bit4,
+                                         Bit5,
+                                         Bits0To3,
+                                         Bits4And5,
+                                         Bits0To5]> {
+  let separator = ",";
+  let printBitEnumPrimaryGroups = 1;
+}
+
 def BitEnum64_None : I64BitEnumAttrCaseNone<"None">;
 def BitEnum64_57   : I64BitEnumAttrCaseBit<"Bit57", 57>;
 def BitEnum64_1    : I64BitEnumAttrCaseBit<"Bit1", 1>;
 def BitEnum64_Test : I64BitEnumAttr<"BitEnum64_Test", "A 64-bit test enum",
                                     [BitEnum64_None, BitEnum64_1, BitEnum64_57]>;
 
 def PrettyIntEnumCase1: I32EnumAttrCase<"Case1", 1, "case_one">;
 def PrettyIntEnumCase2: I32EnumAttrCase<"Case2", 2, "case_two">;
 
 def PrettyIntEnum: I32EnumAttr<"PrettyIntEnum", "A test enum",
                                [PrettyIntEnumCase1, PrettyIntEnumCase2]>;