[llvm] 41379f1 - [ORC] Share ownership of JITDylibs between ExecutionSession and

Lang Hames via llvm-commits llvm-commits at lists.llvm.org
Sun May 10 16:42:49 PDT 2020


Author: Lang Hames
Date: 2020-05-10T16:37:17-07:00
New Revision: 41379f1ec4657860f4840dba914aa643e7938a5c

URL: https://github.com/llvm/llvm-project/commit/41379f1ec4657860f4840dba914aa643e7938a5c
DIFF: https://github.com/llvm/llvm-project/commit/41379f1ec4657860f4840dba914aa643e7938a5c.diff

LOG: [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.

Added: 
    

Modified: 
    llvm/include/llvm/ExecutionEngine/Orc/Core.h
    llvm/lib/ExecutionEngine/Orc/Core.cpp
    llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
    llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ExecutionEngine/Orc/Core.h b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
index 0424d5043cca..a117acefd2d3 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/Core.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
@@ -421,7 +421,7 @@ class MaterializationResponsibility {
 
   /// 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 @@ class MaterializationResponsibility {
 private:
   /// Create a MaterializationResponsibility for the given JITDylib and
   ///        initial symbols.
-  MaterializationResponsibility(JITDylib &JD, SymbolFlagsMap SymbolFlags,
+  MaterializationResponsibility(std::shared_ptr<JITDylib> 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<JITDylib> JD;
   SymbolFlagsMap SymbolFlags;
   SymbolStringPtr InitSymbol;
   VModuleKey K;
@@ -548,6 +550,9 @@ class MaterializationResponsibility {
 /// 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 @@ class MaterializationUnit {
   /// 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 @@ class MaterializationUnit {
 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<JITDylib> 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 @@ class AsynchronousSymbolQuery {
 /// 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<JITDylib> {
   friend class AsynchronousSymbolQuery;
   friend class ExecutionSession;
   friend class Platform;
@@ -1044,6 +1047,7 @@ class JITDylib {
 
   ExecutionSession &ES;
   std::string JITDylibName;
+  bool Open = true;
   SymbolTable Symbols;
   UnmaterializedInfosMap UnmaterializedInfos;
   MaterializingInfosMap MaterializingInfos;
@@ -1090,8 +1094,9 @@ class ExecutionSession {
   using ErrorReporter = std::function<void(Error)>;
 
   /// For dispatching MaterializationUnit::materialize calls.
-  using DispatchMaterializationFunction = std::function<void(
-      JITDylib &JD, std::unique_ptr<MaterializationUnit> MU)>;
+  using DispatchMaterializationFunction =
+      std::function<void(std::unique_ptr<MaterializationUnit> MU,
+                         MaterializationResponsibility MR)>;
 
   /// Construct an ExecutionSession.
   ///
@@ -1243,11 +1248,11 @@ class ExecutionSession {
          SymbolState RequiredState = SymbolState::Ready);
 
   /// Materialize the given unit.
-  void dispatchMaterialization(JITDylib &JD,
-                               std::unique_ptr<MaterializationUnit> MU) {
+  void dispatchMaterialization(std::unique_ptr<MaterializationUnit> 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 @@ class ExecutionSession {
   }
 
   static void
-  materializeOnCurrentThread(JITDylib &JD,
-                             std::unique_ptr<MaterializationUnit> MU) {
-    MU->doMaterialize(JD);
+  materializeOnCurrentThread(std::unique_ptr<MaterializationUnit> MU,
+                             MaterializationResponsibility MR) {
+    MU->materialize(std::move(MR));
   }
 
   void runOutstandingMUs();
@@ -1278,12 +1283,13 @@ class ExecutionSession {
   DispatchMaterializationFunction DispatchMaterialization =
       materializeOnCurrentThread;
 
-  std::vector<std::unique_ptr<JITDylib>> JDs;
+  std::vector<std::shared_ptr<JITDylib>> JDs;
 
   // FIXME: Remove this (and runOutstandingMUs) once the linking layer works
   //        with callbacks from asynchronous queries.
   mutable std::recursive_mutex OutstandingMUsMutex;
-  std::vector<std::pair<JITDylib *, std::unique_ptr<MaterializationUnit>>>
+  std::vector<std::pair<std::unique_ptr<MaterializationUnit>,
+                        MaterializationResponsibility>>
       OutstandingMUs;
 };
 

diff  --git a/llvm/lib/ExecutionEngine/Orc/Core.cpp b/llvm/lib/ExecutionEngine/Orc/Core.cpp
index 7aaf62141c89..bad13cfebbc6 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<MaterializationUnit> 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<AsynchronousSymbolQuery> Q,
   // Add MUs to the OutstandingMUs list.
   {
     std::lock_guard<std::recursive_mutex> 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<JITDylib>(new JITDylib(*this, std::move(Name))));
+        std::shared_ptr<JITDylib>(new JITDylib(*this, std::move(Name))));
     return *JDs.back();
   });
 }
@@ -1972,9 +1978,13 @@ void ExecutionSession::lookup(
   {
     std::lock_guard<std::recursive_mutex> 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<JITDylib *, std::unique_ptr<MaterializationUnit>> JITDylibAndMU;
+    Optional<std::pair<std::unique_ptr<MaterializationUnit>,
+                       MaterializationResponsibility>>
+        JMU;
 
     {
       std::lock_guard<std::recursive_mutex> 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 c8d7b4d2db04..79e502775f79 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<ThreadPool>(hardware_concurrency(S.NumCompileThreads));
     ES->setDispatchMaterialization(
-        [this](JITDylib &JD, std::unique_ptr<MaterializationUnit> MU) {
-          // FIXME: Switch to move capture once we have c++14.
+        [this](std::unique_ptr<MaterializationUnit> MU,
+               MaterializationResponsibility MR) {
+          // FIXME: Switch to move capture once ThreadPool uses unique_function.
           auto SharedMU = std::shared_ptr<MaterializationUnit>(std::move(MU));
-          auto Work = [SharedMU, &JD]() { SharedMU->doMaterialize(JD); };
+          auto SharedMR =
+              std::make_shared<MaterializationResponsibility>(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 a8f85d061bf7..b4e8a8302d3b 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<MaterializationUnit> MU) {
-        if (ExpectNoMoreMaterialization)
-          ADD_FAILURE() << "Unexpected materialization";
-        MU->doMaterialize(JD);
-      });
+  ES.setDispatchMaterialization([&](std::unique_ptr<MaterializationUnit> MU,
+                                    MaterializationResponsibility MR) {
+    if (ExpectNoMoreMaterialization)
+      ADD_FAILURE() << "Unexpected materialization";
+    MU->materialize(std::move(MR));
+  });
 
   auto MU = std::make_unique<SimpleMaterializationUnit>(
       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<MaterializationUnit> MU) {
-        MaterializationThread =
-            std::thread([MU = std::move(MU), &JD] { MU->doMaterialize(JD); });
-      });
+  ES.setDispatchMaterialization([&](std::unique_ptr<MaterializationUnit> MU,
+                                    MaterializationResponsibility MR) {
+    auto SharedMR =
+        std::make_shared<MaterializationResponsibility>(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}})));
 


        


More information about the llvm-commits mailing list