[llvm] 529c091 - [ORC] Simplify JITLinkRedirectableSymbolManager, fix definition locations.
Lang Hames via llvm-commits
llvm-commits at lists.llvm.org
Fri Nov 1 21:33:54 PDT 2024
Author: Lang Hames
Date: 2024-11-02T15:22:02+11:00
New Revision: 529c091381a4d34796b5d0f0f2f5bd17c8a6647e
URL: https://github.com/llvm/llvm-project/commit/529c091381a4d34796b5d0f0f2f5bd17c8a6647e
DIFF: https://github.com/llvm/llvm-project/commit/529c091381a4d34796b5d0f0f2f5bd17c8a6647e.diff
LOG: [ORC] Simplify JITLinkRedirectableSymbolManager, fix definition locations.
Redirectable stubs should be placed in the same JITDylib as their names, with
their lifetimes managed according to the ResourceTracker used when adding them.
The original implementation created a single pool of stubs in the JITDylib
that is passed to the JITLinkRedirectableSymbolManager constructor, but this
may cause the addresses of the redirectable symbols themselves (added to the
JITDylibs passed to createRedirectableSymbols) to appear outside the address
range of their defining JITDylib.
This patch fixes the issue by dropping the pool and emitting a new graph for
each set of requested redirectable symbols. We lose the ability to recycle
stubs, but gain the ability to free them entirely. Since we don't expect stub
reuse to be a common case this is likely the best trade-off.
If in the future we do need to return to a stub pool we should create a
separate one for each JITDylib to ensure that addresses behave as expected.
Added:
Modified:
llvm/include/llvm/ExecutionEngine/JITLink/JITLink.h
llvm/include/llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h
llvm/lib/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.cpp
llvm/unittests/ExecutionEngine/Orc/JITLinkRedirectionManagerTest.cpp
llvm/unittests/ExecutionEngine/Orc/ReOptimizeLayerTest.cpp
Removed:
################################################################################
diff --git a/llvm/include/llvm/ExecutionEngine/JITLink/JITLink.h b/llvm/include/llvm/ExecutionEngine/JITLink/JITLink.h
index dac1b1d8afdc05..9844214c537a06 100644
--- a/llvm/include/llvm/ExecutionEngine/JITLink/JITLink.h
+++ b/llvm/include/llvm/ExecutionEngine/JITLink/JITLink.h
@@ -1077,6 +1077,16 @@ class LinkGraph {
return MutableArrayRef<char>(AllocatedBuffer, SourceStr.size());
}
+ /// Allocate a copy of the given string using the LinkGraph's allocator
+ /// and return it as a StringRef.
+ ///
+ /// This is a convenience wrapper around allocateContent(Twine) that is
+ /// handy when creating new symbol names within the graph.
+ StringRef allocateName(Twine Source) {
+ auto Buf = allocateContent(Source);
+ return {Buf.data(), Buf.size()};
+ }
+
/// Allocate a copy of the given string using the LinkGraph's allocator.
///
/// The allocated string will be terminated with a null character, and the
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h b/llvm/include/llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h
index 8a4740c1dd9cb9..be92f5083969f7 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h
@@ -17,15 +17,16 @@
#include "llvm/ExecutionEngine/Orc/RedirectionManager.h"
#include "llvm/Support/StringSaver.h"
+#include <atomic>
+
namespace llvm {
namespace orc {
-class JITLinkRedirectableSymbolManager : public RedirectableSymbolManager,
- public ResourceManager {
+class JITLinkRedirectableSymbolManager : public RedirectableSymbolManager {
public:
/// Create redirection manager that uses JITLink based implementaion.
static Expected<std::unique_ptr<RedirectableSymbolManager>>
- Create(ObjectLinkingLayer &ObjLinkingLayer, JITDylib &JD) {
+ Create(ObjectLinkingLayer &ObjLinkingLayer) {
auto AnonymousPtrCreator(jitlink::getAnonymousPointerCreator(
ObjLinkingLayer.getExecutionSession().getTargetTriple()));
auto PtrJumpStubCreator(jitlink::getPointerJumpStubCreator(
@@ -35,7 +36,7 @@ class JITLinkRedirectableSymbolManager : public RedirectableSymbolManager,
inconvertibleErrorCode());
return std::unique_ptr<RedirectableSymbolManager>(
new JITLinkRedirectableSymbolManager(
- ObjLinkingLayer, JD, AnonymousPtrCreator, PtrJumpStubCreator));
+ ObjLinkingLayer, AnonymousPtrCreator, PtrJumpStubCreator));
}
void emitRedirectableSymbols(std::unique_ptr<MaterializationResponsibility> R,
@@ -43,61 +44,19 @@ class JITLinkRedirectableSymbolManager : public RedirectableSymbolManager,
Error redirect(JITDylib &TargetJD, const SymbolAddrMap &NewDests) override;
- Error handleRemoveResources(JITDylib &TargetJD, ResourceKey K) override;
-
- void handleTransferResources(JITDylib &TargetJD, ResourceKey DstK,
- ResourceKey SrcK) override;
-
private:
- using StubHandle = unsigned;
- constexpr static unsigned StubBlockSize = 256;
- constexpr static StringRef JumpStubPrefix = "$__IND_JUMP_STUBS";
- constexpr static StringRef StubPtrPrefix = "$IND_JUMP_PTR_";
- constexpr static StringRef JumpStubTableName = "$IND_JUMP_";
- constexpr static StringRef StubPtrTableName = "$__IND_JUMP_PTRS";
-
JITLinkRedirectableSymbolManager(
- ObjectLinkingLayer &ObjLinkingLayer, JITDylib &JD,
+ ObjectLinkingLayer &ObjLinkingLayer,
jitlink::AnonymousPointerCreator &AnonymousPtrCreator,
jitlink::PointerJumpStubCreator &PtrJumpStubCreator)
- : ObjLinkingLayer(ObjLinkingLayer), JD(JD),
+ : ObjLinkingLayer(ObjLinkingLayer),
AnonymousPtrCreator(std::move(AnonymousPtrCreator)),
- PtrJumpStubCreator(std::move(PtrJumpStubCreator)) {
- ObjLinkingLayer.getExecutionSession().registerResourceManager(*this);
- }
-
- ~JITLinkRedirectableSymbolManager() {
- ObjLinkingLayer.getExecutionSession().deregisterResourceManager(*this);
- }
-
- StringRef JumpStubSymbolName(unsigned I) {
- return *ObjLinkingLayer.getExecutionSession().intern(
- (JumpStubPrefix + Twine(I)).str());
- }
-
- StringRef StubPtrSymbolName(unsigned I) {
- return *ObjLinkingLayer.getExecutionSession().intern(
- (StubPtrPrefix + Twine(I)).str());
- }
-
- unsigned GetNumAvailableStubs() const { return AvailableStubs.size(); }
-
- Error redirectInner(JITDylib &TargetJD, const SymbolAddrMap &NewDests);
- Error grow(unsigned Need);
+ PtrJumpStubCreator(std::move(PtrJumpStubCreator)) {}
ObjectLinkingLayer &ObjLinkingLayer;
- JITDylib &JD;
jitlink::AnonymousPointerCreator AnonymousPtrCreator;
jitlink::PointerJumpStubCreator PtrJumpStubCreator;
-
- std::vector<StubHandle> AvailableStubs;
- using SymbolToStubMap = DenseMap<SymbolStringPtr, StubHandle>;
- DenseMap<JITDylib *, SymbolToStubMap> SymbolToStubs;
- std::vector<ExecutorSymbolDef> JumpStubs;
- std::vector<ExecutorSymbolDef> StubPointers;
- DenseMap<ResourceKey, std::vector<SymbolStringPtr>> TrackedResources;
-
- std::mutex Mutex;
+ std::atomic_size_t StubGraphIdx{0};
};
} // namespace orc
diff --git a/llvm/lib/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.cpp b/llvm/lib/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.cpp
index eec70d917237d4..8c6d95374a6158 100644
--- a/llvm/lib/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.cpp
@@ -9,171 +9,88 @@
#include "llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h"
#include "llvm/ExecutionEngine/Orc/Core.h"
+#include "llvm/ExecutionEngine/Orc/DebugUtils.h"
+
#define DEBUG_TYPE "orc"
using namespace llvm;
using namespace llvm::orc;
+namespace {
+constexpr StringRef JumpStubSectionName = "__orc_stubs";
+constexpr StringRef StubPtrSectionName = "__orc_stub_ptrs";
+constexpr StringRef StubSuffix = "$__stub_ptr";
+} // namespace
+
void JITLinkRedirectableSymbolManager::emitRedirectableSymbols(
std::unique_ptr<MaterializationResponsibility> R,
const SymbolAddrMap &InitialDests) {
- auto &ES = ObjLinkingLayer.getExecutionSession();
- std::unique_lock<std::mutex> Lock(Mutex);
- if (GetNumAvailableStubs() < InitialDests.size())
- if (auto Err = grow(InitialDests.size() - GetNumAvailableStubs())) {
- ES.reportError(std::move(Err));
- R->failMaterialization();
- return;
- }
- JITDylib &TargetJD = R->getTargetJITDylib();
- SymbolMap NewSymbolDefs;
- std::vector<SymbolStringPtr> Symbols;
- for (auto &[K, V] : InitialDests) {
- StubHandle StubID = AvailableStubs.back();
- if (SymbolToStubs[&TargetJD].count(K)) {
- ES.reportError(make_error<StringError>(
- "Tried to create duplicate redirectable symbols",
- inconvertibleErrorCode()));
- R->failMaterialization();
- return;
- }
- SymbolToStubs[&TargetJD][K] = StubID;
- NewSymbolDefs[K] = JumpStubs[StubID];
- NewSymbolDefs[K].setFlags(V.getFlags());
- Symbols.push_back(K);
- AvailableStubs.pop_back();
- }
+ auto &ES = ObjLinkingLayer.getExecutionSession();
+ Triple TT = ES.getTargetTriple();
- // FIXME: when this fails we can return stubs to the pool
- if (auto Err = redirectInner(TargetJD, InitialDests)) {
- ES.reportError(std::move(Err));
- R->failMaterialization();
- return;
+ auto G = std::make_unique<jitlink::LinkGraph>(
+ ("<INDIRECT STUBS #" + Twine(++StubGraphIdx) + ">").str(), TT,
+ TT.isArch64Bit() ? 8 : 4,
+ TT.isLittleEndian() ? endianness::little : endianness::big,
+ jitlink::getGenericEdgeKindName);
+ auto &PointerSection =
+ G->createSection(StubPtrSectionName, MemProt::Write | MemProt::Read);
+ auto &StubsSection =
+ G->createSection(JumpStubSectionName, MemProt::Exec | MemProt::Read);
+
+ SymbolFlagsMap NewSymbols;
+ for (auto &[Name, Def] : InitialDests) {
+ jitlink::Symbol *TargetSym = nullptr;
+ if (Def.getAddress())
+ TargetSym = &G->addAbsoluteSymbol(
+ G->allocateName(*Name + "$__init_tgt"), Def.getAddress(), 0,
+ jitlink::Linkage::Strong, jitlink::Scope::Local, false);
+
+ auto PtrName = ES.intern((*Name + StubSuffix).str());
+ auto &Ptr = AnonymousPtrCreator(*G, PointerSection, TargetSym, 0);
+ Ptr.setName(*PtrName);
+ Ptr.setScope(jitlink::Scope::Hidden);
+ auto &Stub = PtrJumpStubCreator(*G, StubsSection, Ptr);
+ Stub.setName(*Name);
+ Stub.setScope(jitlink::Scope::Default);
+ NewSymbols[std::move(PtrName)] = JITSymbolFlags();
}
- // FIXME: return stubs to the pool here too.
- if (auto Err = R->replace(absoluteSymbols(NewSymbolDefs))) {
+ // Try to claim responsibility for the new stub symbols.
+ if (auto Err = R->defineMaterializing(std::move(NewSymbols))) {
ES.reportError(std::move(Err));
- R->failMaterialization();
- return;
+ return R->failMaterialization();
}
- // FIXME: return stubs to the pool here too.
- if (auto Err = R->withResourceKeyDo([&](ResourceKey Key) {
- TrackedResources[Key].insert(TrackedResources[Key].end(),
- Symbols.begin(), Symbols.end());
- })) {
- ES.reportError(std::move(Err));
- R->failMaterialization();
- return;
- }
+ ObjLinkingLayer.emit(std::move(R), std::move(G));
}
Error JITLinkRedirectableSymbolManager::redirect(
JITDylib &TargetJD, const SymbolAddrMap &NewDests) {
- std::unique_lock<std::mutex> Lock(Mutex);
- return redirectInner(TargetJD, NewDests);
-}
+ auto &ES = ObjLinkingLayer.getExecutionSession();
+ SymbolLookupSet LS;
+ DenseMap<NonOwningSymbolStringPtr, SymbolStringPtr> PtrToStub;
+ for (auto &[StubName, Sym] : NewDests) {
+ auto PtrName = ES.intern((*StubName + StubSuffix).str());
+ PtrToStub[NonOwningSymbolStringPtr(PtrName)] = StubName;
+ LS.add(std::move(PtrName));
+ }
+ auto PtrSyms = ES.lookup({{&TargetJD, JITDylibLookupFlags::MatchAllSymbols}},
+ std::move(LS));
+ if (!PtrSyms)
+ return PtrSyms.takeError();
-Error JITLinkRedirectableSymbolManager::redirectInner(
- JITDylib &TargetJD, const SymbolAddrMap &NewDests) {
std::vector<tpctypes::PointerWrite> PtrWrites;
- for (auto &[K, V] : NewDests) {
- if (!SymbolToStubs[&TargetJD].count(K))
- return make_error<StringError>(
- "Tried to redirect non-existent redirectalbe symbol",
- inconvertibleErrorCode());
- StubHandle StubID = SymbolToStubs[&TargetJD].at(K);
- PtrWrites.push_back({StubPointers[StubID].getAddress(), V.getAddress()});
+ for (auto &[PtrName, PtrSym] : *PtrSyms) {
+ auto DestSymI = NewDests.find(PtrToStub[NonOwningSymbolStringPtr(PtrName)]);
+ assert(DestSymI != NewDests.end() && "Bad ptr -> stub mapping");
+ auto &DestSym = DestSymI->second;
+ PtrWrites.push_back({PtrSym.getAddress(), DestSym.getAddress()});
}
+
return ObjLinkingLayer.getExecutionSession()
.getExecutorProcessControl()
.getMemoryAccess()
.writePointers(PtrWrites);
}
-
-Error JITLinkRedirectableSymbolManager::grow(unsigned Need) {
- unsigned OldSize = JumpStubs.size();
- unsigned NumNewStubs = alignTo(Need, StubBlockSize);
- unsigned NewSize = OldSize + NumNewStubs;
-
- JumpStubs.resize(NewSize);
- StubPointers.resize(NewSize);
- AvailableStubs.reserve(NewSize);
-
- SymbolLookupSet LookupSymbols;
- DenseMap<SymbolStringPtr, ExecutorSymbolDef *> NewDefsMap;
-
- auto &ES = ObjLinkingLayer.getExecutionSession();
- Triple TT = ES.getTargetTriple();
- auto G = std::make_unique<jitlink::LinkGraph>(
- "<INDIRECT STUBS>", TT, TT.isArch64Bit() ? 8 : 4,
- TT.isLittleEndian() ? endianness::little : endianness::big,
- jitlink::getGenericEdgeKindName);
- auto &PointerSection =
- G->createSection(StubPtrTableName, MemProt::Write | MemProt::Read);
- auto &StubsSection =
- G->createSection(JumpStubTableName, MemProt::Exec | MemProt::Read);
-
- // FIXME: We can batch the stubs into one block and use address to access them
- for (size_t I = OldSize; I < NewSize; I++) {
- auto &Pointer = AnonymousPtrCreator(*G, PointerSection, nullptr, 0);
-
- StringRef PtrSymName = StubPtrSymbolName(I);
- Pointer.setName(PtrSymName);
- Pointer.setScope(jitlink::Scope::Default);
- LookupSymbols.add(ES.intern(PtrSymName));
- NewDefsMap[ES.intern(PtrSymName)] = &StubPointers[I];
-
- auto &Stub = PtrJumpStubCreator(*G, StubsSection, Pointer);
-
- StringRef JumpStubSymName = JumpStubSymbolName(I);
- Stub.setName(JumpStubSymName);
- Stub.setScope(jitlink::Scope::Default);
- LookupSymbols.add(ES.intern(JumpStubSymName));
- NewDefsMap[ES.intern(JumpStubSymName)] = &JumpStubs[I];
- }
-
- if (auto Err = ObjLinkingLayer.add(JD, std::move(G)))
- return Err;
-
- auto LookupResult = ES.lookup(makeJITDylibSearchOrder(&JD), LookupSymbols);
- if (auto Err = LookupResult.takeError())
- return Err;
-
- for (auto &[K, V] : *LookupResult)
- *NewDefsMap.at(K) = V;
-
- for (size_t I = OldSize; I < NewSize; I++)
- AvailableStubs.push_back(I);
-
- return Error::success();
-}
-
-Error JITLinkRedirectableSymbolManager::handleRemoveResources(
- JITDylib &TargetJD, ResourceKey K) {
- std::unique_lock<std::mutex> Lock(Mutex);
- for (auto &Symbol : TrackedResources[K]) {
- if (!SymbolToStubs[&TargetJD].count(Symbol))
- return make_error<StringError>(
- "Tried to remove non-existent redirectable symbol",
- inconvertibleErrorCode());
- AvailableStubs.push_back(SymbolToStubs[&TargetJD].at(Symbol));
- SymbolToStubs[&TargetJD].erase(Symbol);
- if (SymbolToStubs[&TargetJD].empty())
- SymbolToStubs.erase(&TargetJD);
- }
- TrackedResources.erase(K);
-
- return Error::success();
-}
-
-void JITLinkRedirectableSymbolManager::handleTransferResources(
- JITDylib &TargetJD, ResourceKey DstK, ResourceKey SrcK) {
- std::unique_lock<std::mutex> Lock(Mutex);
- TrackedResources[DstK].insert(TrackedResources[DstK].end(),
- TrackedResources[SrcK].begin(),
- TrackedResources[SrcK].end());
- TrackedResources.erase(SrcK);
-}
diff --git a/llvm/unittests/ExecutionEngine/Orc/JITLinkRedirectionManagerTest.cpp b/llvm/unittests/ExecutionEngine/Orc/JITLinkRedirectionManagerTest.cpp
index ea8d4fa12daaac..2e9733cbb039b7 100644
--- a/llvm/unittests/ExecutionEngine/Orc/JITLinkRedirectionManagerTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/JITLinkRedirectionManagerTest.cpp
@@ -63,51 +63,39 @@ class JITLinkRedirectionManagerTest : public testing::Test {
};
TEST_F(JITLinkRedirectionManagerTest, BasicRedirectionOperation) {
- auto RM = JITLinkRedirectableSymbolManager::Create(*ObjLinkingLayer, *JD);
+ auto RM = JITLinkRedirectableSymbolManager::Create(*ObjLinkingLayer);
// Bail out if we can not create
if (!RM) {
consumeError(RM.takeError());
GTEST_SKIP();
}
- auto DefineTarget = [&](StringRef TargetName, ExecutorAddr Addr) {
- SymbolStringPtr Target = ES->intern(TargetName);
- cantFail(JD->define(std::make_unique<SimpleMaterializationUnit>(
- SymbolFlagsMap({{Target, JITSymbolFlags::Exported}}),
- [&](std::unique_ptr<MaterializationResponsibility> R) -> void {
- // No dependencies registered, can't fail.
- cantFail(
- R->notifyResolved({{Target, {Addr, JITSymbolFlags::Exported}}}));
- cantFail(R->notifyEmitted({}));
- })));
- return cantFail(ES->lookup({JD}, TargetName));
+ auto MakeTarget = [](int (*Fn)()) {
+ return ExecutorSymbolDef(ExecutorAddr::fromPtr(Fn),
+ JITSymbolFlags::Exported |
+ JITSymbolFlags::Callable);
};
- auto InitialTarget =
- DefineTarget("InitialTarget", ExecutorAddr::fromPtr(&initialTarget));
- auto MiddleTarget =
- DefineTarget("MiddleTarget", ExecutorAddr::fromPtr(&middleTarget));
- auto FinalTarget =
- DefineTarget("FinalTarget", ExecutorAddr::fromPtr(&finalTarget));
-
auto RedirectableSymbol = ES->intern("RedirectableTarget");
- EXPECT_THAT_ERROR(
- (*RM)->createRedirectableSymbols(JD->getDefaultResourceTracker(),
- {{RedirectableSymbol, InitialTarget}}),
- Succeeded());
+ EXPECT_THAT_ERROR((*RM)->createRedirectableSymbols(
+ JD->getDefaultResourceTracker(),
+ {{RedirectableSymbol, MakeTarget(initialTarget)}}),
+ Succeeded());
auto RTDef = cantFail(ES->lookup({JD}, RedirectableSymbol));
auto RTPtr = RTDef.getAddress().toPtr<int (*)()>();
auto Result = RTPtr();
EXPECT_EQ(Result, 42) << "Failed to call initial target";
- EXPECT_THAT_ERROR((*RM)->redirect(*JD, {{RedirectableSymbol, MiddleTarget}}),
- Succeeded());
+ EXPECT_THAT_ERROR(
+ (*RM)->redirect(*JD, {{RedirectableSymbol, MakeTarget(middleTarget)}}),
+ Succeeded());
Result = RTPtr();
EXPECT_EQ(Result, 13) << "Failed to call middle redirected target";
- EXPECT_THAT_ERROR((*RM)->redirect(*JD, {{RedirectableSymbol, FinalTarget}}),
- Succeeded());
+ EXPECT_THAT_ERROR(
+ (*RM)->redirect(*JD, {{RedirectableSymbol, MakeTarget(finalTarget)}}),
+ Succeeded());
Result = RTPtr();
EXPECT_EQ(Result, 53) << "Failed to call redirected target";
}
diff --git a/llvm/unittests/ExecutionEngine/Orc/ReOptimizeLayerTest.cpp b/llvm/unittests/ExecutionEngine/Orc/ReOptimizeLayerTest.cpp
index 2b6cd4984c8de1..20db572417d5f6 100644
--- a/llvm/unittests/ExecutionEngine/Orc/ReOptimizeLayerTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/ReOptimizeLayerTest.cpp
@@ -128,7 +128,7 @@ TEST_F(ReOptimizeLayerTest, BasicReOptimization) {
{ExecutorAddr(), JITSymbolFlags::Exported}}})),
Succeeded());
- auto RM = JITLinkRedirectableSymbolManager::Create(*ObjLinkingLayer, *JD);
+ auto RM = JITLinkRedirectableSymbolManager::Create(*ObjLinkingLayer);
EXPECT_THAT_ERROR(RM.takeError(), Succeeded());
ROLayer = std::make_unique<ReOptimizeLayer>(*ES, *DL, *CompileLayer, **RM);
More information about the llvm-commits
mailing list