From 2d10b7b750f97b42055d5b9b08a88c18ff811cd2 Mon Sep 17 00:00:00 2001 From: Lang Hames Date: Fri, 10 Jan 2025 12:09:47 +1100 Subject: [PATCH] Reapply "[ORC][llvm-jitlink] Add SimpleLazyReexportsSpeculator..." with fixes. This reapplies 6d72bf47606, which was reverted in 57447d3ddf to investigate build failures, e.g. https://lab.llvm.org/buildbot/#/builders/3/builds/10114. The original patch contained an invalid unused friend declaration of std::make_shared. This has been removed. --- .../Orc/JITLinkReentryTrampolines.h | 3 +- .../llvm/ExecutionEngine/Orc/LazyReexports.h | 100 +++++++- .../llvm/ExecutionEngine/Orc/TaskDispatch.h | 15 ++ .../Orc/JITLinkReentryTrampolines.cpp | 5 +- .../lib/ExecutionEngine/Orc/LazyReexports.cpp | 225 ++++++++++++++++-- llvm/lib/ExecutionEngine/Orc/TaskDispatch.cpp | 61 +++-- llvm/tools/llvm-jitlink/llvm-jitlink.cpp | 150 +++++++++++- llvm/tools/llvm-jitlink/llvm-jitlink.h | 13 +- 8 files changed, 507 insertions(+), 65 deletions(-) diff --git a/llvm/include/llvm/ExecutionEngine/Orc/JITLinkReentryTrampolines.h b/llvm/include/llvm/ExecutionEngine/Orc/JITLinkReentryTrampolines.h index 673019b748b3b..94d2ff0717eb4 100644 --- a/llvm/include/llvm/ExecutionEngine/Orc/JITLinkReentryTrampolines.h +++ b/llvm/include/llvm/ExecutionEngine/Orc/JITLinkReentryTrampolines.h @@ -65,7 +65,8 @@ class JITLinkReentryTrampolines { Expected> createJITLinkLazyReexportsManager(ObjectLinkingLayer &ObjLinkingLayer, RedirectableSymbolManager &RSMgr, - JITDylib &PlatformJD); + JITDylib &PlatformJD, + LazyReexportsManager::Listener *L = nullptr); } // namespace llvm::orc diff --git a/llvm/include/llvm/ExecutionEngine/Orc/LazyReexports.h b/llvm/include/llvm/ExecutionEngine/Orc/LazyReexports.h index c6b5d08544b1e..635f2b08367e6 100644 --- a/llvm/include/llvm/ExecutionEngine/Orc/LazyReexports.h +++ b/llvm/include/llvm/ExecutionEngine/Orc/LazyReexports.h @@ -179,6 +179,37 @@ class LazyReexportsManager : public ResourceManager { lazyReexports(LazyReexportsManager &, SymbolAliasMap); public: + struct CallThroughInfo { + JITDylibSP JD; + SymbolStringPtr Name; + SymbolStringPtr BodyName; + }; + + class Listener { + public: + using CallThroughInfo = LazyReexportsManager::CallThroughInfo; + + virtual ~Listener(); + + /// Called under the session lock when new lazy reexports are created. + virtual void onLazyReexportsCreated(JITDylib &JD, ResourceKey K, + const SymbolAliasMap &Reexports) = 0; + + /// Called under the session lock when lazy reexports have their ownership + /// transferred to a new ResourceKey. + virtual void onLazyReexportsTransfered(JITDylib &JD, ResourceKey DstK, + ResourceKey SrcK) = 0; + + /// Called under the session lock when lazy reexports are removed. + virtual Error onLazyReexportsRemoved(JITDylib &JD, ResourceKey K) = 0; + + /// Called outside the session lock when a lazy reexport is called. + /// NOTE: Since this is called outside the session lock there is a chance + /// that the reexport referred to has already been removed. Listeners + /// must be prepared to handle requests for stale reexports. + virtual void onLazyReexportCalled(const CallThroughInfo &CTI) = 0; + }; + using OnTrampolinesReadyFn = unique_function> EntryAddrs)>; using EmitTrampolinesFn = @@ -189,7 +220,7 @@ class LazyReexportsManager : public ResourceManager { /// This will work both in-process and out-of-process. static Expected> Create(EmitTrampolinesFn EmitTrampolines, RedirectableSymbolManager &RSMgr, - JITDylib &PlatformJD); + JITDylib &PlatformJD, Listener *L = nullptr); LazyReexportsManager(LazyReexportsManager &&) = delete; LazyReexportsManager &operator=(LazyReexportsManager &&) = delete; @@ -199,12 +230,6 @@ class LazyReexportsManager : public ResourceManager { ResourceKey SrcK) override; private: - struct CallThroughInfo { - SymbolStringPtr Name; - SymbolStringPtr BodyName; - JITDylibSP JD; - }; - class MU; class Plugin; @@ -213,7 +238,7 @@ class LazyReexportsManager : public ResourceManager { LazyReexportsManager(EmitTrampolinesFn EmitTrampolines, RedirectableSymbolManager &RSMgr, JITDylib &PlatformJD, - Error &Err); + Listener *L, Error &Err); std::unique_ptr createLazyReexports(SymbolAliasMap Reexports); @@ -229,6 +254,7 @@ class LazyReexportsManager : public ResourceManager { ExecutionSession &ES; EmitTrampolinesFn EmitTrampolines; RedirectableSymbolManager &RSMgr; + Listener *L; DenseMap> KeyToReentryAddrs; DenseMap CallThroughs; @@ -242,6 +268,64 @@ lazyReexports(LazyReexportsManager &LRM, SymbolAliasMap Reexports) { return LRM.createLazyReexports(std::move(Reexports)); } +class SimpleLazyReexportsSpeculator : public LazyReexportsManager::Listener { +public: + using RecordExecutionFunction = + unique_function; + + static std::shared_ptr + Create(ExecutionSession &ES, RecordExecutionFunction RecordExec = {}) { + class make_shared_helper : public SimpleLazyReexportsSpeculator { + public: + make_shared_helper(ExecutionSession &ES, + RecordExecutionFunction RecordExec) + : SimpleLazyReexportsSpeculator(ES, std::move(RecordExec)) {} + }; + + auto Instance = + std::make_shared(ES, std::move(RecordExec)); + Instance->WeakThis = Instance; + return Instance; + } + + SimpleLazyReexportsSpeculator(SimpleLazyReexportsSpeculator &&) = delete; + SimpleLazyReexportsSpeculator & + operator=(SimpleLazyReexportsSpeculator &&) = delete; + ~SimpleLazyReexportsSpeculator() override; + + void onLazyReexportsCreated(JITDylib &JD, ResourceKey K, + const SymbolAliasMap &Reexports) override; + + void onLazyReexportsTransfered(JITDylib &JD, ResourceKey DstK, + ResourceKey SrcK) override; + + Error onLazyReexportsRemoved(JITDylib &JD, ResourceKey K) override; + + void onLazyReexportCalled(const CallThroughInfo &CTI) override; + + void addSpeculationSuggestions( + std::vector> NewSuggestions); + +private: + SimpleLazyReexportsSpeculator(ExecutionSession &ES, + RecordExecutionFunction RecordExec) + : ES(ES), RecordExec(std::move(RecordExec)) {} + + bool doNextSpeculativeLookup(); + + class SpeculateTask; + + using KeyToFunctionBodiesMap = + DenseMap>; + + ExecutionSession &ES; + RecordExecutionFunction RecordExec; + std::weak_ptr WeakThis; + DenseMap LazyReexports; + std::deque> SpeculateSuggestions; + bool SpeculateTaskActive = false; +}; + } // End namespace orc } // End namespace llvm diff --git a/llvm/include/llvm/ExecutionEngine/Orc/TaskDispatch.h b/llvm/include/llvm/ExecutionEngine/Orc/TaskDispatch.h index d7939864fd609..67164679f9e30 100644 --- a/llvm/include/llvm/ExecutionEngine/Orc/TaskDispatch.h +++ b/llvm/include/llvm/ExecutionEngine/Orc/TaskDispatch.h @@ -92,6 +92,16 @@ makeGenericNamedTask(FnT &&Fn, const char *Desc = nullptr) { Desc); } +/// IdleTask can be used as the basis for low-priority tasks, e.g. speculative +/// lookup. +class IdleTask : public RTTIExtends { +public: + static char ID; + +private: + void anchor() override; +}; + /// Abstract base for classes that dispatch ORC Tasks. class TaskDispatcher { public: @@ -118,9 +128,13 @@ class DynamicThreadPoolTaskDispatcher : public TaskDispatcher { DynamicThreadPoolTaskDispatcher( std::optional MaxMaterializationThreads) : MaxMaterializationThreads(MaxMaterializationThreads) {} + void dispatch(std::unique_ptr T) override; void shutdown() override; private: + bool canRunMaterializationTaskNow(); + bool canRunIdleTaskNow(); + std::mutex DispatchMutex; bool Shutdown = false; size_t Outstanding = 0; @@ -129,6 +143,7 @@ class DynamicThreadPoolTaskDispatcher : public TaskDispatcher { std::optional MaxMaterializationThreads; size_t NumMaterializationThreads = 0; std::deque> MaterializationTaskQueue; + std::deque> IdleTaskQueue; }; #endif // LLVM_ENABLE_THREADS diff --git a/llvm/lib/ExecutionEngine/Orc/JITLinkReentryTrampolines.cpp b/llvm/lib/ExecutionEngine/Orc/JITLinkReentryTrampolines.cpp index be574ef7279c2..d9f859654d41e 100644 --- a/llvm/lib/ExecutionEngine/Orc/JITLinkReentryTrampolines.cpp +++ b/llvm/lib/ExecutionEngine/Orc/JITLinkReentryTrampolines.cpp @@ -173,7 +173,8 @@ void JITLinkReentryTrampolines::emit(ResourceTrackerSP RT, Expected> createJITLinkLazyReexportsManager(ObjectLinkingLayer &ObjLinkingLayer, RedirectableSymbolManager &RSMgr, - JITDylib &PlatformJD) { + JITDylib &PlatformJD, + LazyReexportsManager::Listener *L) { auto JLT = JITLinkReentryTrampolines::Create(ObjLinkingLayer); if (!JLT) return JLT.takeError(); @@ -184,7 +185,7 @@ createJITLinkLazyReexportsManager(ObjectLinkingLayer &ObjLinkingLayer, OnTrampolinesReady) mutable { JLT->emit(std::move(RT), NumTrampolines, std::move(OnTrampolinesReady)); }, - RSMgr, PlatformJD); + RSMgr, PlatformJD, L); } } // namespace llvm::orc diff --git a/llvm/lib/ExecutionEngine/Orc/LazyReexports.cpp b/llvm/lib/ExecutionEngine/Orc/LazyReexports.cpp index df4539ba4329f..7b38621eba824 100644 --- a/llvm/lib/ExecutionEngine/Orc/LazyReexports.cpp +++ b/llvm/lib/ExecutionEngine/Orc/LazyReexports.cpp @@ -280,31 +280,34 @@ class LazyReexportsManager::Plugin : public ObjectLinkingLayer::Plugin { std::mutex M; }; +LazyReexportsManager::Listener::~Listener() = default; + Expected> LazyReexportsManager::Create(EmitTrampolinesFn EmitTrampolines, RedirectableSymbolManager &RSMgr, - JITDylib &PlatformJD) { + JITDylib &PlatformJD, Listener *L) { Error Err = Error::success(); std::unique_ptr LRM(new LazyReexportsManager( - std::move(EmitTrampolines), RSMgr, PlatformJD, Err)); + std::move(EmitTrampolines), RSMgr, PlatformJD, L, Err)); if (Err) return std::move(Err); return std::move(LRM); } Error LazyReexportsManager::handleRemoveResources(JITDylib &JD, ResourceKey K) { - JD.getExecutionSession().runSessionLocked([&]() { + return JD.getExecutionSession().runSessionLocked([&]() -> Error { auto I = KeyToReentryAddrs.find(K); - if (I != KeyToReentryAddrs.end()) { - auto &ReentryAddrs = I->second; - for (auto &ReentryAddr : ReentryAddrs) { - assert(CallThroughs.count(ReentryAddr) && "CallTrhough missing"); - CallThroughs.erase(ReentryAddr); - } - KeyToReentryAddrs.erase(I); + if (I == KeyToReentryAddrs.end()) + return Error::success(); + + auto &ReentryAddrs = I->second; + for (auto &ReentryAddr : ReentryAddrs) { + assert(CallThroughs.count(ReentryAddr) && "CallTrhough missing"); + CallThroughs.erase(ReentryAddr); } + KeyToReentryAddrs.erase(I); + return L ? L->onLazyReexportsRemoved(JD, K) : Error::success(); }); - return Error::success(); } void LazyReexportsManager::handleTransferResources(JITDylib &JD, @@ -323,14 +326,17 @@ void LazyReexportsManager::handleTransferResources(JITDylib &JD, DstAddrs.insert(DstAddrs.end(), SrcAddrs.begin(), SrcAddrs.end()); KeyToReentryAddrs.erase(I); } + if (L) + L->onLazyReexportsTransfered(JD, DstK, SrcK); } } LazyReexportsManager::LazyReexportsManager(EmitTrampolinesFn EmitTrampolines, RedirectableSymbolManager &RSMgr, - JITDylib &PlatformJD, Error &Err) + JITDylib &PlatformJD, Listener *L, + Error &Err) : ES(PlatformJD.getExecutionSession()), - EmitTrampolines(std::move(EmitTrampolines)), RSMgr(RSMgr) { + EmitTrampolines(std::move(EmitTrampolines)), RSMgr(RSMgr), L(L) { using namespace shared; @@ -384,17 +390,22 @@ void LazyReexportsManager::emitRedirectableSymbols( Redirs[Name] = (*ReentryPoints)[I++]; I = 0; - if (auto Err = MR->withResourceKeyDo([&](ResourceKey K) { - for (auto &[Name, AI] : Reexports) { - const auto &ReentryPoint = (*ReentryPoints)[I++]; - CallThroughs[ReentryPoint.getAddress()] = {Name, AI.Aliasee, - &MR->getTargetJITDylib()}; - KeyToReentryAddrs[K].push_back(ReentryPoint.getAddress()); - } - })) { - MR->getExecutionSession().reportError(std::move(Err)); - MR->failMaterialization(); - return; + if (!Reexports.empty()) { + if (auto Err = MR->withResourceKeyDo([&](ResourceKey K) { + auto &JD = MR->getTargetJITDylib(); + auto &ReentryAddrsForK = KeyToReentryAddrs[K]; + for (auto &[Name, AI] : Reexports) { + const auto &ReentryPoint = (*ReentryPoints)[I++]; + CallThroughs[ReentryPoint.getAddress()] = {&JD, Name, AI.Aliasee}; + ReentryAddrsForK.push_back(ReentryPoint.getAddress()); + } + if (L) + L->onLazyReexportsCreated(JD, K, Reexports); + })) { + MR->getExecutionSession().reportError(std::move(Err)); + MR->failMaterialization(); + return; + } } RSMgr.emitRedirectableSymbols(std::move(MR), std::move(Redirs)); @@ -415,6 +426,9 @@ void LazyReexportsManager::resolve(ResolveSendResultFn SendResult, LandingInfo = I->second; }); + if (L) + L->onLazyReexportCalled(LandingInfo); + SymbolInstance LandingSym(LandingInfo.JD, std::move(LandingInfo.BodyName)); LandingSym.lookupAsync([this, JD = std::move(LandingInfo.JD), ReentryName = std::move(LandingInfo.Name), @@ -432,5 +446,168 @@ void LazyReexportsManager::resolve(ResolveSendResultFn SendResult, }); } +class SimpleLazyReexportsSpeculator::SpeculateTask : public IdleTask { +public: + SpeculateTask(std::weak_ptr Speculator) + : Speculator(std::move(Speculator)) {} + + void printDescription(raw_ostream &OS) override { + OS << "Speculative Lookup Task"; + } + + void run() override { + if (auto S = Speculator.lock()) + S->doNextSpeculativeLookup(); + } + +private: + std::weak_ptr Speculator; +}; + +SimpleLazyReexportsSpeculator::~SimpleLazyReexportsSpeculator() { + for (auto &[JD, _] : LazyReexports) + JITDylibSP(JD)->Release(); +} + +void SimpleLazyReexportsSpeculator::onLazyReexportsCreated( + JITDylib &JD, ResourceKey K, const SymbolAliasMap &Reexports) { + if (!LazyReexports.count(&JD)) + JD.Retain(); + auto &BodiesVec = LazyReexports[&JD][K]; + for (auto &[Name, AI] : Reexports) + BodiesVec.push_back(AI.Aliasee); + if (!SpeculateTaskActive) { + SpeculateTaskActive = true; + ES.dispatchTask(std::make_unique(WeakThis)); + } +} + +void SimpleLazyReexportsSpeculator::onLazyReexportsTransfered( + JITDylib &JD, ResourceKey DstK, ResourceKey SrcK) { + + auto I = LazyReexports.find(&JD); + if (I == LazyReexports.end()) + return; + + auto &MapForJD = I->second; + auto J = MapForJD.find(SrcK); + if (J == MapForJD.end()) + return; + + // We have something to transfer. + auto K = MapForJD.find(DstK); + if (K == MapForJD.end()) { + auto Tmp = std::move(J->second); + MapForJD.erase(J); + MapForJD[DstK] = std::move(Tmp); + } else { + auto &SrcNames = J->second; + auto &DstNames = K->second; + DstNames.insert(DstNames.end(), SrcNames.begin(), SrcNames.end()); + MapForJD.erase(J); + } +} + +Error SimpleLazyReexportsSpeculator::onLazyReexportsRemoved(JITDylib &JD, + ResourceKey K) { + + auto I = LazyReexports.find(&JD); + if (I == LazyReexports.end()) + return Error::success(); + + auto &MapForJD = I->second; + MapForJD.erase(K); + + if (MapForJD.empty()) { + LazyReexports.erase(I); + JD.Release(); + } + + return Error::success(); +} + +void SimpleLazyReexportsSpeculator::onLazyReexportCalled( + const CallThroughInfo &CTI) { + if (RecordExec) + RecordExec(CTI); +} + +void SimpleLazyReexportsSpeculator::addSpeculationSuggestions( + std::vector> NewSuggestions) { + ES.runSessionLocked([&]() { + for (auto &[JDName, SymbolName] : NewSuggestions) + SpeculateSuggestions.push_back( + {std::move(JDName), std::move(SymbolName)}); + }); +} + +bool SimpleLazyReexportsSpeculator::doNextSpeculativeLookup() { + // Use existing speculation queue if available, otherwise take the next + // element from LazyReexports. + JITDylibSP SpeculateJD = nullptr; + SymbolStringPtr SpeculateFn; + + auto SpeculateAgain = ES.runSessionLocked([&]() { + while (!SpeculateSuggestions.empty()) { + auto [JDName, SymbolName] = std::move(SpeculateSuggestions.front()); + SpeculateSuggestions.pop_front(); + + if (auto *JD = ES.getJITDylibByName(JDName)) { + SpeculateJD = JD; + SpeculateFn = std::move(SymbolName); + break; + } + } + + if (!SpeculateJD) { + assert(!LazyReexports.empty() && "LazyReexports map is empty"); + auto LRItr = + std::next(LazyReexports.begin(), rand() % LazyReexports.size()); + auto &[JD, KeyToFnBodies] = *LRItr; + + assert(!KeyToFnBodies.empty() && "Key to function bodies map empty"); + auto KeyToFnBodiesItr = + std::next(KeyToFnBodies.begin(), rand() % KeyToFnBodies.size()); + auto &[Key, FnBodies] = *KeyToFnBodiesItr; + + assert(!FnBodies.empty() && "Function bodies list empty"); + auto FnBodyItr = std::next(FnBodies.begin(), rand() % FnBodies.size()); + + SpeculateJD = JITDylibSP(JD); + SpeculateFn = std::move(*FnBodyItr); + + FnBodies.erase(FnBodyItr); + if (FnBodies.empty()) { + KeyToFnBodies.erase(KeyToFnBodiesItr); + if (KeyToFnBodies.empty()) { + LRItr->first->Release(); + LazyReexports.erase(LRItr); + } + } + } + + SpeculateTaskActive = + !SpeculateSuggestions.empty() || !LazyReexports.empty(); + return SpeculateTaskActive; + }); + + LLVM_DEBUG({ + dbgs() << "Issuing speculative lookup for ( " << SpeculateJD->getName() + << ", " << SpeculateFn << " )...\n"; + }); + + ES.lookup( + LookupKind::Static, makeJITDylibSearchOrder(SpeculateJD.get()), + {{std::move(SpeculateFn), SymbolLookupFlags::WeaklyReferencedSymbol}}, + SymbolState::Ready, + [](Expected Result) { consumeError(Result.takeError()); }, + NoDependenciesToRegister); + + if (SpeculateAgain) + ES.dispatchTask(std::make_unique(WeakThis)); + + return false; +} + } // End namespace orc. } // End namespace llvm. diff --git a/llvm/lib/ExecutionEngine/Orc/TaskDispatch.cpp b/llvm/lib/ExecutionEngine/Orc/TaskDispatch.cpp index 1af17e85220db..e87a14f3ea7c4 100644 --- a/llvm/lib/ExecutionEngine/Orc/TaskDispatch.cpp +++ b/llvm/lib/ExecutionEngine/Orc/TaskDispatch.cpp @@ -15,9 +15,13 @@ namespace orc { char Task::ID = 0; char GenericNamedTask::ID = 0; +char IdleTask::ID = 0; + const char *GenericNamedTask::DefaultDescription = "Generic Task"; void Task::anchor() {} +void IdleTask::anchor() {} + TaskDispatcher::~TaskDispatcher() = default; void InPlaceTaskDispatcher::dispatch(std::unique_ptr T) { T->run(); } @@ -26,7 +30,15 @@ void InPlaceTaskDispatcher::shutdown() {} #if LLVM_ENABLE_THREADS void DynamicThreadPoolTaskDispatcher::dispatch(std::unique_ptr T) { - bool IsMaterializationTask = isa(*T); + + enum { Normal, Materialization, Idle } TaskKind; + + if (isa(*T)) + TaskKind = Materialization; + else if (isa(*T)) + TaskKind = Idle; + else + TaskKind = Normal; { std::lock_guard Lock(DispatchMutex); @@ -35,24 +47,24 @@ void DynamicThreadPoolTaskDispatcher::dispatch(std::unique_ptr T) { if (Shutdown) return; - if (IsMaterializationTask) { + if (TaskKind == Materialization) { // If this is a materialization task and there are too many running // already then queue this one up and return early. - if (MaxMaterializationThreads && - NumMaterializationThreads == *MaxMaterializationThreads) { - MaterializationTaskQueue.push_back(std::move(T)); - return; - } + if (!canRunMaterializationTaskNow()) + return MaterializationTaskQueue.push_back(std::move(T)); // Otherwise record that we have a materialization task running. ++NumMaterializationThreads; + } else if (TaskKind == Idle) { + if (!canRunIdleTaskNow()) + return IdleTaskQueue.push_back(std::move(T)); } ++Outstanding; } - std::thread([this, T = std::move(T), IsMaterializationTask]() mutable { + std::thread([this, T = std::move(T), TaskKind]() mutable { while (true) { // Run the task. @@ -67,18 +79,24 @@ void DynamicThreadPoolTaskDispatcher::dispatch(std::unique_ptr T) { // Check the work queue state and either proceed with the next task or // end this thread. std::lock_guard Lock(DispatchMutex); - if (!MaterializationTaskQueue.empty()) { + + if (TaskKind == Materialization) + --NumMaterializationThreads; + --Outstanding; + + if (!MaterializationTaskQueue.empty() && canRunMaterializationTaskNow()) { // If there are any materialization tasks running then steal that work. T = std::move(MaterializationTaskQueue.front()); MaterializationTaskQueue.pop_front(); - if (!IsMaterializationTask) { - ++NumMaterializationThreads; - IsMaterializationTask = true; - } + TaskKind = Materialization; + ++NumMaterializationThreads; + ++Outstanding; + } else if (!IdleTaskQueue.empty() && canRunIdleTaskNow()) { + T = std::move(IdleTaskQueue.front()); + IdleTaskQueue.pop_front(); + TaskKind = Idle; + ++Outstanding; } else { - if (IsMaterializationTask) - --NumMaterializationThreads; - --Outstanding; if (Outstanding == 0) OutstandingCV.notify_all(); return; @@ -92,6 +110,17 @@ void DynamicThreadPoolTaskDispatcher::shutdown() { Shutdown = true; OutstandingCV.wait(Lock, [this]() { return Outstanding == 0; }); } + +bool DynamicThreadPoolTaskDispatcher::canRunMaterializationTaskNow() { + return !MaxMaterializationThreads || + (NumMaterializationThreads < *MaxMaterializationThreads); +} + +bool DynamicThreadPoolTaskDispatcher::canRunIdleTaskNow() { + return !MaxMaterializationThreads || + (Outstanding < *MaxMaterializationThreads); +} + #endif } // namespace orc diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp index 963c36322c8ab..ece8484fe8c7c 100644 --- a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp +++ b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp @@ -83,13 +83,35 @@ using namespace llvm::orc; static cl::OptionCategory JITLinkCategory("JITLink Options"); +static cl::list InputFiles(cl::Positional, cl::OneOrMore, + cl::desc("input files"), + cl::cat(JITLinkCategory)); + static cl::list LazyLink("lazy", cl::desc("Link the following file lazily"), cl::cat(JITLinkCategory)); -static cl::list InputFiles(cl::Positional, cl::OneOrMore, - cl::desc("input files"), - cl::cat(JITLinkCategory)); +enum class SpeculateKind { None, Simple }; + +cl::opt Speculate( + "speculate", cl::desc("Choose speculation scheme"), + cl::init(SpeculateKind::None), + cl::values(clEnumValN(SpeculateKind::None, "none", "No speculation"), + clEnumValN(SpeculateKind::Simple, "simple", + "Simple speculation")), + cl::cat(JITLinkCategory)); + +cl::opt SpeculateOrder( + "speculate-order", + cl::desc("A CSV file containing (JITDylib, Function) pairs to" + "speculatively look up"), + cl::cat(JITLinkCategory)); + +cl::opt RecordLazyExecs( + "record-lazy-execs", + cl::desc("Write lazy-function executions to a CSV file as (JITDylib, " + "function) pairs"), + cl::cat(JITLinkCategory)); static cl::opt MaterializationThreads( "num-threads", cl::desc("Number of materialization threads to use"), @@ -959,17 +981,50 @@ class PhonyExternalsGenerator : public DefinitionGenerator { }; Expected> -createLazyLinkingSupport(ObjectLinkingLayer &OLL, JITDylib &PlatformJD) { - auto RSMgr = JITLinkRedirectableSymbolManager::Create(OLL); +createLazyLinkingSupport(Session &S) { + auto RSMgr = JITLinkRedirectableSymbolManager::Create(S.ObjLayer); if (!RSMgr) return RSMgr.takeError(); - auto LRMgr = createJITLinkLazyReexportsManager(OLL, **RSMgr, PlatformJD); + std::shared_ptr Speculator; + switch (Speculate) { + case SpeculateKind::None: + break; + case SpeculateKind::Simple: + SimpleLazyReexportsSpeculator::RecordExecutionFunction RecordExecs; + + if (!RecordLazyExecs.empty()) + RecordExecs = [&S](const LazyReexportsManager::CallThroughInfo &CTI) { + S.LazyFnExecOrder.push_back({CTI.JD->getName(), CTI.BodyName}); + }; + + Speculator = + SimpleLazyReexportsSpeculator::Create(S.ES, std::move(RecordExecs)); + break; + } + + auto LRMgr = createJITLinkLazyReexportsManager( + S.ObjLayer, **RSMgr, *S.PlatformJD, Speculator.get()); if (!LRMgr) return LRMgr.takeError(); - return std::make_unique(std::move(*RSMgr), - std::move(*LRMgr), OLL); + return std::make_unique( + std::move(*RSMgr), std::move(Speculator), std::move(*LRMgr), S.ObjLayer); +} + +static Error writeLazyExecOrder(Session &S) { + if (RecordLazyExecs.empty()) + return Error::success(); + + std::error_code EC; + raw_fd_ostream ExecOrderOut(RecordLazyExecs, EC); + if (EC) + return createFileError(RecordLazyExecs, EC); + + for (auto &[JDName, FunctionName] : S.LazyFnExecOrder) + ExecOrderOut << JDName << ", " << FunctionName << "\n"; + + return Error::success(); } Expected> Session::Create(Triple TT, @@ -1017,8 +1072,7 @@ Expected> Session::Create(Triple TT, S->Features = std::move(Features); if (lazyLinkingRequested()) { - if (auto LazyLinking = - createLazyLinkingSupport(S->ObjLayer, *S->PlatformJD)) + if (auto LazyLinking = createLazyLinkingSupport(*S)) S->LazyLinking = std::move(*LazyLinking); else return LazyLinking.takeError(); @@ -1028,6 +1082,9 @@ Expected> Session::Create(Triple TT, } Session::~Session() { + if (auto Err = writeLazyExecOrder(*this)) + ES.reportError(std::move(Err)); + if (auto Err = ES.endSession()) ES.reportError(std::move(Err)); } @@ -1702,6 +1759,23 @@ static Error sanitizeArguments(const Triple &TT, const char *ArgV0) { return make_error( "Lazy linking cannot be used with -harness mode", inconvertibleErrorCode()); + } else if (Speculate != SpeculateKind::None) { + errs() << "Warning: -speculate ignored as there are no -lazy inputs\n"; + Speculate = SpeculateKind::None; + } + + if (Speculate == SpeculateKind::None) { + if (!SpeculateOrder.empty()) { + errs() << "Warning: -speculate-order ignored because speculation is " + "disabled\n"; + SpeculateOrder = ""; + } + + if (!RecordLazyExecs.empty()) { + errs() << "Warning: -record-lazy-execs ignored because speculation is " + "disabled\n"; + RecordLazyExecs = ""; + } } return Error::success(); @@ -2267,6 +2341,59 @@ static Error addLibraries(Session &S, return Error::success(); } +static Error addSpeculationOrder(Session &S) { + + if (SpeculateOrder.empty()) + return Error::success(); + + assert(S.LazyLinking && "SpeculateOrder set, but lazy linking not enabled"); + assert(S.LazyLinking->Speculator && "SpeculatoOrder set, but no speculator"); + + auto SpecOrderBuffer = getFile(SpeculateOrder); + if (!SpecOrderBuffer) + return SpecOrderBuffer.takeError(); + + StringRef LineStream((*SpecOrderBuffer)->getBuffer()); + std::vector> SpecOrder; + + size_t LineNumber = 0; + while (!LineStream.empty()) { + ++LineNumber; + + auto MakeSpecOrderErr = [&](StringRef Reason) { + return make_error("Error in speculation order file \"" + + SpeculateOrder + "\" on line " + + Twine(LineNumber) + ": " + Reason, + inconvertibleErrorCode()); + }; + + StringRef CurLine; + std::tie(CurLine, LineStream) = LineStream.split('\n'); + CurLine = CurLine.trim(); + if (CurLine.empty()) + continue; + + auto [JDName, FuncName] = CurLine.split(','); + + if (FuncName.empty()) + return MakeSpecOrderErr("missing ',' separator"); + + JDName = JDName.trim(); + if (JDName.empty()) + return MakeSpecOrderErr("no value for column 1 (JIT Dylib name)"); + + FuncName = FuncName.trim(); + if (FuncName.empty()) + return MakeSpecOrderErr("no value for column 2 (function name)"); + + SpecOrder.push_back({JDName.str(), S.ES.intern(FuncName)}); + } + + S.LazyLinking->Speculator->addSpeculationSuggestions(std::move(SpecOrder)); + + return Error::success(); +} + static Error addSessionInputs(Session &S) { std::map IdxToJD; DenseSet LazyLinkIdxs; @@ -2299,6 +2426,9 @@ static Error addSessionInputs(Session &S) { if (auto Err = addLibraries(S, IdxToJD, LazyLinkIdxs)) return Err; + if (auto Err = addSpeculationOrder(S)) + return Err; + return Error::success(); } diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink.h b/llvm/tools/llvm-jitlink/llvm-jitlink.h index be3710971729a..92d667e797b88 100644 --- a/llvm/tools/llvm-jitlink/llvm-jitlink.h +++ b/llvm/tools/llvm-jitlink/llvm-jitlink.h @@ -33,13 +33,17 @@ namespace llvm { struct Session { struct LazyLinkingSupport { - LazyLinkingSupport(std::unique_ptr RSMgr, - std::unique_ptr LRMgr, - orc::ObjectLinkingLayer &ObjLinkingLayer) - : RSMgr(std::move(RSMgr)), LRMgr(std::move(LRMgr)), + LazyLinkingSupport( + std::unique_ptr RSMgr, + std::shared_ptr Speculator, + std::unique_ptr LRMgr, + orc::ObjectLinkingLayer &ObjLinkingLayer) + : RSMgr(std::move(RSMgr)), Speculator(std::move(Speculator)), + LRMgr(std::move(LRMgr)), LazyObjLinkingLayer(ObjLinkingLayer, *this->LRMgr) {} std::unique_ptr RSMgr; + std::shared_ptr Speculator; std::unique_ptr LRMgr; orc::LazyObjectLinkingLayer LazyObjLinkingLayer; }; @@ -52,6 +56,7 @@ struct Session { std::unique_ptr LazyLinking; orc::JITDylibSearchOrder JDSearchOrder; SubtargetFeatures Features; + std::vector> LazyFnExecOrder; ~Session();