[mlir][Parser] Update DenseElementsAttr to print in hex when the number of elements is over a certain threshold.

[rriddle]

Summary

DenseElementsAttr is used to store tensor data, which in some cases can become extremely large(100s of mb). In these cases it is much more efficient to format the data as a string of hex values instead.

[mehdi_amini]

mlir/lib/IR/AsmPrinter.cpp

66

Can we make it an int and use -1 to disable?

70

I'd lower this to 100 or 500 or something like this: the idea of the threshold to me is that this is the number at which it is unlikely someone really care about reading individual value from the IR text file manually (and they can just round-trip with a higher threshold to get them)

[rriddle]

Resolve comments

mlir/include/mlir/IR/Attributes.h

public:
   /// 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));
   }
 
+  /// 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);
+
   //===--------------------------------------------------------------------===//
   // Iterators
   //===--------------------------------------------------------------------===//
 
   /// A utility iterator that allows walking over the internal Attribute values
   /// of a DenseElementsAttr.
   class AttributeElementIterator
       : public indexed_accessor_iterator<AttributeElementIterator, const void *,
public:
   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 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;
+
   //===--------------------------------------------------------------------===//
   // Mutation Utilities
   //===--------------------------------------------------------------------===//
 
   /// Return a new DenseElementsAttr that has the same data as the current
   /// attribute, but has been reshaped to 'newType'. The new type must have the
   /// same total number of elements as well as element type.
   DenseElementsAttr reshape(ShapedType newType);
 
   /// Generates a new DenseElementsAttr by mapping each int value to a new
   /// underlying APInt. The new values can represent either a integer or float.
   /// This underlying type must be an DenseIntElementsAttr.
   DenseElementsAttr mapValues(Type newElementType,
                               function_ref<APInt(const APInt &)> mapping) const;
 
   /// Generates a new DenseElementsAttr by mapping each float value to a new
   /// underlying APInt. the new values can represent either a integer or float.
   /// This underlying type must be an DenseFPElementsAttr.
   DenseElementsAttr
   mapValues(Type newElementType,
             function_ref<APInt(const APFloat &)> mapping) const;
 
 protected:
-  /// Return the raw storage data held by this attribute.
-  ArrayRef<char> getRawData() const;
-
   /// 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());
   }
 

mlir/lib/IR/AsmPrinter.cpp

 
 /// The OpAsmOpInterface, see OpAsmInterface.td for more details.
 #include "mlir/IR/OpAsmInterface.cpp.inc"
 
 //===----------------------------------------------------------------------===//
 // OpPrintingFlags
 //===----------------------------------------------------------------------===//
 
