Add support for IndexType inside DenseIntElementsAttr.

[silvas]

Summary

This also fixes issues discovered in the parsing/printing path.

[rriddle]

mlir/lib/Parser/Parser.cpp

2047

nit:

Type eltType = type.getElementType();
if (eltType.isIntOrIndex())
  ...

...
auto floatTy = eltType...

2092

nit: Can you just replace this with 1?

[silvas]

Address comments.

[silvas]

Landed in 1b2c7877a4dd3da2ba7462d65a48341f59c6cc61

mlir/include/mlir/IR/StandardTypes.h

 public:
   using Base::Base;
 
   /// Get an instance of the IndexType.
   static IndexType get(MLIRContext *context);
 
   /// Support method to enable LLVM-style type casting.
   static bool kindof(unsigned kind) { return kind == StandardTypes::Index; }
+
+  /// Storage bit width used for IndexType by internal compiler data structures.
+  static constexpr unsigned kInternalStorageBitWidth = 64;
 };
 
 /// Integer types can have arbitrary bitwidth up to a large fixed limit.
 class IntegerType
     : public Type::TypeBase<IntegerType, Type, detail::IntegerTypeStorage> {
 public:
   using Base::Base;
 

mlir/include/mlir/IR/Types.h

#include "DialectSymbolRegistry.def"
 
   /// Return true if this is a signless integer or index type.
   bool isSignlessIntOrIndex();
   /// Return true if this is a signless integer, index, or float type.
   bool isSignlessIntOrIndexOrFloat();
   /// Return true of this is a signless integer or a float type.
   bool isSignlessIntOrFloat();
 
+  /// Return true if this is an integer (of any signedness) or an index type.
+  bool isIntOrIndex();
   /// Return true if this is an integer (of any signedness) or a float type.
   bool isIntOrFloat();
   /// Return true if this is an integer (of any signedness), index, or float
   /// type.
   bool isIntOrIndexOrFloat();
 
   /// Print the current type.
   void print(raw_ostream &os);

mlir/lib/IR/AsmPrinter.cpp

 void ModulePrinter::printDenseElementsAttr(DenseElementsAttr 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().isa<IntegerType>()
+  auto printEltFn = type.getElementType().isIntOrIndex()
                         ? printDenseIntElement
                         : printDenseFloatElement;
 
   // Special case for 0-d and splat tensors.
   if (attr.isSplat()) {
     printEltFn(attr, os, 0, isSigned);
     return;
   }

mlir/lib/IR/AttributeDetail.h

struct TypeAttributeStorage : public AttributeStorage {
 
   Type value;
 };
 
 //===----------------------------------------------------------------------===//
 // Elements Attributes
 //===----------------------------------------------------------------------===//
 
+/// Return the bit width which DenseElementsAttr should use for this type.
+inline size_t getDenseElementBitWidth(Type eltType) {
+  // 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();
+}
+
 /// An attribute representing a reference to a dense vector or tensor object.
 struct DenseElementsAttributeStorage : public AttributeStorage {
   struct KeyTy {
     KeyTy(ShapedType type, ArrayRef<char> data, llvm::hash_code hashCode,
           bool isSplat = false)
         : type(type), data(data), hashCode(hashCode), isSplat(isSplat) {}
 
     /// The type of the dense elements.
struct DenseElementsAttributeStorage : public AttributeStorage {
   bool operator==(const KeyTy &key) const {
     if (key.type != getType())
       return false;
 
     // For boolean splats we need to explicitly check that the first bit is the
     // same. Boolean values are packed at the bit level, and even though a splat
     // is detected the rest of the bits in the first byte may differ from the
     // splat value.
-    if (key.type.getElementTypeBitWidth() == 1) {
+    if (key.type.getElementType().isInteger(1)) {
       if (key.isSplat != isSplat)
         return false;
       if (isSplat)
         return (key.data.front() & 1) == data.front();
     }
 
     // Otherwise, we can default to just checking the data.
     return key.data == data;
static KeyTy getKey(ShapedType ty, ArrayRef<char> data, bool isKnownSplat) {
 
     // Otherwise, we need to check if the data corresponds to a splat or not.
 
     // Handle the simple case of only one element.
     size_t numElements = ty.getNumElements();
     assert(numElements != 1 && "splat of 1 element should already be detected");
 
     // Handle boolean values directly as they are packed to 1-bit.
-    size_t elementWidth = ty.getElementTypeBitWidth();
-    if (elementWidth == 1)
+    if (ty.getElementType().isInteger(1) == 1)
       return getKeyForBoolData(ty, data, numElements);
 
-    // FIXME(b/121118307): using 64 bits for BF16 because it is currently stored
-    // with double semantics.
-    if (ty.getElementType().isBF16())
-      elementWidth = 64;
-
+    size_t elementWidth = getDenseElementBitWidth(ty.getElementType());
     // Non 1-bit dense elements are padded to 8-bits.
     size_t storageSize = llvm::divideCeil(elementWidth, CHAR_BIT);
     assert(((data.size() / storageSize) == numElements) &&
            "data does not hold expected number of elements");
 
     // Create the initial hash value with just the first element.
     auto firstElt = data.take_front(storageSize);
     auto hashVal = llvm::hash_value(firstElt);
construct(AttributeStorageAllocator &allocator, KeyTy key) {
     // 64-bit alignment.
     ArrayRef<char> copy, data = key.data;
     if (!data.empty()) {
       char *rawData = reinterpret_cast<char *>(
           allocator.allocate(data.size(), alignof(uint64_t)));
       std::memcpy(rawData, data.data(), data.size());
 
       // If this is a boolean splat, make sure only the first bit is used.
-      if (key.isSplat && key.type.getElementTypeBitWidth() == 1)
+      if (key.isSplat && key.type.getElementType().isInteger(1))
         rawData[0] &= 1;
       copy = ArrayRef<char>(rawData, data.size());
     }
 
     return new (allocator.allocate<DenseElementsAttributeStorage>())
         DenseElementsAttributeStorage(key.type, copy, key.isSplat);
   }
 

mlir/lib/IR/Attributes.cpp

 
 IntegerAttr IntegerAttr::get(Type type, const APInt &value) {
   return Base::get(type.getContext(), StandardAttributes::Integer, type, value);
 }
 
 IntegerAttr IntegerAttr::get(Type type, int64_t value) {
   // This uses 64 bit APInts by default for index type.
   if (type.isIndex())
-    return get(type, APInt(64, value));
+    return get(type, APInt(IndexType::kInternalStorageBitWidth, value));
 
   auto intType = type.cast<IntegerType>();
   return get(type, APInt(intType.getWidth(), value, intType.isSignedInteger()));
 }
 
 APInt IntegerAttr::getValue() const { return getImpl()->getValue(); }
 
 int64_t IntegerAttr::getInt() const {
uint64_t ElementsAttr::getFlattenedIndex(ArrayRef<uint64_t> index) const {
   }
   return valueIndex;
 }
 
 //===----------------------------------------------------------------------===//
 // DenseElementAttr Utilities
 //===----------------------------------------------------------------------===//
 
-static size_t getDenseElementBitwidth(Type eltType) {
-  // FIXME(b/121118307): using 64 bits for BF16 because it is currently stored
-  // with double semantics.
-  return eltType.isBF16() ? 64 : eltType.getIntOrFloatBitWidth();
-}
-
 /// 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);
 }
 
 /// Set a bit to a specific value.
 static void setBit(char *rawData, size_t bitPos, bool value) {
bool DenseElementsAttr::BoolElementIterator::operator*() const {
   return getBit(getData(), getDataIndex());
 }
 
 /// Constructs a new iterator.
 DenseElementsAttr::IntElementIterator::IntElementIterator(
     DenseElementsAttr attr, size_t dataIndex)
     : DenseElementIndexedIteratorImpl<IntElementIterator, APInt, APInt, APInt>(
           attr.getRawData().data(), attr.isSplat(), dataIndex),
-      bitWidth(getDenseElementBitwidth(attr.getType().getElementType())) {}
+      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);
 }
 
 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); }) {}
 
 //===----------------------------------------------------------------------===//
 // DenseElementsAttr
 //===----------------------------------------------------------------------===//
 
 DenseElementsAttr DenseElementsAttr::get(ShapedType type,
                                          ArrayRef<Attribute> values) {
-  assert(type.getElementType().isIntOrFloat() &&
-         "expected int or float element type");
+  assert(type.getElementType().isIntOrIndexOrFloat() &&
+         "expected int or index or float element type");
   assert(hasSameElementsOrSplat(type, values));
 
   auto eltType = type.getElementType();
-  size_t bitWidth = getDenseElementBitwidth(eltType);
+  size_t bitWidth = getDenseElementBitWidth(eltType);
   size_t storageBitWidth = getDenseElementStorageWidth(bitWidth);
 
   // Compress the attribute values into a character buffer.
   SmallVector<char, 8> data(llvm::divideCeil(storageBitWidth, CHAR_BIT) *
                             values.size());
   APInt intVal;
   for (unsigned i = 0, e = values.size(); i < e; ++i) {
     assert(eltType == values[i].getType() &&
            "expected attribute value to have element type");
 
     switch (eltType.getKind()) {
     case StandardTypes::BF16:
     case StandardTypes::F16:
     case StandardTypes::F32:
     case StandardTypes::F64:
       intVal = values[i].cast<FloatAttr>().getValue().bitcastToAPInt();
       break;
     case StandardTypes::Integer:
+    case StandardTypes::Index:
       intVal = values[i].isa<BoolAttr>()
                    ? APInt(1, values[i].cast<BoolAttr>().getValue() ? 1 : 0)
                    : values[i].cast<IntegerAttr>().getValue();
       break;
     default:
       llvm_unreachable("unexpected element type");
     }
     assert(intVal.getBitWidth() == bitWidth &&
DenseElementsAttr DenseElementsAttr::get(ShapedType type,
   return getRaw(type, buff, /*isSplat=*/(values.size() == 1));
 }
 
 /// 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().isa<IntegerType>());
+  assert(type.getElementType().isIntOrIndex());
   return getRaw(type, 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'.
 DenseElementsAttr DenseElementsAttr::get(ShapedType type,
                                          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'.
 DenseElementsAttr DenseElementsAttr::getRaw(ShapedType type,
                                             ArrayRef<APInt> values) {
   assert(hasSameElementsOrSplat(type, values));
 
-  size_t bitWidth = getDenseElementBitwidth(type.getElementType());
+  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 getRaw(type, elementData, /*isSplat=*/(values.size() == 1));

 }
 
 /// 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,
                               bool isSigned) {
   // Make sure that the data element size is the same as the type element width.
-  if (getDenseElementBitwidth(type.getElementType()) !=
+  if (getDenseElementBitWidth(type.getElementType()) !=
       static_cast<size_t>(dataEltSize * CHAR_BIT))
     return false;
 
-  // Check that the element type is either float or integer.
+  // Check that the element type is either float or integer or index.
   if (!isInt)
     return type.getElementType().isa<FloatType>();
 
+  if (type.getElementType().isIndex())
+    return true;
+
   auto intType = type.getElementType().dyn_cast<IntegerType>();
   if (!intType)
     return false;
 
   // Make sure signedness semantics is consistent.
   if (intType.isSignless())
     return true;
   return intType.isSigned() ? isSigned : !isSigned;
auto DenseElementsAttr::getBoolValues() const
   return {BoolElementIterator(*this, 0),
           BoolElementIterator(*this, getNumElements())};
 }
 
 /// Return the held element values as a range of APInts. The element type of
 /// this attribute must be of integer type.
 auto DenseElementsAttr::getIntValues() const
     -> llvm::iterator_range<IntElementIterator> {
-  assert(getType().getElementType().isa<IntegerType>() &&
-         "expected integer type");
+  assert(getType().getElementType().isIntOrIndex() && "expected integral type");
   return {raw_int_begin(), raw_int_end()};
 }
 auto DenseElementsAttr::int_value_begin() const -> IntElementIterator {
-  assert(getType().getElementType().isa<IntegerType>() &&
-         "expected integer type");
+  assert(getType().getElementType().isIntOrIndex() && "expected integral type");
   return raw_int_begin();
 }
 auto DenseElementsAttr::int_value_end() const -> IntElementIterator {
-  assert(getType().getElementType().isa<IntegerType>() &&
-         "expected integer type");
+  assert(getType().getElementType().isIntOrIndex() && "expected integral type");
   return raw_int_end();
 }
 
 /// 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>();

 //===----------------------------------------------------------------------===//
 // DenseFPElementsAttr
 //===----------------------------------------------------------------------===//
 
 template <typename Fn, typename Attr>
 static ShapedType mappingHelper(Fn mapping, Attr &attr, ShapedType inType,
                                 Type newElementType,
                                 llvm::SmallVectorImpl<char> &data) {
-  size_t bitWidth = getDenseElementBitwidth(newElementType);
+  size_t bitWidth = getDenseElementBitWidth(newElementType);
   size_t storageBitWidth = getDenseElementStorageWidth(bitWidth);
 
   ShapedType newArrayType;
   if (inType.isa<RankedTensorType>())
     newArrayType = RankedTensorType::get(inType.getShape(), newElementType);
   else if (inType.isa<UnrankedTensorType>())
     newArrayType = RankedTensorType::get(inType.getShape(), newElementType);
   else if (inType.isa<VectorType>())
auto newArrayType =
       mappingHelper(mapping, *this, getType(), newElementType, elementData);
 
   return getRaw(newArrayType, elementData, isSplat());
 }
 
 /// 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().isa<IntegerType>();
+         attr.getType().cast<ShapedType>().getElementType().isIntOrIndex();
 }
 
 //===----------------------------------------------------------------------===//
 // OpaqueElementsAttr
 //===----------------------------------------------------------------------===//
 
 OpaqueElementsAttr OpaqueElementsAttr::get(Dialect *dialect, ShapedType type,
                                            StringRef bytes) {

mlir/lib/IR/StandardTypes.cpp

 bool Type::isSignlessIntOrIndexOrFloat() {
   return isa<IndexType>() || isSignlessInteger() || isa<FloatType>();
 }
 
 bool Type::isSignlessIntOrFloat() {
   return isSignlessInteger() || isa<FloatType>();
 }
 
+bool Type::isIntOrIndex() { return isa<IntegerType>() || isIndex(); }
+
 bool Type::isIntOrFloat() { return isa<IntegerType>() || isa<FloatType>(); }
 
 bool Type::isIntOrIndexOrFloat() { return isIntOrFloat() || isIndex(); }
 
 //===----------------------------------------------------------------------===//
 // Integer Type
 //===----------------------------------------------------------------------===//
 

mlir/lib/Parser/Parser.cpp

static Optional<APInt> buildAttributeAPInt(Type type, bool isNegative,
                                            StringRef spelling) {
   // Parse the integer value into an APInt that is big enough to hold the value.
   APInt result;
   bool isHex = spelling.size() > 1 && spelling[1] == 'x';
   if (spelling.getAsInteger(isHex ? 0 : 10, result))
     return llvm::None;
 
   // Extend or truncate the bitwidth to the right size.
-  unsigned width = type.isIndex() ? 64 : type.getIntOrFloatBitWidth();
+  unsigned width = type.isIndex() ? IndexType::kInternalStorageBitWidth
+                                  : type.getIntOrFloatBitWidth();
   if (width > result.getBitWidth()) {
     result = result.zext(width);
   } else if (width < result.getBitWidth()) {
     // The parser can return an unnecessarily wide result with leading zeros.
     // This isn't a problem, but truncating off bits is bad.
     if (result.countLeadingZeros() < result.getBitWidth() - width)
       return llvm::None;
 
case ElementKind::Integer:
       return "'integer'";
     case ElementKind::Float:
       return "'float'";
     }
     llvm_unreachable("unknown element kind");
   }
 
   /// Build a Dense Integer attribute for the given type.
-  DenseElementsAttr getIntAttr(llvm::SMLoc loc, ShapedType type,
-                               IntegerType eltTy);
+  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);
 
   /// Build a Dense attribute with hex data for the given type.
   DenseElementsAttr getHexAttr(llvm::SMLoc loc, ShapedType type);
 
DenseElementsAttr TensorLiteralParser::getAttr(llvm::SMLoc loc,
   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 = type.getElementType().dyn_cast<IntegerType>())
     return getIntAttr(loc, type, intTy);
+  if (auto indexTy = type.getElementType().dyn_cast<IndexType>())
+    return getIntAttr(loc, type, indexTy);
 
   // Otherwise, this must be a floating point type.
   auto floatTy = type.getElementType().dyn_cast<FloatType>();
   if (!floatTy) {
     p.emitError(loc) << "expected floating-point or integer element type, got "
                      << type.getElementType();
     return nullptr;
   }
   return getFloatAttr(loc, type, floatTy);
 }
 
 /// Build a Dense Integer attribute for the given type.
 DenseElementsAttr TensorLiteralParser::getIntAttr(llvm::SMLoc loc,
-                                                  ShapedType type,
-                                                  IntegerType eltTy) {
+                                                  ShapedType type, Type eltTy) {
   std::vector<APInt> intElements;
   intElements.reserve(storage.size());
   auto isUintType = type.getElementType().isUnsignedInteger();
   for (const auto &signAndToken : storage) {
     bool isNegative = signAndToken.first;
     const Token &token = signAndToken.second;
     auto tokenLoc = token.getLoc();
 
for (const auto &signAndToken : storage) {
     }
 
     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)
             << "expected i1 type for 'true' or 'false' values";
-      APInt apInt(eltTy.getWidth(), token.is(Token::kw_true),
+      APInt apInt(eltTy.getIntOrFloatBitWidth(), token.is(Token::kw_true),
                   /*isSigned=*/false);
       intElements.push_back(apInt);
       continue;
     }
 
     // Create APInt values for each element with the correct bitwidth.
     Optional<APInt> apInt =
         buildAttributeAPInt(eltTy, isNegative, token.getSpelling());

mlir/test/IR/parser.mlir

// CHECK: "float32"() {bar = dense<{{\[\[\[}}-5.000000e+00, 6.000000e+00, 1.000000e+00, 2.000000e+00]], {{\[\[}}7.000000e+00, -8.000000e+00, 3.000000e+00, 4.000000e+00]]]> : tensor<2x1x4xf32>} : () -> ()
   "float32"(){bar = dense<[[[-5.0, 6.0, 1.0, 2.0]], [[7.0, -8.0, 3.0, 4.0]]]> : tensor<2x1x4xf32>} : () -> ()
 // CHECK: "float64"() {bar = dense<{{\[\[\[}}-5.000000e+00, 6.000000e+00, 1.000000e+00, 2.000000e+00]], {{\[\[}}7.000000e+00, -8.000000e+00, 3.000000e+00, 4.000000e+00]]]> : tensor<2x1x4xf64>} : () -> ()
   "float64"(){bar = dense<[[[-5.0, 6.0, 1.0, 2.0]], [[7.0, -8.0, 3.0, 4.0]]]> : tensor<2x1x4xf64>} : () -> ()
 
 // 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>} : () -> ()
   return
 }
 
 // CHECK-LABEL: func @densevectorattr
 func @densevectorattr() -> () {
 ^bb0:
 // NOTE: The {{\[\[}} syntax is because "[[" confuses FileCheck.
 // CHECK: "fooi8"() {bar = dense<5> : vector<1x1x1xi8>} : () -> ()