[mlir] Improve BitEnumAttr, update documentation

[krzysz00]

Summary

• Add new operators to BitEnumAttr, mainly not (which only inverts

bits that can be valid bits for the attribute) and xor

• Add new bit enum utility functions: bitEnumClear(bits, bit) and

bitEnumSet(bits, bit, value=true) as they've come up in code I've been
writing that makes use of such enums

• Add rudimentary tests for the enum generator
• Update the OpDefinition documentation to make it contain a correct

example and to have it account for the changes mentioned above.

[antiagainst]

Nice!

[hanchung]

mlir/docs/OpDefinitions.md

1445–1447

The behavior of bitEnumContains changed. Is it intended?

[antiagainst]

mlir/docs/OpDefinitions.md

1445–1447

Good catch about the behavior change! Sorry I missed that.

I think we should not silently changing the behavior. I also see the benefits of supporting the "and" behavior. So maybe change this into two functions, bitEnumContainsAny and bitEnumContainsAll?

[hanchung]

mlir/docs/OpDefinitions.md

1445–1447

Yes, having two methods looks good. (and the naming is more clear to me!) I'll send out a patch.

mlir/docs/OpDefinitions.md

 
 } // namespace Inner
 } // namespace Outer
 ```
 
 Similarly for the following `BitEnumAttr` definition:
 
 ```tablegen
-def None: BitEnumAttrCaseNone<"None">;
-def Bit0: BitEnumAttrCaseBit<"Bit0", 0>;
-def Bit1: BitEnumAttrCaseBit<"Bit1", 1>;
-def Bit2: BitEnumAttrCaseBit<"Bit2", 2>;
-def Bit3: BitEnumAttrCaseBit<"Bit3", 3>;
+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 MyBitEnum: BitEnumAttr<"MyBitEnum", "An example bit enum",
                            [None, Bit0, Bit1, Bit2, Bit3]>;
 ```
 
 We can have:
 
 ```c++
 // An example bit enum
 enum class MyBitEnum : uint32_t {
   None = 0,
   Bit0 = 1,
   Bit1 = 2,
   Bit2 = 4,
   Bit3 = 8,
 };
 
 llvm::Optional<MyBitEnum> symbolizeMyBitEnum(uint32_t);
 std::string stringifyMyBitEnum(MyBitEnum);
 llvm::Optional<MyBitEnum> symbolizeMyBitEnum(llvm::StringRef);