+static llvm::cl::opt<int> printElementsAttrWithHexIfLarger(
+    "mlir-print-elementsattrs-with-hex-if-larger",
+    llvm::cl::desc(
+        "Print DenseElementsAttrs with a hex string that have "
+        "more elements than the given upper limit (use -1 to disable)"),
+    llvm::cl::init(100));
+
 static llvm::cl::opt<unsigned> elideElementsAttrIfLarger(
     "mlir-elide-elementsattrs-if-larger",
     llvm::cl::desc("Elide ElementsAttrs with \"...\" that have "
                    "more elements than the given upper limit"));
 
 static llvm::cl::opt<bool>
     printDebugInfoOpt("mlir-print-debuginfo",
                       llvm::cl::desc("Print debug info in MLIR output"),
public:
   void printIntegerSet(IntegerSet set);
 
 protected:
   void printOptionalAttrDict(ArrayRef<NamedAttribute> attrs,
                              ArrayRef<StringRef> elidedAttrs = {},
                              bool withKeyword = false);
   void printTrailingLocation(Location loc);
   void printLocationInternal(LocationAttr loc, bool pretty = false);
-  void printDenseElementsAttr(DenseElementsAttr attr);
+
+  /// Print a dense elements attribute. If 'allowHex' is true, a hex string is
+  /// used instead of individual elements when the elements attr is large.
+  void printDenseElementsAttr(DenseElementsAttr attr, bool allowHex);
 
   void printDialectAttribute(Attribute attr);
   void printDialectType(Type type);
 
   /// This enum is used to represent the binding strength of the enclosing
   /// context that an AffineExprStorage is being printed in, so we can
   /// intelligently produce parens.
   enum class BindingStrength {
void ModulePrinter::printAttribute(Attribute attr,
   }
   case StandardAttributes::DenseElements: {
     auto eltsAttr = attr.cast<DenseElementsAttr>();
     if (printerFlags.shouldElideElementsAttr(eltsAttr)) {
       printElidedElementsAttr(os);
       break;
     }
     os << "dense<";
-    printDenseElementsAttr(eltsAttr);
+    printDenseElementsAttr(eltsAttr, /*allowHex=*/true);
     os << '>';
     break;
   }
   case StandardAttributes::SparseElements: {
     auto elementsAttr = attr.cast<SparseElementsAttr>();
     if (printerFlags.shouldElideElementsAttr(elementsAttr.getIndices()) ||
         printerFlags.shouldElideElementsAttr(elementsAttr.getValues())) {
       printElidedElementsAttr(os);
       break;
     }
     os << "sparse<";
-    printDenseElementsAttr(elementsAttr.getIndices());
+    printDenseElementsAttr(elementsAttr.getIndices(), /*allowHex=*/false);
     os << ", ";
-    printDenseElementsAttr(elementsAttr.getValues());
+    printDenseElementsAttr(elementsAttr.getValues(), /*allowHex=*/true);
     os << '>';
     break;
   }
 
   // Location attributes.
   case StandardAttributes::CallSiteLocation:
   case StandardAttributes::FileLineColLocation:
   case StandardAttributes::FusedLocation:

 
 /// 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);
 }
 
-void ModulePrinter::printDenseElementsAttr(DenseElementsAttr attr) {
+void ModulePrinter::printDenseElementsAttr(DenseElementsAttr attr,
+                                           bool allowHex) {
   auto type = attr.getType();
   auto shape = type.getShape();
   auto rank = type.getRank();
 
   // The function used to print elements of this attribute.
   auto printEltFn = type.getElementType().isa<IntegerType>()
                         ? printDenseIntElement
                         : printDenseFloatElement;
void ModulePrinter::printDenseElementsAttr(DenseElementsAttr attr,
   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 && printElementsAttrWithHexIfLarger != -1 &&
+      numElements > printElementsAttrWithHexIfLarger) {
+    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;

mlir/lib/IR/Attributes.cpp

DenseElementsAttr DenseElementsAttr::get(ShapedType type,
 
   // 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 getRaw(type, intValues);
 }
 
-// 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'.
+/// 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 getRaw(type, rawBuffer, isSplatBuffer);
+}
+
+/// 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 storageBitWidth = getDenseElementStorageWidth(bitWidth);
   std::vector<char> elementData(llvm::divideCeil(storageBitWidth, CHAR_BIT) *
                                 values.size());

 
 /// A method used to verify specific type invariants that the templatized 'get'
 /// method cannot.
 bool DenseElementsAttr::isValidIntOrFloat(int64_t dataEltSize,
                                           bool isInt) const {
   return ::isValidIntOrFloat(getType(), dataEltSize, isInt);
 }
 
-/// Return the raw storage data held by this attribute.
-ArrayRef<char> DenseElementsAttr::getRawData() const {
-  return static_cast<ImplType *>(impl)->data;
-}
-
 /// Returns if this attribute corresponds to a splat, i.e. if all element
 /// values are the same.
 bool DenseElementsAttr::isSplat() const { return getImpl()->isSplat; }
 
 /// Return the held element values as a range of Attributes.
 auto DenseElementsAttr::getAttributeValues() const
     -> llvm::iterator_range<AttributeElementIterator> {
   return {attr_value_begin(), attr_value_end()};

 }
 auto DenseElementsAttr::float_value_begin() const -> FloatElementIterator {
   return getFloatValues().begin();
 }
 auto DenseElementsAttr::float_value_end() const -> FloatElementIterator {
   return getFloatValues().end();
 }
 
+/// Return the raw storage data held by this attribute.
+ArrayRef<char> DenseElementsAttr::getRawData() const {
+  return static_cast<ImplType *>(impl)->data;
+}
+
 /// Return a new DenseElementsAttr that has the same data as the current
 /// attribute, but has been reshaped to 'newType'. The new type must have the
 /// same total number of elements as well as element type.
 DenseElementsAttr DenseElementsAttr::reshape(ShapedType newType) {
   ShapedType curType = getType();
   if (curType == newType)
     return *this;
 

mlir/lib/Parser/Parser.cpp

Attribute Parser::parseDecOrHexAttr(Type type, bool isNegative) {
   // Otherwise construct an integer attribute.
   if (isNegative ? (int64_t)-val.getValue() >= 0 : (int64_t)val.getValue() < 0)
     return emitError(loc, "integer constant out of range for attribute"),
            nullptr;
 
   return builder.getIntegerAttr(type, isNegative ? -apInt : apInt);
 }
 
+/// Parse elements values stored within a hex etring. On success, the values are
+/// stored into 'result'.
+static ParseResult parseElementAttrHexValues(Parser &parser, Token tok,
+                                             std::string &result) {
+  std::string val = tok.getStringValue();
+  if (val.size() < 2 || val[0] != '0' || val[1] != 'x')
+    return parser.emitError(tok.getLoc(),
+                            "elements hex string should start with '0x'");
+
+  StringRef hexValues = StringRef(val).drop_front(2);
+  if (!llvm::all_of(hexValues, llvm::isHexDigit))
+    return parser.emitError(tok.getLoc(),
+                            "elements hex string only contains hex digits");
+
+  result = llvm::fromHex(hexValues);
+  return success();
+}
+
 /// Parse an opaque elements attribute.
 Attribute Parser::parseOpaqueElementsAttr(Type attrType) {
   consumeToken(Token::kw_opaque);
   if (parseToken(Token::less, "expected '<' after 'opaque'"))
     return nullptr;
 
   if (getToken().isNot(Token::string))
     return (emitError("expected dialect namespace"), nullptr);
 
   auto name = getToken().getStringValue();
   auto *dialect = builder.getContext()->getRegisteredDialect(name);
   // TODO(shpeisman): Allow for having an unknown dialect on an opaque
   // attribute. Otherwise, it can't be roundtripped without having the dialect
   // registered.
   if (!dialect)
     return (emitError("no registered dialect with namespace '" + name + "'"),
             nullptr);
-
   consumeToken(Token::string);
+
   if (parseToken(Token::comma, "expected ','"))
     return nullptr;
 
-  if (getToken().getKind() != Token::string)
-    return (emitError("opaque string should start with '0x'"), nullptr);
-
-  auto val = getToken().getStringValue();
-  if (val.size() < 2 || val[0] != '0' || val[1] != 'x')
-    return (emitError("opaque string should start with '0x'"), nullptr);
-
-  val = val.substr(2);
-  if (!llvm::all_of(val, llvm::isHexDigit))
-    return (emitError("opaque string only contains hex digits"), nullptr);
-
-  consumeToken(Token::string);
-  if (parseToken(Token::greater, "expected '>'"))
+  Token hexTok = getToken();
+  if (parseToken(Token::string, "elements hex string should start with '0x'") ||
+      parseToken(Token::greater, "expected '>'"))
     return nullptr;
-
   auto type = parseElementsLiteralType(attrType);
   if (!type)
     return nullptr;
 
-  return OpaqueElementsAttr::get(dialect, type, llvm::fromHex(val));
+  std::string data;
+  if (parseElementAttrHexValues(*this, hexTok, data))
+    return nullptr;
+  return OpaqueElementsAttr::get(dialect, type, data);
 }
 
 namespace {
 class TensorLiteralParser {
 public:
   TensorLiteralParser(Parser &p) : p(p) {}
 
-  ParseResult parse() {
-    if (p.getToken().is(Token::l_square))
-      return parseList(shape);
-    return parseElement();
-  }
+  /// 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 parse(bool allowHex);
 
   /// 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:
private:
   /// Build a Dense Integer attribute for the given type.
   DenseElementsAttr getIntAttr(llvm::SMLoc loc, ShapedType type,
                                IntegerType 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);
+
   /// Parse a single element, returning failure if it isn't a valid element
   /// literal. For example:
   /// parseElement(1) -> Success, 1
   /// parseElement([1]) -> Failure
   ParseResult parseElement();
 
   /// Parse a list of either lists or elements, returning the dimensions of the
   /// parsed sub-tensors in dims. For example:
   ///   parseList([1, 2, 3]) -> Success, [3]
   ///   parseList([[1, 2], [3, 4]]) -> Success, [2, 2]
   ///   parseList([[1, 2], 3]) -> Failure
   ///   parseList([[1, [2, 3]], [4, [5]]]) -> Failure
   ParseResult parseList(SmallVectorImpl<int64_t> &dims);
 
+  /// Parse a literal that was printed as a hex string.
+  ParseResult parseHexElements();
+
   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) {
+  // If hex is allowed, check for a string literal.
+  if (allowHex && p.getToken().is(Token::string)) {
+    hexStorage = p.getToken();
+    p.consumeToken(Token::string);
+    return success();
+  }
+  // Otherwise, parse a list or an individual element.
+  if (p.getToken().is(Token::l_square))
+    return parseList(shape);
+  return parseElement();
+}
+
 /// Build a dense attribute instance with the parsed elements and the given
 /// shaped type.
 DenseElementsAttr TensorLiteralParser::getAttr(llvm::SMLoc loc,
                                                ShapedType type) {
+  // Check to see if we parsed the literal from a hex string.
+  if (hexStorage.hasValue())
+    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 (!eltTy.isBF16() && !eltTy.isF64()) {
                     &unused);
     }
     floatValues.push_back(apVal);
   }
 
   return DenseElementsAttr::get(type, floatValues);
 }
 
+/// 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.isIntOrFloat()) {
+    p.emitError(loc) << "expected floating-point or integer element type, got "
+                     << elementType;
+    return nullptr;
+  }
+
+  std::string data;
+  if (parseElementAttrHexValues(p, hexStorage.getValue(), data))
+    return nullptr;
+
+  // Check that the size of the hex data correpsonds to the size of the type, or
+  // a splat of the type.
+  if (static_cast<int64_t>(data.size() * CHAR_BIT) !=
+      (type.getNumElements() * elementType.getIntOrFloatBitWidth())) {
+    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);
+}
+
 ParseResult TensorLiteralParser::parseElement() {
   switch (p.getToken().getKind()) {
   // Parse a boolean element.
   case Token::kw_true:
   case Token::kw_false:
   case Token::floatliteral:
   case Token::integer:
     storage.emplace_back(/*isNegative=*/false, p.getToken());

 /// Parse a dense elements attribute.
 Attribute Parser::parseDenseElementsAttr(Type attrType) {
   consumeToken(Token::kw_dense);
   if (parseToken(Token::less, "expected '<' after 'dense'"))
     return nullptr;
 
   // Parse the literal data.
   TensorLiteralParser literalParser(*this);
-  if (literalParser.parse())
+  if (literalParser.parse(/*allowHex=*/true))
     return nullptr;
 
   if (parseToken(Token::greater, "expected '>'"))
     return nullptr;
 
   auto typeLoc = getToken().getLoc();
   auto type = parseElementsLiteralType(attrType);
   if (!type)

 }
 
 /// Parse a sparse elements attribute.
 Attribute Parser::parseSparseElementsAttr(Type attrType) {
   consumeToken(Token::kw_sparse);
   if (parseToken(Token::less, "Expected '<' after 'sparse'"))
     return nullptr;
 
-  /// Parse indices
+  /// Parse the indices. We don't allow hex values here as we may need to use
+  /// the inferred shape.
   auto indicesLoc = getToken().getLoc();
   TensorLiteralParser indiceParser(*this);
-  if (indiceParser.parse())
+  if (indiceParser.parse(/*allowHex=*/false))
     return nullptr;
 
   if (parseToken(Token::comma, "expected ','"))
     return nullptr;
 
-  /// Parse values.
+  /// Parse the values.
   auto valuesLoc = getToken().getLoc();
   TensorLiteralParser valuesParser(*this);
-  if (valuesParser.parse())
+  if (valuesParser.parse(/*allowHex=*/true))
     return nullptr;
 
   if (parseToken(Token::greater, "expected '>'"))
     return nullptr;
 
   auto type = parseElementsLiteralType(attrType);
   if (!type)
     return nullptr;

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

+// RUN: mlir-opt %s -verify-diagnostics -split-input-file -mlir-print-elementsattrs-with-hex-if-larger=1 | FileCheck %s --check-prefix=HEX
+// RUN: mlir-opt %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>} : () -> ()
+
+// -----
+
+// expected-error@+1 {{elements hex string should start with '0x'}}
+"foo.op"() {dense.attr = dense<"00000000000024400000000000001440"> : tensor<2xf64>} : () -> ()
+
+// -----
+
+// expected-error@+1 {{elements hex string only contains hex digits}}
+"foo.op"() {dense.attr = dense<"0x0000000000002440000000000000144X"> : tensor<2xf64>} : () -> ()
+
+// -----
+
+// expected-error@+1 {{expected floating-point or integer element type, got '!unknown<"">'}}
+"foo.op"() {dense.attr = dense<"0x00000000000024400000000000001440"> : tensor<2x!unknown<"">>} : () -> ()
+
+// -----
+
+// expected-error@+1 {{elements hex data size is invalid for provided type}}
+"foo.op"() {dense.attr = dense<"0x00000000000024400000000000001440"> : tensor<4xf64>} : () -> ()

mlir/test/IR/invalid.mlir

 ^bb0:
   "foo"(){bar = opaque<10, "0xQZz123"> : tensor<1xi8>} : () -> () // expected-error {{expected dialect namespace}}
 }
 
 // -----
 
 func @elementsattr_malformed_opaque1() -> () {
 ^bb0:
-  "foo"(){bar = opaque<"", "0xQZz123"> : tensor<1xi8>} : () -> () // expected-error {{opaque string only contains hex digits}}
+  "foo"(){bar = opaque<"", "0xQZz123"> : tensor<1xi8>} : () -> () // expected-error {{elements hex string only contains hex digits}}
 }
 
 // -----
 
 func @elementsattr_malformed_opaque2() -> () {
 ^bb0:
-  "foo"(){bar = opaque<"", "00abc"> : tensor<1xi8>} : () -> () // expected-error {{opaque string should start with '0x'}}
+  "foo"(){bar = opaque<"", "00abc"> : tensor<1xi8>} : () -> () // expected-error {{elements hex string should start with '0x'}}
 }
 
 // -----
 
 func @elementsattr_malformed_opaque3() -> () {
 ^bb0:
   "foo"(){bar = opaque<"t", "0xabc"> : tensor<1xi8>} : () -> () // expected-error {{no registered dialect with namespace 't'}}
 }