[MLIR] Imporve location tracking for parseElementsLiteralType

mlir/lib/AsmParser/AttributeParser.cpp

@@ -951,14 +951,10 @@ Attribute Parser::parseDenseElementsAttr(Type attrType) {
       return nullptr;
   }
 
-  // If the type is specified `parseElementsLiteralType` will not parse a type.
-  // Use the attribute location as the location for error reporting in that
-  // case.
-  auto loc = attrType ? attribLoc : getToken().getLoc();
-  auto type = parseElementsLiteralType(attrType);
+  auto type = parseElementsLiteralType(attribLoc, attrType);
   if (!type)
     return nullptr;
-  return literalParser.getAttr(loc, type);
+  return literalParser.getAttr(attribLoc, type);
 }
 
 Attribute Parser::parseDenseResourceElementsAttr(Type attrType) {
@@ -999,7 +995,7 @@ Attribute Parser::parseDenseResourceElementsAttr(Type attrType) {
 ///   elements-literal-type ::= vector-type | ranked-tensor-type
 ///
 /// This method also checks the type has static shape.
-ShapedType Parser::parseElementsLiteralType(Type type) {
+ShapedType Parser::parseElementsLiteralType(SMLoc loc, Type type) {
   // If the user didn't provide a type, parse the colon type for the literal.
   if (!type) {
     if (parseToken(Token::colon, "expected ':'"))
@@ -1010,12 +1006,14 @@ ShapedType Parser::parseElementsLiteralType(Type type) {
 
   auto sType = dyn_cast<ShapedType>(type);
   if (!sType) {
-    emitError("elements literal must be a shaped type");
+    emitError(loc, "elements literal must be a shaped type");
     return nullptr;
   }
 
-  if (!sType.hasStaticShape())
-    return (emitError("elements literal type must have static shape"), nullptr);
+  if (!sType.hasStaticShape()) {
+    emitError(loc, "elements literal type must have static shape");
+    return nullptr;
+  }
 
   return sType;
 }
@@ -1032,7 +1030,7 @@ Attribute Parser::parseSparseElementsAttr(Type attrType) {
   // of the type.
   Type indiceEltType = builder.getIntegerType(64);
   if (consumeIf(Token::greater)) {
-    ShapedType type = parseElementsLiteralType(attrType);
+    ShapedType type = parseElementsLiteralType(loc, attrType);
     if (!type)
       return nullptr;
 
@@ -1065,7 +1063,7 @@ Attribute Parser::parseSparseElementsAttr(Type attrType) {
   if (parseToken(Token::greater, "expected '>'"))
     return nullptr;
 
-  auto type = parseElementsLiteralType(attrType);
+  auto type = parseElementsLiteralType(loc, attrType);
   if (!type)
     return nullptr;
 

mlir/lib/AsmParser/Parser.h

@@ -288,7 +288,7 @@ class Parser {
 
   /// Parse a dense elements attribute.
   Attribute parseDenseElementsAttr(Type attrType);
-  ShapedType parseElementsLiteralType(Type type);
+  ShapedType parseElementsLiteralType(SMLoc loc, Type type);
 
   /// Parse a dense resource elements attribute.
   Attribute parseDenseResourceElementsAttr(Type attrType);

mlir/test/IR/invalid-builtin-attributes.mlir

@@ -591,3 +591,38 @@ func.func @duplicate_dictionary_attr_key() {
 #attr1 = distinct[0]<43 : i32>
 
 // -----
+
+// Make sure the error is not printed on the return.
+func.func @print_error_on_correct_line() {
+  %0 = arith.constant
+    // expected-error@below {{elements literal must be a shaped type}} 
+    dense<[3]> : i32
+  return
+}
+
+// -----
+
+// Make sure the error is not printed on the return.
+func.func @print_error_on_correct_line() {
+  %0 = arith.constant 
+    // expected-error@below {{elements literal must be a shaped type}}
+    sparse<
+     [
+       [0, 1, 2, 3],
+       [1, 1, 2, 3],
+       [1, 2, 2, 3],
+       [1, 2, 3, 4]
+     ],
+     [1, 1, 1, 1] > : i32
+  return
+}
+
+// -----
+
+// Make sure the error is not printed on the return.
+func.func @print_error_on_correct_line() {
+  %0 = arith.constant 
+    // expected-error@below {{elements literal must be a shaped type}}
+    sparse <> : i32
+  return
+}