Allows deferred location attribute in parseOptionalLocationSpecifier

Before this patch, deferred location in operation like `test.pretty_printed_region` would
break the parser, and so the IR round-trip.

Depends On D117088

Reviewed By: rriddle

Differential Revision: https://reviews.llvm.org/D117413

mlir/include/mlir/IR/OpImplementation.h

@@ -850,10 +850,6 @@ public:
                                               StringRef attrName,
                                               NamedAttrList &attrs) = 0;
 
-  /// Parse a loc(...) specifier if present, filling in result if so.
-  virtual ParseResult
-  parseOptionalLocationSpecifier(Optional<Location> &result) = 0;
-
   //===--------------------------------------------------------------------===//
   // Type Parsing
   //===--------------------------------------------------------------------===//
@@ -1041,6 +1037,13 @@ public:
   using AsmParser::AsmParser;
   ~OpAsmParser() override;
 
+  /// Parse a loc(...) specifier if present, filling in result if so.
+  /// Location for BlockArgument and Operation may be deferred with an alias, in
+  /// which case an OpaqueLoc is set and will be resolved when parsing
+  /// completes.
+  virtual ParseResult
+  parseOptionalLocationSpecifier(Optional<Location> &result) = 0;
+
   /// Return the name of the specified result in the specified syntax, as well
   /// as the sub-element in the name.  It returns an empty string and ~0U for
   /// invalid result numbers.  For example, in this operation:

mlir/lib/Parser/AsmParserImpl.h

@@ -429,22 +429,6 @@ public:
     return success();
   }
 
-  /// Parse a loc(...) specifier if present, filling in result if so.
-  ParseResult
-  parseOptionalLocationSpecifier(Optional<Location> &result) override {
-    // If there is a 'loc' we parse a trailing location.
-    if (!parser.consumeIf(Token::kw_loc))
-      return success();
-    LocationAttr directLoc;
-    if (parser.parseToken(Token::l_paren, "expected '(' in location") ||
-        parser.parseLocationInstance(directLoc) ||
-        parser.parseToken(Token::r_paren, "expected ')' in location"))
-      return failure();
-
-    result = directLoc;
-    return success();
-  }
-
   //===--------------------------------------------------------------------===//
   // Type Parsing
   //===--------------------------------------------------------------------===//

mlir/lib/Parser/Parser.cpp

@@ -341,6 +341,11 @@ public:
   ///
   ParseResult parseTrailingLocationSpecifier(OpOrArgument opOrArgument);
 
+  /// Parse a location alias, that is a sequence looking like: #loc42
+  /// The alias may have already be defined or may be defined later, in which
+  /// case an OpaqueLoc is used a placeholder.
+  ParseResult parseLocationAlias(LocationAttr &loc);
+
   /// This is the structure of a result specifier in the assembly syntax,
   /// including the name, number of results, and location.
   using ResultRecord = std::tuple<StringRef, unsigned, SMLoc>;
@@ -1592,6 +1597,34 @@ public:
     return parser.parseCommaSeparatedListUntil(Token::r_paren, parseElt);
   }
 
+  /// Parse a loc(...) specifier if present, filling in result if so.
+  ParseResult
+  parseOptionalLocationSpecifier(Optional<Location> &result) override {
+    // If there is a 'loc' we parse a trailing location.
+    if (!parser.consumeIf(Token::kw_loc))
+      return success();
+    LocationAttr directLoc;
+    if (parser.parseToken(Token::l_paren, "expected '(' in location"))
+      return failure();
+
+    Token tok = parser.getToken();
+
+    // Check to see if we are parsing a location alias.
+    // Otherwise, we parse the location directly.
+    if (tok.is(Token::hash_identifier)) {
+      if (parser.parseLocationAlias(directLoc))
+        return failure();
+    } else if (parser.parseLocationInstance(directLoc)) {
+      return failure();
+    }
+
+    if (parser.parseToken(Token::r_paren, "expected ')' in location"))
+      return failure();
+
+    result = directLoc;
+    return success();
+  }
+
 private:
   /// Information about the result name specifiers.
   ArrayRef<OperationParser::ResultRecord> resultIDs;
@@ -1719,6 +1752,33 @@ OperationParser::parseCustomOperation(ArrayRef<ResultRecord> resultIDs) {
   return op;
 }
 