-inline MyBitEnum operator|(MyBitEnum lhs, MyBitEnum rhs) {
-  return static_cast<MyBitEnum>(static_cast<uint32_t>(lhs) | static_cast<uint32_t>(rhs));
+
+inline constexpr MyBitEnum operator|(MyBitEnum a, MyBitEnum b) {
+  return static_cast<MyBitEnum>(static_cast<uint32_t>(a) | static_cast<uint32_t>(b));
+}
+inline constexpr MyBitEnum operator&(MyBitEnum a, MyBitEnum b) {
+  return static_cast<MyBitEnum>(static_cast<uint32_t>(a) & static_cast<uint32_t>(b));
+}
+inline constexpr MyBitEnum operator^(MyBitEnum a, MyBitEnum b) {
+  return static_cast<MyBitEnum>(static_cast<uint32_t>(a) ^ static_cast<uint32_t>(b));
+}
+inline constexpr MyBitEnum operator~(MyBitEnum bits) {
+  // Ensure only bits that can be present in the enum are set
+  return static_cast<MyBitEnum>(~static_cast<uint32_t>(bits) & static_cast<uint32_t>(15u));
+}
+inline constexpr bool bitEnumContains(MyBitEnum bits, MyBitEnum bit) {
+  return (bits & bit) == bit;
 }
-inline MyBitEnum operator&(MyBitEnum lhs, MyBitEnum rhs) {
-  return static_cast<MyBitEnum>(static_cast<uint32_t>(lhs) & static_cast<uint32_t>(rhs));
+inline constexpr MyBitEnum bitEnumClear(MyBitEnum bits, MyBitEnum bit) {
+  return bits & ~bit;
 }
-inline bool bitEnumContains(MyBitEnum bits, MyBitEnum bit) {
-  return (static_cast<uint32_t>(bits) & static_cast<uint32_t>(bit)) != 0;
+
+inline std::string stringifyEnum(MyBitEnum enumValue) {
+  return stringifyMyBitEnum(enumValue);
+}
+
+template <typename EnumType>
+::llvm::Optional<EnumType> symbolizeEnum(::llvm::StringRef);
+
+template <>
+inline ::llvm::Optional<MyBitEnum> symbolizeEnum<MyBitEnum>(::llvm::StringRef str) {
+  return symbolizeMyBitEnum(str);
 }
 
 namespace llvm {
 template<> struct DenseMapInfo<::MyBitEnum> {
   using StorageInfo = llvm::DenseMapInfo<uint32_t>;
 
   static inline ::MyBitEnum getEmptyKey() {
     return static_cast<::MyBitEnum>(StorageInfo::getEmptyKey());

 
 ```c++
 std::string stringifyMyBitEnum(MyBitEnum symbol) {
   auto val = static_cast<uint32_t>(symbol);
   assert(15u == (15u | val) && "invalid bits set in bit enum");
   // Special case for all bits unset.
   if (val == 0) return "None";
   llvm::SmallVector<llvm::StringRef, 2> strs;
-  if (1u == (1u & val)) { strs.push_back("Bit0"); }
+  if (1u == (1u & val)) { strs.push_back("tagged"); }
   if (2u == (2u & val)) { strs.push_back("Bit1"); }
   if (4u == (4u & val)) { strs.push_back("Bit2"); }
   if (8u == (8u & val)) { strs.push_back("Bit3"); }
 
   return llvm::join(strs, "|");
 }
 
 llvm::Optional<MyBitEnum> symbolizeMyBitEnum(llvm::StringRef str) {
   // Special case for all bits unset.
   if (str == "None") return MyBitEnum::None;
 
   llvm::SmallVector<llvm::StringRef, 2> symbols;
   str.split(symbols, "|");
 
   uint32_t val = 0;
   for (auto symbol : symbols) {
     auto bit = llvm::StringSwitch<llvm::Optional<uint32_t>>(symbol)
-      .Case("Bit0", 1)
+      .Case("tagged", 1)
       .Case("Bit1", 2)
       .Case("Bit2", 4)
       .Case("Bit3", 8)
       .Default(llvm::None);
     if (bit) { val |= *bit; } else { return llvm::None; }
   }
   return static_cast<MyBitEnum>(val);
 }
 
 llvm::Optional<MyBitEnum> symbolizeMyBitEnum(uint32_t value) {
   // Special case for all bits unset.
   if (value == 0) return MyBitEnum::None;
 
-  if (value & ~(1u | 2u | 4u | 8u)) return llvm::None;
+  if (value & ~static_cast<uint32_t>(15u)) return llvm::None;
   return static_cast<MyBitEnum>(value);
 }
 ```
 
 ## Debugging Tips
 
 ### Run `mlir-tblgen` to see the generated content
 

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 MyBitEnum: I32BitEnumAttr<"MyBitEnum", "An example bit enum",
+                           [None, Bit0, Bit1, Bit2, Bit3]> {
+  let genSpecializedAttr = 0;
+}
+
+// DECL-LABEL: enum class MyBitEnum : uint32_t
+// DECL: None = 0,
+// DECL: Bit0 = 1,
+// DECL: Bit1 = 2,
+// 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);
+
+// 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)

mlir/tools/mlir-tblgen/EnumsGen.cpp

getAllBitsUnsetCase(llvm::ArrayRef<EnumAttrCase> cases) {
   }
   return llvm::None;
 }
 
 // Emits the following inline function for bit enums:
 //
 // inline constexpr <enum-type> operator|(<enum-type> a, <enum-type> b);
 // inline constexpr <enum-type> operator&(<enum-type> a, <enum-type> b);
-// inline constexpr bool bitEnumContains(<enum-type> a, <enum-type> b);
+// inline constexpr <enum-type> operator^(<enum-type> a, <enum-type> b);
+// inline constexpr <enum-type> operator~(<enum-type> bits);
+// inline constexpr bool bitEnumContains(<enum-type> bits, <enum-type> bit);
+// inline constexpr <enum-type> bitEnumClear(<enum-type> bits, <enum-type> bit);
+// inline constexpr <enum-type> bitEnumSet(<enum-type> bits, <enum-type> bit,
+// bool value=true);
 static void emitOperators(const Record &enumDef, raw_ostream &os) {
   EnumAttr enumAttr(enumDef);
   StringRef enumName = enumAttr.getEnumClassName();
   std::string underlyingType = std::string(enumAttr.getUnderlyingType());
-  os << formatv("inline constexpr {0} operator|({0} lhs, {0} rhs) {{\n",
-                enumName)
-     << formatv("  return static_cast<{0}>("
-                "static_cast<{1}>(lhs) | static_cast<{1}>(rhs));\n",
-                enumName, underlyingType)
-     << "}\n";
-  os << formatv("inline constexpr {0} operator&({0} lhs, {0} rhs) {{\n",
-                enumName)
-     << formatv("  return static_cast<{0}>("
-                "static_cast<{1}>(lhs) & static_cast<{1}>(rhs));\n",
-                enumName, underlyingType)
-     << "}\n";
-  os << formatv(
-            "inline constexpr bool bitEnumContains({0} bits, {0} bit) {{\n"
-            "  return (static_cast<{1}>(bits) & static_cast<{1}>(bit)) != 0;\n",
-            enumName, underlyingType)
-     << "}\n";
+  int64_t validBits = enumDef.getValueAsInt("validBits");
+  const char *const operators = R"(
+inline constexpr {0} operator|({0} a, {0} b) {{
+  return static_cast<{0}>(static_cast<{1}>(a) | static_cast<{1}>(b));
+}
+inline constexpr {0} operator&({0} a, {0} b) {{
+  return static_cast<{0}>(static_cast<{1}>(a) & static_cast<{1}>(b));
+}
+inline constexpr {0} operator^({0} a, {0} b) {{
+  return static_cast<{0}>(static_cast<{1}>(a) ^ static_cast<{1}>(b));
+}
+inline constexpr {0} operator~({0} bits) {{
+  // Ensure only bits that can be present in the enum are set
+  return static_cast<{0}>(~static_cast<{1}>(bits) & static_cast<{1}>({2}u));
+}
+inline constexpr bool bitEnumContains({0} bits, {0} bit) {{
+  return (bits & bit) == bit;
+}
+inline constexpr {0} bitEnumClear({0} bits, {0} bit) {{
+  return bits & ~bit;
+}
+inline constexpr {0} bitEnumSet({0} bits, {0} bit, /*optional*/bool value=true) {{
+  return value ? (bits | bit) : bitEnumClear(bits, bit);
+}
+  )";
+  os << formatv(operators, enumName, underlyingType, validBits);
 }
 
 static void emitSymToStrFnForIntEnum(const Record &enumDef, raw_ostream &os) {
   EnumAttr enumAttr(enumDef);
   StringRef enumName = enumAttr.getEnumClassName();
   StringRef symToStrFnName = enumAttr.getSymbolToStringFnName();
   StringRef symToStrFnRetType = enumAttr.getSymbolToStringFnRetType();
   auto enumerants = enumAttr.getAllCases();
static void emitUnderlyingToSymFnForBitEnum(const Record &enumDef,
 
   os << formatv("::llvm::Optional<{0}> {1}({2} value) {{\n", enumName,
                 underlyingToSymFnName, underlyingType);
   if (allBitsUnsetCase) {
     os << "  // Special case for all bits unset.\n";
     os << formatv("  if (value == 0) return {0}::{1};\n\n", enumName,
                   makeIdentifier(allBitsUnsetCase->getSymbol()));
   }
-  llvm::SmallVector<std::string, 8> values;
-  for (const auto &enumerant : enumerants) {
-    if (auto val = enumerant.getValue())
-      values.push_back(std::string(formatv("{0}u", val)));
-  }
-  os << formatv("  if (value & ~static_cast<{0}>({1})) return llvm::None;\n",
-                underlyingType, llvm::join(values, " | "));
+  int64_t validBits = enumDef.getValueAsInt("validBits");
+  os << formatv("  if (value & ~static_cast<{0}>({1}u)) return llvm::None;\n",
+                underlyingType, validBits);
   os << formatv("  return static_cast<{0}>(value);\n", enumName);
   os << "}\n";
 }
 
 static void emitEnumDecl(const Record &enumDef, raw_ostream &os) {
   EnumAttr enumAttr(enumDef);
   StringRef enumName = enumAttr.getEnumClassName();
   StringRef cppNamespace = enumAttr.getCppNamespace();