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 @@ -269,4 +269,24 @@ }]; } +//===----------------------------------------------------------------------===// +// AliasLoc +//===----------------------------------------------------------------------===// + +def FwdAliasLoc : Builtin_LocationAttr<"FwdAliasLoc"> { + let summary = "A location private to the parser, used to resolve aliases."; + let description = [{ + }]; + let parameters = (ins + "StringAttr":$aliasIdentifier, + "const char *":$smLoc + ); + let builders = [ + AttrBuilderWithInferredContext<(ins "StringAttr":$name, "const char *":$smLoc), [{ + return $_get(name.getContext(), name, smLoc); + }]> + ]; +} + + #endif // BUILTIN_LOCATION_ATTRIBUTES_TD diff --git a/mlir/lib/IR/Location.cpp b/mlir/lib/IR/Location.cpp --- a/mlir/lib/IR/Location.cpp +++ b/mlir/lib/IR/Location.cpp @@ -64,8 +64,8 @@ /// Methods for support type inquiry through isa, cast, and dyn_cast. bool LocationAttr::classof(Attribute attr) { - return attr.isa(); + return attr.isa(); } //===----------------------------------------------------------------------===// 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,11 +476,6 @@ /// 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; - /// The builder used when creating parsed operation instances. OpBuilder opBuilder; @@ -537,23 +532,35 @@ // 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); - if (!attr) - return emitError(tokLoc) << "operation location alias was never defined"; - - LocationAttr locAttr = attr.dyn_cast(); + auto getLocation = [&](OpOrArgument opOrArgument, + FwdAliasLoc fwdLoc) -> LocationAttr { + LocationAttr locAttr; + SMLoc tokLoc = SMLoc::getFromPointer(fwdLoc.getSmLoc()); + Attribute attr = attributeAliases.lookup(fwdLoc.getAliasIdentifier()); + if (!attr) { + emitError(tokLoc) << "operation location alias was never defined"; + return locAttr; + } + + locAttr = attr.dyn_cast(); if (!locAttr) - return emitError(tokLoc) - << "expected location, but found '" << attr << "'"; - auto opOrArgument = it.first; - if (auto *op = opOrArgument.dyn_cast()) - op->setLoc(locAttr); - else - opOrArgument.get().setLoc(locAttr); - } + emitError(tokLoc) << "expected location, but found '" << attr << "'"; + return locAttr; + }; + + auto walkRes = topLevelOp->walk([&](Operation *op) { + auto fwdLoc = op->getLoc().dyn_cast(); + if (!fwdLoc) + return WalkResult::advance(); + auto locAttr = getLocation(op, fwdLoc); + if (!locAttr) + return WalkResult::interrupt(); + op->setLoc(locAttr); + // FIXME: block arguments + return WalkResult::advance(); + }); + if (walkRes.wasInterrupted()) + return failure(); // Pop the top level name scope. if (failed(popSSANameScope())) @@ -1731,7 +1738,9 @@ << "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 FwdAliasLoc as a marker. + directLoc = FwdAliasLoc::get(StringAttr::get(getContext(), identifier), + tok.getLoc().getPointer()); } // Otherwise, we parse the location directly.