https://github.com/matthias-springer updated https://github.com/llvm/llvm-project/pull/180397
>From 7b02e6ccea554e67742418f69b1a8c4b15e74aa4 Mon Sep 17 00:00:00 2001 From: Matthias Springer <[email protected]> Date: Sun, 8 Feb 2026 09:07:31 +0000 Subject: [PATCH] [mlir][IR] `DenseElementsAttr`: Remove `i1` dense packing special case --- .../mlir/Bindings/Python/IRAttributes.h | 12 --- mlir/lib/Bindings/Python/IRAttributes.cpp | 98 ++----------------- mlir/lib/IR/AttributeDetail.h | 52 +--------- mlir/lib/IR/BuiltinAttributes.cpp | 60 ++---------- mlir/test/IR/attribute-roundtrip.mlir | 10 -- mlir/test/IR/parse-literal.mlir | 8 +- mlir/test/python/ir/array_attributes.py | 7 +- mlir/unittests/IR/AttributeTest.cpp | 14 --- 8 files changed, 20 insertions(+), 241 deletions(-) delete mode 100644 mlir/test/IR/attribute-roundtrip.mlir diff --git a/mlir/include/mlir/Bindings/Python/IRAttributes.h b/mlir/include/mlir/Bindings/Python/IRAttributes.h index d5d5548602114..6c272f226e5a9 100644 --- a/mlir/include/mlir/Bindings/Python/IRAttributes.h +++ b/mlir/include/mlir/Bindings/Python/IRAttributes.h @@ -442,18 +442,6 @@ class MLIR_PYTHON_API_EXPORTED PyDenseElementsAttribute const std::optional<std::vector<int64_t>> &explicitShape, MlirContext &context); - // There is a complication for boolean numpy arrays, as numpy represents - // them as 8 bits (1 byte) per boolean, whereas MLIR bitpacks them into 8 - // booleans per byte. - static MlirAttribute getBitpackedAttributeFromBooleanBuffer( - Py_buffer &view, std::optional<std::vector<int64_t>> explicitShape, - MlirContext &context); - - // This does the opposite transformation of - // `getBitpackedAttributeFromBooleanBuffer` - std::unique_ptr<nb_buffer_info> - getBooleanBufferFromBitpackedAttribute() const; - template <typename Type> std::unique_ptr<nb_buffer_info> bufferInfo(MlirType shapedType, const char *explicitFormat = nullptr) { diff --git a/mlir/lib/Bindings/Python/IRAttributes.cpp b/mlir/lib/Bindings/Python/IRAttributes.cpp index c271497fcc9f8..16e406c244894 100644 --- a/mlir/lib/Bindings/Python/IRAttributes.cpp +++ b/mlir/lib/Bindings/Python/IRAttributes.cpp @@ -45,16 +45,12 @@ For conversions outside of these types, a `type=` must be explicitly provided and the buffer contents must be bit-castable to the MLIR internal representation: - * Integer types (except for i1): the buffer must be byte aligned to the - next byte boundary. + * Integer types: the buffer must be byte aligned to the next byte boundary. * Floating point types: Must be bit-castable to the given floating point size. - * i1 (bool): Bit packed into 8bit words where the bit pattern matches a - row major ordering. An arbitrary Numpy `bool_` array can be bit packed to - this specification with: `np.packbits(ary, axis=None, bitorder='little')`. + * i1 (bool): Each boolean value is stored as a single byte (0 or 1). -If a single element buffer is passed (or for i1, a single byte with value 0 -or 255), then a splat will be created. +If a single element buffer is passed, then a splat will be created. Args: array: The array or buffer to convert. @@ -63,8 +59,7 @@ or 255), then a splat will be created. type: Skips inference of the MLIR element type and uses this instead. The storage size must be consistent with the actual contents of the buffer. shape: Overrides the shape of the buffer when constructing the MLIR - shaped type. This is needed when the physical and logical shape differ (as - for i1). + shaped type. This is needed when the physical and logical shape differ. context: Explicit context, if not from context manager. Returns: @@ -736,10 +731,7 @@ std::unique_ptr<nb_buffer_info> PyDenseElementsAttribute::accessBuffer() { } else if (mlirTypeIsAInteger(elementType) && mlirIntegerTypeGetWidth(elementType) == 1) { // i1 / bool - // We can not send the buffer directly back to Python, because the i1 - // values are bitpacked within MLIR. We call numpy's unpackbits function - // to convert the bytes. - return getBooleanBufferFromBitpackedAttribute(); + return bufferInfo<bool>(shapedType); } // TODO: Currently crashes the program. @@ -853,9 +845,7 @@ MlirAttribute PyDenseElementsAttribute::getAttributeFromBuffer( bulkLoadElementType = mlirF16TypeGet(context); } else if (format == "?") { // i1 - // The i1 type needs to be bit-packed, so we will handle it separately - return getBitpackedAttributeFromBooleanBuffer(view, explicitShape, - context); + bulkLoadElementType = mlirIntegerTypeGet(context, 1); } else if (isSignedIntegerFormat(format)) { if (view.itemsize == 4) { // i32 @@ -907,82 +897,6 @@ MlirAttribute PyDenseElementsAttribute::getAttributeFromBuffer( return mlirDenseElementsAttrRawBufferGet(type, view.len, view.buf); } -MlirAttribute PyDenseElementsAttribute::getBitpackedAttributeFromBooleanBuffer( - Py_buffer &view, std::optional<std::vector<int64_t>> explicitShape, - MlirContext &context) { - if (llvm::endianness::native != llvm::endianness::little) { - // Given we have no good way of testing the behavior on big-endian - // systems we will throw - throw nb::type_error("Constructing a bit-packed MLIR attribute is " - "unsupported on big-endian systems"); - } - nb::ndarray<uint8_t, nb::numpy, nb::ndim<1>, nb::c_contig> unpackedArray( - /*data=*/static_cast<uint8_t *>(view.buf), - /*shape=*/{static_cast<size_t>(view.len)}); - - nb::module_ numpy = nb::module_::import_("numpy"); - nb::object packbitsFunc = numpy.attr("packbits"); - nb::object packedBooleans = - packbitsFunc(nb::cast(unpackedArray), "bitorder"_a = "little"); - nb_buffer_info pythonBuffer = nb::cast<nb_buffer>(packedBooleans).request(); - - MlirType bitpackedType = getShapedType(mlirIntegerTypeGet(context, 1), - std::move(explicitShape), view); - assert(pythonBuffer.itemsize == 1 && "Packbits must return uint8"); - // Notice that `mlirDenseElementsAttrRawBufferGet` copies the memory of - // packedBooleans, hence the MlirAttribute will remain valid even when - // packedBooleans get reclaimed by the end of the function. - return mlirDenseElementsAttrRawBufferGet(bitpackedType, pythonBuffer.size, - pythonBuffer.ptr); -} - -std::unique_ptr<nb_buffer_info> -PyDenseElementsAttribute::getBooleanBufferFromBitpackedAttribute() const { - if (llvm::endianness::native != llvm::endianness::little) { - // Given we have no good way of testing the behavior on big-endian - // systems we will throw - throw nb::type_error("Constructing a numpy array from a MLIR attribute " - "is unsupported on big-endian systems"); - } - - int64_t numBooleans = mlirElementsAttrGetNumElements(*this); - int64_t numBitpackedBytes = llvm::divideCeil(numBooleans, 8); - uint8_t *bitpackedData = static_cast<uint8_t *>( - const_cast<void *>(mlirDenseElementsAttrGetRawData(*this))); - nb::ndarray<uint8_t, nb::numpy, nb::ndim<1>, nb::c_contig> packedArray( - /*data=*/bitpackedData, - /*shape=*/{static_cast<size_t>(numBitpackedBytes)}); - - nb::module_ numpy = nb::module_::import_("numpy"); - nb::object unpackbitsFunc = numpy.attr("unpackbits"); - nb::object equalFunc = numpy.attr("equal"); - nb::object reshapeFunc = numpy.attr("reshape"); - nb::object unpackedBooleans = - unpackbitsFunc(nb::cast(packedArray), "bitorder"_a = "little"); - - // Unpackbits operates on bytes and gives back a flat 0 / 1 integer array. - // We need to: - // 1. Slice away the padded bits - // 2. Make the boolean array have the correct shape - // 3. Convert the array to a boolean array - unpackedBooleans = unpackedBooleans[nb::slice( - nb::int_(0), nb::int_(numBooleans), nb::int_(1))]; - unpackedBooleans = equalFunc(unpackedBooleans, 1); - - MlirType shapedType = mlirAttributeGetType(*this); - intptr_t rank = mlirShapedTypeGetRank(shapedType); - std::vector<intptr_t> shape(rank); - for (intptr_t i = 0; i < rank; ++i) { - shape[i] = mlirShapedTypeGetDimSize(shapedType, i); - } - unpackedBooleans = reshapeFunc(unpackedBooleans, shape); - - // Make sure the returned nb::buffer_view claims ownership of the data - // in `pythonBuffer` so it remains valid when Python reads it - nb_buffer pythonBuffer = nb::cast<nb_buffer>(unpackedBooleans); - return std::make_unique<nb_buffer_info>(pythonBuffer.request()); -} - PyType_Slot PyDenseElementsAttribute::slots[] = { // Python 3.8 doesn't allow setting the buffer protocol slots from a type spec. #if PY_VERSION_HEX >= 0x03090000 diff --git a/mlir/lib/IR/AttributeDetail.h b/mlir/lib/IR/AttributeDetail.h index 7af5c8cd9191d..c60886bc061ce 100644 --- a/mlir/lib/IR/AttributeDetail.h +++ b/mlir/lib/IR/AttributeDetail.h @@ -33,10 +33,6 @@ namespace detail { /// Return the bit width which DenseElementsAttr should use for this type. inline size_t getDenseElementBitWidth(Type eltType) { - // i1 is stored as a single bit (bit-packed storage). - if (eltType.isInteger(1)) - return 1; - // Check for DenseElementTypeInterface. if (auto denseEltType = llvm::dyn_cast<DenseElementType>(eltType)) return denseEltType.getDenseElementBitSize(); llvm_unreachable("unsupported element type"); @@ -92,10 +88,7 @@ struct DenseIntOrFPElementsAttrStorage : public DenseElementsAttributeStorage { // If the data is already known to be a splat, the key hash value is // directly the data buffer. - 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); } @@ -105,12 +98,8 @@ struct DenseIntOrFPElementsAttrStorage : public DenseElementsAttributeStorage { 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 (isBoolData) - return getKeyForBoolData(ty, data, numElements); - size_t elementWidth = getDenseElementBitWidth(ty.getElementType()); - // Non 1-bit dense elements are padded to 8-bits. + // 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"); @@ -129,45 +118,6 @@ struct DenseIntOrFPElementsAttrStorage : public DenseElementsAttributeStorage { return KeyTy(ty, firstElt, hashVal, /*isSplat=*/true); } - /// 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; - - // 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 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; }) - ? 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); diff --git a/mlir/lib/IR/BuiltinAttributes.cpp b/mlir/lib/IR/BuiltinAttributes.cpp index d9c5fd9acb811..b2f5853269f0a 100644 --- a/mlir/lib/IR/BuiltinAttributes.cpp +++ b/mlir/lib/IR/BuiltinAttributes.cpp @@ -460,7 +460,7 @@ const char DenseIntOrFPElementsAttrStorage::kSplatFalse = 0; /// Get the bitwidth of a dense element type within the buffer. /// DenseElementsAttr requires bitwidths greater than 1 to be aligned by 8. static size_t getDenseElementStorageWidth(size_t origWidth) { - return origWidth == 1 ? origWidth : llvm::alignTo<8>(origWidth); + return llvm::alignTo<8>(origWidth); } static size_t getDenseElementStorageWidth(Type elementType) { return getDenseElementStorageWidth(getDenseElementBitWidth(elementType)); @@ -622,12 +622,6 @@ Attribute DenseElementsAttr::AttributeElementIterator::operator*() const { auto owner = llvm::cast<DenseElementsAttr>(getFromOpaquePointer(base)); Type eltTy = owner.getElementType(); - // Handle i1 (boolean) specially - it's bit-packed and doesn't use interface. - if (eltTy.isInteger(1)) { - bool value = *BoolElementIterator(owner, index); - return IntegerAttr::get(eltTy, APInt(1, value)); - } - // Handle strings specially. if (llvm::isa<DenseStringElementsAttr>(owner)) { ArrayRef<StringRef> vals = owner.getRawStringData(); @@ -654,7 +648,7 @@ DenseElementsAttr::BoolElementIterator::BoolElementIterator( attr.getRawData().data(), attr.isSplat(), dataIndex) {} bool DenseElementsAttr::BoolElementIterator::operator*() const { - return getBit(getData(), getDataIndex()); + return static_cast<bool>(getData()[getDataIndex()]); } //===----------------------------------------------------------------------===// @@ -900,18 +894,8 @@ bool DenseElementsAttr::classof(Attribute attr) { DenseElementsAttr DenseElementsAttr::get(ShapedType type, ArrayRef<Attribute> values) { assert(hasSameNumElementsOrSplat(type, values)); - Type eltType = type.getElementType(); - // Handle i1 (boolean) specially - it's bit-packed. - if (eltType.isInteger(1)) { - SmallVector<bool> boolValues; - boolValues.reserve(values.size()); - for (Attribute attr : values) - boolValues.push_back(llvm::cast<IntegerAttr>(attr).getValue().isOne()); - return get(type, boolValues); - } - // Handle strings specially. if (!llvm::isa<DenseElementType>(eltType)) { SmallVector<StringRef, 8> stringValues; @@ -941,25 +925,9 @@ DenseElementsAttr DenseElementsAttr::get(ShapedType type, ArrayRef<bool> values) { assert(hasSameNumElementsOrSplat(type, values)); assert(type.getElementType().isInteger(1)); - - SmallVector<char> buff(llvm::divideCeil(values.size(), CHAR_BIT)); - - if (!values.empty()) { - bool isSplat = true; - bool firstValue = values[0]; - for (int i = 0, e = values.size(); i != e; ++i) { - isSplat &= values[i] == firstValue; - setBit(buff.data(), i, values[i]); - } - - // Splat of bool is encoded as a byte with all-ones in it. - if (isSplat) { - buff.resize(1); - buff[0] = values[0] ? -1 : 0; - } - } - - return DenseIntOrFPElementsAttr::getRaw(type, buff); + return DenseIntOrFPElementsAttr::getRaw( + type, ArrayRef<char>(reinterpret_cast<const char *>(values.data()), + values.size())); } DenseElementsAttr DenseElementsAttr::get(ShapedType type, @@ -1030,23 +998,7 @@ bool DenseElementsAttr::isValidRawBuffer(ShapedType type, // The initializer is always a splat if the result type has a single element. detectedSplat = numElements == 1; - // Storage width of 1 is special as it is packed by the bit. - if (storageWidth == 1) { - // Check for a splat, or a buffer equal to the number of elements which - // consists of either all 0's or all 1's. - if (rawBuffer.size() == 1) { - auto rawByte = static_cast<uint8_t>(rawBuffer[0]); - if (rawByte == 0 || rawByte == 0xff) { - detectedSplat = true; - return true; - } - } - - // This is a valid non-splat buffer if it has the right size. - return rawBufferWidth == llvm::alignTo<8>(numElements); - } - - // All other types are 8-bit aligned, so we can just check the buffer width + // All types are 8-bit aligned, so we can just check the buffer width // to know if only a single initializer element was passed in. if (rawBufferWidth == storageWidth) { detectedSplat = true; diff --git a/mlir/test/IR/attribute-roundtrip.mlir b/mlir/test/IR/attribute-roundtrip.mlir deleted file mode 100644 index 974dbcae6cf0a..0000000000000 --- a/mlir/test/IR/attribute-roundtrip.mlir +++ /dev/null @@ -1,10 +0,0 @@ -// RUN: mlir-opt -canonicalize %s | mlir-opt | FileCheck %s - -// CHECK-LABEL: @large_i1_tensor_roundtrip -func.func @large_i1_tensor_roundtrip() -> tensor<160xi1> { - %cst_0 = arith.constant dense<"0xFFF00000FF000000FF000000FF000000FF000000"> : tensor<160xi1> - %cst_1 = arith.constant dense<"0xFF000000FF000000FF000000FF000000FF0000F0"> : tensor<160xi1> - // CHECK: dense<"0xFF000000FF000000FF000000FF000000FF000000"> - %0 = arith.andi %cst_0, %cst_1 : tensor<160xi1> - return %0 : tensor<160xi1> -} diff --git a/mlir/test/IR/parse-literal.mlir b/mlir/test/IR/parse-literal.mlir index 71b25e1d86480..36867c56075d0 100644 --- a/mlir/test/IR/parse-literal.mlir +++ b/mlir/test/IR/parse-literal.mlir @@ -36,8 +36,8 @@ func.func @parse_i4_tensor() -> tensor<32xi4> { } // CHECK-LABEL: @parse_i1_tensor -func.func @parse_i1_tensor() -> tensor<256xi1> { - // CHECK: dense<"0x0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F"> : tensor<256xi1> - %0 = arith.constant dense<"0x0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F0F"> : tensor<256xi1> - return %0 : tensor<256xi1> +func.func @parse_i1_tensor() -> tensor<32xi1> { + // CHECK: dense<[true, false, true, false, true, true, true, true, true, true, true, true, true, true, true, true, true, true, true, true, true, true, true, true, true, false, false, false, false, false, false, true]> : tensor<32xi1> + %0 = arith.constant dense<"0x0100010001010101010101010101010101010101010101010100000000000001"> : tensor<32xi1> + return %0 : tensor<32xi1> } diff --git a/mlir/test/python/ir/array_attributes.py b/mlir/test/python/ir/array_attributes.py index 66f7ec8e7fff1..c1e2dd5f5ae5e 100644 --- a/mlir/test/python/ir/array_attributes.py +++ b/mlir/test/python/ir/array_attributes.py @@ -258,11 +258,10 @@ def testGetDenseElementsInteger4(): @run def testGetDenseElementsBool(): with Context(): - bool_array = np.array([[1, 0, 1], [0, 1, 0]], dtype=np.bool_) - array = np.packbits(bool_array, axis=None, bitorder="little") - attr = DenseElementsAttr.get( - array, type=IntegerType.get_signless(1), shape=bool_array.shape + bool_array = np.array( + [[True, False, True], [False, True, False]], dtype=np.bool_ ) + attr = DenseElementsAttr.get(bool_array) # CHECK: dense<{{\[}}[true, false, true], [false, true, false]]> : tensor<2x3xi1> print(attr) diff --git a/mlir/unittests/IR/AttributeTest.cpp b/mlir/unittests/IR/AttributeTest.cpp index fd40404bf3008..404aa8c0dcf3d 100644 --- a/mlir/unittests/IR/AttributeTest.cpp +++ b/mlir/unittests/IR/AttributeTest.cpp @@ -76,20 +76,6 @@ TEST(DenseSplatTest, BoolSplatRawRoundtrip) { EXPECT_EQ(trueSplat, trueSplatFromRaw); } -TEST(DenseSplatTest, BoolSplatSmall) { - MLIRContext context; - Builder builder(&context); - - // Check that splats that don't fill entire byte are handled properly. - auto tensorType = RankedTensorType::get({4}, builder.getI1Type()); - std::vector<char> data{0b00001111}; - auto trueSplatFromRaw = - DenseIntOrFPElementsAttr::getFromRawBuffer(tensorType, data); - EXPECT_TRUE(trueSplatFromRaw.isSplat()); - DenseElementsAttr trueSplat = DenseElementsAttr::get(tensorType, true); - EXPECT_EQ(trueSplat, trueSplatFromRaw); -} - TEST(DenseSplatTest, LargeBoolSplat) { constexpr int64_t boolCount = 56; _______________________________________________ llvm-branch-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits
