[mlir][DenseElementsAttr] Add support for opaque APFloat/APInt complex values.

[rriddle]

Summary

This revision allows for creating DenseElementsAttrs and accessing elements using std::complex<APInt>/std::complex<APFloat>. This allows for opaquely accessing and transforming complex values. This is used by the printer/parser to provide pretty printing for complex values. The form for complex values matches that of std::complex, i.e.:

// `(` element `,` element `)`
dense<(10,10)> : tensor<complex<i64>>

Depends On D79281

[rriddle]

Rebase

[Smit]

Neat!

mlir/unittests/IR/AttributeTest.cpp

30

Also, dump the values?

[kiranchandramohan]

mlir/lib/IR/AsmPrinter.cpp

1548–1558

The code here hits a g++ bug (https://gcc.gnu.org/bugzilla/show_bug.cgi?id=90538) in gcc-9.1/gcc-9.2. This seems to be fixed in 9.3. Will you accept a workaround for this?

/home/user02/open_source/llvm-project/mlir/lib/IR/AsmPrinter.cpp: In instantiation of ‘{anonymous}::ModulePrinter::printDenseIntOrFPElementsAttr(mlir::DenseIntOrFPElementsAttr, bool)::<lambda(auto:2, auto:3, llvm::raw_ostream&, auto:4&& ...)> [with auto:2 = llvm::iterator_range<mlir::DenseElementsAttr::ComplexIntElementIterator>; auto:3 = void (*)(const llvm::APInt&, llvm::raw_ostream&, bool); auto:4 = {bool}]’:
/home/user02/open_source/llvm-project/mlir/lib/IR/AsmPrinter.cpp:1565:77:   required from here
/home/user02/open_source/llvm-project/mlir/lib/IR/AsmPrinter.cpp:1557:42: error: redeclaration of ‘bool& params#0’
 1557 |         printFn(complexValue.imag(), os, params...);
      |                                          ^~~~~~
/home/user02/open_source/llvm-project/mlir/lib/IR/AsmPrinter.cpp:1555:42: note: ‘bool& params#0’ previously declared here
 1555 |         printFn(complexValue.real(), os, params...);
      |                                          ^~~~~~

[mehdi_amini]

mlir/lib/IR/AsmPrinter.cpp

1548–1558

Feel free to send a patch!

mlir/include/mlir/IR/Attributes.h

public:
   static DenseElementsAttr get(ShapedType type, ArrayRef<StringRef> 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'. 'type' must be a vector or tensor with static
   /// shape.
   static DenseElementsAttr get(ShapedType type, ArrayRef<APInt> values);
 
+  /// Constructs a dense complex elements attribute from an array of 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.
+  static DenseElementsAttr get(ShapedType type,
+                               ArrayRef<std::complex<APInt>> values);
+
   /// 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'. 'type' must be a vector or tensor with static
   /// shape.
   static DenseElementsAttr get(ShapedType type, ArrayRef<APFloat> values);
 
+  /// Constructs a dense complex elements attribute from an array of 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.
+  static DenseElementsAttr get(ShapedType type,
+                               ArrayRef<std::complex<APFloat>> values);
+
   /// Construct a dense elements attribute for an initializer_list of 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>
   static DenseElementsAttr get(const ShapedType &type,
                                const std::initializer_list<T> &list) {
     return get(type, ArrayRef<T>(list));
   }
private:
 
     /// Constructs a new iterator.
     IntElementIterator(DenseElementsAttr attr, size_t dataIndex);
 
     /// The bitwidth of the element type.
     size_t bitWidth;
   };
 
+  /// A utility iterator that allows walking over the internal raw complex APInt
+  /// values.
+  class ComplexIntElementIterator
+      : public detail::DenseElementIndexedIteratorImpl<
+            ComplexIntElementIterator, std::complex<APInt>, std::complex<APInt>,
+            std::complex<APInt>> {
+  public:
+    /// Accesses the raw std::complex<APInt> value at this iterator position.
+    std::complex<APInt> operator*() const;
+
+  private:
+    friend DenseElementsAttr;
+
+    /// Constructs a new iterator.
+    ComplexIntElementIterator(DenseElementsAttr attr, size_t dataIndex);
+
+    /// The bitwidth of the element type.
+    size_t bitWidth;
+  };
+
   /// Iterator for walking over APFloat values.
   class FloatElementIterator final
       : public llvm::mapped_iterator<IntElementIterator,
                                      std::function<APFloat(const APInt &)>> {
     friend DenseElementsAttr;
 
     /// Initializes the float element iterator to the specified iterator.
     FloatElementIterator(const llvm::fltSemantics &smt, IntElementIterator it);
 
   public:
     using reference = APFloat;
   };
 
+  /// Iterator for walking over complex APFloat values.
+  class ComplexFloatElementIterator final
+      : public llvm::mapped_iterator<
+            ComplexIntElementIterator,
+            std::function<std::complex<APFloat>(const std::complex<APInt> &)>> {
+    friend DenseElementsAttr;
+
+    /// Initializes the float element iterator to the specified iterator.
+    ComplexFloatElementIterator(const llvm::fltSemantics &smt,
+                                ComplexIntElementIterator it);
+
+  public:
+    using reference = std::complex<APFloat>;
+  };
+
   //===--------------------------------------------------------------------===//
   // Value Querying
   //===--------------------------------------------------------------------===//
 
   /// Returns if this attribute corresponds to a splat, i.e. if all element
   /// values are the same.
   bool isSplat() const;
 
public:
   template <typename T, typename = typename std::enable_if<
                             std::is_same<T, APInt>::value>::type>
   llvm::iterator_range<IntElementIterator> getValues() const {
     return getIntValues();
   }
   IntElementIterator int_value_begin() const;
   IntElementIterator int_value_end() const;
 
+  /// Return the held element values as a range of complex APInts. The element
+  /// type of this attribute must be a complex of integer type.
+  llvm::iterator_range<ComplexIntElementIterator> getComplexIntValues() const;
+  template <typename T, typename = typename std::enable_if<
+                            std::is_same<T, std::complex<APInt>>::value>::type>
+  llvm::iterator_range<ComplexIntElementIterator> getValues() const {
+    return getComplexIntValues();
+  }
+
   /// Return the held element values as a range of APFloat. The element type of
   /// this attribute must be of float type.
   llvm::iterator_range<FloatElementIterator> getFloatValues() const;
   template <typename T, typename = typename std::enable_if<
                             std::is_same<T, APFloat>::value>::type>
   llvm::iterator_range<FloatElementIterator> getValues() const {
     return getFloatValues();
   }
   FloatElementIterator float_value_begin() const;
   FloatElementIterator float_value_end() const;
 
+  /// Return the held element values as a range of complex APFloat. The element
+  /// type of this attribute must be a complex of float type.
+  llvm::iterator_range<ComplexFloatElementIterator>
+  getComplexFloatValues() const;
+  template <typename T, typename = typename std::enable_if<std::is_same<
+                            T, std::complex<APFloat>>::value>::type>
+  llvm::iterator_range<ComplexFloatElementIterator> getValues() const {
+    return getComplexFloatValues();
+  }
+
   /// Return the raw storage data held by this attribute. Users should generally
   /// not use this directly, as the internal storage format is not always in the
   /// form the user might expect.
   ArrayRef<char> getRawData() const;
 
   /// Return the raw StringRef data held by this attribute.
   ArrayRef<StringRef> getRawStringData() const;
 
public:
   /// Method for support type inquiry through isa, cast and dyn_cast.
   static bool kindof(unsigned kind) {
     return kind == StandardAttributes::DenseIntOrFPElements;
   }
 
 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.
+  static DenseElementsAttr getRaw(ShapedType type, size_t storageWidth,
+                                  ArrayRef<APFloat> values, bool isSplat);
+
   /// 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.
-  static DenseElementsAttr getRaw(ShapedType type, ArrayRef<APInt> values);
+  static DenseElementsAttr getRaw(ShapedType type, size_t storageWidth,
+                                  ArrayRef<APInt> values, bool isSplat);
 
   /// 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
getZeroValue() const {
     return getZeroAttr().template cast<T>();
   }
   /// Get a zero for an APInt.
   template <typename T>
   typename std::enable_if<std::is_same<APInt, T>::value, T>::type
   getZeroValue() const {
     return getZeroAPInt();
   }
+  template <typename T>
+  typename std::enable_if<std::is_same<std::complex<APInt>, T>::value, T>::type
+  getZeroValue() const {
+    APInt intZero = getZeroAPInt();
+    return {intZero, intZero};
+  }
   /// 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();
   }
