Fix the MLIR integer attribute parser to be correct in the face of large integer attributes, it was previously artificially limited to 64 bits.

[lattner]

Fix two minor things caught in self-review.

[lattner]

The test failures are unrelated to my change.

[lattner]

Thanks for the review!

mlir/lib/Parser/Parser.cpp

if (apInt != value) {
     return llvm::None;
   }
   return APFloat(type.getFloatSemantics(), apInt);
 }
 
 /// Construct an APint from a parsed value, a known attribute type and
 /// sign.
 static Optional<APInt> buildAttributeAPInt(Type type, bool isNegative,
-                                           uint64_t value) {
-  // We have the integer literal as an uint64_t in val, now convert it into an
-  // APInt and check that we don't overflow.
-  int width = type.isIndex() ? 64 : type.getIntOrFloatBitWidth();
-  APInt apInt(width, value, isNegative);
-  if (apInt != value)
+                                           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();
+  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;
+
+    result = result.trunc(width);
+  }
+
   if (isNegative) {
     // The value is negative, we have an overflow if the sign bit is not set
     // in the negated apInt.
-    apInt.negate();
-    if (!apInt.isSignBitSet())
+    result.negate();
+    if (!result.isSignBitSet())
       return llvm::None;
   } else if ((type.isSignedInteger() || type.isIndex()) &&
-             apInt.isSignBitSet()) {
+             result.isSignBitSet()) {
     // The value is a positive signed integer or index,
     // we have an overflow if the sign bit is set.
     return llvm::None;
   }
 
-  return apInt;
+  return result;
 }
 
 /// Parse a decimal or a hexadecimal literal, which can be either an integer
 /// or a float attribute.
 Attribute Parser::parseDecOrHexAttr(Type type, bool isNegative) {
-  auto val = getToken().getUInt64IntegerValue();
-  if (!val.hasValue())
-    return (emitError("integer constant out of range for attribute"), nullptr);
-
   // Remember if the literal is hexadecimal.
   StringRef spelling = getToken().getSpelling();
   auto loc = state.curToken.getLoc();
   bool isHex = spelling.size() > 1 && spelling[1] == 'x';
 
   consumeToken(Token::integer);
   if (!type) {
     // Default to i64 if not type is specified.
if (isNegative)
              nullptr;
     if (!isHex) {
       emitError(loc, "unexpected decimal integer literal for a float attribute")
               .attachNote()
           << "add a trailing dot to make the literal a float";
       return nullptr;
     }
 
+    auto val = Token::getUInt64IntegerValue(spelling);
+    if (!val.hasValue())
+      return emitError("integer constant out of range for attribute"), nullptr;
+
     // Construct a float attribute bitwise equivalent to the integer literal.
     Optional<APFloat> apVal =
         buildHexadecimalFloatLiteral(this, floatType, *val);
     return apVal ? FloatAttr::get(floatType, *apVal) : Attribute();
   }
 
   if (!type.isa<IntegerType>() && !type.isa<IndexType>())
     return emitError(loc, "integer literal not valid for specified type"),
            nullptr;
 
   if (isNegative && type.isUnsignedInteger()) {
     emitError(loc,
               "negative integer literal not valid for unsigned integer type");
     return nullptr;
   }
 
-  Optional<APInt> apInt = buildAttributeAPInt(type, isNegative, *val);
-
+  Optional<APInt> apInt = buildAttributeAPInt(type, isNegative, spelling);
   if (!apInt)
     return emitError(loc, "integer constant out of range for attribute"),
            nullptr;
   return builder.getIntegerAttr(type, *apInt);
 }
 
 /// Parse elements values stored within a hex etring. On success, the values are
 /// stored into 'result'.
if (token.isAny(Token::kw_true, Token::kw_false)) {
             << "expected i1 type for 'true' or 'false' values";
       APInt apInt(eltTy.getWidth(), token.is(Token::kw_true),
                   /*isSigned=*/false);
       intElements.push_back(apInt);
       continue;
     }
 
     // Create APInt values for each element with the correct bitwidth.
