[mlir][DenseElementsAttr] Add support for ComplexType elements

[rriddle]

Summary

This revision adds support for storing ComplexType elements inside of a DenseElementsAttr. We store complex objects as an array of two elements, matching the definition of std::complex. There is no current attribute storage for ComplexType, but DenseElementsAttr provides API for access/creation using std::complex<>. Given that the internal implementation of DenseElementsAttr is already fairly opaque, the only real complexity here is in the printing/parsing. This revision keeps it simple for now and always uses hex when printing complex elements. A followup will add prettier syntax for this.

[rriddle]

Update

[jpienaar]

mlir/include/mlir/IR/Attributes.h

1057

Comments need update?

1133

Same

mlir/lib/IR/AttributeDetail.h

378

This is a little bit weird. I see some use C64 to refer to 2x64 bit float complex number while XLA consider std::complex<float> as C64 (so 64 is total). So seems the bitwidth is overloaded for complex already, as long as you document it this seems fine then as there doesn't seem to be 1 convention.

[rriddle]

mlir/lib/IR/AttributeDetail.h

378

That is different. We don't suffer from the problematic naming because we explicitly specify the element type. We specify the type exactly the same as C++ complex, i.e. there is no C64 there is only complex<double>/complex<float>/etc. It doesn't seem like there is any ambiguity there?

[rriddle]

Resolve comments

[jpienaar]

Looks good thanks

mlir/lib/IR/AttributeDetail.h

378

You are returning the bitwidth of the storage as element bitwidth which could either be the bitwidth of how it is stored or or the elements. But you are correct in that if one looks at the comment in 375, this would seem to return the storage bitwidth rather than anything intrinsic about the type (e.g., we are storing bf16 as 64 bits :-)). So no ambiguity.

mlir/lib/Parser/Parser.cpp

2068

Now the i vs j debate could heat up :)

mlir/include/mlir/IR/Attributes.h

 //===- Attributes.h - MLIR Attribute Classes --------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_ATTRIBUTES_H
 #define MLIR_IR_ATTRIBUTES_H
 
 #include "mlir/IR/AttributeSupport.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/Sequence.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
+#include <complex>
 
 namespace mlir {
 class AffineMap;
 class Dialect;
 class FunctionType;
 class Identifier;
 class IntegerSet;
 class Location;
protected:
   ptrdiff_t getDataIndex() const {
     bool isSplat = this->base.getInt();
     return isSplat ? 0 : this->index;
   }
 
   /// Return the data base pointer.
   const char *getData() const { return this->base.getPointer(); }
 };
