[mlir] Fix DenseElementsAttr::mapValues(i1, splat).

[wecing]

Summary

Splat of bool is encoded as a byte with all-ones in it [1]. Without this
change, this piece of code:

auto xs = builder.getI32TensorAttr({42, 42, 42, 42});
auto xs2 = xs.mapValues(builder.getI1Type(), [](const llvm::APInt &x) {
  return x.isZero() ? llvm::APInt::getZero(1) : llvm::APInt::getAllOnes(1);
});
xs2.dump();

Prints:

dense<[true, false, false, false]> : tensor<4xi1>

Because only the first bit is set. This applies to both
DenseIntElementsAttr::mapValues() and DenseFPElementsAttr::mapValues().

[1]: https://github.com/llvm/llvm-project/blob/e877b42e2c70813352c1963ea33e992f481d5cba/mlir/lib/IR/BuiltinAttributes.cpp#L984

[rriddle]

Is there a test that you could add? Maybe in mlir/unittests/AttributeTest.cpp?

mlir/lib/IR/BuiltinAttributes.cpp

1524–1525

This could likely be applied within the if (attr.isSplat()) { branch below, after the call to processElt; i.e. you could call and then check the result, fixing if necessary.

[wecing]

Add tests and move special handling code

[wecing]

Include all local changes

[wecing]

Thank you! Since I don't have write access, could you submit the patch for me?

mlir/lib/IR/BuiltinAttributes.cpp

static ShapedType mappingHelper(Fn mapping, Attr &attr, ShapedType inType,
 
   size_t numRawElements = attr.isSplat() ? 1 : newArrayType.getNumElements();
   data.resize(llvm::divideCeil(storageBitWidth * numRawElements, CHAR_BIT));
 
   // Functor used to process a single element value of the attribute.
   auto processElt = [&](decltype(*attr.begin()) value, size_t index) {
     auto newInt = mapping(value);
     assert(newInt.getBitWidth() == bitWidth);
     writeBits(data.data(), index * storageBitWidth, newInt);
   };
 
   // Check for the splat case.
   if (attr.isSplat()) {
+    if (bitWidth == 1) {
+      // Handle the special encoding of splat of bool.
+      data[0] = mapping(*attr.begin()).isZero() ? 0 : -1;
+    } else {
       processElt(*attr.begin(), /*index=*/0);
+    }
     return newArrayType;
   }
 
   // Otherwise, process all of the element values.
   uint64_t elementIdx = 0;
   for (auto value : attr)
     processElt(value, elementIdx++);
   return newArrayType;

mlir/unittests/IR/AttributeTest.cpp

TEST(DenseScalarTest, ExtractZeroRankElement) {
   const int elementValue = 12;
   IntegerType intTy = IntegerType::get(&context, 32);
   Attribute value = IntegerAttr::get(intTy, elementValue);
   RankedTensorType shape = RankedTensorType::get({}, intTy);
 
   auto attr = DenseElementsAttr::get(shape, llvm::makeArrayRef({elementValue}));
   EXPECT_TRUE(attr.getValues<Attribute>()[0] == value);
 }
+
+TEST(DenseSplatMapValuesTest, I32ToTrue) {
+  MLIRContext context;
+  const int elementValue = 12;
+  IntegerType boolTy = IntegerType::get(&context, 1);
+  IntegerType intTy = IntegerType::get(&context, 32);
+  RankedTensorType shape = RankedTensorType::get({4}, intTy);
+
+  auto attr =
+      DenseElementsAttr::get(shape, llvm::makeArrayRef({elementValue}))
+          .mapValues(boolTy, [](const APInt &x) {
+            return x.isZero() ? APInt::getZero(1) : APInt::getAllOnes(1);
+          });
+  EXPECT_EQ(attr.getNumElements(), 4);
+  EXPECT_TRUE(attr.isSplat());
+  EXPECT_TRUE(attr.getSplatValue<BoolAttr>().getValue());
+}
+
+TEST(DenseSplatMapValuesTest, I32ToFalse) {
+  MLIRContext context;
+  const int elementValue = 0;
+  IntegerType boolTy = IntegerType::get(&context, 1);
+  IntegerType intTy = IntegerType::get(&context, 32);
+  RankedTensorType shape = RankedTensorType::get({4}, intTy);
+
+  auto attr =
+      DenseElementsAttr::get(shape, llvm::makeArrayRef({elementValue}))
+          .mapValues(boolTy, [](const APInt &x) {
+            return x.isZero() ? APInt::getZero(1) : APInt::getAllOnes(1);
+          });
+  EXPECT_EQ(attr.getNumElements(), 4);
+  EXPECT_TRUE(attr.isSplat());
+  EXPECT_FALSE(attr.getSplatValue<BoolAttr>().getValue());
+}
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // DenseResourceElementsAttr
 //===----------------------------------------------------------------------===//
 
 template <typename AttrT, typename T>
 static void checkNativeAccess(MLIRContext *ctx, ArrayRef<T> data,