-    auto val = token.getUInt64IntegerValue();
-    if (!val.hasValue()) {
-      p.emitError(tokenLoc, "integer constant out of range for attribute");
-      return nullptr;
-    }
-    Optional<APInt> apInt = buildAttributeAPInt(eltTy, isNegative, *val);
+    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 DenseElementsAttr::get(type, intElements);
 }

mlir/lib/Parser/Token.h

#include "TokenKinds.def"
   // Helpers to decode specific sorts of tokens.
 
   /// For an integer token, return its value as an unsigned.  If it doesn't fit,
   /// return None.
   Optional<unsigned> getUnsignedIntegerValue() const;
 
   /// For an integer token, return its value as an uint64_t.  If it doesn't fit,
   /// return None.
-  Optional<uint64_t> getUInt64IntegerValue() const;
+  static Optional<uint64_t> getUInt64IntegerValue(StringRef spelling);
+  Optional<uint64_t> getUInt64IntegerValue() const {
+    return getUInt64IntegerValue(getSpelling());
+  }
 
   /// For a floatliteral token, return its value as a double. Returns None in
   /// the case of underflow or overflow.
   Optional<double> getFloatingPointValue() const;
 
   /// For an inttype token, return its bitwidth.
   Optional<unsigned> getIntTypeBitwidth() const;
 

mlir/lib/Parser/Token.cpp

Optional<unsigned> Token::getUnsignedIntegerValue() const {
   unsigned result = 0;
   if (spelling.getAsInteger(isHex ? 0 : 10, result))
     return None;
   return result;
 }
 
 /// For an integer token, return its value as a uint64_t.  If it doesn't fit,
 /// return None.
-Optional<uint64_t> Token::getUInt64IntegerValue() const {
+Optional<uint64_t> Token::getUInt64IntegerValue(StringRef spelling) {
   bool isHex = spelling.size() > 1 && spelling[1] == 'x';
 
   uint64_t result = 0;
   if (spelling.getAsInteger(isHex ? 0 : 10, result))
     return None;
   return result;
 }
 

mlir/test/IR/attribute.mlir

"test.in_range_attrs"() {
     attr_17 = -9223372036854775808 : si64,
     // CHECK-SAME: attr_18 = 9223372036854775807 : si64
     attr_18 = 9223372036854775807 : si64,
     // CHECK-SAME: attr_19 = 18446744073709551615 : ui64
     attr_19 = 18446744073709551615 : ui64,
     // CHECK-SAME: attr_20 = 1 : ui1
     attr_20 = 1 : ui1,
     // CHECK-SAME: attr_21 = -1 : si1
-    attr_21 = -1 : si1
-
+    attr_21 = -1 : si1,
+    // CHECK-SAME: attr_22 = 79228162514264337593543950335 : ui96
+    attr_22 = 79228162514264337593543950335 : ui96,
+    // CHECK-SAME: attr_23 = -39614081257132168796771975168 : si96
+    attr_23 = -39614081257132168796771975168 : si96
   } : () -> ()
 
   return
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//

mlir/test/IR/invalid.mlir

 
 func @large_bound() {
   "test.out_of_range_attribute"() {
     // expected-error @+1 {{integer constant out of range for attribute}}
     attr = 18446744073709551616 : ui64
   } : () -> ()
   return
 }
+
+// -----
+
+func @really_large_bound() {
+  "test.out_of_range_attribute"() {
+    // expected-error @+1 {{integer constant out of range for attribute}}
+    attr = 79228162514264337593543950336 : ui96
+  } : () -> ()
+  return
+}
+
+// -----
+
+func @really_large_bound() {
+  "test.out_of_range_attribute"() {
+    // expected-error @+1 {{integer constant out of range for attribute}}
+    attr = 79228162514264337593543950336 : i96
+  } : () -> ()
+  return
+}
+
+// -----
+
+func @really_large_bound() {
+  "test.out_of_range_attribute"() {
+    // expected-error @+1 {{integer constant out of range for attribute}}
+    attr = 39614081257132168796771975168 : si96
+  } : () -> ()
+  return
+}