+
+/// Type trait detector that checks if a given type T is a complex type.
+template <typename T> struct is_complex_t : public std::false_type {};
+template <typename T>
+struct is_complex_t<std::complex<T>> : public std::true_type {};
 } // namespace detail
 
 /// An attribute that represents a reference to a dense vector or tensor object.
 ///
 class DenseElementsAttr : public ElementsAttr {
 public:
   using ElementsAttr::ElementsAttr;
 
static DenseElementsAttr get(const ShapedType &type, ArrayRef<T> values) {
     return getRawIntOrFloat(
         type, ArrayRef<char>(data, values.size() * sizeof(T)), sizeof(T),
         std::numeric_limits<T>::is_integer, std::numeric_limits<T>::is_signed);
   }
 
   /// Constructs a dense integer elements attribute from a single element.
   template <typename T, typename = typename std::enable_if<
                             std::numeric_limits<T>::is_integer ||
-                            llvm::is_one_of<T, float, double>::value>::type>
+                            llvm::is_one_of<T, float, double>::value ||
+                            detail::is_complex_t<T>::value>::type>
   static DenseElementsAttr get(const ShapedType &type, T value) {
     return get(type, llvm::makeArrayRef(value));
   }
 
+  /// Constructs a dense complex elements attribute from an array of complex
+  /// values. Each value is expected to be the same bitwidth of the element type
+  /// of 'type'. 'type' must be a vector or tensor with static shape.
+  template <typename T, typename ElementT = typename T::value_type,
+            typename = typename std::enable_if<
+                detail::is_complex_t<T>::value &&
+                (std::numeric_limits<ElementT>::is_integer ||
+                 llvm::is_one_of<ElementT, float, double>::value)>::type>
+  static DenseElementsAttr get(const ShapedType &type, ArrayRef<T> values) {
+    const char *data = reinterpret_cast<const char *>(values.data());
+    return getRawComplex(type, ArrayRef<char>(data, values.size() * sizeof(T)),
+                         sizeof(T), std::numeric_limits<ElementT>::is_integer,
+                         std::numeric_limits<ElementT>::is_signed);
+  }
+
   /// Overload of the above 'get' method that is specialized for boolean values.
   static DenseElementsAttr get(ShapedType type, ArrayRef<bool> values);
 
   /// Overload of the above 'get' method that is specialized for StringRef
   /// values.
   static DenseElementsAttr get(ShapedType type, ArrayRef<StringRef> values);
 
   /// Constructs a dense integer elements attribute from an array of APInt
public:
 
   /// 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.
   static DenseElementsAttr getFromRawBuffer(ShapedType type,
                                             ArrayRef<char> rawBuffer,
                                             bool isSplatBuffer);
 
+  /// 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.
+  static bool isValidRawBuffer(ShapedType type, ArrayRef<char> rawBuffer,
+                               bool &detectedSplat);
+
   //===--------------------------------------------------------------------===//
   // Iterators
   //===--------------------------------------------------------------------===//
 
   /// A utility iterator that allows walking over the internal Attribute values
   /// of a DenseElementsAttr.
   class AttributeElementIterator
       : public llvm::indexed_accessor_iterator<AttributeElementIterator,
public:
   /// values.
   template <typename T, typename = typename std::enable_if<
                             (!std::is_same<T, bool>::value &&
                              std::numeric_limits<T>::is_integer) ||
                             llvm::is_one_of<T, float, double>::value>::type>
   llvm::iterator_range<ElementIterator<T>> getValues() const {
     assert(isValidIntOrFloat(sizeof(T), std::numeric_limits<T>::is_integer,
                              std::numeric_limits<T>::is_signed));
-    auto rawData = getRawData().data();
+    const char *rawData = getRawData().data();
+    bool splat = isSplat();
+    return {ElementIterator<T>(rawData, splat, 0),
+            ElementIterator<T>(rawData, splat, getNumElements())};
+  }
+
+  /// Return the held element values as a range of std::complex.
+  template <typename T, typename ElementT = typename T::value_type,
+            typename = typename std::enable_if<
+                detail::is_complex_t<T>::value &&
+                (std::numeric_limits<ElementT>::is_integer ||
+                 llvm::is_one_of<ElementT, float, double>::value)>::type>
+  llvm::iterator_range<ElementIterator<T>> getValues() const {
+    assert(isValidComplex(sizeof(T), std::numeric_limits<ElementT>::is_integer,
+                          std::numeric_limits<ElementT>::is_signed));
+    const char *rawData = getRawData().data();
     bool splat = isSplat();
     return {ElementIterator<T>(rawData, splat, 0),
             ElementIterator<T>(rawData, splat, getNumElements())};
   }
 
+  /// Return the held element values as a range of StringRef.
   template <typename T, typename = typename std::enable_if<
                             std::is_same<T, StringRef>::value>::type>
   llvm::iterator_range<ElementIterator<StringRef>> getValues() const {
     auto stringRefs = getRawStringData();
     const char *ptr = reinterpret_cast<const char *>(stringRefs.data());
     bool splat = isSplat();
     return {ElementIterator<StringRef>(ptr, splat, 0),
             ElementIterator<StringRef>(ptr, splat, getNumElements())};
protected:
   /// Get iterators to the raw APInt values for each element in this attribute.
   IntElementIterator raw_int_begin() const {
     return IntElementIterator(*this, 0);
   }
   IntElementIterator raw_int_end() const {
     return IntElementIterator(*this, getNumElements());
   }
 
   /// 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);
 
   /// Check the information for a C++ data type, check if this type is valid for
   /// the current attribute. This method is used to verify specific type
   /// invariants that the templatized 'getValues' method cannot.
   bool isValidIntOrFloat(int64_t dataEltSize, bool isInt, bool isSigned) const;
+
+  /// Check the information for a C++ data type, check if this type is valid for
+  /// the current attribute. This method is used to verify specific type
+  /// invariants that the templatized 'getValues' method cannot.
+  bool isValidComplex(int64_t dataEltSize, bool isInt, bool isSigned) const;
 };
 
 /// An attribute class for representing dense arrays of strings. The structure
 /// storing and querying a list of densely packed strings.
 class DenseStringElementsAttr
     : public Attribute::AttrBase<DenseStringElementsAttr, DenseElementsAttr,
                                  detail::DenseStringElementsAttributeStorage> {
 
protected:
   /// of 'type'. 'type' must be a vector or tensor with static shape.
   static DenseElementsAttr getRaw(ShapedType type, 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);
 
   /// 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);
 };
 
private:
   }
   /// Get a zero for an APFloat.
   template <typename T>
   typename std::enable_if<std::is_same<APFloat, T>::value, T>::type
   getZeroValue() const {
     return getZeroAPFloat();
   }
 