+  template <typename T>
+  typename std::enable_if<std::is_same<std::complex<APFloat>, T>::value,
+                          T>::type
+  getZeroValue() const {
+    APFloat floatZero = getZeroAPFloat();
+    return {floatZero, floatZero};
+  }
 
   /// 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, StringRef>::value ||
-          detail::is_complex_t<T>::value,
+          (detail::is_complex_t<T>::value &&
+           !llvm::is_one_of<T, std::complex<APInt>,
+                            std::complex<APFloat>>::value),
       T>::type
   getZeroValue() const {
     return T();
   }
 
   /// Flatten, and return, all of the sparse indices in this attribute in
   /// row-major order.
   std::vector<ptrdiff_t> getFlattenedSparseIndices() const;

mlir/lib/IR/AsmPrinter.cpp

void ModulePrinter::printAttribute(Attribute attr,
 
   // Don't print the type if we must elide it, or if it is a None type.
   if (typeElision != AttrTypeElision::Must && !attrType.isa<NoneType>()) {
     os << " : ";
     printType(attrType);
   }
 }
 
-/// Print the integer element of the given DenseElementsAttr at 'index'.
-static void printDenseIntElement(DenseElementsAttr attr, raw_ostream &os,
-                                 unsigned index, bool isSigned) {
-  APInt value = *std::next(attr.int_value_begin(), index);
+/// Print the integer element of a DenseElementsAttr.
+static void printDenseIntElement(const APInt &value, raw_ostream &os,
+                                 bool isSigned) {
   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) {
-  APFloat value = *std::next(attr.float_value_begin(), index);
-  printFloatValue(value, os);
-}
-
 static void
 printDenseElementsAttrImpl(bool isSplat, ShapedType type, raw_ostream &os,
                            function_ref<void(unsigned)> printEltFn) {
   // Special case for 0-d and splat tensors.
   if (isSplat)
     return printEltFn(0);
 
   // 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 (type.getElementType().isa<ComplexType>() ||
-      (!attr.isSplat() && allowHex &&
-       shouldPrintElementsAttrWithHex(numElements))) {
+  if (!attr.isSplat() && allowHex &&
+      shouldPrintElementsAttrWithHex(numElements)) {
     ArrayRef<char> rawData = attr.getRawData();
     os << '"' << "0x" << llvm::toHex(StringRef(rawData.data(), rawData.size()))
        << "\"";
     return;
   }
 
-  if (elementType.isIntOrIndex()) {
+  if (ComplexType complexTy = elementType.dyn_cast<ComplexType>()) {
+    auto printComplexValue = [&](auto complexValues, auto printFn,
+                                 raw_ostream &os, auto &&... params) {
+      printDenseElementsAttrImpl(attr.isSplat(), type, os, [&](unsigned index) {
+        auto complexValue = *(complexValues.begin() + index);
+        os << "(";
+        printFn(complexValue.real(), os, params...);
+        os << ",";
+        printFn(complexValue.imag(), os, params...);
+        os << ")";
+      });
+    };
+
+    Type complexElementType = complexTy.getElementType();
+    if (complexElementType.isa<IntegerType>())
+      printComplexValue(attr.getComplexIntValues(), printDenseIntElement, os,
+                        /*isSigned=*/!complexElementType.isUnsignedInteger());
+    else
+      printComplexValue(attr.getComplexFloatValues(), printFloatValue, os);
+  } else if (elementType.isIntOrIndex()) {
     bool isSigned = !elementType.isUnsignedInteger();
+    auto intValues = attr.getIntValues();
     printDenseElementsAttrImpl(attr.isSplat(), type, os, [&](unsigned index) {
-      printDenseIntElement(attr, os, index, isSigned);
+      printDenseIntElement(*(intValues.begin() + index), os, isSigned);
     });
   } else {
     assert(elementType.isa<FloatType>() && "unexpected element type");
+    auto floatValues = attr.getFloatValues();
     printDenseElementsAttrImpl(attr.isSplat(), type, os, [&](unsigned index) {
-      printDenseFloatElement(attr, os, index);
+      printFloatValue(*(floatValues.begin() + index), os);
     });
   }
 }
 
 void ModulePrinter::printDenseStringElementsAttr(DenseStringElementsAttr attr) {
   ArrayRef<StringRef> data = attr.getRawStringData();
   auto printFn = [&](unsigned index) {
     os << "\"";

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;
+  // Align the width for complex to 8 to make storage and interpretation easier.
+  if (ComplexType comp = eltType.dyn_cast<ComplexType>())
+    return llvm::alignTo<8>(getDenseElementBitWidth(comp.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

 // DenseElementAttr Utilities
 //===----------------------------------------------------------------------===//
 
 /// 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);
 }
+static size_t getDenseElementStorageWidth(Type elementType) {
+  return getDenseElementStorageWidth(getDenseElementBitWidth(elementType));
+}
 
 /// Set a bit to a specific value.
 static void setBit(char *rawData, size_t bitPos, bool value) {
   if (value)
     rawData[bitPos / CHAR_BIT] |= (1 << (bitPos % CHAR_BIT));
   else
     rawData[bitPos / CHAR_BIT] &= ~(1 << (bitPos % CHAR_BIT));
 }
static bool hasSameElementsOrSplat(ShapedType type, const Values &values) {
   return (values.size() == 1) ||
          (type.getNumElements() == static_cast<int64_t>(values.size()));
 }
 
 //===----------------------------------------------------------------------===//
 // DenseElementAttr Iterators
 //===----------------------------------------------------------------------===//
 
-/// Constructs a new iterator.
+//===----------------------------------------------------------------------===//
+// AttributeElementIterator
+
 DenseElementsAttr::AttributeElementIterator::AttributeElementIterator(
     DenseElementsAttr attr, size_t index)
     : llvm::indexed_accessor_iterator<AttributeElementIterator, const void *,
                                       Attribute, Attribute, Attribute>(
           attr.getAsOpaquePointer(), index) {}
 
-/// Accesses the Attribute value at this iterator position.
 Attribute DenseElementsAttr::AttributeElementIterator::operator*() const {
   auto owner = getFromOpaquePointer(base).cast<DenseElementsAttr>();
   Type eltTy = owner.getType().getElementType();
   if (auto intEltTy = eltTy.dyn_cast<IntegerType>()) {
     if (intEltTy.getWidth() == 1)
       return BoolAttr::get((*IntElementIterator(owner, index)).isOneValue(),
                            owner.getContext());
     return IntegerAttr::get(eltTy, *IntElementIterator(owner, index));
   }
   if (auto floatEltTy = eltTy.dyn_cast<FloatType>()) {
     IntElementIterator intIt(owner, index);
     FloatElementIterator floatIt(floatEltTy.getFloatSemantics(), intIt);
     return FloatAttr::get(eltTy, *floatIt);
   }
   llvm_unreachable("unexpected element type");
 }
 
-/// Constructs a new iterator.
+//===----------------------------------------------------------------------===//
+// BoolElementIterator
+
 DenseElementsAttr::BoolElementIterator::BoolElementIterator(
     DenseElementsAttr attr, size_t dataIndex)
     : DenseElementIndexedIteratorImpl<BoolElementIterator, bool, bool, bool>(
           attr.getRawData().data(), attr.isSplat(), dataIndex) {}
 
-/// Accesses the bool value at this iterator position.
 bool DenseElementsAttr::BoolElementIterator::operator*() const {
   return getBit(getData(), getDataIndex());
 }
 
-/// Constructs a new iterator.
+//===----------------------------------------------------------------------===//
+// IntElementIterator
+
 DenseElementsAttr::IntElementIterator::IntElementIterator(
     DenseElementsAttr attr, size_t dataIndex)
     : DenseElementIndexedIteratorImpl<IntElementIterator, APInt, APInt, APInt>(
           attr.getRawData().data(), attr.isSplat(), dataIndex),
       bitWidth(getDenseElementBitWidth(attr.getType().getElementType())) {}
 
-/// Accesses the raw APInt value at this iterator position.
 APInt DenseElementsAttr::IntElementIterator::operator*() const {
   return readBits(getData(),
                   getDataIndex() * getDenseElementStorageWidth(bitWidth),
                   bitWidth);
 }
 
+//===----------------------------------------------------------------------===//
+// ComplexIntElementIterator
+
+DenseElementsAttr::ComplexIntElementIterator::ComplexIntElementIterator(
+    DenseElementsAttr attr, size_t dataIndex)
+    : DenseElementIndexedIteratorImpl<ComplexIntElementIterator,
+                                      std::complex<APInt>, std::complex<APInt>,
+                                      std::complex<APInt>>(
+          attr.getRawData().data(), attr.isSplat(), dataIndex) {
+  auto complexType = attr.getType().getElementType().cast<ComplexType>();
+  bitWidth = getDenseElementBitWidth(complexType.getElementType());
+}
+
+std::complex<APInt>
+DenseElementsAttr::ComplexIntElementIterator::operator*() const {
+  size_t storageWidth = getDenseElementStorageWidth(bitWidth);
+  size_t offset = getDataIndex() * storageWidth * 2;
+  return {readBits(getData(), offset, bitWidth),
+          readBits(getData(), offset + storageWidth, bitWidth)};
+}
+
+//===----------------------------------------------------------------------===//
+// FloatElementIterator
+
 DenseElementsAttr::FloatElementIterator::FloatElementIterator(
     const llvm::fltSemantics &smt, IntElementIterator it)
     : llvm::mapped_iterator<IntElementIterator,
                             std::function<APFloat(const APInt &)>>(
           it, [&](const APInt &val) { return APFloat(smt, val); }) {}
 
 //===----------------------------------------------------------------------===//
+// ComplexFloatElementIterator
+
+DenseElementsAttr::ComplexFloatElementIterator::ComplexFloatElementIterator(
+    const llvm::fltSemantics &smt, ComplexIntElementIterator it)
+    : llvm::mapped_iterator<
+          ComplexIntElementIterator,
+          std::function<std::complex<APFloat>(const std::complex<APInt> &)>>(
+          it, [&](const std::complex<APInt> &val) -> std::complex<APFloat> {
+            return {APFloat(smt, val.real()), APFloat(smt, val.imag())};
+          }) {}
+
+//===----------------------------------------------------------------------===//
 // DenseElementsAttr
 //===----------------------------------------------------------------------===//
 
 DenseElementsAttr DenseElementsAttr::get(ShapedType type,
                                          ArrayRef<Attribute> values) {
   assert(type.getElementType().isIntOrIndexOrFloat() &&
          "expected int or index or float element type");
   assert(hasSameElementsOrSplat(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());
-  return DenseIntOrFPElementsAttr::getRaw(type, values);
+  assert(hasSameElementsOrSplat(type, values));
+  size_t storageBitWidth = getDenseElementStorageWidth(type.getElementType());
+  return DenseIntOrFPElementsAttr::getRaw(type, storageBitWidth, values,
+                                          /*isSplat=*/(values.size() == 1));
+}
+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));
 }
 
 // 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>());
-
-  // Convert the APFloat values to APInt and create a dense elements attribute.
-  std::vector<APInt> intValues(values.size());
-  for (unsigned i = 0, e = values.size(); i != e; ++i)
-    intValues[i] = values[i].bitcastToAPInt();
-  return DenseIntOrFPElementsAttr::getRaw(type, intValues);
+  assert(hasSameElementsOrSplat(type, values));
+  size_t storageBitWidth = getDenseElementStorageWidth(type.getElementType());
+  return DenseIntOrFPElementsAttr::getRaw(type, storageBitWidth, values,
+                                          /*isSplat=*/(values.size() == 1));
+}
+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));
 }
 
 /// 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);
 }
 
 /// 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 storageWidth = getDenseElementStorageWidth(type.getElementType());
   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());

 auto DenseElementsAttr::int_value_begin() const -> IntElementIterator {
   assert(getType().getElementType().isIntOrIndex() && "expected integral type");
   return raw_int_begin();
 }
 auto DenseElementsAttr::int_value_end() const -> IntElementIterator {
   assert(getType().getElementType().isIntOrIndex() && "expected integral type");
   return raw_int_end();
 }
