Allows deferred location attribute in `parseOptionalLocationSpecifier`

[mehdi_amini]

Summary

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

[rriddle]

mlir/lib/Parser/Parser.cpp

344–347

1757–1758

1803–1804

Is this ever false now?

[mehdi_amini]

Rebase

[mehdi_amini]

mlir/lib/Parser/Parser.cpp

1803–1804

I don't think so! This is also in the parent revision, I fixed it there

[rriddle]

LG

mlir/include/mlir/IR/OpImplementation.h

public:
   }
 
   /// Parse an optional @-identifier and store it (without the '@' symbol) in a
   /// string attribute named 'attrName'.
   virtual ParseResult parseOptionalSymbolName(StringAttr &result,
                                               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
   //===--------------------------------------------------------------------===//
 
   /// Parse a type.
   virtual ParseResult parseType(Type &result) = 0;
 
   /// Parse a custom type with the provided callback, unless the next

 /// parse the indices, then calling `parseColonTypeList` to parse the result
 /// type.
 ///
 class OpAsmParser : public AsmParser {
 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:
   ///
   ///  %x, %y:2, %z = foo.op
   ///
   ///    getResultName(0) == {"x", 0 }
   ///    getResultName(1) == {"y", 0 }

mlir/lib/Parser/AsmParserImpl.h

ParseResult parseOptionalSymbolName(StringAttr &result, StringRef attrName,
     // reference.
     if (parser.getState().asmState) {
       parser.getState().asmState->addUses(SymbolRefAttr::get(result),
                                           atToken.getLocRange());
     }
     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
   //===--------------------------------------------------------------------===//
 
   /// Parse a type.
   ParseResult parseType(Type &result) override {
     return failure(!(result = parser.parseType()));
   }

mlir/lib/Parser/Parser.cpp

public:
 
   /// Parse an optional trailing location and add it to the specifier Operation
   /// or `OperandType` if present.
   ///
   ///   trailing-location ::= (`loc` (`(` location `)` | attribute-alias))?
   ///
   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>;
 
   /// Parse an operation instance that is in the op-defined custom form.
   /// resultInfo specifies information about the "%name =" specifiers.
   Operation *parseCustomOperation(ArrayRef<ResultRecord> resultIDs);
 
auto parseElt = [&]() -> ParseResult {
       lhs.push_back(regionArg);
       rhs.push_back(operand);
       types.push_back(type);
       return success();
     };
     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;
 
   /// The abstract information of the operation.
   function_ref<ParseResult(OpAsmParser &, OperationState &)> parseAssembly;
   bool isIsolatedFromAbove;
   StringRef opName;
OperationParser::parseCustomOperation(ArrayRef<ResultRecord> resultIDs) {
 
   // Otherwise, create the operation and try to parse a location for it.
   Operation *op = opBuilder.createOperation(opState);
   if (parseTrailingLocationSpecifier(op))
     return nullptr;
   return op;
 }
 
-ParseResult
-OperationParser::parseTrailingLocationSpecifier(OpOrArgument opOrArgument) {
-  // If there is a 'loc' we parse a trailing location.
-  if (!consumeIf(Token::kw_loc))
-    return success();
-  if (parseToken(Token::l_paren, "expected '(' in location"))
-    return failure();
+ParseResult OperationParser::parseLocationAlias(LocationAttr &loc) {
   Token tok = getToken();
-
-  // Check to see if we are parsing a location alias.
-  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 << "'";
-    }
+  consumeToken(Token::hash_identifier);
+           << "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});
-    }
+    if (!(loc = attr.dyn_cast<LocationAttr>()))
+    loc = OpaqueLoc::get(deferredLocsReferences.size(),
+    deferredLocsReferences.push_back(DeferredLocInfo{tok.getLoc(), identifier});
+  return success();
+}
+
+ParseResult
+OperationParser::parseTrailingLocationSpecifier(OpOrArgument opOrArgument) {
+  // If there is a 'loc' we parse a trailing location.
+  if (!consumeIf(Token::kw_loc))
+    return success();
+  if (parseToken(Token::l_paren, "expected '(' in location"))
+    return failure();
+  Token tok = getToken();
 
-    // Otherwise, we parse the location directly.
+  // Check to see if we are parsing a location alias.
+  LocationAttr directLoc;
+  if (tok.is(Token::hash_identifier)) {
+    if (parseLocationAlias(directLoc))
+      return failure();
   } else if (parseLocationInstance(directLoc)) {
     return failure();
   }
 
   if (parseToken(Token::r_paren, "expected ')' in location"))
     return failure();
 
   if (auto *op = opOrArgument.dyn_cast<Operation *>())
     op->setLoc(directLoc);
   else
     opOrArgument.get<BlockArgument>().setLoc(directLoc);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Region Parsing

mlir/test/IR/pretty_printed_region_op.mlir

-// 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
 
 // -----
 
 func @pretty_printed_region_op(%arg0 : f32, %arg1 : f32) -> (f32) {
 // CHECK-CUSTOM:  test.pretty_printed_region %arg1, %arg0 start special.op end : (f32, f32) -> f32
 // CHECK-GENERIC: "test.pretty_printed_region"(%arg1, %arg0)
 // CHECK-GENERIC:   ^bb0(%arg[[x:[0-9]+]]: f32, %arg[[y:[0-9]+]]: f32
 // CHECK-GENERIC:     %[[RES:.*]] = "special.op"(%arg[[x]], %arg[[y]]) : (f32, f32) -> f32
// CHECK:          : (f32, f32) -> f32
   %0 = test.pretty_printed_region %arg1, %arg0 ( {
     ^bb0(%arg2: f32, %arg3: f32):
       %1 = "non.special.op"(%arg2, %arg3) : (f32, f32) -> f32
       "test.return"(%1) : (f32) -> ()
     }) : (f32, 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
+}