-  /// Get a zero for a StringRef.
-  template <typename T>
-  typename std::enable_if<std::is_same<StringRef, T>::value, T>::type
-  getZeroValue() const {
-    return StringRef();
-  }
-
-  /// Get a zero for an C++ integer or float type.
+  /// Get a zero for an C++ integer, float, StringRef, or complex type.
   template <typename T>
-  typename std::enable_if<std::numeric_limits<T>::is_integer ||
-                              llvm::is_one_of<T, float, double>::value,
-                          T>::type
+  typename std::enable_if<
+      std::numeric_limits<T>::is_integer ||
+          llvm::is_one_of<T, float, double, StringRef>::value ||
+          detail::is_complex_t<T>::value,
   getZeroValue() const {
-    return T(0);
+    return T();
   }
 
   /// Flatten, and return, all of the sparse indices in this attribute in
   /// row-major order.
   std::vector<ptrdiff_t> getFlattenedSparseIndices() const;
 };
 
 /// An attribute that represents a reference to a splat vector or tensor

mlir/lib/IR/AsmPrinter.cpp

static void printDenseIntElement(DenseElementsAttr attr, raw_ostream &os,
   if (value.getBitWidth() == 1)
     os << (value.getBoolValue() ? "true" : "false");
   else
     value.print(os, isSigned);
 }
 
 /// Print the float element of the given DenseElementsAttr at 'index'.
 static void printDenseFloatElement(DenseElementsAttr attr, raw_ostream &os,
-                                   unsigned index, bool isSigned) {
-  assert(isSigned && "floating point values are always signed");
+                                   unsigned index) {
   APFloat value = *std::next(attr.float_value_begin(), index);
   printFloatValue(value, os);
 }
 