+auto DenseElementsAttr::getComplexIntValues() const
+    -> llvm::iterator_range<ComplexIntElementIterator> {
+  Type eltTy = getType().getElementType().cast<ComplexType>().getElementType();
+  (void)eltTy;
+  assert(eltTy.isa<IntegerType>() && "expected complex integral type");
+  return {ComplexIntElementIterator(*this, 0),
+          ComplexIntElementIterator(*this, getNumElements())};
+}
 
 /// Return the held element values as a range of APFloat. The element type of
 /// this attribute must be of float type.
 auto DenseElementsAttr::getFloatValues() const
     -> llvm::iterator_range<FloatElementIterator> {
   auto elementType = getType().getElementType().cast<FloatType>();
-  assert(elementType.isa<FloatType>() && "expected float type");
   const auto &elementSemantics = elementType.getFloatSemantics();
   return {FloatElementIterator(elementSemantics, raw_int_begin()),
           FloatElementIterator(elementSemantics, raw_int_end())};
 }
 auto DenseElementsAttr::float_value_begin() const -> FloatElementIterator {
   return getFloatValues().begin();
 }
 auto DenseElementsAttr::float_value_end() const -> FloatElementIterator {
   return getFloatValues().end();
 }
+auto DenseElementsAttr::getComplexFloatValues() const
+    -> llvm::iterator_range<ComplexFloatElementIterator> {
+  Type eltTy = getType().getElementType().cast<ComplexType>().getElementType();
+  assert(eltTy.isa<FloatType>() && "expected complex float type");
+  const auto &semantics = eltTy.cast<FloatType>().getFloatSemantics();
+  return {{semantics, {*this, 0}},
+          {semantics, {*this, static_cast<size_t>(getNumElements())}}};
+}
 
 /// Return the raw storage data held by this attribute.
 ArrayRef<char> DenseElementsAttr::getRawData() const {
   return static_cast<DenseIntOrFPElementsAttributeStorage *>(impl)->data;
 }
 
 ArrayRef<StringRef> DenseElementsAttr::getRawStringData() const {
   return static_cast<DenseStringElementsAttributeStorage *>(impl)->data;
DenseStringElementsAttr::get(ShapedType type, ArrayRef<StringRef> values) {
   return Base::get(type.getContext(), StandardAttributes::DenseStringElements,
                    type, values, (values.size() == 1));
 }
 
 //===----------------------------------------------------------------------===//
 // 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, CHAR_BIT) * llvm::size(values));
+  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);
+  }
+}
+
+/// 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) {
+  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);
+}
+
 /// 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,
