[Mlir-commits] [mlir] a440b04 - Use Opaque location for handling deferred references to aliases internally to the parser (NFC)

[Mehdi Amini]

Author: Mehdi Amini
Date: 2022-01-18T20:43:36Z
New Revision: a440b0404f8625984389db9ee102ba308754747a

URL: https://github.com/llvm/llvm-project/commit/a440b0404f8625984389db9ee102ba308754747a
DIFF: https://github.com/llvm/llvm-project/commit/a440b0404f8625984389db9ee102ba308754747a.diff

LOG: Use Opaque location for handling deferred references to aliases internally to the parser (NFC)

This will allow to return to the client of `parseLocationInstance` a location that is resolved later.

Reviewed By: rriddle

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

Added: 
    

Modified: 
    mlir/include/mlir/IR/BuiltinLocationAttributes.td
    mlir/include/mlir/IR/Location.h
    mlir/lib/Parser/Parser.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/BuiltinLocationAttributes.td b/mlir/include/mlir/IR/BuiltinLocationAttributes.td
index 535016b9912ad..0bb3b7a8a59a9 100644
--- a/mlir/include/mlir/IR/BuiltinLocationAttributes.td
+++ b/mlir/include/mlir/IR/BuiltinLocationAttributes.td
@@ -202,12 +202,12 @@ def OpaqueLoc : Builtin_LocationAttr<"OpaqueLoc"> {
     /// Returns an instance of opaque location which contains a given pointer to
     /// an object. The corresponding MLIR location is set to UnknownLoc.
     template <typename T>
-    static Location get(T underlyingLocation, MLIRContext *context);
+    static OpaqueLoc get(T underlyingLocation, MLIRContext *context);
 
     /// Returns an instance of opaque location which contains a given pointer to
     /// an object and an additional MLIR location.
     template <typename T>
-    static Location get(T underlyingLocation, Location fallbackLocation) {
+    static OpaqueLoc get(T underlyingLocation, Location fallbackLocation) {
       return get(reinterpret_cast<uintptr_t>(underlyingLocation),
                  TypeID::get<T>(), fallbackLocation);
     }

diff  --git a/mlir/include/mlir/IR/Location.h b/mlir/include/mlir/IR/Location.h
index 4a3f22e8eea53..16cd3889e5763 100644
--- a/mlir/include/mlir/IR/Location.h
+++ b/mlir/include/mlir/IR/Location.h
@@ -115,7 +115,7 @@ namespace mlir {
 /// Returns an instance of opaque location which contains a given pointer to
 /// an object. The corresponding MLIR location is set to UnknownLoc.
 template <typename T>
-inline Location OpaqueLoc::get(T underlyingLocation, MLIRContext *context) {
+inline OpaqueLoc OpaqueLoc::get(T underlyingLocation, MLIRContext *context) {
   return get(reinterpret_cast<uintptr_t>(underlyingLocation), TypeID::get<T>(),
              UnknownLoc::get(context));
 }

diff  --git a/mlir/lib/Parser/Parser.cpp b/mlir/lib/Parser/Parser.cpp
index da765b61243ba..e8a99f3b2bc78 100644
--- a/mlir/lib/Parser/Parser.cpp
+++ b/mlir/lib/Parser/Parser.cpp
@@ -476,10 +476,14 @@ class OperationParser : public Parser {
   /// their first reference, to allow checking for use of undefined values.
   DenseMap<Value, SMLoc> forwardRefPlaceholders;
 
-  /// A set of operations whose locations reference aliases that have yet to
-  /// be resolved.
-  SmallVector<std::pair<OpOrArgument, Token>, 8>
-      opsAndArgumentsWithDeferredLocs;
+  /// Deffered locations: when parsing `loc(#loc42)` we add an entry to this
+  /// map. After parsing the definition `#loc42 = ...` we'll patch back users
+  /// of this location.
+  struct DeferredLocInfo {
+    SMLoc loc;
+    StringRef identifier;
+  };
+  std::vector<DeferredLocInfo> deferredLocsReferences;
 
   /// The builder used when creating parsed operation instances.
   OpBuilder opBuilder;
@@ -537,23 +541,36 @@ ParseResult OperationParser::finalize() {
 
   // Resolve the locations of any deferred operations.
   auto &attributeAliases = state.symbols.attributeAliasDefinitions;
-  for (std::pair<OpOrArgument, Token> &it : opsAndArgumentsWithDeferredLocs) {
-    llvm::SMLoc tokLoc = it.second.getLoc();
-    StringRef identifier = it.second.getSpelling().drop_front();
-    Attribute attr = attributeAliases.lookup(identifier);
+  auto locID = TypeID::get<DeferredLocInfo *>();
+  auto resolveLocation = [&](auto &opOrArgument) -> LogicalResult {
+    auto fwdLoc = opOrArgument.getLoc().template dyn_cast<OpaqueLoc>();
+    if (!fwdLoc || fwdLoc.getUnderlyingTypeID() != locID)
+      return success();
+    auto locInfo = deferredLocsReferences[fwdLoc.getUnderlyingLocation()];
+    Attribute attr = attributeAliases.lookup(locInfo.identifier);
     if (!attr)
-      return emitError(tokLoc) << "operation location alias was never defined";
-
-    LocationAttr locAttr = attr.dyn_cast<LocationAttr>();
+      return emitError(locInfo.loc)
+             << "operation location alias was never defined";
+    auto locAttr = attr.dyn_cast<LocationAttr>();
     if (!locAttr)
-      return emitError(tokLoc)
+      return emitError(locInfo.loc)
              << "expected location, but found '" << attr << "'";
-    auto opOrArgument = it.first;
-    if (auto *op = opOrArgument.dyn_cast<Operation *>())
-      op->setLoc(locAttr);
-    else
-      opOrArgument.get<BlockArgument>().setLoc(locAttr);
-  }
+    opOrArgument.setLoc(locAttr);
+    return success();
+  };
+
+  auto walkRes = topLevelOp->walk([&](Operation *op) {
+    if (failed(resolveLocation(*op)))
+      return WalkResult::interrupt();
+    for (Region &region : op->getRegions())
+      for (Block &block : region.getBlocks())
+        for (BlockArgument arg : block.getArguments())
+          if (failed(resolveLocation(arg)))
+            return WalkResult::interrupt();
+    return WalkResult::advance();
+  });
+  if (walkRes.wasInterrupted())
+    return failure();
 
   // Pop the top level name scope.
   if (failed(popSSANameScope()))
@@ -1731,7 +1748,12 @@ OperationParser::parseTrailingLocationSpecifier(OpOrArgument opOrArgument) {
                << "expected location, but found '" << attr << "'";
     } else {
       // Otherwise, remember this operation and resolve its location later.
-      opsAndArgumentsWithDeferredLocs.emplace_back(opOrArgument, tok);
+      // 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.
@@ -1742,12 +1764,10 @@ OperationParser::parseTrailingLocationSpecifier(OpOrArgument opOrArgument) {
   if (parseToken(Token::r_paren, "expected ')' in location"))
     return failure();
 
-  if (directLoc) {
-    if (auto *op = opOrArgument.dyn_cast<Operation *>())
-      op->setLoc(directLoc);
-    else
-      opOrArgument.get<BlockArgument>().setLoc(directLoc);
-  }
+  if (auto *op = opOrArgument.dyn_cast<Operation *>())
+    op->setLoc(directLoc);
+  else
+    opOrArgument.get<BlockArgument>().setLoc(directLoc);
   return success();
 }