[mlir] Optimize the parsing of ElementsAttr hex strings

[rriddle]

Summary

This revision optimizes the parsing of hex strings by using the checked variant of llvm::fromHex, and adding a specialized method to Token for extracting hex strings. This leads a large decrease in compile time when parsing large hex constants (one example: 2.6 seconds -> 370 miliseconds)

Depends On D90265

[rdzhabarov]

Pretty cool improvement!

mlir/lib/Parser/AttributeParser.cpp

413

nit: can you fix typo "etring".

mlir/lib/Parser/Token.cpp

128

i see that we put comments in both .h and cpp files. It's a bit redundant, imo, but as far as i can tell it's used for all other methods.

do we stick to this convention?

[rdzhabarov]

mlir/lib/Parser/Token.cpp

131

Is this to drop quotes from begin and end? (might be trivial, but i'm not too familiar with codebase).

135

nit: hex will be moved anyway, no need to make it explicit.

[mehdi_amini]

mlir/lib/Parser/Token.cpp

131

(Worth a comment I think)

135

I believe you actually need to be explicit here because the type of the returned value does not match the type of hex.

[rriddle]

mlir/lib/Parser/Token.cpp

135

That is what I would have expected as well, but I checked with godbolt and clang/gcc seem to do the right thing with llvm::Optional.

[rriddle]

Resolve comments

mlir/lib/Parser/Token.cpp

128

i see that we put comments in both .h and cpp files. It's a bit redundant, imo, but as far as i can tell it's used for all other methods.

do we stick to this convention?

Generally, yeah. I don't do it if the declaration is already in the same file though(that is way to redundant for my taste).

mlir/lib/Parser/AttributeParser.cpp

return emitError(loc, "integer constant out of range for attribute"),
            nullptr;
   return builder.getIntegerAttr(type, *apInt);
 }
 
 //===----------------------------------------------------------------------===//
 // TensorLiteralParser
 //===----------------------------------------------------------------------===//
 
-/// Parse elements values stored within a hex etring. On success, the values are
+/// Parse elements values stored within a hex string. 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);
+  if (Optional<std::string> value = tok.getHexStringValue()) {
+    result = std::move(*value);
     return success();
   }
+  return parser.emitError(
+      tok.getLoc(), "expected string containing hex digits starting with `0x`");
+}
 
 namespace {
 /// This class implements a parser for TensorLiterals. A tensor literal is
 /// either a single element (e.g, 5) or a multi-dimensional list of elements
 /// (e.g., [[5, 5]]).
 class TensorLiteralParser {
 public:
   TensorLiteralParser(Parser &p) : p(p) {}

mlir/lib/Parser/Token.h

#include "TokenKinds.def"
   /// the identifier, returning None if it is a named identifier like #x or
   /// if the integer doesn't fit.
   Optional<unsigned> getHashIdentifierNumber() const;
 
   /// Given a token containing a string literal, return its value, including
   /// removing the quote characters and unescaping the contents of the string.
   std::string getStringValue() const;
 
+  /// Given a token containing a hex string literal, return its value or None if
+  /// the token does not contain a valid hex string. A hex string literal is a
+  /// string starting with `0x` and only containing hex digits.
+  Optional<std::string> getHexStringValue() const;
+
   /// Given a token containing a symbol reference, return the unescaped string
   /// value.
   std::string getSymbolReference() const;
 
   // Location processing.
   llvm::SMLoc getLoc() const;
   llvm::SMLoc getEndLoc() const;
   llvm::SMRange getLocRange() const;

mlir/lib/Parser/Token.cpp

for (unsigned i = 0, e = bytes.size(); i != e;) {
 
     assert(llvm::isHexDigit(c1) && llvm::isHexDigit(c2) && "invalid escape");
     result.push_back((llvm::hexDigitValue(c1) << 4) | llvm::hexDigitValue(c2));
   }
 
   return result;
 }
 
+/// Given a token containing a hex string literal, return its value or None if
+/// the token does not contain a valid hex string.
+Optional<std::string> Token::getHexStringValue() const {
+  assert(getKind() == string);
+
+  // Get the internal string data, without the quotes.
+  StringRef bytes = getSpelling().drop_front().drop_back();
+
+  // Try to extract the binary data from the hex string.
+  std::string hex;
+  if (!bytes.consume_front("0x") || !llvm::tryGetFromHex(bytes, hex))
+    return llvm::None;
+  return hex;
+}
+
 /// Given a token containing a symbol reference, return the unescaped string
 /// value.
 std::string Token::getSymbolReference() const {
   assert(is(Token::at_identifier) && "expected valid @-identifier");
   StringRef nameStr = getSpelling().drop_front();
 
   // Check to see if the reference is a string literal, or a bare identifier.
   if (nameStr.front() == '"')

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

 // 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<"0x2041A040"> : tensor<2xbf16>} : () -> ()
 
 // -----
 
-// expected-error@+1 {{elements hex string should start with '0x'}}
+// expected-error@+1 {{expected string containing hex digits starting with `0x`}}
 "foo.op"() {dense.attr = dense<"00000000000024400000000000001440"> : tensor<2xf64>} : () -> ()
 
 // -----
 
-// expected-error@+1 {{elements hex string only contains hex digits}}
+// expected-error@+1 {{expected string containing hex digits starting with `0x`}}
 "foo.op"() {dense.attr = dense<"0x0000000000002440000000000000144X"> : tensor<2xf64>} : () -> ()
 
 // -----
 
 // 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 {{elements hex string only contains hex digits}}
+  "foo"(){bar = opaque<"", "0xQZz123"> : tensor<1xi8>} : () -> () // expected-error {{expected string containing hex digits starting with `0x`}}
 }
 
 // -----
 
 func @elementsattr_malformed_opaque2() -> () {
 ^bb0:
-  "foo"(){bar = opaque<"", "00abc"> : tensor<1xi8>} : () -> () // expected-error {{elements hex string should start with '0x'}}
+  "foo"(){bar = opaque<"", "00abc"> : tensor<1xi8>} : () -> () // expected-error {{expected string containing hex digits starting with `0x`}}
 }
 
 // -----
 
 func @elementsattr_malformed_opaque3() -> () {
 ^bb0:
   "foo"(){bar = opaque<"t", "0xabc"> : tensor<1xi8>} : () -> () // expected-error {{no registered dialect with namespace 't'}}
 }