[llvm-branch-commits] [mlir] [mlir][IR] `DenseElementsAttr`: Remove `i1` dense packing special case (PR #180397)
Matthias Springer via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Sun Feb 8 02:06:14 PST 2026
https://github.com/matthias-springer updated https://github.com/llvm/llvm-project/pull/180397
>From de936681e3d4183e10fe5b0d1aa391e6736f43fb Mon Sep 17 00:00:00 2001
From: Matthias Springer <me at m-sp.org>
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 5ff9afd0875f1..0772dad057a03 100644
--- a/mlir/include/mlir/Bindings/Python/IRAttributes.h
+++ b/mlir/include/mlir/Bindings/Python/IRAttributes.h
@@ -441,18 +441,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 05c0c5e825df3..011db9c147d8e 100644
--- a/mlir/lib/Bindings/Python/IRAttributes.cpp
+++ b/mlir/lib/Bindings/Python/IRAttributes.cpp
@@ -46,16 +46,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.
@@ -64,8 +60,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:
@@ -739,10 +734,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.
@@ -856,9 +848,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
@@ -910,82 +900,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;
More information about the llvm-branch-commits
mailing list