[DenseElementAttr] Simplify the public API for creating these.

[lattner]

Summary

Instead of requiring the client to compute the "isSplat" bit,
compute it internally. This makes the logic more consistent
and defines away a lot of "elements.size()==1" in the clients.

This addresses Issue #55185

[rriddle]

mlir/include/mlir/IR/BuiltinAttributes.td

240

249

257

318

This feels unrelated.

mlir/lib/IR/BuiltinAttributes.cpp

739

Please remove the trailing comment and prefer one placed before the line. This style of this file doesn't use those types of comments.

1084–1085

[lattner]

Thank you for the review!

mlir/include/mlir/IR/BuiltinAttributes.td

318

removed.

mlir/lib/IR/BuiltinAttributes.cpp

1084–1085

Your suggestion isn't correct. I'm not a fan of "isValidRawBuffer", but this is calling it to get the "isSplat" value back out by-reference.

mlir/include/mlir/IR/BuiltinAttributes.h

public:
   ///
   /// The format of the raw buffer is a densely packed array of values that
   /// can be bitcast to the storage format of the element type specified.
   /// Types that are not byte aligned will be:
   ///   - For bitwidth > 1: Rounded up to the next byte.
   ///   - For bitwidth = 1: Packed into 8bit bytes with bits corresponding to
   ///     the linear order of the shape type from MSB to LSB, padded to on the
   ///     right.
-  ///
-  /// If `isSplatBuffer` is true, then the raw buffer should contain a
-  /// single element (or for the case of 1-bit, a single byte of 0 or 255),
-  /// which will be used to construct a splat.
   static DenseElementsAttr getFromRawBuffer(ShapedType type,
-                                            ArrayRef<char> rawBuffer,
-                                            bool isSplatBuffer);
+                                            ArrayRef<char> rawBuffer);
 
   /// Returns true if the given buffer is a valid raw buffer for the given type.
   /// `detectedSplat` is set if the buffer is valid and represents a splat
   /// buffer. The definition may be expanded over time, but currently, a
   /// splat buffer is detected if:
   ///   - For >1bit: The buffer consists of a single element.
   ///   - For 1bit: The buffer consists of a single byte with value 0 or 255.
   ///

mlir/include/mlir/IR/BuiltinAttributes.td

def Builtin_DenseIntOrFPElementsAttr : Builtin_Attr<
 
   protected:
     friend DenseElementsAttr;
 
     /// Constructs a dense elements attribute from an array of raw APFloat
     /// values. Each APFloat value is expected to have the same bitwidth as the
     /// element type of 'type'. 'type' must be a vector or tensor with static
     /// shape.
+    ///
+    /// If the `values` array only has a single element, then this  constructs
+    /// splat of that value.
     static DenseElementsAttr getRaw(ShapedType type, size_t storageWidth,
-                                    ArrayRef<APFloat> values, bool isSplat);
+                                    ArrayRef<APFloat> values);
 
     /// Constructs a dense elements attribute from an array of raw APInt values.
     /// Each APInt value is expected to have the same bitwidth as the element
     /// type of 'type'. 'type' must be a vector or tensor with static shape.
+    ///
+    /// If the `values` array only has a single element, then this  constructs
+    /// splat of that value.
     static DenseElementsAttr getRaw(ShapedType type, size_t storageWidth,
-                                    ArrayRef<APInt> values, bool isSplat);
+                                    ArrayRef<APInt> values);
 
     /// Get or create a new dense elements attribute instance with the given raw
     /// data buffer. 'type' must be a vector or tensor with static shape.
