Fix raw buffer size when creating a DenseElementsAttr from an array of attributes.

PiperOrigin-RevId: 229973134

diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp
index 1938ddce668e831e8bb854e0196249f00f7f48db..6b03041257ba74969327314facf5b58af3c58228 100644 (file)
--- a/mlir/lib/IR/MLIRContext.cpp
+++ b/mlir/lib/IR/MLIRContext.cpp
@@ -1078,6 +1078,8 @@ DenseElementsAttr DenseElementsAttr::get(VectorOrTensorType type,
                                          ArrayRef<Attribute> values) {
   assert(type.getElementType().isIntOrFloat() &&
          "expected int or float element type");
+  assert(values.size() == type.getNumElements() &&
+         "expected 'values' to contain the same number of elements as 'type'");
 
   // FIXME: using 64 bits for BF16 because it is currently stored with double
   // semantics.
@@ -1085,7 +1087,7 @@ DenseElementsAttr DenseElementsAttr::get(VectorOrTensorType type,
   size_t bitWidth = eltType.isBF16() ? 64 : eltType.getIntOrFloatBitWidth();
 
   // Compress the attribute values into a character buffer.
-  SmallVector<char, 8> data(type.getSizeInBits() * 8L);
+  SmallVector<char, 8> data(APInt::getNumWords(bitWidth * values.size()) * 8L);
   for (unsigned i = 0, e = values.size(); i < e; ++i) {
     unsigned bitPos = i * bitWidth;