[compiler-rt] [llvm] [ORC][Runtime] Enhancing ELF Platform with Push-Request Model for Initializers (PR #102846)

via llvm-commits llvm-commits at lists.llvm.org
Sun Aug 11 20:59:23 PDT 2024


https://github.com/SahilPatidar created https://github.com/llvm/llvm-project/pull/102846

With the help of @lhames, this PR introduces a push-request model to manage list initializers in the runtime state for each `JITDylib`, similar to how `MachO` and `COFF` handle it. Previously, the ELF runtime lacked the ability to register, deregister, or retain initializers, causing issues when re-running `dlopen`. Initializers were erased after `rt_getInitializers` was used.

To address this, we introduce `__orc_rt_elfnix_register_init_sections` and `__orc_rt_elfnix_register_jitdylib` (and corresponding deregister functions). This model allows the runtime to request the push of initializers for the `JITDylib`, with the platform handling this via a record method. Additionally, we leverage `MachO`'s `RecordSectionsTracker` to store initializers, ensuring only newly updated initializers are executed using `processNewSections`, improving the runtime's efficiency and reliability.

>From 51ce29878dc2f414b4eaadf9ecac5ba85fd6a0fd Mon Sep 17 00:00:00 2001
From: SahilPatidar <patidarsahil2001 at gmail.com>
Date: Fri, 9 Aug 2024 11:43:48 +0530
Subject: [PATCH 1/3] [ORC][Runtime][ELF] Improve ELF runtime with experimental
 enhancements

---
 compiler-rt/lib/orc/common.h                  |  97 ++++
 compiler-rt/lib/orc/elfnix_platform.cpp       | 432 +++++++++++++-----
 compiler-rt/lib/orc/elfnix_platform.h         |  56 +--
 .../llvm/ExecutionEngine/Orc/ELFNixPlatform.h | 126 ++---
 .../ExecutionEngine/Orc/ELFNixPlatform.cpp    | 431 +++++++++--------
 5 files changed, 682 insertions(+), 460 deletions(-)

diff --git a/compiler-rt/lib/orc/common.h b/compiler-rt/lib/orc/common.h
index 73c5c4a2bd8d47..a03ae0bddb893d 100644
--- a/compiler-rt/lib/orc/common.h
+++ b/compiler-rt/lib/orc/common.h
@@ -17,6 +17,17 @@
 #include "orc_rt/c_api.h"
 #include <type_traits>
 
+#include "bitmask_enum.h"
+#include "debug.h"
+#include "error.h"
+#include "executor_address.h"
+#include <algorithm>
+#include <ios>
+#include <map>
+#include <mutex>
+#include <sstream>
+#include <vector>
+
 /// This macro should be used to define tags that will be associated with
 /// handlers in the JIT process, and call can be used to define tags f
 #define ORC_RT_JIT_DISPATCH_TAG(X) \
@@ -45,4 +56,90 @@ ORC_RT_IMPORT orc_rt_CWrapperFunctionResult
 __orc_rt_jit_dispatch(__orc_rt_Opaque *DispatchCtx, const void *FnTag,
                       const char *Data, size_t Size) ORC_RT_WEAK_IMPORT;
 
+/// Used to manage sections of fixed-sized metadata records (e.g. pointer
+/// sections, selector refs, etc.)
+template <typename RecordElement> class RecordSectionsTracker {
+public:
+/// Add a section to the "new" list.
+void add(__orc_rt::span<RecordElement> Sec) { New.push_back(std::move(Sec)); }
+
+/// Returns true if there are new sections to process.
+bool hasNewSections() const { return !New.empty(); }
+
+/// Returns the number of new sections to process.
+size_t numNewSections() const { return New.size(); }
+
+/// Process all new sections.
+template <typename ProcessSectionFunc>
+std::enable_if_t<std::is_void_v<
+    std::invoke_result_t<ProcessSectionFunc, __orc_rt::span<RecordElement>>>>
+processNewSections(ProcessSectionFunc &&ProcessSection) {
+  for (auto &Sec : New)
+    ProcessSection(Sec);
+  moveNewToProcessed();
+}
+
+/// Proces all new sections with a fallible handler.
+///
+/// Successfully handled sections will be moved to the Processed
+/// list.
+template <typename ProcessSectionFunc>
+std::enable_if_t<
+    std::is_same_v<__orc_rt::Error, std::invoke_result_t<ProcessSectionFunc,
+                                                __orc_rt::span<RecordElement>>>,
+    __orc_rt::Error>
+processNewSections(ProcessSectionFunc &&ProcessSection) {
+  for (size_t I = 0; I != New.size(); ++I) {
+    if (auto Err = ProcessSection(New[I])) {
+      for (size_t J = 0; J != I; ++J)
+        Processed.push_back(New[J]);
+      New.erase(New.begin(), New.begin() + I);
+      return Err;
+    }
+  }
+  moveNewToProcessed();
+  return __orc_rt::Error::success();
+}
+
+/// Move all sections back to New for reprocessing.
+void reset() {
+  moveNewToProcessed();
+  New = std::move(Processed);
+}
+
+/// Remove the section with the given range.
+bool removeIfPresent(__orc_rt::ExecutorAddrRange R) {
+  if (removeIfPresent(New, R))
+    return true;
+  return removeIfPresent(Processed, R);
+}
+
+private:
+void moveNewToProcessed() {
+  if (Processed.empty())
+    Processed = std::move(New);
+  else {
+    Processed.reserve(Processed.size() + New.size());
+    std::copy(New.begin(), New.end(), std::back_inserter(Processed));
+    New.clear();
+  }
+}
+
+bool removeIfPresent(std::vector<__orc_rt::span<RecordElement>> &V,
+                     __orc_rt::ExecutorAddrRange R) {
+  auto RI = std::find_if(
+    V.rbegin(), V.rend(),
+    [RS = R.toSpan<RecordElement>()](const __orc_rt::span<RecordElement> &E) {
+        return E.data() == RS.data();
+    });
+  if (RI != V.rend()) {
+    V.erase(std::next(RI).base());
+    return true;
+  }
+  return false;
+}
+
+std::vector<__orc_rt::span<RecordElement>> Processed;
+std::vector<__orc_rt::span<RecordElement>> New;
+};
 #endif // ORC_RT_COMMON_H
diff --git a/compiler-rt/lib/orc/elfnix_platform.cpp b/compiler-rt/lib/orc/elfnix_platform.cpp
index c087e71038f950..29582ac6867871 100644
--- a/compiler-rt/lib/orc/elfnix_platform.cpp
+++ b/compiler-rt/lib/orc/elfnix_platform.cpp
@@ -29,6 +29,7 @@ using namespace __orc_rt::elfnix;
 
 // Declare function tags for functions in the JIT process.
 ORC_RT_JIT_DISPATCH_TAG(__orc_rt_elfnix_get_initializers_tag)
+ORC_RT_JIT_DISPATCH_TAG(__orc_rt_elfnix_push_initializers_tag)
 ORC_RT_JIT_DISPATCH_TAG(__orc_rt_elfnix_get_deinitializers_tag)
 ORC_RT_JIT_DISPATCH_TAG(__orc_rt_elfnix_symbol_lookup_tag)
 
@@ -56,21 +57,6 @@ Error validatePointerSectionExtent(const char *SectionName,
   return Error::success();
 }
 
-Error runInitArray(const std::vector<ExecutorAddrRange> &InitArraySections,
-                   const ELFNixJITDylibInitializers &MOJDIs) {
-
-  for (const auto &ModInits : InitArraySections) {
-    if (auto Err = validatePointerSectionExtent(".init_array", ModInits))
-      return Err;
-
-    using InitFunc = void (*)();
-    for (auto *Init : ModInits.toSpan<InitFunc>())
-      (*Init)();
-  }
-
-  return Error::success();
-}
-
 struct TLSInfoEntry {
   unsigned long Key = 0;
   unsigned long DataAddress = 0;
@@ -91,10 +77,18 @@ class ELFNixPlatformRuntimeState {
   using AtExitsVector = std::vector<AtExitEntry>;
 
   struct PerJITDylibState {
+    std::string Name;
     void *Header = nullptr;
     size_t RefCount = 0;
+    size_t LinkedAgainstRefCount = 0;
     bool AllowReinitialization = false;
     AtExitsVector AtExits;
+    std::vector<PerJITDylibState *> Deps;
+    RecordSectionsTracker<void (*)()> RecordedInits;
+
+    bool referenced() const {
+      return LinkedAgainstRefCount != 0 || RefCount != 0;
+    }
   };
 
 public:
@@ -112,6 +106,12 @@ class ELFNixPlatformRuntimeState {
   ELFNixPlatformRuntimeState &operator=(ELFNixPlatformRuntimeState &&) = delete;
 
   Error registerObjectSections(ELFNixPerObjectSectionsToRegister POSR);
+  Error registerJITDylib(std::string &Name, void *Handle);
+  Error deregisterJITDylib(void *Handle);
+  Error registerInits(
+    ExecutorAddr HeaderAddr, std::vector<ExecutorAddrRange> Inits);
+  Error deregisterInits(
+    ExecutorAddr HeaderAddr, std::vector<ExecutorAddrRange> Inits);
   Error deregisterObjectSections(ELFNixPerObjectSectionsToRegister POSR);
 
   const char *dlerror();
@@ -121,6 +121,8 @@ class ELFNixPlatformRuntimeState {
 
   int registerAtExit(void (*F)(void *), void *Arg, void *DSOHandle);
   void runAtExits(void *DSOHandle);
+  void runAtExits(std::unique_lock<std::recursive_mutex> &JDStateLock,
+                  PerJITDylibState &JDS);
 
   /// Returns the base address of the section containing ThreadData.
   Expected<std::pair<const char *, size_t>>
@@ -131,18 +133,23 @@ class ELFNixPlatformRuntimeState {
 private:
   PerJITDylibState *getJITDylibStateByHeaderAddr(void *DSOHandle);
   PerJITDylibState *getJITDylibStateByName(std::string_view Path);
-  PerJITDylibState &
-  getOrCreateJITDylibState(ELFNixJITDylibInitializers &MOJDIs);
 
   Error registerThreadDataSection(span<const char> ThreadDataSection);
 
   Expected<ExecutorAddr> lookupSymbolInJITDylib(void *DSOHandle,
                                                 std::string_view Symbol);
 
-  Expected<ELFNixJITDylibInitializerSequence>
-  getJITDylibInitializersByName(std::string_view Path);
-  Expected<void *> dlopenInitialize(std::string_view Path, int Mode);
-  Error initializeJITDylib(ELFNixJITDylibInitializers &MOJDIs);
+  Error runInits(
+    std::unique_lock<std::recursive_mutex> &JDStatesLock, PerJITDylibState &JDS);
+  Expected<void *> dlopenImpl(std::string_view Path, int Mode);
+  Error dlopenFull(std::unique_lock<std::recursive_mutex> &JDStatesLock,
+                   PerJITDylibState &JDS);
+  Error dlopenInitialize(
+    std::unique_lock<std::recursive_mutex> &JDStatesLock,
+    PerJITDylibState &JDS, ELFNixJITDylibDepInfoMap &DepInfo);
+  Error dlcloseImpl(void *DSOHandle);
+  Error dlcloseInitialize(
+    std::unique_lock<std::recursive_mutex> &JDStatesLock, PerJITDylibState &JDS);
 
   static ELFNixPlatformRuntimeState *MOPS;
 
@@ -206,6 +213,90 @@ Error ELFNixPlatformRuntimeState::registerObjectSections(
   return Error::success();
 }
 
+Error ELFNixPlatformRuntimeState::registerJITDylib(std::string &Name, void *Handle) {
+  std::lock_guard<std::recursive_mutex> Lock(JDStatesMutex);
+
+  if (JDStates.count(Handle)) {
+    std::ostringstream ErrStream;
+    ErrStream << "Duplicate JITDylib registration for header " << Handle
+              << " (name = " << Name << ")";
+    return make_error<StringError>(ErrStream.str());
+  }
+
+  if (JDNameToHeader.count(Name)) {
+    std::ostringstream ErrStream;
+    ErrStream << "Duplicate JITDylib registration for header " << Handle
+              << " (header = " << Handle << ")";
+    return make_error<StringError>(ErrStream.str());
+  }
+
+  auto &JD = JDStates[Handle];
+  JD.Header = Handle;
+  JD.Name = std::move(Name);
+  JDNameToHeader[JD.Name] = Handle;
+  return Error::success();
+}
+
+Error ELFNixPlatformRuntimeState::deregisterJITDylib(void *Handle) {
+  std::lock_guard<std::recursive_mutex> Lock(JDStatesMutex);
+
+  auto I = JDStates.find(Handle);
+  if (I == JDStates.end()) {
+    std::ostringstream ErrStream;
+    ErrStream << "Attempted to deregister unrecognized header " << Handle;
+    return make_error<StringError>(ErrStream.str());
+  }
+
+  auto J = JDNameToHeader.find(
+      std::string(I->second.Name.data(), I->second.Name.size()));
+  assert(J != JDNameToHeader.end() &&
+        "Missing JDNameToHeader entry for JITDylib");
+  JDNameToHeader.erase(J);
+  JDStates.erase(I);
+  return Error::success();
+}
+
+Error ELFNixPlatformRuntimeState::registerInits(
+  ExecutorAddr HeaderAddr, std::vector<ExecutorAddrRange> Inits) {
+  std::lock_guard<std::recursive_mutex> Lock(JDStatesMutex);
+  PerJITDylibState *JDS = getJITDylibStateByHeaderAddr(HeaderAddr.toPtr<void *>());
+
+  if (!JDS) {
+    std::ostringstream ErrStream;
+    ErrStream << "Could not register object platform sections for "
+                 "unrecognized header "
+              << HeaderAddr.toPtr<void *>();
+    return make_error<StringError>(ErrStream.str());
+  }
+
+  for (auto &I : Inits) {
+    JDS->RecordedInits.add(I.toSpan<void (*)()>());
+  }
+
+  return Error::success();
+}
+
+Error ELFNixPlatformRuntimeState::deregisterInits(
+    ExecutorAddr HeaderAddr, std::vector<ExecutorAddrRange> Inits) {
+
+  std::lock_guard<std::recursive_mutex> Lock(JDStatesMutex);
+  PerJITDylibState *JDS = getJITDylibStateByHeaderAddr(HeaderAddr.toPtr<void *>());
+
+  if (!JDS) {
+    std::ostringstream ErrStream;
+    ErrStream << "Could not register object platform sections for unrecognized "
+                 "header "
+              << HeaderAddr.toPtr<void *>();
+    return make_error<StringError>(ErrStream.str());
+  }
+
+  for (auto &I : Inits) {
+    JDS->RecordedInits.removeIfPresent(I);
+  }
+
+  return Error::success();
+}
+
 Error ELFNixPlatformRuntimeState::deregisterObjectSections(
     ELFNixPerObjectSectionsToRegister POSR) {
   if (POSR.EHFrameSection.Start)
@@ -217,28 +308,26 @@ Error ELFNixPlatformRuntimeState::deregisterObjectSections(
 const char *ELFNixPlatformRuntimeState::dlerror() { return DLFcnError.c_str(); }
 
 void *ELFNixPlatformRuntimeState::dlopen(std::string_view Path, int Mode) {
-  std::lock_guard<std::recursive_mutex> Lock(JDStatesMutex);
+  // std::lock_guard<std::recursive_mutex> Lock(JDStatesMutex);
 
-  // Use fast path if all JITDylibs are already loaded and don't require
-  // re-running initializers.
-  if (auto *JDS = getJITDylibStateByName(Path)) {
-    if (!JDS->AllowReinitialization) {
-      ++JDS->RefCount;
-      return JDS->Header;
-    }
-  }
 
-  auto H = dlopenInitialize(Path, Mode);
-  if (!H) {
+
+  if (auto H = dlopenImpl(Path, Mode))
+    return *H;
+  else {
+    // FIXME: Make dlerror thread safe.
     DLFcnError = toString(H.takeError());
     return nullptr;
   }
-
-  return *H;
 }
 
 int ELFNixPlatformRuntimeState::dlclose(void *DSOHandle) {
-  runAtExits(DSOHandle);
+
+
+  if (auto Err = dlcloseImpl(DSOHandle)) {
+      DLFcnError = toString(std::move(Err));
+    return -1;
+  }
   return 0;
 }
 
@@ -264,15 +353,16 @@ int ELFNixPlatformRuntimeState::registerAtExit(void (*F)(void *), void *Arg,
 }
 
 void ELFNixPlatformRuntimeState::runAtExits(void *DSOHandle) {
-  // FIXME: Should atexits be allowed to run concurrently with access to
-  // JDState?
-  AtExitsVector V;
-  {
-    std::lock_guard<std::recursive_mutex> Lock(JDStatesMutex);
-    auto *JDS = getJITDylibStateByHeaderAddr(DSOHandle);
-    assert(JDS && "JITDlybi state not initialized");
-    std::swap(V, JDS->AtExits);
-  }
+  std::unique_lock<std::recursive_mutex> Lock(JDStatesMutex);
+  PerJITDylibState *JDS = getJITDylibStateByHeaderAddr(DSOHandle);
+
+  if (JDS)
+    runAtExits(Lock, *JDS);
+}
+
+void ELFNixPlatformRuntimeState::runAtExits(
+  std::unique_lock<std::recursive_mutex> &JDStateLock, PerJITDylibState &JDS) {
+  AtExitsVector V = std::move(JDS.AtExits);
 
   while (!V.empty()) {
     auto &AE = V.back();
@@ -297,8 +387,13 @@ ELFNixPlatformRuntimeState::getThreadDataSectionFor(const char *ThreadData) {
 ELFNixPlatformRuntimeState::PerJITDylibState *
 ELFNixPlatformRuntimeState::getJITDylibStateByHeaderAddr(void *DSOHandle) {
   auto I = JDStates.find(DSOHandle);
-  if (I == JDStates.end())
-    return nullptr;
+ // if (I == JDStates.end())
+ //   return nullptr;
+  if (I == JDStates.end()) {
+    I = JDStates.insert(std::make_pair(DSOHandle, PerJITDylibState())).first;
+    I->second.Header = DSOHandle;
+  }
+
   return &I->second;
 }
 
@@ -315,24 +410,6 @@ ELFNixPlatformRuntimeState::getJITDylibStateByName(std::string_view Name) {
   return &J->second;
 }
 
-ELFNixPlatformRuntimeState::PerJITDylibState &
-ELFNixPlatformRuntimeState::getOrCreateJITDylibState(
-    ELFNixJITDylibInitializers &MOJDIs) {
-  void *Header = MOJDIs.DSOHandleAddress.toPtr<void *>();
-
-  auto &JDS = JDStates[Header];
-
-  // If this entry hasn't been created yet.
-  if (!JDS.Header) {
-    assert(!JDNameToHeader.count(MOJDIs.Name) &&
-           "JITDylib has header map entry but no name map entry");
-    JDNameToHeader[MOJDIs.Name] = Header;
-    JDS.Header = Header;
-  }
-
-  return JDS;
-}
-
 Error ELFNixPlatformRuntimeState::registerThreadDataSection(
     span<const char> ThreadDataSection) {
   std::lock_guard<std::mutex> Lock(ThreadDataSectionsMutex);
@@ -360,73 +437,139 @@ ELFNixPlatformRuntimeState::lookupSymbolInJITDylib(void *DSOHandle,
   return Result;
 }
 
-Expected<ELFNixJITDylibInitializerSequence>
-ELFNixPlatformRuntimeState::getJITDylibInitializersByName(
-    std::string_view Path) {
-  Expected<ELFNixJITDylibInitializerSequence> Result(
-      (ELFNixJITDylibInitializerSequence()));
-  std::string PathStr(Path.data(), Path.size());
-  if (auto Err =
-          WrapperFunction<SPSExpected<SPSELFNixJITDylibInitializerSequence>(
-              SPSString)>::call(&__orc_rt_elfnix_get_initializers_tag, Result,
-                                Path))
-    return std::move(Err);
-  return Result;
+Error ELFNixPlatformRuntimeState::runInits(
+  std::unique_lock<std::recursive_mutex> &JDStatesLock, PerJITDylibState &JDS) {
+  std::vector<span<void (*)()>> InitSections;
+  InitSections.reserve(JDS.RecordedInits.numNewSections());
+
+  JDS.RecordedInits.processNewSections(
+    [&](span<void (*)()> Inits) { InitSections.push_back(Inits); });
+
+  JDStatesLock.unlock();
+  for(auto Sec : InitSections)
+    for (auto *Init : Sec)
+      Init();
+
+  JDStatesLock.lock();
+
+  return Error::success();
 }
 
-Expected<void *>
-ELFNixPlatformRuntimeState::dlopenInitialize(std::string_view Path, int Mode) {
-  // Either our JITDylib wasn't loaded, or it or one of its dependencies allows
-  // reinitialization. We need to call in to the JIT to see if there's any new
-  // work pending.
-  auto InitSeq = getJITDylibInitializersByName(Path);
-  if (!InitSeq)
-    return InitSeq.takeError();
-
-  // Init sequences should be non-empty.
-  if (InitSeq->empty())
-    return make_error<StringError>(
-        "__orc_rt_elfnix_get_initializers returned an "
-        "empty init sequence");
-
-  // Otherwise register and run initializers for each JITDylib.
-  for (auto &MOJDIs : *InitSeq)
-    if (auto Err = initializeJITDylib(MOJDIs))
-      return std::move(Err);
-
-  // Return the header for the last item in the list.
-  auto *JDS = getJITDylibStateByHeaderAddr(
-      InitSeq->back().DSOHandleAddress.toPtr<void *>());
-  assert(JDS && "Missing state entry for JD");
+Expected<void *> ELFNixPlatformRuntimeState::dlopenImpl(std::string_view Path,
+                                                        int Mode) {
+  std::unique_lock<std::recursive_mutex> Lock(JDStatesMutex);
+  PerJITDylibState *JDS = getJITDylibStateByName(Path);
+
+  if (!JDS)
+    return make_error<StringError>("No registered JTIDylib for path " +
+                                   std::string(Path.data(), Path.size()));
+
+  if (auto Err = dlopenFull(Lock, *JDS))
+    return std::move(Err);
+
+  ++JDS->RefCount;
+
   return JDS->Header;
 }
 
-long getPriority(const std::string &name) {
-  auto pos = name.find_last_not_of("0123456789");
-  if (pos == name.size() - 1)
-    return 65535;
-  else
-    return std::strtol(name.c_str() + pos + 1, nullptr, 10);
+Error ELFNixPlatformRuntimeState::dlopenFull(
+  std::unique_lock<std::recursive_mutex> &JDStateLock, PerJITDylibState &JDS) {
+  Expected<ELFNixJITDylibDepInfoMap> DepInfo((ELFNixJITDylibDepInfoMap()));
+  JDStateLock.unlock();
+  if (auto Err = WrapperFunction<SPSExpected<SPSELFNixJITDylibDepInfoMap>(
+        SPSExecutorAddr)>::call(&__orc_rt_elfnix_push_initializers_tag,
+                                DepInfo, ExecutorAddr::fromPtr(JDS.Header)))
+    return Err;
+  JDStateLock.lock();
+
+  if (!DepInfo)
+    return DepInfo.takeError();
+
+  if (auto Err = dlopenInitialize(JDStateLock, JDS, *DepInfo))
+    return Err;
+
+  if (!DepInfo->empty()) {
+    std::ostringstream ErrStream;
+    ErrStream << "Encountered unrecognized dep-info key headers "
+                 "while processing dlopen of "
+              << JDS.Name;
+    return make_error<StringError>(ErrStream.str());
+  }
+
+  return Error::success();
 }
 
-Error ELFNixPlatformRuntimeState::initializeJITDylib(
-    ELFNixJITDylibInitializers &MOJDIs) {
+Error ELFNixPlatformRuntimeState::dlopenInitialize(
+  std::unique_lock<std::recursive_mutex> &JDStatesLock, PerJITDylibState &JDS,
+  ELFNixJITDylibDepInfoMap &DepInfo) {
+
+  auto I = DepInfo.find(ExecutorAddr::fromPtr(JDS.Header));
+  if (I == DepInfo.end())
+    return Error::success();
+
+  auto Deps = std::move(I->second);
+  DepInfo.erase(I);
+
+  std::vector<PerJITDylibState *> OldDeps;
+  std::swap(JDS.Deps, OldDeps);
+  JDS.Deps.reserve(Deps.size());
+  for (auto H : Deps) {
+    PerJITDylibState *DepJDS = getJITDylibStateByHeaderAddr(H.toPtr<void *>());
+    if (!DepJDS) {
+      std::ostringstream ErrStream;
+      ErrStream << "Encountered unrecognized dep header "
+                << H.toPtr<void *>() << " while initializing "
+                << JDS.Name;
+      return make_error<StringError>(ErrStream.str());
+    }
+    ++DepJDS->LinkedAgainstRefCount;
+    if (auto Err = dlopenInitialize(JDStatesLock, *DepJDS, DepInfo))
+      return Err;
+  }
 
-  auto &JDS = getOrCreateJITDylibState(MOJDIs);
-  ++JDS.RefCount;
+  if (auto Err = runInits(JDStatesLock, JDS))
+    return Err;
 
-  using SectionList = std::vector<ExecutorAddrRange>;
-  std::sort(MOJDIs.InitSections.begin(), MOJDIs.InitSections.end(),
-            [](const std::pair<std::string, SectionList> &LHS,
-               const std::pair<std::string, SectionList> &RHS) -> bool {
-              return getPriority(LHS.first) < getPriority(RHS.first);
-            });
-  for (auto &Entry : MOJDIs.InitSections)
-    if (auto Err = runInitArray(Entry.second, MOJDIs))
-      return Err;
+  for (auto *DepJDS : OldDeps) {
+    --DepJDS->LinkedAgainstRefCount;
+    if (!DepJDS->referenced())
+      if (auto Err = dlcloseInitialize(JDStatesLock, *DepJDS))
+        return Err;
+  }
+  return Error::success();
+}
+
+Error ELFNixPlatformRuntimeState::dlcloseImpl(void *DSOHandle) {
+
+  std::unique_lock<std::recursive_mutex> Lock(JDStatesMutex);
+  PerJITDylibState *JDS = getJITDylibStateByHeaderAddr(DSOHandle);
+
+  if (!JDS) {
+    std::ostringstream ErrStream;
+    ErrStream << "No registered JITDylib for " << DSOHandle;
+    return make_error<StringError>(ErrStream.str());
+  }
+
+  --JDS->RefCount;
+
+  if (!JDS->referenced())
+    return dlcloseInitialize(Lock, *JDS);
+
+  return Error::success();
+}
+
+Error ELFNixPlatformRuntimeState::dlcloseInitialize(
+  std::unique_lock<std::recursive_mutex> &JDStatesLock, PerJITDylibState &JDS) {
+  runAtExits(JDStatesLock, JDS);
+  JDS.RecordedInits.reset();
+  for (auto *DepJDS : JDS.Deps)
+    if (!JDS.referenced())
+      if (auto Err = dlcloseInitialize(JDStatesLock, *DepJDS))
+        return Err;
 
   return Error::success();
 }
+
 class ELFNixPlatformRuntimeTLVManager {
 public:
   void *getInstance(const char *ThreadData);
@@ -472,11 +615,12 @@ void destroyELFNixTLVMgr(void *ELFNixTLVMgr) {
 
 ORC_RT_INTERFACE orc_rt_CWrapperFunctionResult
 __orc_rt_elfnix_platform_bootstrap(char *ArgData, size_t ArgSize) {
-  return WrapperFunction<void(uint64_t)>::handle(
+  return WrapperFunction<SPSError(SPSExecutorAddr)>::handle(
              ArgData, ArgSize,
-             [](uint64_t &DSOHandle) {
+             [](ExecutorAddr DSOHandle) {
                ELFNixPlatformRuntimeState::initialize(
-                   reinterpret_cast<void *>(DSOHandle));
+                   DSOHandle.toPtr<void *>());
+                   return Error::success();
              })
       .release();
 }
@@ -487,6 +631,50 @@ __orc_rt_elfnix_platform_shutdown(char *ArgData, size_t ArgSize) {
   return WrapperFunctionResult().release();
 }
 
+ORC_RT_INTERFACE orc_rt_CWrapperFunctionResult
+__orc_rt_elfnix_register_jitdylib(char *ArgData, size_t ArgSize) {
+  return WrapperFunction<SPSError(SPSString, SPSExecutorAddr)>::
+    handle(ArgData, ArgSize,
+            [](std::string &JDName, ExecutorAddr HeaderAddr) {
+              return ELFNixPlatformRuntimeState::get().registerJITDylib(
+                JDName, HeaderAddr.toPtr<void *>());
+            }).release();
+}
+
+ORC_RT_INTERFACE orc_rt_CWrapperFunctionResult
+__orc_rt_elfnix_deregister_jitdylib(char *ArgData, size_t ArgSize) {
+  return WrapperFunction<SPSError(SPSExecutorAddr)>::
+    handle(ArgData, ArgSize,
+            [](ExecutorAddr HeaderAddr) {
+              return ELFNixPlatformRuntimeState::get().deregisterJITDylib(
+                HeaderAddr.toPtr<void *>());
+            }).release();
+}
+
+ORC_RT_INTERFACE orc_rt_CWrapperFunctionResult
+__orc_rt_elfnix_register_init_sections(char *ArgData, size_t ArgSize) {
+  return WrapperFunction<SPSError(SPSExecutorAddr,
+                                  SPSSequence<SPSExecutorAddrRange>)>::
+    handle(ArgData, ArgSize,
+            [](ExecutorAddr HeaderAddr, std::vector<ExecutorAddrRange> &Inits) {
+              return ELFNixPlatformRuntimeState::get().registerInits(
+                  HeaderAddr, std::move(Inits));
+            })
+        .release();
+}
+
+ORC_RT_INTERFACE orc_rt_CWrapperFunctionResult
+__orc_rt_elfnix_deregister_init_sections(char *ArgData, size_t ArgSize) {
+  return WrapperFunction<SPSError(SPSExecutorAddr,
+                                  SPSSequence<SPSExecutorAddrRange>)>::
+    handle(ArgData, ArgSize,
+            [](ExecutorAddr HeaderAddr, std::vector<ExecutorAddrRange> &Inits) {
+              return ELFNixPlatformRuntimeState::get().deregisterInits(
+                  HeaderAddr, std::move(Inits));
+            })
+        .release();
+}
+
 /// Wrapper function for registering metadata on a per-object basis.
 ORC_RT_INTERFACE orc_rt_CWrapperFunctionResult
 __orc_rt_elfnix_register_object_sections(char *ArgData, size_t ArgSize) {
diff --git a/compiler-rt/lib/orc/elfnix_platform.h b/compiler-rt/lib/orc/elfnix_platform.h
index e0ee9591dfc622..133be30726664b 100644
--- a/compiler-rt/lib/orc/elfnix_platform.h
+++ b/compiler-rt/lib/orc/elfnix_platform.h
@@ -37,26 +37,10 @@ struct ELFNixPerObjectSectionsToRegister {
   ExecutorAddrRange ThreadDataSection;
 };
 
-struct ELFNixJITDylibInitializers {
-  using SectionList = std::vector<ExecutorAddrRange>;
+using ELFNixJITDylibDepInfo = std::vector<ExecutorAddr>;
 
-  ELFNixJITDylibInitializers() = default;
-  ELFNixJITDylibInitializers(std::string Name, ExecutorAddr DSOHandleAddress)
-      : Name(std::move(Name)), DSOHandleAddress(std::move(DSOHandleAddress)) {}
-
-  std::string Name;
-  ExecutorAddr DSOHandleAddress;
-
-  std::vector<std::pair<std::string, SectionList>> InitSections;
-};
-
-class ELFNixJITDylibDeinitializers {};
-
-using ELFNixJITDylibInitializerSequence =
-    std::vector<ELFNixJITDylibInitializers>;
-
-using ELFNixJITDylibDeinitializerSequence =
-    std::vector<ELFNixJITDylibDeinitializers>;
+using ELFNixJITDylibDepInfoMap =
+    std::unordered_map<ExecutorAddr, ELFNixJITDylibDepInfo>;
 
 enum dlopen_mode : int {
   ORC_RT_RTLD_LAZY = 0x1,
@@ -94,37 +78,9 @@ class SPSSerializationTraits<SPSELFNixPerObjectSectionsToRegister,
   }
 };
 
-using SPSNamedExecutorAddrRangeSequenceMap =
-    SPSSequence<SPSTuple<SPSString, SPSExecutorAddrRangeSequence>>;
-
-using SPSELFNixJITDylibInitializers =
-    SPSTuple<SPSString, SPSExecutorAddr, SPSNamedExecutorAddrRangeSequenceMap>;
-
-using SPSELFNixJITDylibInitializerSequence =
-    SPSSequence<SPSELFNixJITDylibInitializers>;
-
-/// Serialization traits for ELFNixJITDylibInitializers.
-template <>
-class SPSSerializationTraits<SPSELFNixJITDylibInitializers,
-                             elfnix::ELFNixJITDylibInitializers> {
-public:
-  static size_t size(const elfnix::ELFNixJITDylibInitializers &MOJDIs) {
-    return SPSELFNixJITDylibInitializers::AsArgList::size(
-        MOJDIs.Name, MOJDIs.DSOHandleAddress, MOJDIs.InitSections);
-  }
-
-  static bool serialize(SPSOutputBuffer &OB,
-                        const elfnix::ELFNixJITDylibInitializers &MOJDIs) {
-    return SPSELFNixJITDylibInitializers::AsArgList::serialize(
-        OB, MOJDIs.Name, MOJDIs.DSOHandleAddress, MOJDIs.InitSections);
-  }
-
-  static bool deserialize(SPSInputBuffer &IB,
-                          elfnix::ELFNixJITDylibInitializers &MOJDIs) {
-    return SPSELFNixJITDylibInitializers::AsArgList::deserialize(
-        IB, MOJDIs.Name, MOJDIs.DSOHandleAddress, MOJDIs.InitSections);
-  }
-};
+using SPSELFNixJITDylibDepInfo = SPSSequence<SPSExecutorAddr>;
+using SPSELFNixJITDylibDepInfoMap =
+    SPSSequence<SPSTuple<SPSExecutorAddr, SPSELFNixJITDylibDepInfo>>;
 
 } // end namespace __orc_rt
 
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h b/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h
index 9b635064fc6a54..addf5248ab40dc 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h
@@ -31,25 +31,9 @@ struct ELFPerObjectSectionsToRegister {
   ExecutorAddrRange ThreadDataSection;
 };
 
-struct ELFNixJITDylibInitializers {
-  using SectionList = std::vector<ExecutorAddrRange>;
-
-  ELFNixJITDylibInitializers(std::string Name, ExecutorAddr DSOHandleAddress)
-      : Name(std::move(Name)), DSOHandleAddress(std::move(DSOHandleAddress)) {}
-
-  std::string Name;
-  ExecutorAddr DSOHandleAddress;
-
-  StringMap<SectionList> InitSections;
-};
-
-class ELFNixJITDylibDeinitializers {};
-
-using ELFNixJITDylibInitializerSequence =
-    std::vector<ELFNixJITDylibInitializers>;
-
-using ELFNixJITDylibDeinitializerSequence =
-    std::vector<ELFNixJITDylibDeinitializers>;
+using ELFNixJITDylibDepInfo = std::vector<ExecutorAddr>;
+using ELFNixJITDylibDepInfoMap =
+    std::vector<std::pair<ExecutorAddr, ELFNixJITDylibDepInfo>>;
 
 /// Mediates between ELFNix initialization and ExecutionSession state.
 class ELFNixPlatform : public Platform {
@@ -161,7 +145,8 @@ class ELFNixPlatform : public Platform {
                                      jitlink::PassConfiguration &Config);
 
     void addDSOHandleSupportPasses(MaterializationResponsibility &MR,
-                                   jitlink::PassConfiguration &Config);
+                                   jitlink::PassConfiguration &Config,
+                                   bool IsBootstraping);
 
     void addEHAndTLVSupportPasses(MaterializationResponsibility &MR,
                                   jitlink::PassConfiguration &Config);
@@ -169,7 +154,8 @@ class ELFNixPlatform : public Platform {
     Error preserveInitSections(jitlink::LinkGraph &G,
                                MaterializationResponsibility &MR);
 
-    Error registerInitSections(jitlink::LinkGraph &G, JITDylib &JD);
+    Error registerInitSections(jitlink::LinkGraph &G, JITDylib &JD,
+                               bool IsBootstraping);
 
     Error fixTLVSectionsAndEdges(jitlink::LinkGraph &G, JITDylib &JD);
 
@@ -178,11 +164,8 @@ class ELFNixPlatform : public Platform {
     InitSymbolDepMap InitSymbolDeps;
   };
 
-  using SendInitializerSequenceFn =
-      unique_function<void(Expected<ELFNixJITDylibInitializerSequence>)>;
-
-  using SendDeinitializerSequenceFn =
-      unique_function<void(Expected<ELFNixJITDylibDeinitializerSequence>)>;
+  using PushInitializersSendResultFn =
+      unique_function<void(Expected<ELFNixJITDylibDepInfoMap>)>;
 
   using SendSymbolAddressFn = unique_function<void(Expected<ExecutorAddr>)>;
 
@@ -196,18 +179,10 @@ class ELFNixPlatform : public Platform {
   // Associate ELFNixPlatform JIT-side runtime support functions with handlers.
   Error associateRuntimeSupportFunctions(JITDylib &PlatformJD);
 
-  void getInitializersBuildSequencePhase(SendInitializerSequenceFn SendResult,
-                                         JITDylib &JD,
-                                         std::vector<JITDylibSP> DFSLinkOrder);
-
-  void getInitializersLookupPhase(SendInitializerSequenceFn SendResult,
-                                  JITDylib &JD);
-
-  void rt_getInitializers(SendInitializerSequenceFn SendResult,
-                          StringRef JDName);
+  void pushInitializersLoop(PushInitializersSendResultFn SendResult, JITDylibSP JD);
 
-  void rt_getDeinitializers(SendDeinitializerSequenceFn SendResult,
-                            ExecutorAddr Handle);
+  void rt_recordInitializers(PushInitializersSendResultFn SendResult,
+                             ExecutorAddr JDHeader);
 
   void rt_lookupSymbol(SendSymbolAddressFn SendResult, ExecutorAddr Handle,
                        StringRef SymbolName);
@@ -215,22 +190,34 @@ class ELFNixPlatform : public Platform {
   // Records the addresses of runtime symbols used by the platform.
   Error bootstrapELFNixRuntime(JITDylib &PlatformJD);
 
-  Error registerInitInfo(JITDylib &JD,
-                         ArrayRef<jitlink::Section *> InitSections);
-
   Error registerPerObjectSections(const ELFPerObjectSectionsToRegister &POSR);
 
   Expected<uint64_t> createPThreadKey();
 
+  struct JDBootstrapState {
+    JITDylib *JD = nullptr;
+    std::string JDName;
+    ExecutorAddr HeaderAddr;
+    SmallVector<ExecutorAddrRange> Initializers;
+  };
+
+  std::map<JITDylib *, JDBootstrapState> JDBootstrapStates;
+
   ExecutionSession &ES;
   ObjectLinkingLayer &ObjLinkingLayer;
 
   SymbolStringPtr DSOHandleSymbol;
+  ExecutorAddr DSOHandleAddr;
   std::atomic<bool> RuntimeBootstrapped{false};
 
   ExecutorAddr orc_rt_elfnix_platform_bootstrap;
   ExecutorAddr orc_rt_elfnix_platform_shutdown;
+  ExecutorAddr orc_rt_elfnix_register_jitdylib;
+  ExecutorAddr orc_rt_elfnix_deregister_jitdylib;
+  ExecutorAddr orc_rt_elfnix_register_init_sections;
+  ExecutorAddr orc_rt_elfnix_deregister_init_sections;
   ExecutorAddr orc_rt_elfnix_register_object_sections;
+  ExecutorAddr orc_rt_elfnix_deregister_object_sections;
   ExecutorAddr orc_rt_elfnix_create_pthread_key;
 
   DenseMap<JITDylib *, SymbolLookupSet> RegisteredInitSymbols;
@@ -238,10 +225,10 @@ class ELFNixPlatform : public Platform {
   // InitSeqs gets its own mutex to avoid locking the whole session when
   // aggregating data from the jitlink.
   std::mutex PlatformMutex;
-  DenseMap<JITDylib *, ELFNixJITDylibInitializers> InitSeqs;
   std::vector<ELFPerObjectSectionsToRegister> BootstrapPOSRs;
 
   DenseMap<ExecutorAddr, JITDylib *> HandleAddrToJITDylib;
+  DenseMap<JITDylib *, ExecutorAddr> JITDylibToHandleAddr;
   DenseMap<JITDylib *, uint64_t> JITDylibToPThreadKey;
 };
 
@@ -273,63 +260,12 @@ class SPSSerializationTraits<SPSELFPerObjectSectionsToRegister,
   }
 };
 
-using SPSNamedExecutorAddrRangeSequenceMap =
-    SPSSequence<SPSTuple<SPSString, SPSExecutorAddrRangeSequence>>;
-
-using SPSELFNixJITDylibInitializers =
-    SPSTuple<SPSString, SPSExecutorAddr, SPSNamedExecutorAddrRangeSequenceMap>;
-
-using SPSELFNixJITDylibInitializerSequence =
-    SPSSequence<SPSELFNixJITDylibInitializers>;
-
-/// Serialization traits for ELFNixJITDylibInitializers.
-template <>
-class SPSSerializationTraits<SPSELFNixJITDylibInitializers,
-                             ELFNixJITDylibInitializers> {
-public:
-  static size_t size(const ELFNixJITDylibInitializers &MOJDIs) {
-    return SPSELFNixJITDylibInitializers::AsArgList::size(
-        MOJDIs.Name, MOJDIs.DSOHandleAddress, MOJDIs.InitSections);
-  }
-
-  static bool serialize(SPSOutputBuffer &OB,
-                        const ELFNixJITDylibInitializers &MOJDIs) {
-    return SPSELFNixJITDylibInitializers::AsArgList::serialize(
-        OB, MOJDIs.Name, MOJDIs.DSOHandleAddress, MOJDIs.InitSections);
-  }
-
-  static bool deserialize(SPSInputBuffer &IB,
-                          ELFNixJITDylibInitializers &MOJDIs) {
-    return SPSELFNixJITDylibInitializers::AsArgList::deserialize(
-        IB, MOJDIs.Name, MOJDIs.DSOHandleAddress, MOJDIs.InitSections);
-  }
-};
-
-using SPSELFJITDylibDeinitializers = SPSEmpty;
-
-using SPSELFJITDylibDeinitializerSequence =
-    SPSSequence<SPSELFJITDylibDeinitializers>;
-
-template <>
-class SPSSerializationTraits<SPSELFJITDylibDeinitializers,
-                             ELFNixJITDylibDeinitializers> {
-public:
-  static size_t size(const ELFNixJITDylibDeinitializers &MOJDDs) { return 0; }
-
-  static bool serialize(SPSOutputBuffer &OB,
-                        const ELFNixJITDylibDeinitializers &MOJDDs) {
-    return true;
-  }
-
-  static bool deserialize(SPSInputBuffer &IB,
-                          ELFNixJITDylibDeinitializers &MOJDDs) {
-    MOJDDs = ELFNixJITDylibDeinitializers();
-    return true;
-  }
-};
+using SPSELFNixJITDylibDepInfoMap =
+    SPSSequence<SPSTuple<SPSExecutorAddr, SPSSequence<SPSExecutorAddr>>>;
 
 } // end namespace shared
 } // end namespace orc
 } // end namespace llvm
 
 #endif // LLVM_EXECUTIONENGINE_ORC_ELFNIXPLATFORM_H
+
diff --git a/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp b/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
index 2b6c4b9e7f4316..0be3d778ed1f03 100644
--- a/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
@@ -1,4 +1,4 @@
-//===------ ELFNixPlatform.cpp - Utilities for executing MachO in Orc -----===//
+//===------ ELFNixPlatform.cpp - Utilities for executing ELFNix in Orc -----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -15,6 +15,7 @@
 #include "llvm/ExecutionEngine/JITLink/x86_64.h"
 #include "llvm/ExecutionEngine/Orc/DebugUtils.h"
 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
+#include "llvm/ExecutionEngine/Orc/LookupAndRecordAddrs.h"
 #include "llvm/ExecutionEngine/Orc/Shared/ObjectFormats.h"
 #include "llvm/Support/BinaryByteStream.h"
 #include "llvm/Support/Debug.h"
@@ -174,16 +175,28 @@ ELFNixPlatform::Create(ExecutionSession &ES,
 }
 
 Error ELFNixPlatform::setupJITDylib(JITDylib &JD) {
-  return JD.define(
-      std::make_unique<DSOHandleMaterializationUnit>(*this, DSOHandleSymbol));
+  if (auto Err = JD.define(std::make_unique<DSOHandleMaterializationUnit>(
+          *this, DSOHandleSymbol)))
+    return Err;
+
+  return ES.lookup({&JD}, DSOHandleSymbol).takeError();
 }
 
 Error ELFNixPlatform::teardownJITDylib(JITDylib &JD) {
+  std::lock_guard<std::mutex> Lock(PlatformMutex);
+  auto I = JITDylibToHandleAddr.find(&JD);
+  if (I != JITDylibToHandleAddr.end()) {
+    assert(HandleAddrToJITDylib.count(I->second) &&
+           "HandleAddrToJITDylib missing entry");
+    HandleAddrToJITDylib.erase(I->second);
+    JITDylibToHandleAddr.erase(I);
+  }
   return Error::success();
 }
 
 Error ELFNixPlatform::notifyAdding(ResourceTracker &RT,
                                    const MaterializationUnit &MU) {
+
   auto &JD = RT.getJITDylib();
   const auto &InitSym = MU.getInitializerSymbol();
   if (!InitSym)
@@ -265,20 +278,27 @@ ELFNixPlatform::ELFNixPlatform(
     : ES(ES), ObjLinkingLayer(ObjLinkingLayer),
       DSOHandleSymbol(ES.intern("__dso_handle")) {
   ErrorAsOutParameter _(&Err);
-
   ObjLinkingLayer.addPlugin(std::make_unique<ELFNixPlatformPlugin>(*this));
 
   PlatformJD.addGenerator(std::move(OrcRuntimeGenerator));
 
   // PlatformJD hasn't been 'set-up' by the platform yet (since we're creating
   // the platform now), so set it up.
-  if (auto E2 = setupJITDylib(PlatformJD)) {
+  if (auto E2 =
+          PlatformJD.define(std::make_unique<DSOHandleMaterializationUnit>(
+              *this, DSOHandleSymbol))) {
     Err = std::move(E2);
     return;
   }
 
-  RegisteredInitSymbols[&PlatformJD].add(
-      DSOHandleSymbol, SymbolLookupFlags::WeaklyReferencedSymbol);
+  auto E = ES.lookup({&PlatformJD}, DSOHandleSymbol);
+  if (auto E2 = E.takeError()) {
+    Err = std::move(E2);
+    return;
+  }
+  DSOHandleAddr = E->getAddress();
+  // RegisteredInitSymbols[&PlatformJD].add(
+  //   DSOHandleSymbol, SymbolLookupFlags::WeaklyReferencedSymbol);
 
   // Associate wrapper function tags with JIT-side function implementations.
   if (auto E2 = associateRuntimeSupportFunctions(PlatformJD)) {
@@ -293,22 +313,18 @@ ELFNixPlatform::ELFNixPlatform(
     Err = std::move(E2);
     return;
   }
+
+  JDBootstrapStates.clear();
 }
 
 Error ELFNixPlatform::associateRuntimeSupportFunctions(JITDylib &PlatformJD) {
   ExecutionSession::JITDispatchHandlerAssociationMap WFs;
 
-  using GetInitializersSPSSig =
-      SPSExpected<SPSELFNixJITDylibInitializerSequence>(SPSString);
-  WFs[ES.intern("__orc_rt_elfnix_get_initializers_tag")] =
-      ES.wrapAsyncWithSPS<GetInitializersSPSSig>(
-          this, &ELFNixPlatform::rt_getInitializers);
-
-  using GetDeinitializersSPSSig =
-      SPSExpected<SPSELFJITDylibDeinitializerSequence>(SPSExecutorAddr);
-  WFs[ES.intern("__orc_rt_elfnix_get_deinitializers_tag")] =
-      ES.wrapAsyncWithSPS<GetDeinitializersSPSSig>(
-          this, &ELFNixPlatform::rt_getDeinitializers);
+  using RecordInitializersSPSSig =
+      SPSExpected<SPSELFNixJITDylibDepInfoMap>(SPSExecutorAddr);
+  WFs[ES.intern("__orc_rt_elfnix_push_initializers_tag")] =
+      ES.wrapAsyncWithSPS<RecordInitializersSPSSig>(
+          this, &ELFNixPlatform::rt_recordInitializers);
 
   using LookupSymbolSPSSig =
       SPSExpected<SPSExecutorAddr>(SPSExecutorAddr, SPSString);
@@ -319,110 +335,120 @@ Error ELFNixPlatform::associateRuntimeSupportFunctions(JITDylib &PlatformJD) {
   return ES.registerJITDispatchHandlers(PlatformJD, std::move(WFs));
 }
 
-void ELFNixPlatform::getInitializersBuildSequencePhase(
-    SendInitializerSequenceFn SendResult, JITDylib &JD,
-    std::vector<JITDylibSP> DFSLinkOrder) {
-  ELFNixJITDylibInitializerSequence FullInitSeq;
-  {
-    std::lock_guard<std::mutex> Lock(PlatformMutex);
-    for (auto &InitJD : reverse(DFSLinkOrder)) {
-      LLVM_DEBUG({
-        dbgs() << "ELFNixPlatform: Appending inits for \"" << InitJD->getName()
-               << "\" to sequence\n";
-      });
-      auto ISItr = InitSeqs.find(InitJD.get());
-      if (ISItr != InitSeqs.end()) {
-        FullInitSeq.emplace_back(std::move(ISItr->second));
-        InitSeqs.erase(ISItr);
-      }
-    }
-  }
-
-  SendResult(std::move(FullInitSeq));
-}
-
-void ELFNixPlatform::getInitializersLookupPhase(
-    SendInitializerSequenceFn SendResult, JITDylib &JD) {
-
-  auto DFSLinkOrder = JD.getDFSLinkOrder();
-  if (!DFSLinkOrder) {
-    SendResult(DFSLinkOrder.takeError());
-    return;
-  }
-
+void ELFNixPlatform::pushInitializersLoop(
+    PushInitializersSendResultFn SendResult, JITDylibSP JD) {
   DenseMap<JITDylib *, SymbolLookupSet> NewInitSymbols;
+  DenseMap<JITDylib *, SmallVector<JITDylib *>> JDDepMap;
+  SmallVector<JITDylib *, 16> Worklist({JD.get()});
+
   ES.runSessionLocked([&]() {
-    for (auto &InitJD : *DFSLinkOrder) {
-      auto RISItr = RegisteredInitSymbols.find(InitJD.get());
+    while (!Worklist.empty()) {
+      // FIXME: Check for defunct dylibs.
+
+      auto DepJD = Worklist.back();
+      Worklist.pop_back();
+
+      // If we've already visited this JITDylib on this iteration then continue.
+      if (JDDepMap.count(DepJD))
+        continue;
+
+      // Add dep info.
+      auto &DM = JDDepMap[DepJD];
+      DepJD->withLinkOrderDo([&](const JITDylibSearchOrder &O) {
+        for (auto &KV : O) {
+          if (KV.first == DepJD)
+            continue;
+          DM.push_back(KV.first);
+          Worklist.push_back(KV.first);
+        }
+      });
+
+      // Add any registered init symbols.
+      auto RISItr = RegisteredInitSymbols.find(DepJD);
       if (RISItr != RegisteredInitSymbols.end()) {
-        NewInitSymbols[InitJD.get()] = std::move(RISItr->second);
+        NewInitSymbols[DepJD] = std::move(RISItr->second);
         RegisteredInitSymbols.erase(RISItr);
       }
     }
   });
 
-  // If there are no further init symbols to look up then move on to the next
-  // phase.
+  // If there are no further init symbols to look up then send the link order
+  // (as a list of header addresses) to the caller.
   if (NewInitSymbols.empty()) {
-    getInitializersBuildSequencePhase(std::move(SendResult), JD,
-                                      std::move(*DFSLinkOrder));
+
+    // To make the list intelligible to the runtime we need to convert all
+    // JITDylib pointers to their header addresses. Only include JITDylibs
+    // that appear in the JITDylibToHandleAddr map (i.e. those that have been
+    // through setupJITDylib) -- bare JITDylibs aren't managed by the platform.
+    DenseMap<JITDylib *, ExecutorAddr> HeaderAddrs;
+    HeaderAddrs.reserve(JDDepMap.size());
+    {
+      std::lock_guard<std::mutex> Lock(PlatformMutex);
+      for (auto &KV : JDDepMap) {
+        auto I = JITDylibToHandleAddr.find(KV.first);
+        if (I != JITDylibToHandleAddr.end())
+          HeaderAddrs[KV.first] = I->second;
+      }
+    }
+
+    // Build the dep info map to return.
+    ELFNixJITDylibDepInfoMap DIM;
+    DIM.reserve(JDDepMap.size());
+    for (auto &KV : JDDepMap) {
+      auto HI = HeaderAddrs.find(KV.first);
+      // Skip unmanaged JITDylibs.
+      if (HI == HeaderAddrs.end())
+        continue;
+      auto H = HI->second;
+      ELFNixJITDylibDepInfo DepInfo;
+      for (auto &Dep : KV.second) {
+        auto HJ = HeaderAddrs.find(Dep);
+        if (HJ != HeaderAddrs.end())
+          DepInfo.push_back(HJ->second);
+      }
+      DIM.push_back(std::make_pair(H, std::move(DepInfo)));
+    }
+    SendResult(DIM);
     return;
   }
 
   // Otherwise issue a lookup and re-run this phase when it completes.
   lookupInitSymbolsAsync(
-      [this, SendResult = std::move(SendResult), &JD](Error Err) mutable {
+      [this, SendResult = std::move(SendResult), JD](Error Err) mutable {
         if (Err)
           SendResult(std::move(Err));
         else
-          getInitializersLookupPhase(std::move(SendResult), JD);
+          pushInitializersLoop(std::move(SendResult), JD);
       },
       ES, std::move(NewInitSymbols));
 }
 
-void ELFNixPlatform::rt_getInitializers(SendInitializerSequenceFn SendResult,
-                                        StringRef JDName) {
-  LLVM_DEBUG({
-    dbgs() << "ELFNixPlatform::rt_getInitializers(\"" << JDName << "\")\n";
-  });
-
-  JITDylib *JD = ES.getJITDylibByName(JDName);
-  if (!JD) {
-    LLVM_DEBUG({
-      dbgs() << "  No such JITDylib \"" << JDName << "\". Sending error.\n";
-    });
-    SendResult(make_error<StringError>("No JITDylib named " + JDName,
-                                       inconvertibleErrorCode()));
-    return;
-  }
-
-  getInitializersLookupPhase(std::move(SendResult), *JD);
-}
-
-void ELFNixPlatform::rt_getDeinitializers(
-    SendDeinitializerSequenceFn SendResult, ExecutorAddr Handle) {
-  LLVM_DEBUG({
-    dbgs() << "ELFNixPlatform::rt_getDeinitializers(\"" << Handle << "\")\n";
-  });
-
-  JITDylib *JD = nullptr;
-
+void ELFNixPlatform::rt_recordInitializers(
+    PushInitializersSendResultFn SendResult, ExecutorAddr JDHeaderAddr) {
+  JITDylibSP JD;
   {
     std::lock_guard<std::mutex> Lock(PlatformMutex);
-    auto I = HandleAddrToJITDylib.find(Handle);
+    auto I = HandleAddrToJITDylib.find(JDHeaderAddr);
     if (I != HandleAddrToJITDylib.end())
       JD = I->second;
   }
 
+  LLVM_DEBUG({
+    dbgs() << "ELFNixPlatform::rt_pushInitializers(" << JDHeaderAddr << ") ";
+    if (JD)
+      dbgs() << "pushing initializers for " << JD->getName() << "\n";
+    else
+      dbgs() << "No JITDylib for header address.\n";
+  });
+
   if (!JD) {
-    LLVM_DEBUG(dbgs() << "  No JITDylib for handle " << Handle << "\n");
-    SendResult(make_error<StringError>("No JITDylib associated with handle " +
-                                           formatv("{0:x}", Handle),
+    SendResult(make_error<StringError>("No JITDylib with header addr " +
+                                           formatv("{0:x}", JDHeaderAddr),
                                        inconvertibleErrorCode()));
     return;
   }
 
-  SendResult(ELFNixJITDylibDeinitializerSequence());
+  pushInitializersLoop(std::move(SendResult), JD);
 }
 
 void ELFNixPlatform::rt_lookupSymbol(SendSymbolAddressFn SendResult,
@@ -474,42 +500,31 @@ void ELFNixPlatform::rt_lookupSymbol(SendSymbolAddressFn SendResult,
 }
 
 Error ELFNixPlatform::bootstrapELFNixRuntime(JITDylib &PlatformJD) {
+  if (auto Err = lookupAndRecordAddrs(
+          ES, LookupKind::Static, makeJITDylibSearchOrder(&PlatformJD),
+          {
+              {ES.intern("__orc_rt_elfnix_platform_bootstrap"),
+               &orc_rt_elfnix_platform_bootstrap},
+              {ES.intern("__orc_rt_elfnix_platform_shutdown"),
+               &orc_rt_elfnix_platform_shutdown},
+              {ES.intern("__orc_rt_elfnix_register_jitdylib"),
+               &orc_rt_elfnix_register_jitdylib},
+              {ES.intern("__orc_rt_elfnix_deregister_jitdylib"),
+               &orc_rt_elfnix_deregister_jitdylib},
+              {ES.intern("__orc_rt_elfnix_register_init_sections"),
+               &orc_rt_elfnix_register_init_sections},
+              {ES.intern("__orc_rt_elfnix_deregister_init_sections"),
+               &orc_rt_elfnix_deregister_init_sections},
+              {ES.intern("__orc_rt_elfnix_register_object_sections"),
+               &orc_rt_elfnix_register_object_sections},
+              {ES.intern("__orc_rt_elfnix_deregister_object_sections"),
+               &orc_rt_elfnix_deregister_object_sections},
+          }))
+    return Err;
 
-  std::pair<const char *, ExecutorAddr *> Symbols[] = {
-      {"__orc_rt_elfnix_platform_bootstrap", &orc_rt_elfnix_platform_bootstrap},
-      {"__orc_rt_elfnix_platform_shutdown", &orc_rt_elfnix_platform_shutdown},
-      {"__orc_rt_elfnix_register_object_sections",
-       &orc_rt_elfnix_register_object_sections},
-      {"__orc_rt_elfnix_create_pthread_key",
-       &orc_rt_elfnix_create_pthread_key}};
-
-  SymbolLookupSet RuntimeSymbols;
-  std::vector<std::pair<SymbolStringPtr, ExecutorAddr *>> AddrsToRecord;
-  for (const auto &KV : Symbols) {
-    auto Name = ES.intern(KV.first);
-    RuntimeSymbols.add(Name);
-    AddrsToRecord.push_back({std::move(Name), KV.second});
-  }
-
-  auto RuntimeSymbolAddrs = ES.lookup(
-      {{&PlatformJD, JITDylibLookupFlags::MatchAllSymbols}}, RuntimeSymbols);
-  if (!RuntimeSymbolAddrs)
-    return RuntimeSymbolAddrs.takeError();
-
-  for (const auto &KV : AddrsToRecord) {
-    auto &Name = KV.first;
-    assert(RuntimeSymbolAddrs->count(Name) && "Missing runtime symbol?");
-    *KV.second = (*RuntimeSymbolAddrs)[Name].getAddress();
-  }
-
-  auto PJDDSOHandle = ES.lookup(
-      {{&PlatformJD, JITDylibLookupFlags::MatchAllSymbols}}, DSOHandleSymbol);
-  if (!PJDDSOHandle)
-    return PJDDSOHandle.takeError();
 
-  if (auto Err = ES.callSPSWrapper<void(uint64_t)>(
-          orc_rt_elfnix_platform_bootstrap,
-          PJDDSOHandle->getAddress().getValue()))
+  if (auto Err = ES.callSPSWrapper<void(SPSExecutorAddr)>(
+          orc_rt_elfnix_platform_bootstrap, DSOHandleAddr))
     return Err;
 
   // FIXME: Ordering is fuzzy here. We're probably best off saying
@@ -522,45 +537,28 @@ Error ELFNixPlatform::bootstrapELFNixRuntime(JITDylib &PlatformJD) {
     DeferredPOSRs = std::move(BootstrapPOSRs);
   }
 
+  for (auto KV : JDBootstrapStates) {
+    auto &JDBState = KV.second;
+    if (auto Err = ES.callSPSWrapper<void(SPSString, SPSExecutorAddr)>(
+            orc_rt_elfnix_register_jitdylib, JDBState.JDName,
+            JDBState.HeaderAddr))
+      return Err;
+  }
+
   for (auto &D : DeferredPOSRs)
     if (auto Err = registerPerObjectSections(D))
       return Err;
 
-  return Error::success();
-}
-
-Error ELFNixPlatform::registerInitInfo(
-    JITDylib &JD, ArrayRef<jitlink::Section *> InitSections) {
-
-  std::unique_lock<std::mutex> Lock(PlatformMutex);
-
-  ELFNixJITDylibInitializers *InitSeq = nullptr;
-  {
-    auto I = InitSeqs.find(&JD);
-    if (I == InitSeqs.end()) {
-      // If there's no init sequence entry yet then we need to look up the
-      // header symbol to force creation of one.
-      Lock.unlock();
-
-      auto SearchOrder =
-          JD.withLinkOrderDo([](const JITDylibSearchOrder &SO) { return SO; });
-      if (auto Err = ES.lookup(SearchOrder, DSOHandleSymbol).takeError())
-        return Err;
-
-      Lock.lock();
-      I = InitSeqs.find(&JD);
-      assert(I != InitSeqs.end() &&
-             "Entry missing after header symbol lookup?");
-    }
-    InitSeq = &I->second;
-  }
-
-  for (auto *Sec : InitSections) {
-    // FIXME: Avoid copy here.
-    jitlink::SectionRange R(*Sec);
-    InitSeq->InitSections[Sec->getName()].push_back(R.getRange());
+  for (auto KV : JDBootstrapStates) {
+    auto &JDBState = KV.second;
+    if (JDBState.Initializers.empty())
+      continue;
+    if (auto Err = ES.callSPSWrapper<void(SPSExecutorAddr,
+                                          SPSSequence<SPSExecutorAddrRange>)>(
+            orc_rt_elfnix_register_init_sections, JDBState.HeaderAddr,
+            JDBState.Initializers))
+      return Err;
   }
-
   return Error::success();
 }
 
@@ -599,21 +597,32 @@ void ELFNixPlatform::ELFNixPlatformPlugin::modifyPassConfig(
     MaterializationResponsibility &MR, jitlink::LinkGraph &LG,
     jitlink::PassConfiguration &Config) {
 
-  // If the initializer symbol is the __dso_handle symbol then just add
-  // the DSO handle support passes.
-  if (MR.getInitializerSymbol() == MP.DSOHandleSymbol) {
-    addDSOHandleSupportPasses(MR, Config);
-    // The DSOHandle materialization unit doesn't require any other
-    // support, so we can bail out early.
-    return;
-  }
+  bool IsBootstrapping = !MP.RuntimeBootstrapped.load();
 
-  // If the object contains initializers then add passes to record them.
-  if (MR.getInitializerSymbol())
-    addInitializerSupportPasses(MR, Config);
+  if (auto InitializerSymbol = MR.getInitializerSymbol()) {
+    // If the initializer symbol is the __dso_handle symbol then just add
+    // the DSO handle support passes.
+    if (MR.getInitializerSymbol() == MP.DSOHandleSymbol) {
+      addDSOHandleSupportPasses(MR, Config, IsBootstrapping);
+      // The DSOHandle materialization unit doesn't require any other
+      // support, so we can bail out early.
+      return;
+    }
+    Config.PrePrunePasses.push_back(
+        [this, &MR](jitlink::LinkGraph &G) -> Error {
+          if (auto Err = preserveInitSections(G, MR))
+            return Err;
+          return Error::success();
+        });
+  }
 
   // Add passes for eh-frame and TLV support.
   addEHAndTLVSupportPasses(MR, Config);
+
+  Config.PostFixupPasses.push_back([this, &JD = MR.getTargetJITDylib(),
+                                    IsBootstrapping](jitlink::LinkGraph &G) {
+    return registerInitSections(G, JD, IsBootstrapping);
+  });
 }
 
 ObjectLinkingLayer::Plugin::SyntheticSymbolDependenciesMap
@@ -630,26 +639,12 @@ ELFNixPlatform::ELFNixPlatformPlugin::getSyntheticSymbolDependencies(
   return SyntheticSymbolDependenciesMap();
 }
 
-void ELFNixPlatform::ELFNixPlatformPlugin::addInitializerSupportPasses(
-    MaterializationResponsibility &MR, jitlink::PassConfiguration &Config) {
-
-  /// Preserve init sections.
-  Config.PrePrunePasses.push_back([this, &MR](jitlink::LinkGraph &G) -> Error {
-    if (auto Err = preserveInitSections(G, MR))
-      return Err;
-    return Error::success();
-  });
-
-  Config.PostFixupPasses.push_back(
-      [this, &JD = MR.getTargetJITDylib()](jitlink::LinkGraph &G) {
-        return registerInitSections(G, JD);
-      });
-}
-
 void ELFNixPlatform::ELFNixPlatformPlugin::addDSOHandleSupportPasses(
-    MaterializationResponsibility &MR, jitlink::PassConfiguration &Config) {
+    MaterializationResponsibility &MR, jitlink::PassConfiguration &Config,
+    bool IsBootstraping) {
 
-  Config.PostAllocationPasses.push_back([this, &JD = MR.getTargetJITDylib()](
+  Config.PostAllocationPasses.push_back([this, &JD = MR.getTargetJITDylib(),
+                                         IsBootstraping](
                                             jitlink::LinkGraph &G) -> Error {
     auto I = llvm::find_if(G.defined_symbols(), [this](jitlink::Symbol *Sym) {
       return Sym->getName() == *MP.DSOHandleSymbol;
@@ -659,9 +654,26 @@ void ELFNixPlatform::ELFNixPlatformPlugin::addDSOHandleSupportPasses(
       std::lock_guard<std::mutex> Lock(MP.PlatformMutex);
       auto HandleAddr = (*I)->getAddress();
       MP.HandleAddrToJITDylib[HandleAddr] = &JD;
-      assert(!MP.InitSeqs.count(&JD) && "InitSeq entry for JD already exists");
-      MP.InitSeqs.insert(std::make_pair(
-          &JD, ELFNixJITDylibInitializers(JD.getName(), HandleAddr)));
+      MP.JITDylibToHandleAddr[&JD] = HandleAddr;
+
+      if (!IsBootstraping) {
+        G.allocActions().push_back(
+            {cantFail(WrapperFunctionCall::Create<
+                      SPSArgList<SPSString, SPSExecutorAddr>>(
+                 MP.orc_rt_elfnix_register_jitdylib, JD.getName(), HandleAddr)),
+             cantFail(WrapperFunctionCall::Create<SPSArgList<SPSExecutorAddr>>(
+                 MP.orc_rt_elfnix_deregister_jitdylib, HandleAddr))});
+      } else {
+        // G.allocActions().push_back(
+        //   {{},
+        //    cantFail(WrapperFunctionCall::Create<SPSArgList<SPSExecutorAddr>>(
+        //         MP.orc_rt_elfnix_deregister_jitdylib, HandleAddr))});
+        JDBootstrapState BState;
+        BState.JD = &JD;
+        BState.JDName = JD.getName();
+        BState.HeaderAddr = HandleAddr;
+        MP.JDBootstrapStates.emplace(&JD, BState);
+      }
     }
     return Error::success();
   });
@@ -771,8 +783,8 @@ Error ELFNixPlatform::ELFNixPlatformPlugin::preserveInitSections(
 }
 
 Error ELFNixPlatform::ELFNixPlatformPlugin::registerInitSections(
-    jitlink::LinkGraph &G, JITDylib &JD) {
-
+    jitlink::LinkGraph &G, JITDylib &JD, bool IsBootstraping) {
+  SmallVector<ExecutorAddrRange> ELFNixPlatformSecs;
   SmallVector<jitlink::Section *> InitSections;
 
   LLVM_DEBUG(dbgs() << "ELFNixPlatform::registerInitSections\n");
@@ -780,6 +792,8 @@ Error ELFNixPlatform::ELFNixPlatformPlugin::registerInitSections(
   for (auto &Sec : G.sections()) {
     if (isELFInitializerSection(Sec.getName())) {
       InitSections.push_back(&Sec);
+      jitlink::SectionRange R(Sec);
+      ELFNixPlatformSecs.push_back(R.getRange());
     }
   }
 
@@ -791,8 +805,39 @@ Error ELFNixPlatform::ELFNixPlatformPlugin::registerInitSections(
       dbgs() << "  " << Sec->getName() << ": " << R.getRange() << "\n";
     }
   });
+  using SPSRegisterInitSectionsArgs =
+      SPSArgList<SPSExecutorAddr, SPSSequence<SPSExecutorAddrRange>>;
 
-  return MP.registerInitInfo(JD, InitSections);
+  ExecutorAddr HeaderAddr;
+  {
+    std::lock_guard<std::mutex> Lock(MP.PlatformMutex);
+    auto I = MP.JITDylibToHandleAddr.find(&JD);
+    assert(I != MP.JITDylibToHandleAddr.end() && "No header registered for JD");
+    assert(I->second && "Null header registered for JD");
+    HeaderAddr = I->second;
+  }
+
+  if (IsBootstraping) {
+    auto &JBS = MP.JDBootstrapStates[&JD];
+    for (auto &I : ELFNixPlatformSecs)
+      JBS.Initializers.push_back(I);
+    // G.allocActions().push_back(
+    //    {{},
+    //     cantFail(
+    //       WrapperFunctionCall::Create<SPSRegisterInitSectionsArgs>(
+    //           MP.orc_rt_elfnix_deregister_init_sections, HeaderAddr,
+    //           ELFNixPlatformSecs))});
+  } else {
+    G.allocActions().push_back(
+        {cantFail(WrapperFunctionCall::Create<SPSRegisterInitSectionsArgs>(
+             MP.orc_rt_elfnix_register_init_sections, HeaderAddr,
+             ELFNixPlatformSecs)),
+         cantFail(WrapperFunctionCall::Create<SPSRegisterInitSectionsArgs>(
+             MP.orc_rt_elfnix_deregister_init_sections, HeaderAddr,
+             ELFNixPlatformSecs))});
+  }
+
+  return Error::success();
 }
 
 Error ELFNixPlatform::ELFNixPlatformPlugin::fixTLVSectionsAndEdges(

>From 7bd1d51b9aa6cc3cbd2cc54d3d79ba337b253867 Mon Sep 17 00:00:00 2001
From: SahilPatidar <patidarsahil2001 at gmail.com>
Date: Fri, 9 Aug 2024 17:39:05 +0530
Subject: [PATCH 2/3] Replace direct call to
 `orc_rt_elfnix_register_object_sections` with `AllocActions`.

---
 compiler-rt/lib/orc/elfnix_platform.cpp       | 16 +++++++-------
 .../llvm/ExecutionEngine/Orc/ELFNixPlatform.h |  3 ++-
 .../ExecutionEngine/Orc/ELFNixPlatform.cpp    | 22 +++++++++++--------
 3 files changed, 23 insertions(+), 18 deletions(-)

diff --git a/compiler-rt/lib/orc/elfnix_platform.cpp b/compiler-rt/lib/orc/elfnix_platform.cpp
index 29582ac6867871..45b17cef0570e6 100644
--- a/compiler-rt/lib/orc/elfnix_platform.cpp
+++ b/compiler-rt/lib/orc/elfnix_platform.cpp
@@ -213,6 +213,14 @@ Error ELFNixPlatformRuntimeState::registerObjectSections(
   return Error::success();
 }
 
+Error ELFNixPlatformRuntimeState::deregisterObjectSections(
+    ELFNixPerObjectSectionsToRegister POSR) {
+  if (POSR.EHFrameSection.Start)
+    deregisterEHFrameSection(POSR.EHFrameSection.Start.toPtr<const char *>());
+
+  return Error::success();
+}
+
 Error ELFNixPlatformRuntimeState::registerJITDylib(std::string &Name, void *Handle) {
   std::lock_guard<std::recursive_mutex> Lock(JDStatesMutex);
 
@@ -297,14 +305,6 @@ Error ELFNixPlatformRuntimeState::deregisterInits(
   return Error::success();
 }
 
-Error ELFNixPlatformRuntimeState::deregisterObjectSections(
-    ELFNixPerObjectSectionsToRegister POSR) {
-  if (POSR.EHFrameSection.Start)
-    deregisterEHFrameSection(POSR.EHFrameSection.Start.toPtr<const char *>());
-
-  return Error::success();
-}
-
 const char *ELFNixPlatformRuntimeState::dlerror() { return DLFcnError.c_str(); }
 
 void *ELFNixPlatformRuntimeState::dlopen(std::string_view Path, int Mode) {
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h b/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h
index addf5248ab40dc..4fbbd2b484d345 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h
@@ -190,7 +190,8 @@ class ELFNixPlatform : public Platform {
   // Records the addresses of runtime symbols used by the platform.
   Error bootstrapELFNixRuntime(JITDylib &PlatformJD);
 
-  Error registerPerObjectSections(const ELFPerObjectSectionsToRegister &POSR);
+  Error registerPerObjectSections(
+      jitlink::LinkGraph &G,const ELFPerObjectSectionsToRegister &POSR);
 
   Expected<uint64_t> createPThreadKey();
 
diff --git a/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp b/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
index 0be3d778ed1f03..b233588b7a5d4f 100644
--- a/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
@@ -546,7 +546,9 @@ Error ELFNixPlatform::bootstrapELFNixRuntime(JITDylib &PlatformJD) {
   }
 
   for (auto &D : DeferredPOSRs)
-    if (auto Err = registerPerObjectSections(D))
+    if (auto Err = ES.callSPSWrapper<void(
+                    SPSELFPerObjectSectionsToRegister)>(
+         orc_rt_elfnix_register_object_sections, D))
       return Err;
 
   for (auto KV : JDBootstrapStates) {
@@ -563,7 +565,7 @@ Error ELFNixPlatform::bootstrapELFNixRuntime(JITDylib &PlatformJD) {
 }
 
 Error ELFNixPlatform::registerPerObjectSections(
-    const ELFPerObjectSectionsToRegister &POSR) {
+    jitlink::LinkGraph &G, const ELFPerObjectSectionsToRegister &POSR) {
 
   if (!orc_rt_elfnix_register_object_sections)
     return make_error<StringError>("Attempting to register per-object "
@@ -571,12 +573,14 @@ Error ELFNixPlatform::registerPerObjectSections(
                                    "been loaded yet",
                                    inconvertibleErrorCode());
 
-  Error ErrResult = Error::success();
-  if (auto Err = ES.callSPSWrapper<shared::SPSError(
-                     SPSELFPerObjectSectionsToRegister)>(
-          orc_rt_elfnix_register_object_sections, ErrResult, POSR))
-    return Err;
-  return ErrResult;
+  using SPSRegisterObjSectionsArgs = SPSArgList<SPSELFPerObjectSectionsToRegister>;
+  G.allocActions().push_back(
+        {cantFail(WrapperFunctionCall::Create<SPSRegisterObjSectionsArgs>(
+             orc_rt_elfnix_register_object_sections, POSR)),
+         cantFail(WrapperFunctionCall::Create<SPSRegisterObjSectionsArgs>(
+             orc_rt_elfnix_deregister_object_sections, POSR))});
+
+  return Error::success();
 }
 
 Expected<uint64_t> ELFNixPlatform::createPThreadKey() {
@@ -738,7 +742,7 @@ void ELFNixPlatform::ELFNixPlatformPlugin::addEHAndTLVSupportPasses(
       }
 
       // Otherwise register it immediately.
-      if (auto Err = MP.registerPerObjectSections(POSR))
+      if (auto Err = MP.registerPerObjectSections(G,POSR))
         return Err;
     }
 

>From 1e0785cb2b934f0955103dbed8723570506591ed Mon Sep 17 00:00:00 2001
From: SahilPatidar <patidarsahil2001 at gmail.com>
Date: Sun, 11 Aug 2024 16:24:42 +0530
Subject: [PATCH 3/3] Update and fix bootstrap process logic

---
 compiler-rt/lib/orc/common.h                  | 160 +++++++++---------
 compiler-rt/lib/orc/elfnix_platform.cpp       | 136 ++++++++-------
 .../llvm/ExecutionEngine/Orc/ELFNixPlatform.h |   5 +-
 .../ExecutionEngine/Orc/ELFNixPlatform.cpp    |  92 +++++-----
 4 files changed, 185 insertions(+), 208 deletions(-)

diff --git a/compiler-rt/lib/orc/common.h b/compiler-rt/lib/orc/common.h
index a03ae0bddb893d..0db540c4dfde8f 100644
--- a/compiler-rt/lib/orc/common.h
+++ b/compiler-rt/lib/orc/common.h
@@ -14,18 +14,11 @@
 #define ORC_RT_COMMON_H
 
 #include "compiler.h"
-#include "orc_rt/c_api.h"
-#include <type_traits>
-
-#include "bitmask_enum.h"
-#include "debug.h"
 #include "error.h"
 #include "executor_address.h"
+#include "orc_rt/c_api.h"
+#include <type_traits>
 #include <algorithm>
-#include <ios>
-#include <map>
-#include <mutex>
-#include <sstream>
 #include <vector>
 
 /// This macro should be used to define tags that will be associated with
@@ -60,86 +53,87 @@ __orc_rt_jit_dispatch(__orc_rt_Opaque *DispatchCtx, const void *FnTag,
 /// sections, selector refs, etc.)
 template <typename RecordElement> class RecordSectionsTracker {
 public:
-/// Add a section to the "new" list.
-void add(__orc_rt::span<RecordElement> Sec) { New.push_back(std::move(Sec)); }
-
-/// Returns true if there are new sections to process.
-bool hasNewSections() const { return !New.empty(); }
-
-/// Returns the number of new sections to process.
-size_t numNewSections() const { return New.size(); }
-
-/// Process all new sections.
-template <typename ProcessSectionFunc>
-std::enable_if_t<std::is_void_v<
-    std::invoke_result_t<ProcessSectionFunc, __orc_rt::span<RecordElement>>>>
-processNewSections(ProcessSectionFunc &&ProcessSection) {
-  for (auto &Sec : New)
-    ProcessSection(Sec);
-  moveNewToProcessed();
-}
-
-/// Proces all new sections with a fallible handler.
-///
-/// Successfully handled sections will be moved to the Processed
-/// list.
-template <typename ProcessSectionFunc>
-std::enable_if_t<
-    std::is_same_v<__orc_rt::Error, std::invoke_result_t<ProcessSectionFunc,
-                                                __orc_rt::span<RecordElement>>>,
-    __orc_rt::Error>
-processNewSections(ProcessSectionFunc &&ProcessSection) {
-  for (size_t I = 0; I != New.size(); ++I) {
-    if (auto Err = ProcessSection(New[I])) {
-      for (size_t J = 0; J != I; ++J)
-        Processed.push_back(New[J]);
-      New.erase(New.begin(), New.begin() + I);
-      return Err;
+  /// Add a section to the "new" list.
+  void add(__orc_rt::span<RecordElement> Sec) { New.push_back(std::move(Sec)); }
+
+  /// Returns true if there are new sections to process.
+  bool hasNewSections() const { return !New.empty(); }
+
+  /// Returns the number of new sections to process.
+  size_t numNewSections() const { return New.size(); }
+
+  /// Process all new sections.
+  template <typename ProcessSectionFunc>
+  std::enable_if_t<std::is_void_v<
+      std::invoke_result_t<ProcessSectionFunc, __orc_rt::span<RecordElement>>>>
+  processNewSections(ProcessSectionFunc &&ProcessSection) {
+    for (auto &Sec : New)
+      ProcessSection(Sec);
+    moveNewToProcessed();
+  }
+
+  /// Proces all new sections with a fallible handler.
+  ///
+  /// Successfully handled sections will be moved to the Processed
+  /// list.
+  template <typename ProcessSectionFunc>
+  std::enable_if_t<
+      std::is_same_v<__orc_rt::Error,
+                     std::invoke_result_t<ProcessSectionFunc,
+                                          __orc_rt::span<RecordElement>>>,
+      __orc_rt::Error>
+  processNewSections(ProcessSectionFunc &&ProcessSection) {
+    for (size_t I = 0; I != New.size(); ++I) {
+      if (auto Err = ProcessSection(New[I])) {
+        for (size_t J = 0; J != I; ++J)
+          Processed.push_back(New[J]);
+        New.erase(New.begin(), New.begin() + I);
+        return Err;
+      }
     }
+    moveNewToProcessed();
+    return __orc_rt::Error::success();
+  }
+
+  /// Move all sections back to New for reprocessing.
+  void reset() {
+    moveNewToProcessed();
+    New = std::move(Processed);
+  }
+
+  /// Remove the section with the given range.
+  bool removeIfPresent(__orc_rt::ExecutorAddrRange R) {
+    if (removeIfPresent(New, R))
+      return true;
+    return removeIfPresent(Processed, R);
   }
-  moveNewToProcessed();
-  return __orc_rt::Error::success();
-}
-
-/// Move all sections back to New for reprocessing.
-void reset() {
-  moveNewToProcessed();
-  New = std::move(Processed);
-}
-
-/// Remove the section with the given range.
-bool removeIfPresent(__orc_rt::ExecutorAddrRange R) {
-  if (removeIfPresent(New, R))
-    return true;
-  return removeIfPresent(Processed, R);
-}
 
 private:
-void moveNewToProcessed() {
-  if (Processed.empty())
-    Processed = std::move(New);
-  else {
-    Processed.reserve(Processed.size() + New.size());
-    std::copy(New.begin(), New.end(), std::back_inserter(Processed));
-    New.clear();
+  void moveNewToProcessed() {
+    if (Processed.empty())
+      Processed = std::move(New);
+    else {
+      Processed.reserve(Processed.size() + New.size());
+      std::copy(New.begin(), New.end(), std::back_inserter(Processed));
+      New.clear();
+    }
   }
-}
-
-bool removeIfPresent(std::vector<__orc_rt::span<RecordElement>> &V,
-                     __orc_rt::ExecutorAddrRange R) {
-  auto RI = std::find_if(
-    V.rbegin(), V.rend(),
-    [RS = R.toSpan<RecordElement>()](const __orc_rt::span<RecordElement> &E) {
-        return E.data() == RS.data();
-    });
-  if (RI != V.rend()) {
-    V.erase(std::next(RI).base());
-    return true;
+
+  bool removeIfPresent(std::vector<__orc_rt::span<RecordElement>> &V,
+                       __orc_rt::ExecutorAddrRange R) {
+    auto RI = std::find_if(V.rbegin(), V.rend(),
+                           [RS = R.toSpan<RecordElement>()](
+                               const __orc_rt::span<RecordElement> &E) {
+                             return E.data() == RS.data();
+                           });
+    if (RI != V.rend()) {
+      V.erase(std::next(RI).base());
+      return true;
+    }
+    return false;
   }
-  return false;
-}
 
-std::vector<__orc_rt::span<RecordElement>> Processed;
-std::vector<__orc_rt::span<RecordElement>> New;
+  std::vector<__orc_rt::span<RecordElement>> Processed;
+  std::vector<__orc_rt::span<RecordElement>> New;
 };
 #endif // ORC_RT_COMMON_H
diff --git a/compiler-rt/lib/orc/elfnix_platform.cpp b/compiler-rt/lib/orc/elfnix_platform.cpp
index 45b17cef0570e6..ff56db5f26d5aa 100644
--- a/compiler-rt/lib/orc/elfnix_platform.cpp
+++ b/compiler-rt/lib/orc/elfnix_platform.cpp
@@ -28,9 +28,7 @@ using namespace __orc_rt;
 using namespace __orc_rt::elfnix;
 
 // Declare function tags for functions in the JIT process.
-ORC_RT_JIT_DISPATCH_TAG(__orc_rt_elfnix_get_initializers_tag)
 ORC_RT_JIT_DISPATCH_TAG(__orc_rt_elfnix_push_initializers_tag)
-ORC_RT_JIT_DISPATCH_TAG(__orc_rt_elfnix_get_deinitializers_tag)
 ORC_RT_JIT_DISPATCH_TAG(__orc_rt_elfnix_symbol_lookup_tag)
 
 // eh-frame registration functions, made available via aliases
@@ -108,10 +106,10 @@ class ELFNixPlatformRuntimeState {
   Error registerObjectSections(ELFNixPerObjectSectionsToRegister POSR);
   Error registerJITDylib(std::string &Name, void *Handle);
   Error deregisterJITDylib(void *Handle);
-  Error registerInits(
-    ExecutorAddr HeaderAddr, std::vector<ExecutorAddrRange> Inits);
-  Error deregisterInits(
-    ExecutorAddr HeaderAddr, std::vector<ExecutorAddrRange> Inits);
+  Error registerInits(ExecutorAddr HeaderAddr,
+                      std::vector<ExecutorAddrRange> Inits);
+  Error deregisterInits(ExecutorAddr HeaderAddr,
+                        std::vector<ExecutorAddrRange> Inits);
   Error deregisterObjectSections(ELFNixPerObjectSectionsToRegister POSR);
 
   const char *dlerror();
@@ -139,17 +137,17 @@ class ELFNixPlatformRuntimeState {
   Expected<ExecutorAddr> lookupSymbolInJITDylib(void *DSOHandle,
                                                 std::string_view Symbol);
 
-  Error runInits(
-    std::unique_lock<std::recursive_mutex> &JDStatesLock, PerJITDylibState &JDS);
+  Error runInits(std::unique_lock<std::recursive_mutex> &JDStatesLock,
+                 PerJITDylibState &JDS);
   Expected<void *> dlopenImpl(std::string_view Path, int Mode);
   Error dlopenFull(std::unique_lock<std::recursive_mutex> &JDStatesLock,
                    PerJITDylibState &JDS);
-  Error dlopenInitialize(
-    std::unique_lock<std::recursive_mutex> &JDStatesLock,
-    PerJITDylibState &JDS, ELFNixJITDylibDepInfoMap &DepInfo);
+  Error dlopenInitialize(std::unique_lock<std::recursive_mutex> &JDStatesLock,
+                         PerJITDylibState &JDS,
+                         ELFNixJITDylibDepInfoMap &DepInfo);
   Error dlcloseImpl(void *DSOHandle);
-  Error dlcloseInitialize(
-    std::unique_lock<std::recursive_mutex> &JDStatesLock, PerJITDylibState &JDS);
+  Error dlcloseInitialize(std::unique_lock<std::recursive_mutex> &JDStatesLock,
+                          PerJITDylibState &JDS);
 
   static ELFNixPlatformRuntimeState *MOPS;
 
@@ -221,7 +219,8 @@ Error ELFNixPlatformRuntimeState::deregisterObjectSections(
   return Error::success();
 }
 
-Error ELFNixPlatformRuntimeState::registerJITDylib(std::string &Name, void *Handle) {
+Error ELFNixPlatformRuntimeState::registerJITDylib(std::string &Name,
+                                                   void *Handle) {
   std::lock_guard<std::recursive_mutex> Lock(JDStatesMutex);
 
   if (JDStates.count(Handle)) {
@@ -258,16 +257,17 @@ Error ELFNixPlatformRuntimeState::deregisterJITDylib(void *Handle) {
   auto J = JDNameToHeader.find(
       std::string(I->second.Name.data(), I->second.Name.size()));
   assert(J != JDNameToHeader.end() &&
-        "Missing JDNameToHeader entry for JITDylib");
+         "Missing JDNameToHeader entry for JITDylib");
   JDNameToHeader.erase(J);
   JDStates.erase(I);
   return Error::success();
 }
 
 Error ELFNixPlatformRuntimeState::registerInits(
-  ExecutorAddr HeaderAddr, std::vector<ExecutorAddrRange> Inits) {
+    ExecutorAddr HeaderAddr, std::vector<ExecutorAddrRange> Inits) {
   std::lock_guard<std::recursive_mutex> Lock(JDStatesMutex);
-  PerJITDylibState *JDS = getJITDylibStateByHeaderAddr(HeaderAddr.toPtr<void *>());
+  PerJITDylibState *JDS =
+      getJITDylibStateByHeaderAddr(HeaderAddr.toPtr<void *>());
 
   if (!JDS) {
     std::ostringstream ErrStream;
@@ -288,7 +288,8 @@ Error ELFNixPlatformRuntimeState::deregisterInits(
     ExecutorAddr HeaderAddr, std::vector<ExecutorAddrRange> Inits) {
 
   std::lock_guard<std::recursive_mutex> Lock(JDStatesMutex);
-  PerJITDylibState *JDS = getJITDylibStateByHeaderAddr(HeaderAddr.toPtr<void *>());
+  PerJITDylibState *JDS =
+      getJITDylibStateByHeaderAddr(HeaderAddr.toPtr<void *>());
 
   if (!JDS) {
     std::ostringstream ErrStream;
@@ -308,10 +309,6 @@ Error ELFNixPlatformRuntimeState::deregisterInits(
 const char *ELFNixPlatformRuntimeState::dlerror() { return DLFcnError.c_str(); }
 
 void *ELFNixPlatformRuntimeState::dlopen(std::string_view Path, int Mode) {
-  // std::lock_guard<std::recursive_mutex> Lock(JDStatesMutex);
-
-
-
   if (auto H = dlopenImpl(Path, Mode))
     return *H;
   else {
@@ -322,10 +319,8 @@ void *ELFNixPlatformRuntimeState::dlopen(std::string_view Path, int Mode) {
 }
 
 int ELFNixPlatformRuntimeState::dlclose(void *DSOHandle) {
-
-
   if (auto Err = dlcloseImpl(DSOHandle)) {
-      DLFcnError = toString(std::move(Err));
+    DLFcnError = toString(std::move(Err));
     return -1;
   }
   return 0;
@@ -361,7 +356,8 @@ void ELFNixPlatformRuntimeState::runAtExits(void *DSOHandle) {
 }
 
 void ELFNixPlatformRuntimeState::runAtExits(
-  std::unique_lock<std::recursive_mutex> &JDStateLock, PerJITDylibState &JDS) {
+    std::unique_lock<std::recursive_mutex> &JDStateLock,
+    PerJITDylibState &JDS) {
   AtExitsVector V = std::move(JDS.AtExits);
 
   while (!V.empty()) {
@@ -387,12 +383,8 @@ ELFNixPlatformRuntimeState::getThreadDataSectionFor(const char *ThreadData) {
 ELFNixPlatformRuntimeState::PerJITDylibState *
 ELFNixPlatformRuntimeState::getJITDylibStateByHeaderAddr(void *DSOHandle) {
   auto I = JDStates.find(DSOHandle);
- // if (I == JDStates.end())
- //   return nullptr;
-  if (I == JDStates.end()) {
-    I = JDStates.insert(std::make_pair(DSOHandle, PerJITDylibState())).first;
-    I->second.Header = DSOHandle;
-  }
+  if (I == JDStates.end())
+    return nullptr;
 
   return &I->second;
 }
@@ -438,15 +430,16 @@ ELFNixPlatformRuntimeState::lookupSymbolInJITDylib(void *DSOHandle,
 }
 
 Error ELFNixPlatformRuntimeState::runInits(
-  std::unique_lock<std::recursive_mutex> &JDStatesLock, PerJITDylibState &JDS) {
+    std::unique_lock<std::recursive_mutex> &JDStatesLock,
+    PerJITDylibState &JDS) {
   std::vector<span<void (*)()>> InitSections;
   InitSections.reserve(JDS.RecordedInits.numNewSections());
 
   JDS.RecordedInits.processNewSections(
-    [&](span<void (*)()> Inits) { InitSections.push_back(Inits); });
+      [&](span<void (*)()> Inits) { InitSections.push_back(Inits); });
 
   JDStatesLock.unlock();
-  for(auto Sec : InitSections)
+  for (auto Sec : InitSections)
     for (auto *Init : Sec)
       Init();
 
@@ -473,12 +466,13 @@ Expected<void *> ELFNixPlatformRuntimeState::dlopenImpl(std::string_view Path,
 }
 
 Error ELFNixPlatformRuntimeState::dlopenFull(
-  std::unique_lock<std::recursive_mutex> &JDStateLock, PerJITDylibState &JDS) {
+    std::unique_lock<std::recursive_mutex> &JDStateLock,
+    PerJITDylibState &JDS) {
   Expected<ELFNixJITDylibDepInfoMap> DepInfo((ELFNixJITDylibDepInfoMap()));
   JDStateLock.unlock();
   if (auto Err = WrapperFunction<SPSExpected<SPSELFNixJITDylibDepInfoMap>(
-        SPSExecutorAddr)>::call(&__orc_rt_elfnix_push_initializers_tag,
-                                DepInfo, ExecutorAddr::fromPtr(JDS.Header)))
+          SPSExecutorAddr)>::call(&__orc_rt_elfnix_push_initializers_tag,
+                                  DepInfo, ExecutorAddr::fromPtr(JDS.Header)))
     return Err;
   JDStateLock.lock();
 
@@ -500,8 +494,8 @@ Error ELFNixPlatformRuntimeState::dlopenFull(
 }
 
 Error ELFNixPlatformRuntimeState::dlopenInitialize(
-  std::unique_lock<std::recursive_mutex> &JDStatesLock, PerJITDylibState &JDS,
-  ELFNixJITDylibDepInfoMap &DepInfo) {
+    std::unique_lock<std::recursive_mutex> &JDStatesLock, PerJITDylibState &JDS,
+    ELFNixJITDylibDepInfoMap &DepInfo) {
 
   auto I = DepInfo.find(ExecutorAddr::fromPtr(JDS.Header));
   if (I == DepInfo.end())
@@ -517,9 +511,8 @@ Error ELFNixPlatformRuntimeState::dlopenInitialize(
     PerJITDylibState *DepJDS = getJITDylibStateByHeaderAddr(H.toPtr<void *>());
     if (!DepJDS) {
       std::ostringstream ErrStream;
-      ErrStream << "Encountered unrecognized dep header "
-                << H.toPtr<void *>() << " while initializing "
-                << JDS.Name;
+      ErrStream << "Encountered unrecognized dep header " << H.toPtr<void *>()
+                << " while initializing " << JDS.Name;
       return make_error<StringError>(ErrStream.str());
     }
     ++DepJDS->LinkedAgainstRefCount;
@@ -559,7 +552,8 @@ Error ELFNixPlatformRuntimeState::dlcloseImpl(void *DSOHandle) {
 }
 
 Error ELFNixPlatformRuntimeState::dlcloseInitialize(
-  std::unique_lock<std::recursive_mutex> &JDStatesLock, PerJITDylibState &JDS) {
+    std::unique_lock<std::recursive_mutex> &JDStatesLock,
+    PerJITDylibState &JDS) {
   runAtExits(JDStatesLock, JDS);
   JDS.RecordedInits.reset();
   for (auto *DepJDS : JDS.Deps)
@@ -620,7 +614,7 @@ __orc_rt_elfnix_platform_bootstrap(char *ArgData, size_t ArgSize) {
              [](ExecutorAddr DSOHandle) {
                ELFNixPlatformRuntimeState::initialize(
                    DSOHandle.toPtr<void *>());
-                   return Error::success();
+               return Error::success();
              })
       .release();
 }
@@ -633,46 +627,50 @@ __orc_rt_elfnix_platform_shutdown(char *ArgData, size_t ArgSize) {
 
 ORC_RT_INTERFACE orc_rt_CWrapperFunctionResult
 __orc_rt_elfnix_register_jitdylib(char *ArgData, size_t ArgSize) {
-  return WrapperFunction<SPSError(SPSString, SPSExecutorAddr)>::
-    handle(ArgData, ArgSize,
-            [](std::string &JDName, ExecutorAddr HeaderAddr) {
-              return ELFNixPlatformRuntimeState::get().registerJITDylib(
-                JDName, HeaderAddr.toPtr<void *>());
-            }).release();
+  return WrapperFunction<SPSError(SPSString, SPSExecutorAddr)>::handle(
+             ArgData, ArgSize,
+             [](std::string &JDName, ExecutorAddr HeaderAddr) {
+               return ELFNixPlatformRuntimeState::get().registerJITDylib(
+                   JDName, HeaderAddr.toPtr<void *>());
+             })
+      .release();
 }
 
 ORC_RT_INTERFACE orc_rt_CWrapperFunctionResult
 __orc_rt_elfnix_deregister_jitdylib(char *ArgData, size_t ArgSize) {
-  return WrapperFunction<SPSError(SPSExecutorAddr)>::
-    handle(ArgData, ArgSize,
-            [](ExecutorAddr HeaderAddr) {
-              return ELFNixPlatformRuntimeState::get().deregisterJITDylib(
-                HeaderAddr.toPtr<void *>());
-            }).release();
+  return WrapperFunction<SPSError(SPSExecutorAddr)>::handle(
+             ArgData, ArgSize,
+             [](ExecutorAddr HeaderAddr) {
+               return ELFNixPlatformRuntimeState::get().deregisterJITDylib(
+                   HeaderAddr.toPtr<void *>());
+             })
+      .release();
 }
 
 ORC_RT_INTERFACE orc_rt_CWrapperFunctionResult
 __orc_rt_elfnix_register_init_sections(char *ArgData, size_t ArgSize) {
   return WrapperFunction<SPSError(SPSExecutorAddr,
                                   SPSSequence<SPSExecutorAddrRange>)>::
-    handle(ArgData, ArgSize,
-            [](ExecutorAddr HeaderAddr, std::vector<ExecutorAddrRange> &Inits) {
-              return ELFNixPlatformRuntimeState::get().registerInits(
-                  HeaderAddr, std::move(Inits));
-            })
-        .release();
+      handle(ArgData, ArgSize,
+             [](ExecutorAddr HeaderAddr,
+                std::vector<ExecutorAddrRange> &Inits) {
+               return ELFNixPlatformRuntimeState::get().registerInits(
+                   HeaderAddr, std::move(Inits));
+             })
+          .release();
 }
 
 ORC_RT_INTERFACE orc_rt_CWrapperFunctionResult
 __orc_rt_elfnix_deregister_init_sections(char *ArgData, size_t ArgSize) {
   return WrapperFunction<SPSError(SPSExecutorAddr,
                                   SPSSequence<SPSExecutorAddrRange>)>::
-    handle(ArgData, ArgSize,
-            [](ExecutorAddr HeaderAddr, std::vector<ExecutorAddrRange> &Inits) {
-              return ELFNixPlatformRuntimeState::get().deregisterInits(
-                  HeaderAddr, std::move(Inits));
-            })
-        .release();
+      handle(ArgData, ArgSize,
+             [](ExecutorAddr HeaderAddr,
+                std::vector<ExecutorAddrRange> &Inits) {
+               return ELFNixPlatformRuntimeState::get().deregisterInits(
+                   HeaderAddr, std::move(Inits));
+             })
+          .release();
 }
 
 /// Wrapper function for registering metadata on a per-object basis.
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h b/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h
index 4fbbd2b484d345..ea417ebdc04dbd 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h
@@ -146,7 +146,7 @@ class ELFNixPlatform : public Platform {
 
     void addDSOHandleSupportPasses(MaterializationResponsibility &MR,
                                    jitlink::PassConfiguration &Config,
-                                   bool IsBootstraping);
+                                   bool IsBootstrapping);
 
     void addEHAndTLVSupportPasses(MaterializationResponsibility &MR,
                                   jitlink::PassConfiguration &Config);
@@ -155,7 +155,7 @@ class ELFNixPlatform : public Platform {
                                MaterializationResponsibility &MR);
 
     Error registerInitSections(jitlink::LinkGraph &G, JITDylib &JD,
-                               bool IsBootstraping);
+                               bool IsBootstrapping);
 
     Error fixTLVSectionsAndEdges(jitlink::LinkGraph &G, JITDylib &JD);
 
@@ -208,7 +208,6 @@ class ELFNixPlatform : public Platform {
   ObjectLinkingLayer &ObjLinkingLayer;
 
   SymbolStringPtr DSOHandleSymbol;
-  ExecutorAddr DSOHandleAddr;
   std::atomic<bool> RuntimeBootstrapped{false};
 
   ExecutorAddr orc_rt_elfnix_platform_bootstrap;
diff --git a/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp b/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
index b233588b7a5d4f..e4b05cdc75113f 100644
--- a/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
@@ -284,22 +284,11 @@ ELFNixPlatform::ELFNixPlatform(
 
   // PlatformJD hasn't been 'set-up' by the platform yet (since we're creating
   // the platform now), so set it up.
-  if (auto E2 =
-          PlatformJD.define(std::make_unique<DSOHandleMaterializationUnit>(
-              *this, DSOHandleSymbol))) {
+  if (auto E2 = setupJITDylib(PlatformJD)) {
     Err = std::move(E2);
     return;
   }
 
-  auto E = ES.lookup({&PlatformJD}, DSOHandleSymbol);
-  if (auto E2 = E.takeError()) {
-    Err = std::move(E2);
-    return;
-  }
-  DSOHandleAddr = E->getAddress();
-  // RegisteredInitSymbols[&PlatformJD].add(
-  //   DSOHandleSymbol, SymbolLookupFlags::WeaklyReferencedSymbol);
-
   // Associate wrapper function tags with JIT-side function implementations.
   if (auto E2 = associateRuntimeSupportFunctions(PlatformJD)) {
     Err = std::move(E2);
@@ -313,8 +302,6 @@ ELFNixPlatform::ELFNixPlatform(
     Err = std::move(E2);
     return;
   }
-
-  JDBootstrapStates.clear();
 }
 
 Error ELFNixPlatform::associateRuntimeSupportFunctions(JITDylib &PlatformJD) {
@@ -434,7 +421,7 @@ void ELFNixPlatform::rt_recordInitializers(
   }
 
   LLVM_DEBUG({
-    dbgs() << "ELFNixPlatform::rt_pushInitializers(" << JDHeaderAddr << ") ";
+    dbgs() << "ELFNixPlatform::rt_recordInitializers(" << JDHeaderAddr << ") ";
     if (JD)
       dbgs() << "pushing initializers for " << JD->getName() << "\n";
     else
@@ -500,9 +487,11 @@ void ELFNixPlatform::rt_lookupSymbol(SendSymbolAddressFn SendResult,
 }
 
 Error ELFNixPlatform::bootstrapELFNixRuntime(JITDylib &PlatformJD) {
+  ExecutorAddr DSOHandleAddr;
   if (auto Err = lookupAndRecordAddrs(
           ES, LookupKind::Static, makeJITDylibSearchOrder(&PlatformJD),
           {
+              {DSOHandleSymbol, &DSOHandleAddr},
               {ES.intern("__orc_rt_elfnix_platform_bootstrap"),
                &orc_rt_elfnix_platform_bootstrap},
               {ES.intern("__orc_rt_elfnix_platform_shutdown"),
@@ -522,11 +511,18 @@ Error ELFNixPlatform::bootstrapELFNixRuntime(JITDylib &PlatformJD) {
           }))
     return Err;
 
-
   if (auto Err = ES.callSPSWrapper<void(SPSExecutorAddr)>(
           orc_rt_elfnix_platform_bootstrap, DSOHandleAddr))
     return Err;
 
+  for (auto KV : JDBootstrapStates) {
+    auto &JDBState = KV.second;
+    if (auto Err = ES.callSPSWrapper<void(SPSString, SPSExecutorAddr)>(
+            orc_rt_elfnix_register_jitdylib, JDBState.JDName,
+            JDBState.HeaderAddr))
+      return Err;
+  }
+
   // FIXME: Ordering is fuzzy here. We're probably best off saying
   // "behavior is undefined if code that uses the runtime is added before
   // the platform constructor returns", then move all this to the constructor.
@@ -537,18 +533,9 @@ Error ELFNixPlatform::bootstrapELFNixRuntime(JITDylib &PlatformJD) {
     DeferredPOSRs = std::move(BootstrapPOSRs);
   }
 
-  for (auto KV : JDBootstrapStates) {
-    auto &JDBState = KV.second;
-    if (auto Err = ES.callSPSWrapper<void(SPSString, SPSExecutorAddr)>(
-            orc_rt_elfnix_register_jitdylib, JDBState.JDName,
-            JDBState.HeaderAddr))
-      return Err;
-  }
-
   for (auto &D : DeferredPOSRs)
-    if (auto Err = ES.callSPSWrapper<void(
-                    SPSELFPerObjectSectionsToRegister)>(
-         orc_rt_elfnix_register_object_sections, D))
+    if (auto Err = ES.callSPSWrapper<void(SPSELFPerObjectSectionsToRegister)>(
+            orc_rt_elfnix_register_object_sections, D))
       return Err;
 
   for (auto KV : JDBootstrapStates) {
@@ -561,6 +548,7 @@ Error ELFNixPlatform::bootstrapELFNixRuntime(JITDylib &PlatformJD) {
             JDBState.Initializers))
       return Err;
   }
+
   return Error::success();
 }
 
@@ -573,12 +561,13 @@ Error ELFNixPlatform::registerPerObjectSections(
                                    "been loaded yet",
                                    inconvertibleErrorCode());
 
-  using SPSRegisterObjSectionsArgs = SPSArgList<SPSELFPerObjectSectionsToRegister>;
+  using SPSRegisterObjSectionsArgs =
+      SPSArgList<SPSELFPerObjectSectionsToRegister>;
   G.allocActions().push_back(
-        {cantFail(WrapperFunctionCall::Create<SPSRegisterObjSectionsArgs>(
-             orc_rt_elfnix_register_object_sections, POSR)),
-         cantFail(WrapperFunctionCall::Create<SPSRegisterObjSectionsArgs>(
-             orc_rt_elfnix_deregister_object_sections, POSR))});
+      {cantFail(WrapperFunctionCall::Create<SPSRegisterObjSectionsArgs>(
+           orc_rt_elfnix_register_object_sections, POSR)),
+       cantFail(WrapperFunctionCall::Create<SPSRegisterObjSectionsArgs>(
+           orc_rt_elfnix_deregister_object_sections, POSR))});
 
   return Error::success();
 }
@@ -645,10 +634,10 @@ ELFNixPlatform::ELFNixPlatformPlugin::getSyntheticSymbolDependencies(
 
 void ELFNixPlatform::ELFNixPlatformPlugin::addDSOHandleSupportPasses(
     MaterializationResponsibility &MR, jitlink::PassConfiguration &Config,
-    bool IsBootstraping) {
+    bool IsBootstrapping) {
 
   Config.PostAllocationPasses.push_back([this, &JD = MR.getTargetJITDylib(),
-                                         IsBootstraping](
+                                         IsBootstrapping](
                                             jitlink::LinkGraph &G) -> Error {
     auto I = llvm::find_if(G.defined_symbols(), [this](jitlink::Symbol *Sym) {
       return Sym->getName() == *MP.DSOHandleSymbol;
@@ -660,7 +649,7 @@ void ELFNixPlatform::ELFNixPlatformPlugin::addDSOHandleSupportPasses(
       MP.HandleAddrToJITDylib[HandleAddr] = &JD;
       MP.JITDylibToHandleAddr[&JD] = HandleAddr;
 
-      if (!IsBootstraping) {
+      if (!IsBootstrapping) {
         G.allocActions().push_back(
             {cantFail(WrapperFunctionCall::Create<
                       SPSArgList<SPSString, SPSExecutorAddr>>(
@@ -668,10 +657,10 @@ void ELFNixPlatform::ELFNixPlatformPlugin::addDSOHandleSupportPasses(
              cantFail(WrapperFunctionCall::Create<SPSArgList<SPSExecutorAddr>>(
                  MP.orc_rt_elfnix_deregister_jitdylib, HandleAddr))});
       } else {
-        // G.allocActions().push_back(
-        //   {{},
-        //    cantFail(WrapperFunctionCall::Create<SPSArgList<SPSExecutorAddr>>(
-        //         MP.orc_rt_elfnix_deregister_jitdylib, HandleAddr))});
+        /*G.allocActions().push_back(
+              {{},
+               cantFail(WrapperFunctionCall::Create<SPSArgList<SPSExecutorAddr>>(
+                   MP.orc_rt_elfnix_deregister_jitdylib, HandleAddr))});*/
         JDBootstrapState BState;
         BState.JD = &JD;
         BState.JDName = JD.getName();
@@ -742,7 +731,7 @@ void ELFNixPlatform::ELFNixPlatformPlugin::addEHAndTLVSupportPasses(
       }
 
       // Otherwise register it immediately.
-      if (auto Err = MP.registerPerObjectSections(G,POSR))
+      if (auto Err = MP.registerPerObjectSections(G, POSR))
         return Err;
     }
 
@@ -787,15 +776,13 @@ Error ELFNixPlatform::ELFNixPlatformPlugin::preserveInitSections(
 }
 
 Error ELFNixPlatform::ELFNixPlatformPlugin::registerInitSections(
-    jitlink::LinkGraph &G, JITDylib &JD, bool IsBootstraping) {
+    jitlink::LinkGraph &G, JITDylib &JD, bool IsBootstrapping) {
   SmallVector<ExecutorAddrRange> ELFNixPlatformSecs;
-  SmallVector<jitlink::Section *> InitSections;
 
   LLVM_DEBUG(dbgs() << "ELFNixPlatform::registerInitSections\n");
 
   for (auto &Sec : G.sections()) {
     if (isELFInitializerSection(Sec.getName())) {
-      InitSections.push_back(&Sec);
       jitlink::SectionRange R(Sec);
       ELFNixPlatformSecs.push_back(R.getRange());
     }
@@ -804,9 +791,9 @@ Error ELFNixPlatform::ELFNixPlatformPlugin::registerInitSections(
   // Dump the scraped inits.
   LLVM_DEBUG({
     dbgs() << "ELFNixPlatform: Scraped " << G.getName() << " init sections:\n";
-    for (auto *Sec : InitSections) {
-      jitlink::SectionRange R(*Sec);
-      dbgs() << "  " << Sec->getName() << ": " << R.getRange() << "\n";
+    for (auto &Sec : G.sections()) {
+      jitlink::SectionRange R(Sec);
+      dbgs() << "  " << Sec.getName() << ": " << R.getRange() << "\n";
     }
   });
   using SPSRegisterInitSectionsArgs =
@@ -821,16 +808,15 @@ Error ELFNixPlatform::ELFNixPlatformPlugin::registerInitSections(
     HeaderAddr = I->second;
   }
 
-  if (IsBootstraping) {
+  if (IsBootstrapping) {
     auto &JBS = MP.JDBootstrapStates[&JD];
     for (auto &I : ELFNixPlatformSecs)
       JBS.Initializers.push_back(I);
-    // G.allocActions().push_back(
-    //    {{},
-    //     cantFail(
-    //       WrapperFunctionCall::Create<SPSRegisterInitSectionsArgs>(
-    //           MP.orc_rt_elfnix_deregister_init_sections, HeaderAddr,
-    //           ELFNixPlatformSecs))});
+     /*G.allocActions().push_back(
+           {{},
+            cantFail(WrapperFunctionCall::Create<SPSRegisterInitSectionsArgs>(
+                MP.orc_rt_elfnix_deregister_init_sections, HeaderAddr,
+                ELFNixPlatformSecs))});*/
   } else {
     G.allocActions().push_back(
         {cantFail(WrapperFunctionCall::Create<SPSRegisterInitSectionsArgs>(



More information about the llvm-commits mailing list