+ParseResult OperationParser::parseLocationAlias(LocationAttr &loc) {
+  Token tok = getToken();
+  consumeToken(Token::hash_identifier);
+  StringRef identifier = tok.getSpelling().drop_front();
+  if (identifier.contains('.')) {
+    return emitError(tok.getLoc())
+           << "expected location, but found dialect attribute: '#" << identifier
+           << "'";
+  }
+
+  // If this alias can be resolved, do it now.
+  Attribute attr = state.symbols.attributeAliasDefinitions.lookup(identifier);
+  if (attr) {
+    if (!(loc = attr.dyn_cast<LocationAttr>()))
+      return emitError(tok.getLoc())
+             << "expected location, but found '" << attr << "'";
+  } else {
+    // Otherwise, remember this operation and resolve its location later.
+    // In the meantime, use a special OpaqueLoc as a marker.
+    loc = OpaqueLoc::get(deferredLocsReferences.size(),
+                         TypeID::get<DeferredLocInfo *>(),
+                         UnknownLoc::get(getContext()));
+    deferredLocsReferences.push_back(DeferredLocInfo{tok.getLoc(), identifier});
+  }
+  return success();
+}
+
 ParseResult
 OperationParser::parseTrailingLocationSpecifier(OpOrArgument opOrArgument) {
   // If there is a 'loc' we parse a trailing location.
@@ -1729,34 +1789,11 @@ OperationParser::parseTrailingLocationSpecifier(OpOrArgument opOrArgument) {
   Token tok = getToken();
 
   // Check to see if we are parsing a location alias.
+  // Otherwise, we parse the location directly.
   LocationAttr directLoc;
   if (tok.is(Token::hash_identifier)) {
-    consumeToken();
-
-    StringRef identifier = tok.getSpelling().drop_front();
-    if (identifier.contains('.')) {
-      return emitError(tok.getLoc())
-             << "expected location, but found dialect attribute: '#"
-             << identifier << "'";
-    }
-
-    // If this alias can be resolved, do it now.
-    Attribute attr = state.symbols.attributeAliasDefinitions.lookup(identifier);
-    if (attr) {
-      if (!(directLoc = attr.dyn_cast<LocationAttr>()))
-        return emitError(tok.getLoc())
-               << "expected location, but found '" << attr << "'";
-    } else {
-      // Otherwise, remember this operation and resolve its location later.
-      // In the meantime, use a special OpaqueLoc as a marker.
-      directLoc = OpaqueLoc::get(deferredLocsReferences.size(),
-                                 TypeID::get<DeferredLocInfo *>(),
-                                 UnknownLoc::get(getContext()));
-      deferredLocsReferences.push_back(
-          DeferredLocInfo{tok.getLoc(), identifier});
-    }
-
-    // Otherwise, we parse the location directly.
+    if (parseLocationAlias(directLoc))
+      return failure();
   } else if (parseLocationInstance(directLoc)) {
     return failure();
   }

mlir/test/IR/pretty_printed_region_op.mlir

@@ -1,5 +1,6 @@
-// RUN: mlir-opt -allow-unregistered-dialect -split-input-file  %s | FileCheck %s --check-prefixes=CHECK-CUSTOM,CHECK
+// RUN: mlir-opt -allow-unregistered-dialect -split-input-file %s | FileCheck %s --check-prefixes=CHECK-CUSTOM,CHECK
 // RUN: mlir-opt -allow-unregistered-dialect -mlir-print-op-generic -split-input-file   %s | FileCheck %s --check-prefixes=CHECK,CHECK-GENERIC
+// RUN: mlir-opt -allow-unregistered-dialect -split-input-file --mlir-print-op-generic --mlir-print-debuginfo -mlir-print-local-scope  %s | FileCheck %s --check-prefixes=CHECK-LOCATION
 
 // -----
 
@@ -33,3 +34,16 @@ func @pretty_printed_region_op(%arg0 : f32, %arg1 : f32) -> (f32) {
     return %0 : f32
 }
 
+// -----
+
+
+func @pretty_printed_region_op_deferred_loc(%arg0 : f32, %arg1 : f32) -> (f32) {
+// CHECK-LOCATION: "test.pretty_printed_region"(%arg1, %arg0)
+// CHECK-LOCATION:   ^bb0(%arg[[x:[0-9]+]]: f32 loc("foo"), %arg[[y:[0-9]+]]: f32 loc("foo")
+// CHECK-LOCATION:     %[[RES:.*]] = "special.op"(%arg[[x]], %arg[[y]]) : (f32, f32) -> f32
+// CHECK-LOCATION:     "test.return"(%[[RES]]) : (f32) -> ()
+// CHECK-LOCATION:  : (f32, f32) -> f32
+
+  %res = test.pretty_printed_region %arg1, %arg0 start special.op end : (f32, f32) -> (f32) loc("foo")
+  return %res : f32
+}