diff --git a/mlir/include/mlir/IR/BuiltinLocationAttributes.td b/mlir/include/mlir/IR/BuiltinLocationAttributes.td --- a/mlir/include/mlir/IR/BuiltinLocationAttributes.td +++ b/mlir/include/mlir/IR/BuiltinLocationAttributes.td @@ -202,12 +202,12 @@ /// Returns an instance of opaque location which contains a given pointer to /// an object. The corresponding MLIR location is set to UnknownLoc. template - 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 - static Location get(T underlyingLocation, Location fallbackLocation) { + static OpaqueLoc get(T underlyingLocation, Location fallbackLocation) { return get(reinterpret_cast(underlyingLocation), TypeID::get(), fallbackLocation); } diff --git a/mlir/include/mlir/IR/Location.h b/mlir/include/mlir/IR/Location.h --- a/mlir/include/mlir/IR/Location.h +++ b/mlir/include/mlir/IR/Location.h @@ -115,7 +115,7 @@ /// Returns an instance of opaque location which contains a given pointer to /// an object. The corresponding MLIR location is set to UnknownLoc. template -inline Location OpaqueLoc::get(T underlyingLocation, MLIRContext *context) { +inline OpaqueLoc OpaqueLoc::get(T underlyingLocation, MLIRContext *context) { return get(reinterpret_cast(underlyingLocation), TypeID::get(), UnknownLoc::get(context)); } diff --git a/mlir/lib/Parser/Parser.cpp b/mlir/lib/Parser/Parser.cpp --- a/mlir/lib/Parser/Parser.cpp +++ b/mlir/lib/Parser/Parser.cpp @@ -476,10 +476,14 @@ /// their first reference, to allow checking for use of undefined values. DenseMap forwardRefPlaceholders; - /// A set of operations whose locations reference aliases that have yet to - /// be resolved. - SmallVector, 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 deferredLocsReferences; /// The builder used when creating parsed operation instances. OpBuilder opBuilder; @@ -537,23 +541,36 @@ // Resolve the locations of any deferred operations. auto &attributeAliases = state.symbols.attributeAliasDefinitions; - for (std::pair &it : opsAndArgumentsWithDeferredLocs) { - llvm::SMLoc tokLoc = it.second.getLoc(); - StringRef identifier = it.second.getSpelling().drop_front(); - Attribute attr = attributeAliases.lookup(identifier); + auto locID = TypeID::get(); + auto resolveLocation = [&](auto &opOrArgument) -> LogicalResult { + auto fwdLoc = opOrArgument.getLoc().template dyn_cast(); + 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(); + return emitError(locInfo.loc) + << "operation location alias was never defined"; + auto locAttr = attr.dyn_cast(); if (!locAttr) - return emitError(tokLoc) + return emitError(locInfo.loc) << "expected location, but found '" << attr << "'"; - auto opOrArgument = it.first; - if (auto *op = opOrArgument.dyn_cast()) - op->setLoc(locAttr); - else - opOrArgument.get().setLoc(locAttr); - } + opOrArgument.setLoc(locAttr); + return success(); + }; + + auto walkRes = topLevelOp->walk([&](Operation *op) { + if (failed(resolveLocation(*op))) + return WalkResult::interrupt(); + for (Region ®ion : 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 @@ << "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(), + UnknownLoc::get(getContext())); + deferredLocsReferences.push_back( + DeferredLocInfo{tok.getLoc(), identifier}); } // Otherwise, we parse the location directly. @@ -1742,12 +1764,10 @@ if (parseToken(Token::r_paren, "expected ')' in location")) return failure(); - if (directLoc) { - if (auto *op = opOrArgument.dyn_cast()) - op->setLoc(directLoc); - else - opOrArgument.get().setLoc(directLoc); - } + if (auto *op = opOrArgument.dyn_cast()) + op->setLoc(directLoc); + else + opOrArgument.get().setLoc(directLoc); return success(); }