[mlir] Improve bitEnumContains methods.

[hanchung]

Summary

https://github.com/llvm/llvm-project/commit/839b436c93604e042f74050cf2adadd75f30e898
changes the behavior. Based on the discussion, we also want to support
"and" behavior. The revision changes it into two functions, bitEnumContainsAny
and bitEnumContainsAll.

[krzysz00]

This patch seems incorrect, though I'm not unhappy with the idea

mlir/docs/OpDefinitions.md

1449

I don't think this is the check - this just checks that there's a bit in either input that's set. I think this should and bits and bit together, then cast to the underlying type and check if the result isn't zero.

mlir/tools/mlir-tblgen/EnumsGen.cpp

165

Also, this should be generic to the enum using {0} and {1}

[hanchung]

mlir/docs/OpDefinitions.md

1449

Sorry that I copied wrong line. You're right on it.

[hanchung]

address issues

[hanchung]

fix EnumsGen.cpp

[hanchung]

fix EnumsGen.cpp

[krzysz00]

LGTM

Could we also have a bitEnumContains that asserts the bit only has one bit set, for backwards compatibility? That might be a separate patch.

[hanchung]

update existing uses

[hanchung]

update tests

mlir/docs/OpDefinitions.md

 }
 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) {
+inline constexpr bool bitEnumContainsAll(MyBitEnum bits, MyBitEnum bit) {
   return (bits & bit) == bit;
 }
+inline constexpr bool bitEnumContainsAny(MyBitEnum bits, MyBitEnum bit) {
+  return (static_cast<uint32_t>(bits) & static_cast<uint32_t>(bit)) != 0;
+}
 inline constexpr MyBitEnum bitEnumClear(MyBitEnum bits, MyBitEnum bit) {
   return bits & ~bit;
 }
 
 inline std::string stringifyEnum(MyBitEnum enumValue) {
   return stringifyMyBitEnum(enumValue);
 }
 

mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp

if (parser.parseOptionalLSquare()) {
     return success();
   }
 
   spirv::MemoryAccess memoryAccessAttr;
   if (parseEnumStrAttr<spirv::MemoryAccessAttr>(memoryAccessAttr, parser, state,
                                                 kMemoryAccessAttrName))
     return failure();
 
-  if (spirv::bitEnumContains(memoryAccessAttr, spirv::MemoryAccess::Aligned)) {
+  if (spirv::bitEnumContainsAll(memoryAccessAttr,
+                                spirv::MemoryAccess::Aligned)) {
     // Parse integer attribute for alignment.
     Attribute alignmentAttr;
     Type i32Type = parser.getBuilder().getIntegerType(32);
     if (parser.parseComma() ||
         parser.parseAttribute(alignmentAttr, i32Type, kAlignmentAttrName,
                               state.attributes)) {
       return failure();
     }
if (parser.parseOptionalLSquare()) {
     return success();
   }
 
   spirv::MemoryAccess memoryAccessAttr;
   if (parseEnumStrAttr<spirv::MemoryAccessAttr>(memoryAccessAttr, parser, state,
                                                 kSourceMemoryAccessAttrName))
     return failure();
 
-  if (spirv::bitEnumContains(memoryAccessAttr, spirv::MemoryAccess::Aligned)) {
+  if (spirv::bitEnumContainsAll(memoryAccessAttr,
+                                spirv::MemoryAccess::Aligned)) {
     // Parse integer attribute for alignment.
     Attribute alignmentAttr;
     Type i32Type = parser.getBuilder().getIntegerType(32);
     if (parser.parseComma() ||
         parser.parseAttribute(alignmentAttr, i32Type, kSourceAlignmentAttrName,
                               state.attributes)) {
       return failure();
     }
static void printMemoryAccessAttribute(
     Optional<uint32_t> alignmentAttrValue = None) {
   // Print optional memory access attribute.
   if (auto memAccess = (memoryAccessAtrrValue ? memoryAccessAtrrValue
                                               : memoryOp.memory_access())) {
     elidedAttrs.push_back(kMemoryAccessAttrName);
 
     printer << " [\"" << stringifyMemoryAccess(*memAccess) << "\"";
 
-    if (spirv::bitEnumContains(*memAccess, spirv::MemoryAccess::Aligned)) {
+    if (spirv::bitEnumContainsAll(*memAccess, spirv::MemoryAccess::Aligned)) {
       // Print integer alignment attribute.
       if (auto alignment = (alignmentAttrValue ? alignmentAttrValue
                                                : memoryOp.alignment())) {
         elidedAttrs.push_back(kAlignmentAttrName);
         printer << ", " << alignment;
       }
     }
     printer << "]";
static void printSourceMemoryAccessAttribute(
 
   // Print optional memory access attribute.
   if (auto memAccess = (memoryAccessAtrrValue ? memoryAccessAtrrValue
                                               : memoryOp.memory_access())) {
     elidedAttrs.push_back(kSourceMemoryAccessAttrName);
 
     printer << " [\"" << stringifyMemoryAccess(*memAccess) << "\"";
 
-    if (spirv::bitEnumContains(*memAccess, spirv::MemoryAccess::Aligned)) {
+    if (spirv::bitEnumContainsAll(*memAccess, spirv::MemoryAccess::Aligned)) {
       // Print integer alignment attribute.
       if (auto alignment = (alignmentAttrValue ? alignmentAttrValue
                                                : memoryOp.alignment())) {
         elidedAttrs.push_back(kSourceAlignmentAttrName);
         printer << ", " << alignment;
       }
     }
     printer << "]";
spirv::ImageOperands noSupportOperands =
       spirv::ImageOperands::Bias | spirv::ImageOperands::Lod |
       spirv::ImageOperands::Grad | spirv::ImageOperands::ConstOffset |
       spirv::ImageOperands::Offset | spirv::ImageOperands::ConstOffsets |
       spirv::ImageOperands::Sample | spirv::ImageOperands::MinLod |
       spirv::ImageOperands::MakeTexelAvailable |
       spirv::ImageOperands::MakeTexelVisible |
       spirv::ImageOperands::SignExtend | spirv::ImageOperands::ZeroExtend;
 
-  if (spirv::bitEnumContains(attr.getValue(), noSupportOperands))
+  if (spirv::bitEnumContainsAll(attr.getValue(), noSupportOperands))
     llvm_unreachable("unimplemented operands of Image Operands");
 
   return success();
 }
 
 static LogicalResult verifyCastOp(Operation *op,
                                   bool requireSameBitWidth = true,
                                   bool skipBitWidthCheck = false) {
static LogicalResult verifyMemoryAccessAttribute(MemoryOpTy memoryOp) {
 
   auto memAccess = memAccessAttr.template cast<spirv::MemoryAccessAttr>();
 
   if (!memAccess) {
     return memoryOp.emitOpError("invalid memory access specifier: ")
            << memAccessAttr;
   }
 
-  if (spirv::bitEnumContains(memAccess.getValue(),
+  if (spirv::bitEnumContainsAll(memAccess.getValue(),
                                 spirv::MemoryAccess::Aligned)) {
     if (!op->getAttr(kAlignmentAttrName)) {
       return memoryOp.emitOpError("missing alignment value");
     }
   } else {
     if (op->getAttr(kAlignmentAttrName)) {
       return memoryOp.emitOpError(
           "invalid alignment specification with non-aligned memory access "
           "specification");
static LogicalResult verifySourceMemoryAccessAttribute(MemoryOpTy memoryOp) {
 
   auto memAccess = memAccessAttr.template cast<spirv::MemoryAccessAttr>();
 
   if (!memAccess) {
     return memoryOp.emitOpError("invalid memory access specifier: ")
            << memAccess;
   }
 
-  if (spirv::bitEnumContains(memAccess.getValue(),
+  if (spirv::bitEnumContainsAll(memAccess.getValue(),
                                 spirv::MemoryAccess::Aligned)) {
     if (!op->getAttr(kSourceAlignmentAttrName)) {
       return memoryOp.emitOpError("missing alignment value");
     }
   } else {
     if (op->getAttr(kSourceAlignmentAttrName)) {
       return memoryOp.emitOpError(
           "invalid alignment specification with non-aligned memory access "
           "specification");

mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp

const std::pair<FastmathFlags, FuncT> handlers[] = {
       {FastmathFlags::contract, &llvmFMF::setAllowContract},
       {FastmathFlags::afn,      &llvmFMF::setApproxFunc},
       {FastmathFlags::reassoc,  &llvmFMF::setAllowReassoc},
       // clang-format on
   };
   llvm::FastMathFlags ret;
   auto fmf = op.getFastmathFlags();
   for (auto it : handlers)
-    if (bitEnumContains(fmf, it.first))
+    if (bitEnumContainsAll(fmf, it.first))
       (ret.*(it.second))(true);
   return ret;
 }
 
 /// Returns an LLVM metadata node corresponding to a loop option. This metadata
 /// is attached to an llvm.loop node.
 static llvm::MDNode *getLoopOptionMetadata(llvm::LLVMContext &ctx,
                                            LoopOptionCase option,

mlir/tools/mlir-tblgen/EnumsGen.cpp

 }
 
 // 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 <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 bool bitEnumContainsAll(<enum-type> bits, <enum-type> bit);
+// inline constexpr bool bitEnumContainsAny(<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());
   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) {{
+inline constexpr bool bitEnumContainsAll({0} bits, {0} bit) {{
   return (bits & bit) == bit;
 }
+inline constexpr bool bitEnumContainsAny({0} bits, {0} bit) {{
+  return (static_cast<{1}>(bits) & static_cast<{1}>(bit)) != 0;
+}
 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);

mlir/unittests/TableGen/EnumsGenTest.cpp

TEST(EnumsGenTest, GeneratedSymbolToStringFnForPrimaryGroupBitEnum) {
   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,
+  EXPECT_TRUE(bitEnumContainsAll(BitEnumWithNone::Bit0 | BitEnumWithNone::Bit3,
                                  BitEnumWithNone::Bit0));
-  EXPECT_FALSE(bitEnumContains(BitEnumWithNone::Bit0 & BitEnumWithNone::Bit3,
+  EXPECT_FALSE(bitEnumContainsAll(BitEnumWithNone::Bit0 & BitEnumWithNone::Bit3,
                                   BitEnumWithNone::Bit0));
 }
 
 TEST(EnumsGenTest, GeneratedSymbolToCustomStringFn) {
   EXPECT_EQ(stringifyPrettyIntEnum(PrettyIntEnum::Case1), "case_one");
   EXPECT_EQ(stringifyPrettyIntEnum(PrettyIntEnum::Case2), "case_two");
 }
 
 TEST(EnumsGenTest, GeneratedCustomStringToSymbolFn) {