From 314f044878920439c5e1d76b6488deaa0189e545 Mon Sep 17 00:00:00 2001
From: lorenzo chelini <lchelini@nvidia.com>
Date: Thu, 20 Feb 2025 12:05:31 +0100
Subject: [PATCH] [MLIR] Imporve location tracking for
 `parseElementsLiteralType`

This commit improves line location tracking in case of error reporting to the
user in `parseElementsLiteralType`. There are two cases: the type is already
parsed [1] or not yet parsed [2]. We print the error at the attribute's location
in both cases to ensure consistency.

Case 1)
```mlir
memref<i32> = dense<[3]>
              ^
```

Case 2)
```mlir
dense<[3]> :  memref<i32>
^
```
---
 mlir/lib/AsmParser/AttributeParser.cpp       | 22 ++++++------
 mlir/lib/AsmParser/Parser.h                  |  2 +-
 mlir/test/IR/invalid-builtin-attributes.mlir | 35 ++++++++++++++++++++
 3 files changed, 46 insertions(+), 13 deletions(-)

diff --git a/mlir/lib/AsmParser/AttributeParser.cpp b/mlir/lib/AsmParser/AttributeParser.cpp
index ff616dac9625b..32c68e093a0b0 100644
--- a/mlir/lib/AsmParser/AttributeParser.cpp
+++ b/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;
 
diff --git a/mlir/lib/AsmParser/Parser.h b/mlir/lib/AsmParser/Parser.h
index 1b8aa7c9dce6f..ecc128cf767b3 100644
--- a/mlir/lib/AsmParser/Parser.h
+++ b/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);
diff --git a/mlir/test/IR/invalid-builtin-attributes.mlir b/mlir/test/IR/invalid-builtin-attributes.mlir
index 5098fe751fd01..64201950b2306 100644
--- a/mlir/test/IR/invalid-builtin-attributes.mlir
+++ b/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
+}