-static void printDenseStringElement(DenseStringElementsAttr attr,
-                                    raw_ostream &os, unsigned index) {
-  os << "\"";
-  printEscapedString(attr.getRawStringData()[index], os);
-  os << "\"";
-}
-
-void ModulePrinter::printDenseElementsAttr(DenseElementsAttr attr,
-                                           bool allowHex) {
-  if (auto stringAttr = attr.dyn_cast<DenseStringElementsAttr>()) {
-    printDenseStringElementsAttr(stringAttr);
-    return;
-  }
-
-  printDenseIntOrFPElementsAttr(attr.cast<DenseIntOrFPElementsAttr>(),
-                                allowHex);
-}
-
-void ModulePrinter::printDenseIntOrFPElementsAttr(DenseIntOrFPElementsAttr attr,
-                                                  bool allowHex) {
-  auto type = attr.getType();
-  auto shape = type.getShape();
-  auto rank = type.getRank();
-  bool isSigned = !type.getElementType().isUnsignedInteger();
-
-  // The function used to print elements of this attribute.
-  auto printEltFn = type.getElementType().isIntOrIndex()
-                        ? printDenseIntElement
-                        : printDenseFloatElement;
-
+static void
+printDenseElementsAttrImpl(bool isSplat, ShapedType type, raw_ostream &os,
+                           function_ref<void(unsigned)> printEltFn) {
   // Special case for 0-d and splat tensors.
-  if (attr.isSplat()) {
-    printEltFn(attr, os, 0, isSigned);
-    return;
-  }
+  if (isSplat)
+    return printEltFn(0);
 
   // Special case for degenerate tensors.
   auto numElements = type.getNumElements();
+  int64_t rank = type.getRank();
   if (numElements == 0) {
     for (int i = 0; i < rank; ++i)
       os << '[';
     for (int i = 0; i < rank; ++i)
       os << ']';
     return;
   }
 
-  // Check to see if we should format this attribute as a hex string.
-  if (allowHex && shouldPrintElementsAttrWithHex(numElements)) {
-    ArrayRef<char> rawData = attr.getRawData();
-    os << '"' << "0x" << llvm::toHex(StringRef(rawData.data(), rawData.size()))
-       << "\"";
-    return;
-  }
-
   // We use a mixed-radix counter to iterate through the shape. When we bump a
   // non-least-significant digit, we emit a close bracket. When we next emit an
   // element we re-open all closed brackets.
 
   // The mixed-radix counter, with radices in 'shape'.
   SmallVector<unsigned, 4> counter(rank, 0);
   // The number of brackets that have been opened and not closed.
   unsigned openBrackets = 0;
 
-  auto bumpCounter = [&]() {
+  auto shape = type.getShape();
+  auto bumpCounter = [&] {
     // Bump the least significant digit.
     ++counter[rank - 1];
     // Iterate backwards bubbling back the increment.
     for (unsigned i = rank - 1; i > 0; --i)
       if (counter[i] >= shape[i]) {
         // Index 'i' is rolled over. Bump (i-1) and close a bracket.
         counter[i] = 0;
         ++counter[i - 1];
         --openBrackets;
         os << ']';
       }
   };
 
   for (unsigned idx = 0, e = numElements; idx != e; ++idx) {
     if (idx != 0)
       os << ", ";
     while (openBrackets++ < rank)
       os << '[';
     openBrackets = rank;
-    printEltFn(attr, os, idx, isSigned);
+    printEltFn(idx);
     bumpCounter();
   }
   while (openBrackets-- > 0)
     os << ']';
 }
 
-void ModulePrinter::printDenseStringElementsAttr(DenseStringElementsAttr attr) {
-  auto type = attr.getType();
-  auto shape = type.getShape();
-  auto rank = type.getRank();
+void ModulePrinter::printDenseElementsAttr(DenseElementsAttr attr,
+                                           bool allowHex) {
+  if (auto stringAttr = attr.dyn_cast<DenseStringElementsAttr>())
+    return printDenseStringElementsAttr(stringAttr);
 
-  // Special case for 0-d and splat tensors.
-  if (attr.isSplat()) {
-    printDenseStringElement(attr, os, 0);
-    return;
-  }
+  printDenseIntOrFPElementsAttr(attr.cast<DenseIntOrFPElementsAttr>(),
+                                allowHex);
 
-  // Special case for degenerate tensors.
+void ModulePrinter::printDenseIntOrFPElementsAttr(DenseIntOrFPElementsAttr attr,
+                                                  bool allowHex) {
+  auto type = attr.getType();
+  auto elementType = type.getElementType();
+
+  // Check to see if we should format this attribute as a hex string.
+  // TODO: Add support for formatting complex elements nicely.
   auto numElements = type.getNumElements();
-  if (numElements == 0) {
-    for (int i = 0; i < rank; ++i)
-      os << '[';
-    for (int i = 0; i < rank; ++i)
-      os << ']';
+  if (type.getElementType().isa<ComplexType>() ||
+      (!attr.isSplat() && allowHex &&
+       shouldPrintElementsAttrWithHex(numElements))) {
+    ArrayRef<char> rawData = attr.getRawData();
+    os << '"' << "0x" << llvm::toHex(StringRef(rawData.data(), rawData.size()))
+       << "\"";
     return;
   }
 
-  // We use a mixed-radix counter to iterate through the shape. When we bump a
-  // non-least-significant digit, we emit a close bracket. When we next emit an
-  // element we re-open all closed brackets.
-
-  // The mixed-radix counter, with radices in 'shape'.
-  SmallVector<unsigned, 4> counter(rank, 0);
-  // The number of brackets that have been opened and not closed.
-  unsigned openBrackets = 0;
-
-  auto bumpCounter = [&]() {
-    // Bump the least significant digit.
-    ++counter[rank - 1];
-    // Iterate backwards bubbling back the increment.
-    for (unsigned i = rank - 1; i > 0; --i)
-      if (counter[i] >= shape[i]) {
-        // Index 'i' is rolled over. Bump (i-1) and close a bracket.
-        counter[i] = 0;
-        ++counter[i - 1];
-        --openBrackets;
-        os << ']';
-      }
-  };
-
-  for (unsigned idx = 0, e = numElements; idx != e; ++idx) {
-    if (idx != 0)
-      os << ", ";
-    while (openBrackets++ < rank)
-      os << '[';
-    openBrackets = rank;
-    printDenseStringElement(attr, os, idx);
-    bumpCounter();
-  }
-  while (openBrackets-- > 0)
-    os << ']';
+  if (elementType.isIntOrIndex()) {
+    bool isSigned = !elementType.isUnsignedInteger();
+    printDenseElementsAttrImpl(attr.isSplat(), type, os, [&](unsigned index) {
+      printDenseIntElement(attr, os, index, isSigned);
+    });
+  } else {
+    assert(elementType.isa<FloatType>() && "unexpected element type");
+    printDenseElementsAttrImpl(attr.isSplat(), type, os, [&](unsigned index) {
+      printDenseFloatElement(attr, os, index);
+    });
+
+void ModulePrinter::printDenseStringElementsAttr(DenseStringElementsAttr attr) {
+  ArrayRef<StringRef> data = attr.getRawStringData();
+  auto printFn = [&](unsigned index) {
+    os << "\"";
+    printEscapedString(data[index], os);
+    os << "\"";
+  };
+  printDenseElementsAttrImpl(attr.isSplat(), attr.getType(), os, printFn);
 }
 
 void ModulePrinter::printType(Type type) {
   if (!type) {
     os << "<<NULL TYPE>>";
     return;
   }
 

mlir/lib/IR/AttributeDetail.h

 };
 
 //===----------------------------------------------------------------------===//
 // Elements Attributes
 //===----------------------------------------------------------------------===//
 
 /// Return the bit width which DenseElementsAttr should use for this type.
 inline size_t getDenseElementBitWidth(Type eltType) {
+  if (ComplexType complex = eltType.dyn_cast<ComplexType>())
+    return getDenseElementBitWidth(complex.getElementType()) * 2;
   // FIXME(b/121118307): using 64 bits for BF16 because it is currently stored
   // with double semantics.
   if (eltType.isBF16())
     return 64;
   if (eltType.isIndex())
     return IndexType::kInternalStorageBitWidth;
   return eltType.getIntOrFloatBitWidth();
 }

mlir/lib/IR/Attributes.cpp

 /// 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);
 }
 
+/// 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 elementWidth = getDenseElementBitWidth(type.getElementType());
+  size_t storageWidth = getDenseElementStorageWidth(elementWidth);
+  size_t rawBufferWidth = rawBuffer.size() * CHAR_BIT;
+
+  // 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.
+    if ((detectedSplat = rawBuffer.size() == 1))
+      return true;
+    return rawBufferWidth == llvm::alignTo<8>(type.getNumElements());
+  }
+  // All other types are 8-bit aligned.
+  if ((detectedSplat = rawBufferWidth == storageWidth))
+    return true;
+  return rawBufferWidth == (storageWidth * type.getNumElements());
+}
+
 /// Check the information for a C++ data type, check if this type is valid for
 /// the current attribute. This method is used to verify specific type
 /// invariants that the templatized 'getValues' method cannot.
-static bool isValidIntOrFloat(ShapedType type, int64_t dataEltSize, bool isInt,
+static bool isValidIntOrFloat(Type type, int64_t dataEltSize, bool isInt,
                               bool isSigned) {
   // Make sure that the data element size is the same as the type element width.
-  if (getDenseElementBitWidth(type.getElementType()) !=
+  if (getDenseElementBitWidth(type) !=
       static_cast<size_t>(dataEltSize * CHAR_BIT))
     return false;
 
   // Check that the element type is either float or integer or index.
   if (!isInt)
-    return type.getElementType().isa<FloatType>();
-
-  if (type.getElementType().isIndex())
+    return type.isa<FloatType>();
+  if (type.isIndex())
     return true;
 
-  auto intType = type.getElementType().dyn_cast<IntegerType>();
+  auto intType = type.dyn_cast<IntegerType>();
   if (!intType)
     return false;
 
   // Make sure signedness semantics is consistent.
   if (intType.isSignless())
     return true;
   return intType.isSigned() ? isSigned : !isSigned;
 }
 
 /// Defaults down the subclass implementation.
+DenseElementsAttr DenseElementsAttr::getRawComplex(ShapedType type,
+                                                   ArrayRef<char> data,
+                                                   int64_t dataEltSize,
+                                                   bool isInt, bool isSigned) {
+  return DenseIntOrFPElementsAttr::getRawComplex(type, data, dataEltSize, isInt,
+                                                 isSigned);
+}
 DenseElementsAttr DenseElementsAttr::getRawIntOrFloat(ShapedType type,
                                                       ArrayRef<char> data,
                                                       int64_t dataEltSize,
                                                       bool isInt,
                                                       bool isSigned) {
   return DenseIntOrFPElementsAttr::getRawIntOrFloat(type, data, dataEltSize,
                                                     isInt, isSigned);
 }
 
 /// A method used to verify specific type invariants that the templatized 'get'
 /// method cannot.
 bool DenseElementsAttr::isValidIntOrFloat(int64_t dataEltSize, bool isInt,
                                           bool isSigned) const {
-  return ::isValidIntOrFloat(getType(), dataEltSize, isInt, isSigned);
+  return ::isValidIntOrFloat(getType().getElementType(), dataEltSize, isInt,
+                             isSigned);
+}
+
+/// Check the information for a C++ data type, check if this type is valid for
+/// the current attribute.
+bool DenseElementsAttr::isValidComplex(int64_t dataEltSize, bool isInt,
+                                       bool isSigned) const {
+  return ::isValidIntOrFloat(
+      getType().getElementType().cast<ComplexType>().getElementType(),
+      dataEltSize / 2, isInt, isSigned);
 }
 
 /// Returns if this attribute corresponds to a splat, i.e. if all element
 /// values are the same.
 bool DenseElementsAttr::isSplat() const {
   return static_cast<DenseElementsAttributeStorage *>(impl)->isSplat;
 }
 
DenseElementsAttr DenseIntOrFPElementsAttr::getRaw(ShapedType type,
                                                    bool isSplat) {
   assert((type.isa<RankedTensorType>() || type.isa<VectorType>()) &&
          "type must be ranked tensor or vector");
   assert(type.hasStaticShape() && "type must have static shape");
   return Base::get(type.getContext(), StandardAttributes::DenseIntOrFPElements,
                    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);
+}
+
 /// 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, dataEltSize, isInt, 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);
 }
 
 //===----------------------------------------------------------------------===//
 // DenseFPElementsAttr

mlir/lib/Parser/Parser.cpp

 
 /// Build a dense attribute instance with the parsed elements and the given
 /// shaped type.
 DenseElementsAttr TensorLiteralParser::getAttr(llvm::SMLoc loc,
                                                ShapedType type) {
   Type eltType = type.getElementType();
 
   // Check to see if we parse the literal from a hex string.
-  if (hexStorage.hasValue() && eltType.isIntOrFloat())
+  if (hexStorage.hasValue() &&
+      (eltType.isIntOrFloat() || eltType.isa<ComplexType>()))
     return getHexAttr(loc, type);
 
   // Check that the parsed storage size has the same number of elements to the
   // type, or is a known splat.
   if (!shape.empty() && getShape() != type.getShape()) {
     p.emitError(loc) << "inferred shape of elements literal ([" << getShape()
                      << "]) does not match type ([" << type.getShape() << "])";
     return nullptr;
   }
 
   // If the type is an integer, build a set of APInt values from the storage
   // with the correct bitwidth.
   if (auto intTy = eltType.dyn_cast<IntegerType>())
     return getIntAttr(loc, type, intTy);
   if (auto indexTy = eltType.dyn_cast<IndexType>())
     return getIntAttr(loc, type, indexTy);
 
   // If parsing a floating point type.
   if (auto floatTy = eltType.dyn_cast<FloatType>())
     return getFloatAttr(loc, type, floatTy);
 
+  // If parsing a complex type.
+  // TODO: Support complex elements with pretty element printing.
+  if (eltType.isa<ComplexType>()) {
+    p.emitError(loc) << "complex elements only support hex formatting";
+    return nullptr;
+  }
+
   // Other types are assumed to be string representations.
   return getStringAttr(loc, type, type.getElementType());
 }
 
 /// Build a Dense Integer attribute for the given type.
 DenseElementsAttr TensorLiteralParser::getIntAttr(llvm::SMLoc loc,
                                                   ShapedType type, Type eltTy) {
   std::vector<APInt> intElements;
DenseElementsAttr TensorLiteralParser::getStringAttr(llvm::SMLoc loc,
 
   return DenseStringElementsAttr::get(type, stringRefValues);
 }
 
 /// Build a Dense attribute with hex data for the given type.
 DenseElementsAttr TensorLiteralParser::getHexAttr(llvm::SMLoc loc,
                                                   ShapedType type) {
   Type elementType = type.getElementType();
-  if (!elementType.isa<FloatType>() && !elementType.isa<IntegerType>()) {
-    p.emitError(loc) << "expected floating-point or integer element type, got "
-                     << elementType;
+  if (!elementType.isIntOrIndexOrFloat() && !elementType.isa<ComplexType>()) {
+    p.emitError(loc)
+        << "expected floating-point, integer, or complex element type, got "
     return nullptr;
   }
 
   std::string data;
   if (parseElementAttrHexValues(p, hexStorage.getValue(), data))
     return nullptr;
 
-  // Check that the size of the hex data corresponds to the size of the type, or
-  // a splat of the type.
-  // TODO: bf16 is currently stored as a double, this should be removed when
-  // APFloat properly supports it.
-  int64_t elementWidth =
-      elementType.isBF16() ? 64 : elementType.getIntOrFloatBitWidth();
-  if (static_cast<int64_t>(data.size() * CHAR_BIT) !=
-      (type.getNumElements() * elementWidth)) {
+  ArrayRef<char> rawData(data.data(), data.size());
+  bool detectedSplat = false;
+  if (!DenseElementsAttr::isValidRawBuffer(type, rawData, detectedSplat)) {
     p.emitError(loc) << "elements hex data size is invalid for provided type: "
                      << type;
     return nullptr;
   }
 
-  return DenseElementsAttr::getFromRawBuffer(
-      type, ArrayRef<char>(data.data(), data.size()), /*isSplatBuffer=*/false);
+  return DenseElementsAttr::getFromRawBuffer(type, rawData, detectedSplat);
 }
 
 ParseResult TensorLiteralParser::parseElement() {
   switch (p.getToken().getKind()) {
   // Parse a boolean element.
   case Token::kw_true:
   case Token::kw_false:
   case Token::floatliteral:

mlir/test/IR/dense-elements-hex.mlir

 // RUN: mlir-opt -allow-unregistered-dialect %s -verify-diagnostics -split-input-file -mlir-print-elementsattrs-with-hex-if-larger=1 | FileCheck %s --check-prefix=HEX
 // RUN: mlir-opt -allow-unregistered-dialect %s -verify-diagnostics -split-input-file | FileCheck %s
 
 // HEX: dense<"0x00000000000024400000000000001440"> : tensor<2xf64>
 "foo.op"() {dense.attr = dense<[10.0, 5.0]> : tensor<2xf64>} : () -> ()
 
 // CHECK: dense<[1.000000e+01, 5.000000e+00]> : tensor<2xf64>
 "foo.op"() {dense.attr = dense<"0x00000000000024400000000000001440"> : tensor<2xf64>} : () -> ()
 
+// CHECK: dense<"0x00000000000024400000000000001440"> : tensor<1xcomplex<f64>>
+"foo.op"() {dense.attr = dense<"0x00000000000024400000000000001440"> : tensor<1xcomplex<f64>>} : () -> ()
+
+// CHECK: dense<"0x00000000000024400000000000001440"> : tensor<10xcomplex<f64>>
+"foo.op"() {dense.attr = dense<"0x00000000000024400000000000001440"> : tensor<10xcomplex<f64>>} : () -> ()
+
 // CHECK: dense<[1.000000e+01, 5.000000e+00]> : tensor<2xbf16>
 "foo.op"() {dense.attr = dense<"0x00000000000024400000000000001440"> : tensor<2xbf16>} : () -> ()
 
 // -----
 
 // expected-error@+1 {{elements hex string should start with '0x'}}
 "foo.op"() {dense.attr = dense<"00000000000024400000000000001440"> : tensor<2xf64>} : () -> ()
 

mlir/unittests/IR/AttributeTest.cpp

static void testSplat(Type eltType, const EltTy &splatElt) {
 
   // Check that the generated splat is the same for 1 element and N elements.
   DenseElementsAttr splat = DenseElementsAttr::get(shape, splatElt);
   EXPECT_TRUE(splat.isSplat());
 
   auto detectedSplat =
       DenseElementsAttr::get(shape, llvm::makeArrayRef({splatElt, splatElt}));
   EXPECT_EQ(detectedSplat, splat);
+
+  for (auto newValue : detectedSplat.template getValues<EltTy>())
+    EXPECT_EQ(newValue, splatElt);
 }
 
 namespace {
 TEST(DenseSplatTest, BoolSplat) {
   MLIRContext context;
   IntegerType boolTy = IntegerType::get(1, &context);
   RankedTensorType shape = RankedTensorType::get({2, 2}, boolTy);
 

 TEST(DenseSplatTest, StringAttrSplat) {
   MLIRContext context;
   Type stringType =
       OpaqueType::get(Identifier::get("test", &context), "string", &context);
   Attribute stringAttr = StringAttr::get("test-string", stringType);
   testSplat(stringType, stringAttr);
 }
 
+TEST(DenseComplexTest, ComplexFloatSplat) {
+  MLIRContext context;
+  ComplexType complexType = ComplexType::get(FloatType::getF32(&context));
+  std::complex<float> value(10.0, 15.0);
+  testSplat(complexType, value);
+}
+
+TEST(DenseComplexTest, ComplexIntSplat) {
+  MLIRContext context;
+  ComplexType complexType = ComplexType::get(IntegerType::get(64, &context));
+  std::complex<int64_t> value(10, 15);
+  testSplat(complexType, value);
+}
+
 } // end namespace