-                                                   ArrayRef<APInt> values) {
-  assert(hasSameElementsOrSplat(type, values));
-
-  size_t bitWidth = getDenseElementBitWidth(type.getElementType());
-  size_t storageBitWidth = getDenseElementStorageWidth(bitWidth);
-  std::vector<char> elementData(llvm::divideCeil(storageBitWidth, CHAR_BIT) *
-                                values.size());
-  for (unsigned i = 0, e = values.size(); i != e; ++i) {
-    assert(values[i].getBitWidth() == bitWidth);
-    writeBits(elementData.data(), i * storageBitWidth, values[i]);
-  }
-  return DenseIntOrFPElementsAttr::getRaw(type, elementData,
-                                          /*isSplat=*/(values.size() == 1));
+                                                   size_t storageWidth,
+                                                   ArrayRef<APInt> values,
+                                                   bool isSplat) {
+  std::vector<char> data;
+  writeAPIntsToBuffer(storageWidth, data, values);
+  return DenseIntOrFPElementsAttr::getRaw(type, data, isSplat);
 }
 
 DenseElementsAttr DenseIntOrFPElementsAttr::getRaw(ShapedType type,
                                                    ArrayRef<char> data,
                                                    bool isSplat) {
   assert((type.isa<RankedTensorType>() || type.isa<VectorType>()) &&
          "type must be ranked tensor or vector");
   assert(type.hasStaticShape() && "type must have static shape");

mlir/lib/Parser/Parser.cpp

public:
 
   /// Build a dense attribute instance with the parsed elements and the given
   /// shaped type.
   DenseElementsAttr getAttr(llvm::SMLoc loc, ShapedType type);
 
   ArrayRef<int64_t> getShape() const { return shape; }
 
 private:
-  enum class ElementKind { Boolean, Integer, Float, String };
-
-  /// Return a string to represent the given element kind.
-  const char *getElementKindStr(ElementKind kind) {
-    switch (kind) {
-    case ElementKind::Boolean:
-      return "'boolean'";
-    case ElementKind::Integer:
-      return "'integer'";
-    case ElementKind::Float:
-      return "'float'";
-    case ElementKind::String:
-      return "'string'";
-    }
-    llvm_unreachable("unknown element kind");
-  }
-
-  /// Build a Dense Integer attribute for the given type.
-  DenseElementsAttr getIntAttr(llvm::SMLoc loc, ShapedType type, Type eltTy);
-
-  /// Build a Dense Float attribute for the given type.
-  DenseElementsAttr getFloatAttr(llvm::SMLoc loc, ShapedType type,
-                                 FloatType eltTy);
+  /// Get the parsed elements for an integer attribute.
+  ParseResult getIntAttrElements(llvm::SMLoc loc, Type eltTy,
+                                 std::vector<APInt> &intValues);
+
+  /// Get the parsed elements for a float attribute.
+  ParseResult getFloatAttrElements(llvm::SMLoc loc, FloatType eltTy,
+                                   std::vector<APFloat> &floatValues);
 
   /// Build a Dense String attribute for the given type.
   DenseElementsAttr getStringAttr(llvm::SMLoc loc, ShapedType type, Type eltTy);
 
   /// Build a Dense attribute with hex data for the given type.
   DenseElementsAttr getHexAttr(llvm::SMLoc loc, ShapedType type);
 
   /// Parse a single element, returning failure if it isn't a valid element
private:
   Parser &p;
 
   /// The shape inferred from the parsed elements.
   SmallVector<int64_t, 4> shape;
 
   /// Storage used when parsing elements, this is a pair of <is_negated, token>.
   std::vector<std::pair<bool, Token>> storage;
 
-  /// A flag that indicates the type of elements that have been parsed.
-  Optional<ElementKind> knownEltKind;
-
   /// Storage used when parsing elements that were stored as hex values.
   Optional<Token> hexStorage;
 };
 } // namespace
 
 /// Parse the elements of a tensor literal. If 'allowHex' is true, the parser
 /// may also parse a tensor literal that is store as a hex string.
 ParseResult TensorLiteralParser::parse(bool allowHex) {
DenseElementsAttr TensorLiteralParser::getAttr(llvm::SMLoc loc,
   // 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";
+  // Handle complex types in the specific element type cases below.
+  bool isComplex = false;
+  if (ComplexType complexTy = eltType.dyn_cast<ComplexType>()) {
+    eltType = complexTy.getElementType();
+    isComplex = true;
+  }
+
+  // Handle integer and index types.
+  if (eltType.isIntOrIndex()) {
+    std::vector<APInt> intValues;
+    if (failed(getIntAttrElements(loc, eltType, intValues)))
+      return nullptr;
+    if (isComplex) {
+      // If this is a complex, treat the parsed values as complex values.
+      auto complexData = llvm::makeArrayRef(
+          reinterpret_cast<std::complex<APInt> *>(intValues.data()),
+          intValues.size() / 2);
+      return DenseElementsAttr::get(type, complexData);
+    }
+    return DenseElementsAttr::get(type, intValues);
+  }
+  // Handle floating point types.
+  if (FloatType floatTy = eltType.dyn_cast<FloatType>()) {
+    std::vector<APFloat> floatValues;
+    if (failed(getFloatAttrElements(loc, floatTy, floatValues)))
       return nullptr;
+    if (isComplex) {
+      // If this is a complex, treat the parsed values as complex values.
+      auto complexData = llvm::makeArrayRef(
+          reinterpret_cast<std::complex<APFloat> *>(floatValues.data()),
+          floatValues.size() / 2);
+      return DenseElementsAttr::get(type, complexData);
+    }
+    return DenseElementsAttr::get(type, floatValues);
   }
 
   // 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;
-  intElements.reserve(storage.size());
-  auto isUintType = type.getElementType().isUnsignedInteger();
+ParseResult
+TensorLiteralParser::getIntAttrElements(llvm::SMLoc loc, Type eltTy,
+                                        std::vector<APInt> &intValues) {
+  intValues.reserve(storage.size());
+  bool isUintType = eltTy.isUnsignedInteger();
   for (const auto &signAndToken : storage) {
     bool isNegative = signAndToken.first;
     const Token &token = signAndToken.second;
     auto tokenLoc = token.getLoc();
 
     if (isNegative && isUintType) {
-      p.emitError(tokenLoc)
+      return p.emitError(tokenLoc)
              << "expected unsigned integer elements, but parsed negative value";
-      return nullptr;
     }
 
     // Check to see if floating point values were parsed.
     if (token.is(Token::floatliteral)) {
-      p.emitError(tokenLoc)
+      return p.emitError(tokenLoc)
              << "expected integer elements, but parsed floating-point";
-      return nullptr;
     }
 
     assert(token.isAny(Token::integer, Token::kw_true, Token::kw_false) &&
            "unexpected token type");
     if (token.isAny(Token::kw_true, Token::kw_false)) {
       if (!eltTy.isInteger(1)) {
-        p.emitError(tokenLoc)
+        return p.emitError(tokenLoc)
                << "expected i1 type for 'true' or 'false' values";
-        return nullptr;
       }
       APInt apInt(1, token.is(Token::kw_true), /*isSigned=*/false);
-      intElements.push_back(apInt);
+      intValues.push_back(apInt);
       continue;
     }
 
     // Create APInt values for each element with the correct bitwidth.
     Optional<APInt> apInt =
         buildAttributeAPInt(eltTy, isNegative, token.getSpelling());
     if (!apInt)
-      return (p.emitError(tokenLoc, "integer constant out of range for type"),
-              nullptr);
-    intElements.push_back(*apInt);
+      return p.emitError(tokenLoc, "integer constant out of range for type");
+    intValues.push_back(*apInt);
   }
-
-  return DenseElementsAttr::get(type, intElements);
+  return success();
 }
 
 /// Build a Dense Float attribute for the given type.
-DenseElementsAttr TensorLiteralParser::getFloatAttr(llvm::SMLoc loc,
-                                                    ShapedType type,
-                                                    FloatType eltTy) {
-  std::vector<APFloat> floatValues;
+ParseResult
+TensorLiteralParser::getFloatAttrElements(llvm::SMLoc loc, FloatType eltTy,
+                                          std::vector<APFloat> &floatValues) {
   floatValues.reserve(storage.size());
   for (const auto &signAndToken : storage) {
     bool isNegative = signAndToken.first;
     const Token &token = signAndToken.second;
 
     // Handle hexadecimal float literals.
     if (token.is(Token::integer) && token.getSpelling().startswith("0x")) {
       if (isNegative) {
-        p.emitError(token.getLoc())
+        return p.emitError(token.getLoc())
                << "hexadecimal float literal should not have a leading minus";
-        return nullptr;
       }
       auto val = token.getUInt64IntegerValue();
       if (!val.hasValue()) {
-        p.emitError("hexadecimal float constant out of range for attribute");
-        return nullptr;
+        return p.emitError(
+            "hexadecimal float constant out of range for attribute");
       }
       Optional<APFloat> apVal = buildHexadecimalFloatLiteral(&p, eltTy, *val);
       if (!apVal)
-        return nullptr;
+        return failure();
       floatValues.push_back(*apVal);
       continue;
     }
 
     // Check to see if any decimal integers or booleans were parsed.
-    if (!token.is(Token::floatliteral)) {
-      p.emitError() << "expected floating-point elements, but parsed integer";
-      return nullptr;
-    }
+    if (!token.is(Token::floatliteral))
+      return p.emitError()
+             << "expected floating-point elements, but parsed integer";
 
     // Build the float values from tokens.
     auto val = token.getFloatingPointValue();
-    if (!val.hasValue()) {
-      p.emitError("floating point value too large for attribute");
-      return nullptr;
-    }
+    if (!val.hasValue())
+      return p.emitError("floating point value too large for attribute");
+
     // Treat BF16 as double because it is not supported in LLVM's APFloat.
     APFloat apVal(isNegative ? -*val : *val);
     if (!eltTy.isBF16() && !eltTy.isF64()) {
       bool unused;
       apVal.convert(eltTy.getFloatSemantics(), APFloat::rmNearestTiesToEven,
                     &unused);
     }
     floatValues.push_back(apVal);
   }
-
-  return DenseElementsAttr::get(type, floatValues);
+  return success();
 }
 
 /// Build a Dense String attribute for the given type.
 DenseElementsAttr TensorLiteralParser::getStringAttr(llvm::SMLoc loc,
                                                      ShapedType type,
                                                      Type eltTy) {
   if (hexStorage.hasValue()) {
     auto stringValue = hexStorage.getValue().getStringValue();
case Token::minus:
     storage.emplace_back(/*isNegative=*/true, p.getToken());
     p.consumeToken();
     break;
 
   case Token::string:
     storage.emplace_back(/*isNegative=*/ false, p.getToken());
     p.consumeToken();
     break;
+
+  // Parse a complex element of the form '(' element ',' element ')'.
+  case Token::l_paren:
+    p.consumeToken(Token::l_paren);
+    if (parseElement() ||
+        p.parseToken(Token::comma, "expected ',' between complex elements") ||
+        parseElement() ||
+        p.parseToken(Token::r_paren, "expected ')' after complex elements"))
+      return failure();
+    break;
+
   default:
     return p.emitError("expected element literal of primitive type");
   }
 
   return success();
 }
 
 /// Parse a list of either lists or elements, returning the dimensions of the

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<2xcomplex<f64>>
+"foo.op"() {dense.attr = dense<"0x0000000000002440000000000000144000000000000024400000000000001440"> : tensor<2xcomplex<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/test/IR/invalid.mlir

 
 func @elementsattr_toolarge2() -> () {
 ^bb0:
   "foo"(){bar = dense<[-777]> : tensor<1xi8>} : () -> () // expected-error {{integer constant out of range}}
 }
 
 // -----
 
+"foo"(){bar = dense<[()]> : tensor<complex<i64>>} : () -> () // expected-error {{expected element literal of primitive type}}
+
+// -----
+
+"foo"(){bar = dense<[(10)]> : tensor<complex<i64>>} : () -> () // expected-error {{expected ',' between complex elements}}
+
+// -----
+
+"foo"(){bar = dense<[(10,)]> : tensor<complex<i64>>} : () -> () // expected-error {{expected element literal of primitive type}}
+
+// -----
+
+"foo"(){bar = dense<[(10,10]> : tensor<complex<i64>>} : () -> () // expected-error {{expected ')' after complex elements}}
+
+// -----
+
 func @elementsattr_malformed_opaque() -> () {
 ^bb0:
   "foo"(){bar = opaque<10, "0xQZz123"> : tensor<1xi8>} : () -> () // expected-error {{expected dialect namespace}}
 }
 
 // -----
 
 func @elementsattr_malformed_opaque1() -> () {

mlir/test/IR/parser.mlir

// CHECK: "intscalar"() {bar = dense<1> : tensor<i32>} : () -> ()
   "intscalar"(){bar = dense<1> : tensor<i32>} : () -> ()
 // CHECK: "floatscalar"() {bar = dense<5.000000e+00> : tensor<f32>} : () -> ()
   "floatscalar"(){bar = dense<5.0> : tensor<f32>} : () -> ()
 
 // CHECK: "index"() {bar = dense<1> : tensor<index>} : () -> ()
   "index"(){bar = dense<1> : tensor<index>} : () -> ()
 // CHECK: "index"() {bar = dense<[1, 2]> : tensor<2xindex>} : () -> ()
   "index"(){bar = dense<[1, 2]> : tensor<2xindex>} : () -> ()
+
+  // CHECK: dense<(1,1)> : tensor<complex<i64>>
+  "complex_attr"(){bar = dense<(1,1)> : tensor<complex<i64>>} : () -> ()
+  // CHECK: dense<[(1,1), (2,2)]> : tensor<2xcomplex<i64>>
+  "complex_attr"(){bar = dense<[(1,1), (2,2)]> : tensor<2xcomplex<i64>>} : () -> ()
+  // CHECK: dense<(1.000000e+00,0.000000e+00)> : tensor<complex<f32>>
+  "complex_attr"(){bar = dense<(1.000000e+00,0.000000e+00)> : tensor<complex<f32>>} : () -> ()
+  // CHECK: dense<[(1.000000e+00,0.000000e+00), (2.000000e+00,2.000000e+00)]> : tensor<2xcomplex<f32>>
+  "complex_attr"(){bar = dense<[(1.000000e+00,0.000000e+00), (2.000000e+00,2.000000e+00)]> : tensor<2xcomplex<f32>>} : () -> ()
   return
 }
 
 // CHECK-LABEL: func @densevectorattr
 func @densevectorattr() -> () {
 ^bb0:
 // NOTE: The {{\[\[}} syntax is because "[[" confuses FileCheck.
 // CHECK: "fooi8"() {bar = dense<5> : vector<1x1x1xi8>} : () -> ()

mlir/unittests/IR/AttributeTest.cpp

static void testSplat(Type eltType, const EltTy &splatElt) {
   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);
+    EXPECT_TRUE(newValue == splatElt);
 }
 
 namespace {
 TEST(DenseSplatTest, BoolSplat) {
   MLIRContext context;
   IntegerType boolTy = IntegerType::get(1, &context);
   RankedTensorType shape = RankedTensorType::get({2, 2}, boolTy);
 

 
 TEST(DenseComplexTest, ComplexIntSplat) {
   MLIRContext context;
   ComplexType complexType = ComplexType::get(IntegerType::get(64, &context));
   std::complex<int64_t> value(10, 15);
   testSplat(complexType, value);
 }
 
+TEST(DenseComplexTest, ComplexAPFloatSplat) {
+  MLIRContext context;
+  ComplexType complexType = ComplexType::get(FloatType::getF32(&context));
+  std::complex<APFloat> value(APFloat(10.0f), APFloat(15.0f));
+  testSplat(complexType, value);
+}
+
+TEST(DenseComplexTest, ComplexAPIntSplat) {
+  MLIRContext context;
+  ComplexType complexType = ComplexType::get(IntegerType::get(64, &context));
+  std::complex<APInt> value(APInt(64, 10), APInt(64, 15));
+  testSplat(complexType, value);
+}
+
 } // end namespace