From 41379f1ec4657860f4840dba914aa643e7938a5c Mon Sep 17 00:00:00 2001 From: Lang Hames Date: Sun, 10 May 2020 11:34:04 -0700 Subject: [PATCH] [ORC] Share ownership of JITDylibs between ExecutionSession and MaterializationResponsibility. MaterializationResponsibility objects provide a connection between a materialization process (compiler, jit linker, etc.) and the JIT state held in the ExecutionSession and JITDylib objects. Switching to shared ownership extends the lifetime of JITDylibs to ensure they remain accessible until all materializers targeting them have completed. This will allow (in a follow-up patch) JITDylibs to be removed from the ExecutionSession and placed in a pending-destruction state while they are kept alive to communicate errors to/from any still-runnning materialization processes. The intent is to enable JITDylibs to be safely removed even if they have running compiles targeting them. --- llvm/include/llvm/ExecutionEngine/Orc/Core.h | 60 +++++++++-------- llvm/lib/ExecutionEngine/Orc/Core.cpp | 75 +++++++++++++--------- llvm/lib/ExecutionEngine/Orc/LLJIT.cpp | 11 +++- .../unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp | 26 ++++---- 4 files changed, 99 insertions(+), 73 deletions(-) diff --git a/llvm/include/llvm/ExecutionEngine/Orc/Core.h b/llvm/include/llvm/ExecutionEngine/Orc/Core.h index 0424d50..a117acef 100644 --- a/llvm/include/llvm/ExecutionEngine/Orc/Core.h +++ b/llvm/include/llvm/ExecutionEngine/Orc/Core.h @@ -421,7 +421,7 @@ public: /// Returns the target JITDylib that these symbols are being materialized /// into. - JITDylib &getTargetJITDylib() const { return JD; } + JITDylib &getTargetJITDylib() const { return *JD; } /// Returns the VModuleKey for this instance. VModuleKey getVModuleKey() const { return K; } @@ -526,14 +526,16 @@ public: private: /// Create a MaterializationResponsibility for the given JITDylib and /// initial symbols. - MaterializationResponsibility(JITDylib &JD, SymbolFlagsMap SymbolFlags, + MaterializationResponsibility(std::shared_ptr JD, + SymbolFlagsMap SymbolFlags, SymbolStringPtr InitSymbol, VModuleKey K) - : JD(JD), SymbolFlags(std::move(SymbolFlags)), + : JD(std::move(JD)), SymbolFlags(std::move(SymbolFlags)), InitSymbol(std::move(InitSymbol)), K(std::move(K)) { + assert(this->JD && "Cannot initialize with null JD"); assert(!this->SymbolFlags.empty() && "Materializing nothing?"); } - JITDylib &JD; + std::shared_ptr JD; SymbolFlagsMap SymbolFlags; SymbolStringPtr InitSymbol; VModuleKey K; @@ -548,6 +550,9 @@ private: /// is requested via the lookup method. The JITDylib will call discard if a /// stronger definition is added or already present. class MaterializationUnit { + friend class ExecutionSession; + friend class JITDylib; + public: MaterializationUnit(SymbolFlagsMap InitalSymbolFlags, SymbolStringPtr InitSymbol, VModuleKey K) @@ -569,13 +574,10 @@ public: /// Returns the initialization symbol for this MaterializationUnit (if any). const SymbolStringPtr &getInitializerSymbol() const { return InitSymbol; } - /// Called by materialization dispatchers (see - /// ExecutionSession::DispatchMaterializationFunction) to trigger - /// materialization of this MaterializationUnit. - void doMaterialize(JITDylib &JD) { - materialize(MaterializationResponsibility( - JD, std::move(SymbolFlags), std::move(InitSymbol), std::move(K))); - } + /// Implementations of this method should materialize all symbols + /// in the materialzation unit, except for those that have been + /// previously discarded. + virtual void materialize(MaterializationResponsibility R) = 0; /// Called by JITDylibs to notify MaterializationUnits that the given symbol /// has been overridden. @@ -592,10 +594,11 @@ protected: private: virtual void anchor(); - /// Implementations of this method should materialize all symbols - /// in the materialzation unit, except for those that have been - /// previously discarded. - virtual void materialize(MaterializationResponsibility R) = 0; + MaterializationResponsibility + createMaterializationResponsibility(std::shared_ptr JD) { + return MaterializationResponsibility(std::move(JD), std::move(SymbolFlags), + std::move(InitSymbol), K); + } /// Implementations of this method should discard the given symbol /// from the source (e.g. if the source is an LLVM IR Module and the @@ -773,7 +776,7 @@ private: /// their addresses may be used as keys for resource management. /// JITDylib state changes must be made via an ExecutionSession to guarantee /// that they are synchronized with respect to other JITDylib operations. -class JITDylib { +class JITDylib : public std::enable_shared_from_this { friend class AsynchronousSymbolQuery; friend class ExecutionSession; friend class Platform; @@ -1044,6 +1047,7 @@ private: ExecutionSession &ES; std::string JITDylibName; + bool Open = true; SymbolTable Symbols; UnmaterializedInfosMap UnmaterializedInfos; MaterializingInfosMap MaterializingInfos; @@ -1090,8 +1094,9 @@ public: using ErrorReporter = std::function; /// For dispatching MaterializationUnit::materialize calls. - using DispatchMaterializationFunction = std::function MU)>; + using DispatchMaterializationFunction = + std::function MU, + MaterializationResponsibility MR)>; /// Construct an ExecutionSession. /// @@ -1243,11 +1248,11 @@ public: SymbolState RequiredState = SymbolState::Ready); /// Materialize the given unit. - void dispatchMaterialization(JITDylib &JD, - std::unique_ptr MU) { + void dispatchMaterialization(std::unique_ptr MU, + MaterializationResponsibility MR) { assert(MU && "MU must be non-null"); - DEBUG_WITH_TYPE("orc", dumpDispatchInfo(JD, *MU)); - DispatchMaterialization(JD, std::move(MU)); + DEBUG_WITH_TYPE("orc", dumpDispatchInfo(MR.getTargetJITDylib(), *MU)); + DispatchMaterialization(std::move(MU), std::move(MR)); } /// Dump the state of all the JITDylibs in this session. @@ -1259,9 +1264,9 @@ private: } static void - materializeOnCurrentThread(JITDylib &JD, - std::unique_ptr MU) { - MU->doMaterialize(JD); + materializeOnCurrentThread(std::unique_ptr MU, + MaterializationResponsibility MR) { + MU->materialize(std::move(MR)); } void runOutstandingMUs(); @@ -1278,12 +1283,13 @@ private: DispatchMaterializationFunction DispatchMaterialization = materializeOnCurrentThread; - std::vector> JDs; + std::vector> JDs; // FIXME: Remove this (and runOutstandingMUs) once the linking layer works // with callbacks from asynchronous queries. mutable std::recursive_mutex OutstandingMUsMutex; - std::vector>> + std::vector, + MaterializationResponsibility>> OutstandingMUs; }; diff --git a/llvm/lib/ExecutionEngine/Orc/Core.cpp b/llvm/lib/ExecutionEngine/Orc/Core.cpp index 7aaf621..bad13cf 100644 --- a/llvm/lib/ExecutionEngine/Orc/Core.cpp +++ b/llvm/lib/ExecutionEngine/Orc/Core.cpp @@ -187,12 +187,12 @@ MaterializationResponsibility::~MaterializationResponsibility() { } SymbolNameSet MaterializationResponsibility::getRequestedSymbols() const { - return JD.getRequestedSymbols(SymbolFlags); + return JD->getRequestedSymbols(SymbolFlags); } Error MaterializationResponsibility::notifyResolved(const SymbolMap &Symbols) { LLVM_DEBUG({ - dbgs() << "In " << JD.getName() << " resolving " << Symbols << "\n"; + dbgs() << "In " << JD->getName() << " resolving " << Symbols << "\n"; }); #ifndef NDEBUG for (auto &KV : Symbols) { @@ -207,16 +207,16 @@ Error MaterializationResponsibility::notifyResolved(const SymbolMap &Symbols) { } #endif - return JD.resolve(Symbols); + return JD->resolve(Symbols); } Error MaterializationResponsibility::notifyEmitted() { LLVM_DEBUG({ - dbgs() << "In " << JD.getName() << " emitting " << SymbolFlags << "\n"; + dbgs() << "In " << JD->getName() << " emitting " << SymbolFlags << "\n"; }); - if (auto Err = JD.emit(SymbolFlags)) + if (auto Err = JD->emit(SymbolFlags)) return Err; SymbolFlags.clear(); @@ -227,10 +227,10 @@ Error MaterializationResponsibility::defineMaterializing( SymbolFlagsMap NewSymbolFlags) { LLVM_DEBUG({ - dbgs() << "In " << JD.getName() << " defining materializing symbols " - << NewSymbolFlags << "\n"; - }); - if (auto AcceptedDefs = JD.defineMaterializing(std::move(NewSymbolFlags))) { + dbgs() << "In " << JD->getName() << " defining materializing symbols " + << NewSymbolFlags << "\n"; + }); + if (auto AcceptedDefs = JD->defineMaterializing(std::move(NewSymbolFlags))) { // Add all newly accepted symbols to this responsibility object. for (auto &KV : *AcceptedDefs) SymbolFlags.insert(KV); @@ -242,17 +242,17 @@ Error MaterializationResponsibility::defineMaterializing( void MaterializationResponsibility::failMaterialization() { LLVM_DEBUG({ - dbgs() << "In " << JD.getName() << " failing materialization for " + dbgs() << "In " << JD->getName() << " failing materialization for " << SymbolFlags << "\n"; }); JITDylib::FailedSymbolsWorklist Worklist; for (auto &KV : SymbolFlags) - Worklist.push_back(std::make_pair(&JD, KV.first)); + Worklist.push_back(std::make_pair(JD.get(), KV.first)); SymbolFlags.clear(); - JD.notifyFailed(std::move(Worklist)); + JD->notifyFailed(std::move(Worklist)); } void MaterializationResponsibility::replace( @@ -271,12 +271,12 @@ void MaterializationResponsibility::replace( if (MU->getInitializerSymbol() == InitSymbol) InitSymbol = nullptr; - LLVM_DEBUG(JD.getExecutionSession().runSessionLocked([&]() { - dbgs() << "In " << JD.getName() << " replacing symbols with " << *MU + LLVM_DEBUG(JD->getExecutionSession().runSessionLocked([&]() { + dbgs() << "In " << JD->getName() << " replacing symbols with " << *MU << "\n"; });); - JD.replace(std::move(MU)); + JD->replace(std::move(MU)); } MaterializationResponsibility @@ -315,7 +315,7 @@ void MaterializationResponsibility::addDependencies( }); assert(SymbolFlags.count(Name) && "Symbol not covered by this MaterializationResponsibility instance"); - JD.addDependencies(Name, Dependencies); + JD->addDependencies(Name, Dependencies); } void MaterializationResponsibility::addDependenciesForAll( @@ -325,7 +325,7 @@ void MaterializationResponsibility::addDependenciesForAll( << Dependencies << "\n"; }); for (auto &KV : SymbolFlags) - JD.addDependencies(KV.first, Dependencies); + JD->addDependencies(KV.first, Dependencies); } AbsoluteSymbolsMaterializationUnit::AbsoluteSymbolsMaterializationUnit( @@ -703,8 +703,11 @@ void JITDylib::replace(std::unique_ptr MU) { return nullptr; }); - if (MustRunMU) - ES.dispatchMaterialization(*this, std::move(MustRunMU)); + if (MustRunMU) { + auto MR = + MustRunMU->createMaterializationResponsibility(shared_from_this()); + ES.dispatchMaterialization(std::move(MustRunMU), std::move(MR)); + } } SymbolNameSet @@ -1448,8 +1451,11 @@ JITDylib::legacyLookup(std::shared_ptr Q, // Add MUs to the OutstandingMUs list. { std::lock_guard Lock(ES.OutstandingMUsMutex); - for (auto &MU : MUs) - ES.OutstandingMUs.push_back(make_pair(this, std::move(MU))); + auto ThisJD = shared_from_this(); + for (auto &MU : MUs) { + auto MR = MU->createMaterializationResponsibility(ThisJD); + ES.OutstandingMUs.push_back(make_pair(std::move(MU), std::move(MR))); + } } ES.runOutstandingMUs(); @@ -1783,7 +1789,7 @@ JITDylib &ExecutionSession::createBareJITDylib(std::string Name) { assert(!getJITDylibByName(Name) && "JITDylib with that name already exists"); return runSessionLocked([&, this]() -> JITDylib & { JDs.push_back( - std::unique_ptr(new JITDylib(*this, std::move(Name)))); + std::shared_ptr(new JITDylib(*this, std::move(Name)))); return *JDs.back(); }); } @@ -1972,9 +1978,13 @@ void ExecutionSession::lookup( { std::lock_guard Lock(OutstandingMUsMutex); - for (auto &KV : CollectedMUsMap) - for (auto &MU : KV.second) - OutstandingMUs.push_back(std::make_pair(KV.first, std::move(MU))); + for (auto &KV : CollectedMUsMap) { + auto JD = KV.first->shared_from_this(); + for (auto &MU : KV.second) { + auto MR = MU->createMaterializationResponsibility(JD); + OutstandingMUs.push_back(std::make_pair(std::move(MU), std::move(MR))); + } + } } runOutstandingMUs(); @@ -2069,22 +2079,23 @@ void ExecutionSession::dump(raw_ostream &OS) { void ExecutionSession::runOutstandingMUs() { while (1) { - std::pair> JITDylibAndMU; + Optional, + MaterializationResponsibility>> + JMU; { std::lock_guard Lock(OutstandingMUsMutex); if (!OutstandingMUs.empty()) { - JITDylibAndMU = std::move(OutstandingMUs.back()); + JMU.emplace(std::move(OutstandingMUs.back())); OutstandingMUs.pop_back(); } } - if (JITDylibAndMU.first) { - assert(JITDylibAndMU.second && "JITDylib, but no MU?"); - dispatchMaterialization(*JITDylibAndMU.first, - std::move(JITDylibAndMU.second)); - } else + if (!JMU) break; + + assert(JMU->first && "No MU?"); + dispatchMaterialization(std::move(JMU->first), std::move(JMU->second)); } } diff --git a/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp b/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp index c8d7b4d..79e5027 100644 --- a/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp +++ b/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp @@ -1074,10 +1074,15 @@ LLJIT::LLJIT(LLJITBuilderState &S, Error &Err) CompileThreads = std::make_unique(hardware_concurrency(S.NumCompileThreads)); ES->setDispatchMaterialization( - [this](JITDylib &JD, std::unique_ptr MU) { - // FIXME: Switch to move capture once we have c++14. + [this](std::unique_ptr MU, + MaterializationResponsibility MR) { + // FIXME: Switch to move capture once ThreadPool uses unique_function. auto SharedMU = std::shared_ptr(std::move(MU)); - auto Work = [SharedMU, &JD]() { SharedMU->doMaterialize(JD); }; + auto SharedMR = + std::make_shared(std::move(MR)); + auto Work = [SharedMU, SharedMR]() mutable { + SharedMU->materialize(std::move(*SharedMR)); + }; CompileThreads->async(std::move(Work)); }); } diff --git a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp index a8f85d0..b4e8a83 100644 --- a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp +++ b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp @@ -1008,12 +1008,12 @@ TEST_F(CoreAPIsStandardTest, TestBasicWeakSymbolMaterialization) { TEST_F(CoreAPIsStandardTest, DefineMaterializingSymbol) { bool ExpectNoMoreMaterialization = false; - ES.setDispatchMaterialization( - [&](JITDylib &JD, std::unique_ptr MU) { - if (ExpectNoMoreMaterialization) - ADD_FAILURE() << "Unexpected materialization"; - MU->doMaterialize(JD); - }); + ES.setDispatchMaterialization([&](std::unique_ptr MU, + MaterializationResponsibility MR) { + if (ExpectNoMoreMaterialization) + ADD_FAILURE() << "Unexpected materialization"; + MU->materialize(std::move(MR)); + }); auto MU = std::make_unique( SymbolFlagsMap({{Foo, FooSym.getFlags()}}), @@ -1186,11 +1186,15 @@ TEST_F(CoreAPIsStandardTest, TestLookupWithThreadedMaterialization) { #if LLVM_ENABLE_THREADS std::thread MaterializationThread; - ES.setDispatchMaterialization( - [&](JITDylib &JD, std::unique_ptr MU) { - MaterializationThread = - std::thread([MU = std::move(MU), &JD] { MU->doMaterialize(JD); }); - }); + ES.setDispatchMaterialization([&](std::unique_ptr MU, + MaterializationResponsibility MR) { + auto SharedMR = + std::make_shared(std::move(MR)); + MaterializationThread = + std::thread([MU = std::move(MU), MR = std::move(SharedMR)] { + MU->materialize(std::move(*MR)); + }); + }); cantFail(JD.define(absoluteSymbols({{Foo, FooSym}}))); -- 2.7.4