-    static DenseElementsAttr getRaw(ShapedType type, ArrayRef<char> data,
-                                    bool isSplat);
+    ///
+    /// If the `values` array only has a single element, then this  constructs
+    /// splat of that value.
+    static DenseElementsAttr getRaw(ShapedType type, ArrayRef<char> data);
 
     /// Overload of the raw 'get' method that asserts that the given type is of
     /// complex type. This method is used to verify type invariants that the
     /// templatized 'get' method cannot.
     static DenseElementsAttr getRawComplex(ShapedType type, ArrayRef<char> data,
                                            int64_t dataEltSize, bool isInt,
                                            bool isSigned);
 
     /// Overload of the raw 'get' method that asserts that the given type is of
     /// integer or floating-point type. This method is used to verify type
     /// invariants that the templatized 'get' method cannot.
     static DenseElementsAttr getRawIntOrFloat(ShapedType type,
                                               ArrayRef<char> data,
                                               int64_t dataEltSize, bool isInt,
                                               bool isSigned);
-
   public:
   }];
   let genAccessors = 0;
   let genStorageClass = 0;
   let skipDefaultBuilders = 1;
 }
 
 //===----------------------------------------------------------------------===//
let description = [{
     dense<["example1", "example2"]> : tensor<2x!foo.string>
     ```
   }];
   let parameters = (ins AttributeSelfTypeParameter<"", "ShapedType">:$type,
                         "ArrayRef<StringRef>":$value);
   let builders = [
     AttrBuilderWithInferredContext<(ins "ShapedType":$type,
                                         "ArrayRef<StringRef>":$values), [{
-      return $_get(type.getContext(), type, values,
+      return $_get(type.getContext(), type, values, 
                    /* isSplat */(values.size() == 1));
     }]>,
   ];
   let extraClassDeclaration = [{
     using DenseElementsAttr::empty;
     using DenseElementsAttr::getNumElements;
     using DenseElementsAttr::getElementType;
     using DenseElementsAttr::getValues;

mlir/lib/CAPI/IR/BuiltinAttributes.cpp

 MlirAttribute mlirDenseElementsAttrRawBufferGet(MlirType shapedType,
                                                 size_t rawBufferSize,
                                                 const void *rawBuffer) {
   auto shapedTypeCpp = unwrap(shapedType).cast<ShapedType>();
   ArrayRef<char> rawBufferCpp(static_cast<const char *>(rawBuffer),
                               rawBufferSize);
   bool isSplat = false;
   if (!DenseElementsAttr::isValidRawBuffer(shapedTypeCpp, rawBufferCpp,
-                                           isSplat)) {
+                                           isSplat))
     return mlirAttributeGetNull();
-  }
-  return wrap(DenseElementsAttr::getFromRawBuffer(shapedTypeCpp, rawBufferCpp,
-                                                  isSplat));
+  return wrap(DenseElementsAttr::getFromRawBuffer(shapedTypeCpp, rawBufferCpp));
 }
 
 MlirAttribute mlirDenseElementsAttrSplatGet(MlirType shapedType,
                                             MlirAttribute element) {
   return wrap(DenseElementsAttr::get(unwrap(shapedType).cast<ShapedType>(),
                                      unwrap(element)));
 }
 MlirAttribute mlirDenseElementsAttrBoolSplatGet(MlirType shapedType,

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

struct FoldReshapeWithConstant : OpRewritePattern<TensorReshapeOp> {
   LogicalResult matchAndRewrite(TensorReshapeOp reshapeOp,
                                 PatternRewriter &rewriter) const override {
     DenseElementsAttr attr;
     if (!matchPattern(reshapeOp.src(), m_Constant(&attr)))
       return failure();
     if (!attr || !attr.isSplat())
       return failure();
     DenseElementsAttr newAttr = DenseElementsAttr::getFromRawBuffer(
-        reshapeOp.getResultType(), attr.getRawData(), true);
+        reshapeOp.getResultType(), attr.getRawData());
     rewriter.replaceOpWithNewOp<arith::ConstantOp>(reshapeOp, newAttr);
     return success();
   }
 };
 
 /// Reshape of a FromElements can be replaced with a FromElements of the result
 /// type
 template <typename TensorReshapeOp>

mlir/lib/IR/BuiltinAttributes.cpp

else if (eltType.isa<IntegerType, IndexType>())
       intVal = values[i].cast<IntegerAttr>().getValue();
     else
       llvm_unreachable("unexpected element type");
 
     assert(intVal.getBitWidth() == bitWidth &&
            "expected value to have same bitwidth as element type");
     writeBits(data.data(), i * storageBitWidth, intVal);
   }
-  return DenseIntOrFPElementsAttr::getRaw(type, data,
-                                          /*isSplat=*/(values.size() == 1));
+
+  // Handle the special encoding of splat of bool.
+  if (values.size() == 1 && values[0].getType().isInteger(1))
+    data[0] = data[0] ? -1 : 0;
+
+  return DenseIntOrFPElementsAttr::getRaw(type, data);
 }
 
 DenseElementsAttr DenseElementsAttr::get(ShapedType type,
                                          ArrayRef<bool> values) {
   assert(hasSameElementsOrSplat(type, values));
   assert(type.getElementType().isInteger(1));
 
   std::vector<char> buff(llvm::divideCeil(values.size(), CHAR_BIT));
-  for (int i = 0, e = values.size(); i != e; ++i)
+
+  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]);
-  return DenseIntOrFPElementsAttr::getRaw(type, buff,
-                                          /*isSplat=*/(values.size() == 1));
+    }
+
+    if (isSplat) { // special encoding for splat.
+      buff.resize(1);
+      buff[0] = values[0] ? -1 : 0;
+    }
+  }
+
+  return DenseIntOrFPElementsAttr::getRaw(type, buff);
 }
 
 DenseElementsAttr DenseElementsAttr::get(ShapedType type,
                                          ArrayRef<StringRef> values) {
   assert(!type.getElementType().isIntOrFloat());
   return DenseStringElementsAttr::get(type, values);
 }
 
 /// Constructs a dense integer elements attribute from an array of APInt
 /// values. Each APInt value is expected to have the same bitwidth as the
 /// element type of 'type'.
 DenseElementsAttr DenseElementsAttr::get(ShapedType type,
                                          ArrayRef<APInt> values) {
   assert(type.getElementType().isIntOrIndex());
   assert(hasSameElementsOrSplat(type, values));
   size_t storageBitWidth = getDenseElementStorageWidth(type.getElementType());
-  return DenseIntOrFPElementsAttr::getRaw(type, storageBitWidth, values,
-                                          /*isSplat=*/(values.size() == 1));
+  return DenseIntOrFPElementsAttr::getRaw(type, storageBitWidth, values);
 }
 DenseElementsAttr DenseElementsAttr::get(ShapedType type,
                                          ArrayRef<std::complex<APInt>> values) {
   ComplexType complex = type.getElementType().cast<ComplexType>();
   assert(complex.getElementType().isa<IntegerType>());
   assert(hasSameElementsOrSplat(type, values));
   size_t storageBitWidth = getDenseElementStorageWidth(complex) / 2;
   ArrayRef<APInt> intVals(reinterpret_cast<const APInt *>(values.data()),
                           values.size() * 2);
-  return DenseIntOrFPElementsAttr::getRaw(type, storageBitWidth, intVals,
-                                          /*isSplat=*/(values.size() == 1));
+  return DenseIntOrFPElementsAttr::getRaw(type, storageBitWidth, intVals);
 }
 
 // Constructs a dense float elements attribute from an array of APFloat
 // values. Each APFloat value is expected to have the same bitwidth as the
 // element type of 'type'.
 DenseElementsAttr DenseElementsAttr::get(ShapedType type,
                                          ArrayRef<APFloat> values) {
   assert(type.getElementType().isa<FloatType>());
   assert(hasSameElementsOrSplat(type, values));
   size_t storageBitWidth = getDenseElementStorageWidth(type.getElementType());
-  return DenseIntOrFPElementsAttr::getRaw(type, storageBitWidth, values,
-                                          /*isSplat=*/(values.size() == 1));
+  return DenseIntOrFPElementsAttr::getRaw(type, storageBitWidth, values);
 }
 DenseElementsAttr
 DenseElementsAttr::get(ShapedType type,
                        ArrayRef<std::complex<APFloat>> values) {
   ComplexType complex = type.getElementType().cast<ComplexType>();
   assert(complex.getElementType().isa<FloatType>());
   assert(hasSameElementsOrSplat(type, values));
   ArrayRef<APFloat> apVals(reinterpret_cast<const APFloat *>(values.data()),
                            values.size() * 2);
   size_t storageBitWidth = getDenseElementStorageWidth(complex) / 2;
-  return DenseIntOrFPElementsAttr::getRaw(type, storageBitWidth, apVals,
-                                          /*isSplat=*/(values.size() == 1));
+  return DenseIntOrFPElementsAttr::getRaw(type, storageBitWidth, apVals);
 }
 
 /// Construct a dense elements attribute from a raw buffer representing the
 /// data for this attribute. Users should generally not use this methods as
 /// the expected buffer format may not be a form the user expects.
-DenseElementsAttr DenseElementsAttr::getFromRawBuffer(ShapedType type,
-                                                      ArrayRef<char> rawBuffer,
-                                                      bool isSplatBuffer) {
-  return DenseIntOrFPElementsAttr::getRaw(type, rawBuffer, isSplatBuffer);
+DenseElementsAttr
+DenseElementsAttr::getFromRawBuffer(ShapedType type, ArrayRef<char> rawBuffer) {
+  return DenseIntOrFPElementsAttr::getRaw(type, rawBuffer);
 }
 
 /// Returns true if the given buffer is a valid raw buffer for the given type.
 bool DenseElementsAttr::isValidRawBuffer(ShapedType type,
                                          ArrayRef<char> rawBuffer,
                                          bool &detectedSplat) {
   size_t storageWidth = getDenseElementStorageWidth(type.getElementType());
   size_t rawBufferWidth = rawBuffer.size() * CHAR_BIT;
DenseElementsAttr DenseElementsAttr::reshape(ShapedType newType) {
   ShapedType curType = getType();
   if (curType == newType)
     return *this;
 
   assert(newType.getElementType() == curType.getElementType() &&
          "expected the same element type");
   assert(newType.getNumElements() == curType.getNumElements() &&
          "expected the same number of elements");
-  return DenseIntOrFPElementsAttr::getRaw(newType, getRawData(), isSplat());
+  return DenseIntOrFPElementsAttr::getRaw(newType, getRawData());
 }
 
 DenseElementsAttr DenseElementsAttr::resizeSplat(ShapedType newType) {
   assert(isSplat() && "expected a splat type");
 
   ShapedType curType = getType();
   if (curType == newType)
     return *this;
 
   assert(newType.getElementType() == curType.getElementType() &&
          "expected the same element type");
-  return DenseIntOrFPElementsAttr::getRaw(newType, getRawData(), true);
+  return DenseIntOrFPElementsAttr::getRaw(newType, getRawData());
 }
 
 /// Return a new DenseElementsAttr that has the same data as the current
 /// attribute, but has bitcast elements such that it is now 'newType'. The new
 /// type must have the same shape and element types of the same bitwidth as the
 /// current type.
 DenseElementsAttr DenseElementsAttr::bitcast(Type newElType) {
   ShapedType curType = getType();
   Type curElType = curType.getElementType();
   if (curElType == newElType)
     return *this;
 
   assert(getDenseElementBitWidth(newElType) ==
              getDenseElementBitWidth(curElType) &&
          "expected element types with the same bitwidth");
   return DenseIntOrFPElementsAttr::getRaw(curType.clone(newElType),
-                                          getRawData(), isSplat());
+                                          getRawData());
 }
 
 DenseElementsAttr
 DenseElementsAttr::mapValues(Type newElementType,
                              function_ref<APInt(const APInt &)> mapping) const {
   return cast<DenseIntElementsAttr>().mapValues(newElementType, mapping);
 }
 

 //===----------------------------------------------------------------------===//
 // DenseIntOrFPElementsAttr
 //===----------------------------------------------------------------------===//
 
 /// Utility method to write a range of APInt values to a buffer.
 template <typename APRangeT>
 static void writeAPIntsToBuffer(size_t storageWidth, std::vector<char> &data,
                                 APRangeT &&values) {
-  data.resize(llvm::divideCeil(storageWidth * llvm::size(values), CHAR_BIT));
+  size_t numValues = llvm::size(values);
+  data.resize(llvm::divideCeil(storageWidth * numValues, CHAR_BIT));
   size_t offset = 0;
   for (auto it = values.begin(), e = values.end(); it != e;
        ++it, offset += storageWidth) {
     assert((*it).getBitWidth() <= storageWidth);
     writeBits(data.data(), offset, *it);
   }
+
+  // Handle the special encoding of splat of a boolean.
+  if (numValues == 1 && (*values.begin()).getBitWidth() == 1)
+    data[0] = data[0] ? -1 : 0;
 }
 
 /// Constructs a dense elements attribute from an array of raw APFloat values.
 /// Each APFloat value is expected to have the same bitwidth as the element
 /// type of 'type'. 'type' must be a vector or tensor with static shape.
 DenseElementsAttr DenseIntOrFPElementsAttr::getRaw(ShapedType type,
                                                    size_t storageWidth,
-                                                   ArrayRef<APFloat> values,
-                                                   bool isSplat) {
+                                                   ArrayRef<APFloat> values) {
   std::vector<char> data;
   auto unwrapFloat = [](const APFloat &val) { return val.bitcastToAPInt(); };
   writeAPIntsToBuffer(storageWidth, data, llvm::map_range(values, unwrapFloat));
-  return DenseIntOrFPElementsAttr::getRaw(type, data, isSplat);
+  return DenseIntOrFPElementsAttr::getRaw(type, data);
 }
 
 /// Constructs a dense elements attribute from an array of raw APInt values.
 /// Each APInt value is expected to have the same bitwidth as the element type
 /// of 'type'.
 DenseElementsAttr DenseIntOrFPElementsAttr::getRaw(ShapedType type,
                                                    size_t storageWidth,
-                                                   ArrayRef<APInt> values,
-                                                   bool isSplat) {
+                                                   ArrayRef<APInt> values) {
   std::vector<char> data;
   writeAPIntsToBuffer(storageWidth, data, values);
-  return DenseIntOrFPElementsAttr::getRaw(type, data, isSplat);
+  return DenseIntOrFPElementsAttr::getRaw(type, data);
 }
 
 DenseElementsAttr DenseIntOrFPElementsAttr::getRaw(ShapedType type,
-                                                   ArrayRef<char> data,
-                                                   bool isSplat) {
+                                                   ArrayRef<char> data) {
   assert((type.isa<RankedTensorType, VectorType>()) &&
          "type must be ranked tensor or vector");
   assert(type.hasStaticShape() && "type must have static shape");
+  bool isSplat = false;
+  bool isValid = isValidRawBuffer(type, data, isSplat);
+  assert(isValid);
+  (void)isValid;
   return Base::get(type.getContext(), type, data, isSplat);
 }
 
 /// Overload of the raw 'get' method that asserts that the given type is of
 /// complex type. This method is used to verify type invariants that the
 /// templatized 'get' method cannot.
 DenseElementsAttr DenseIntOrFPElementsAttr::getRawComplex(ShapedType type,
                                                           ArrayRef<char> data,
                                                           int64_t dataEltSize,
                                                           bool isInt,
                                                           bool isSigned) {
   assert(::isValidIntOrFloat(
       type.getElementType().cast<ComplexType>().getElementType(),
       dataEltSize / 2, isInt, isSigned));
 
   int64_t numElements = data.size() / dataEltSize;
   assert(numElements == 1 || numElements == type.getNumElements());
-  return getRaw(type, data, /*isSplat=*/numElements == 1);
+  return getRaw(type, data);
 }
 
 /// Overload of the 'getRaw' method that asserts that the given type is of
 /// integer type. This method is used to verify type invariants that the
 /// templatized 'get' method cannot.
 DenseElementsAttr
 DenseIntOrFPElementsAttr::getRawIntOrFloat(ShapedType type, ArrayRef<char> data,
                                            int64_t dataEltSize, bool isInt,
                                            bool isSigned) {
   assert(
       ::isValidIntOrFloat(type.getElementType(), dataEltSize, isInt, isSigned));
 
   int64_t numElements = data.size() / dataEltSize;
   assert(numElements == 1 || numElements == type.getNumElements());
-  return getRaw(type, data, /*isSplat=*/numElements == 1);
+  (void)numElements;
+  return getRaw(type, data);
 }
 
 void DenseIntOrFPElementsAttr::convertEndianOfCharForBEmachine(
     const char *inRawData, char *outRawData, size_t elementBitWidth,
     size_t numElements) {
   using llvm::support::ulittle16_t;
   using llvm::support::ulittle32_t;
   using llvm::support::ulittle64_t;

 }
 
 DenseElementsAttr DenseFPElementsAttr::mapValues(
     Type newElementType, function_ref<APInt(const APFloat &)> mapping) const {
   llvm::SmallVector<char, 8> elementData;
   auto newArrayType =
       mappingHelper(mapping, *this, getType(), newElementType, elementData);
 
-  return getRaw(newArrayType, elementData, isSplat());
+  return getRaw(newArrayType, elementData);
 }
 
 /// Method for supporting type inquiry through isa, cast and dyn_cast.
 bool DenseFPElementsAttr::classof(Attribute attr) {
   return attr.isa<DenseElementsAttr>() &&
          attr.getType().cast<ShapedType>().getElementType().isa<FloatType>();
 }
 
 //===----------------------------------------------------------------------===//
 // DenseIntElementsAttr
 //===----------------------------------------------------------------------===//
 
 DenseElementsAttr DenseIntElementsAttr::mapValues(
     Type newElementType, function_ref<APInt(const APInt &)> mapping) const {
   llvm::SmallVector<char, 8> elementData;
   auto newArrayType =
       mappingHelper(mapping, *this, getType(), newElementType, elementData);
-
-  return getRaw(newArrayType, elementData, isSplat());
+  return getRaw(newArrayType, elementData);
 }
 
 /// Method for supporting type inquiry through isa, cast and dyn_cast.
 bool DenseIntElementsAttr::classof(Attribute attr) {
   return attr.isa<DenseElementsAttr>() &&
          attr.getType().cast<ShapedType>().getElementType().isIntOrIndex();
 }
 

mlir/lib/Parser/AttributeParser.cpp

if (llvm::support::endian::system_endianness() ==
     // Convert endianess in big-endian(BE) machines. `rawData` is
     // little-endian(LE) because HEX in raw data of dense element attribute
     // is always LE format. It is converted into BE here to be used in BE
     // machines.
     SmallVector<char, 64> outDataVec(rawData.size());
     MutableArrayRef<char> convRawData(outDataVec);
     DenseIntOrFPElementsAttr::convertEndianOfArrayRefForBEmachine(
         rawData, convRawData, type);
-    return DenseElementsAttr::getFromRawBuffer(type, convRawData,
-                                               detectedSplat);
+    return DenseElementsAttr::getFromRawBuffer(type, convRawData);
   }
 
-  return DenseElementsAttr::getFromRawBuffer(type, rawData, detectedSplat);
+  return DenseElementsAttr::getFromRawBuffer(type, rawData);
 }
 
 ParseResult TensorLiteralParser::parseElement() {
   switch (p.getToken().getKind()) {
   // Parse a boolean element.
   case Token::kw_true:
   case Token::kw_false:
   case Token::floatliteral: