[mlir] Fix DenseElementsAttr treatment of bool splat of "true"

[rriddle]

Summary

Boolean splats currently can't roundtrip via the "raw" DenseElementsAttr
API. This is because internally we treat true splats in some cases as "1"(one bit set)
and in other cases as "0xFF"(all bits set). This commit cleans up this handling to
consistently use 0xFF (all bits set) as the value for a splat of true.

Depends on D133460

[jpienaar]

Roundtripping catching bugs :)

mlir/lib/IR/AttributeDetail.h

struct KeyTy {
     llvm::hash_code hashCode;
 
     /// A boolean that indicates if this data is a splat or not.
     bool isSplat;
   };
 
   /// Compare this storage instance with the provided key.
   bool operator==(const KeyTy &key) const {
-    if (key.type != type)
-      return false;
-
-    // For boolean splats we need to explicitly check that the first bit is the
-    // same. Boolean values are packed at the bit level, and even though a splat
-    // is detected the rest of the bits in the first byte may differ from the
-    // splat value.
-    if (key.type.getElementType().isInteger(1)) {
-      if (key.isSplat != isSplat)
-        return false;
-      if (isSplat)
-        return (key.data.front() & 1) == data.front();
-    }
-
-    // Otherwise, we can default to just checking the data.
-    return key.data == data;
+    return key.type == type && key.data == data;
   }
 
   /// Construct a key from a shaped type, raw data buffer, and a flag that
   /// signals if the data is already known to be a splat. Callers to this
   /// function are expected to tag preknown splat values when possible, e.g. one
   /// element shapes.
   static KeyTy getKey(ShapedType ty, ArrayRef<char> data, bool isKnownSplat) {
     // Handle an empty storage instance.
     if (data.empty())
       return KeyTy(ty, data, 0);
 
     // If the data is already known to be a splat, the key hash value is
     // directly the data buffer.
-    if (isKnownSplat)
+    bool isBoolData = ty.getElementType().isInteger(1);
+    if (isKnownSplat) {
+      if (isBoolData)
+        return getKeyForSplatBoolData(ty, data[0] != 0);
       return KeyTy(ty, data, llvm::hash_value(data), isKnownSplat);
+    }
 
     // Otherwise, we need to check if the data corresponds to a splat or not.
 
     // Handle the simple case of only one element.
     size_t numElements = ty.getNumElements();
     assert(numElements != 1 && "splat of 1 element should already be detected");
 
     // Handle boolean values directly as they are packed to 1-bit.
-    if (ty.getElementType().isInteger(1) == 1)
+    if (isBoolData)
       return getKeyForBoolData(ty, data, numElements);
 
     size_t elementWidth = getDenseElementBitWidth(ty.getElementType());
     // Non 1-bit dense elements are padded to 8-bits.
     size_t storageSize = llvm::divideCeil(elementWidth, CHAR_BIT);
     assert(((data.size() / storageSize) == numElements) &&
            "data does not hold expected number of elements");
 
struct DenseIntOrFPElementsAttrStorage : public DenseElementsAttributeStorage {
   }
 
   /// Construct a key with a set of boolean data.
   static KeyTy getKeyForBoolData(ShapedType ty, ArrayRef<char> data,
                                  size_t numElements) {
     ArrayRef<char> splatData = data;
     bool splatValue = splatData.front() & 1;
 
-    // Helper functor to generate a KeyTy for a boolean splat value.
-    auto generateSplatKey = [=] {
-      return KeyTy(ty, data.take_front(1),
-                   llvm::hash_value(ArrayRef<char>(splatValue ? 1 : 0)),
-                   /*isSplat=*/true);
-    };
+    // Check the simple case where the data matches the known splat value.
+    if (splatData == ArrayRef<char>(splatValue ? kSplatTrue : kSplatFalse))
+      return getKeyForSplatBoolData(ty, splatValue);
 
     // Handle the case where the potential splat value is 1 and the number of
     // elements is non 8-bit aligned.
     size_t numOddElements = numElements % CHAR_BIT;
     if (splatValue && numOddElements != 0) {
       // Check that all bits are set in the last value.
       char lastElt = splatData.back();
       if (lastElt != llvm::maskTrailingOnes<unsigned char>(numOddElements))
         return KeyTy(ty, data, llvm::hash_value(data));
 
       // If this is the only element, the data is known to be a splat.
       if (splatData.size() == 1)
-        return generateSplatKey();
+        return getKeyForSplatBoolData(ty, splatValue);
       splatData = splatData.drop_back();
     }
 
     // Check that the data buffer corresponds to a splat of the proper mask.
     char mask = splatValue ? ~0 : 0;
     return llvm::all_of(splatData, [mask](char c) { return c == mask; })
-               ? generateSplatKey()
+               ? getKeyForSplatBoolData(ty, splatValue)
                : KeyTy(ty, data, llvm::hash_value(data));
   }
 
+  /// Return a key to use for a boolean splat of the given value.
+  static KeyTy getKeyForSplatBoolData(ShapedType type, bool splatValue) {
+    const char &splatData = splatValue ? kSplatTrue : kSplatFalse;
+    return KeyTy(type, splatData, llvm::hash_value(splatData),
+                 /*isSplat=*/true);
+  }
+
   /// Hash the key for the storage.
   static llvm::hash_code hashKey(const KeyTy &key) {
     return llvm::hash_combine(key.type, key.hashCode);
   }
 
   /// Construct a new storage instance.
   static DenseIntOrFPElementsAttrStorage *
   construct(AttributeStorageAllocator &allocator, KeyTy key) {
     // If the data buffer is non-empty, we copy it into the allocator with a
     // 64-bit alignment.
     ArrayRef<char> copy, data = key.data;
     if (!data.empty()) {
       char *rawData = reinterpret_cast<char *>(
           allocator.allocate(data.size(), alignof(uint64_t)));
       std::memcpy(rawData, data.data(), data.size());
-
-      // If this is a boolean splat, make sure only the first bit is used.
-      if (key.isSplat && key.type.getElementType().isInteger(1))
-        rawData[0] &= 1;
       copy = ArrayRef<char>(rawData, data.size());
     }
 
     return new (allocator.allocate<DenseIntOrFPElementsAttrStorage>())
         DenseIntOrFPElementsAttrStorage(key.type, copy, key.isSplat);
   }
 
   ArrayRef<char> data;
+
+  /// The values used to denote a boolean splat value.
+  static constexpr char kSplatTrue = ~0;
+  static constexpr char kSplatFalse = 0;
 };
 
 /// An attribute representing a reference to a dense vector or tensor object
 /// containing strings.
 struct DenseStringElementsAttrStorage : public DenseElementsAttributeStorage {
   DenseStringElementsAttrStorage(ShapedType ty, ArrayRef<StringRef> data,
                                  bool isSplat = false)
       : DenseElementsAttributeStorage(ty, isSplat), data(data) {}

mlir/unittests/IR/AttributeTest.cpp

TEST(DenseSplatTest, BoolSplat) {
   /// Detect and handle splat within 8 elements (bool values are bit-packed).
   /// True.
   auto detectedSplat = DenseElementsAttr::get(shape, {true, true, true, true});
   EXPECT_EQ(detectedSplat, trueSplat);
   /// False.
   detectedSplat = DenseElementsAttr::get(shape, {false, false, false, false});
   EXPECT_EQ(detectedSplat, falseSplat);
 }
+TEST(DenseSplatTest, BoolSplatRawRoundtrip) {
+  MLIRContext context;
+  IntegerType boolTy = IntegerType::get(&context, 1);
+  RankedTensorType shape = RankedTensorType::get({2, 2}, boolTy);
+
+  // Check that splat booleans properly round trip via the raw API.
+  DenseElementsAttr trueSplat = DenseElementsAttr::get(shape, true);
+  EXPECT_TRUE(trueSplat.isSplat());
+  DenseElementsAttr trueSplatFromRaw =
+      DenseElementsAttr::getFromRawBuffer(shape, trueSplat.getRawData());
+  EXPECT_TRUE(trueSplatFromRaw.isSplat());
+
+  EXPECT_EQ(trueSplat, trueSplatFromRaw);
+}
 
 TEST(DenseSplatTest, LargeBoolSplat) {
   constexpr int64_t boolCount = 56;
 
   MLIRContext context;
   IntegerType boolTy = IntegerType::get(&context, 1);
   RankedTensorType shape = RankedTensorType::get({boolCount}, boolTy);