[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 Sep 22 22:41:42 PDT 2024


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

>From 48fe06651e0a28abffeacf1d9023f6e9ce57e997 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/9] [ORC][Runtime][ELF] Improve ELF runtime with experimental
 enhancements

---
 compiler-rt/lib/orc/common.h                  |  97 ++++
 compiler-rt/lib/orc/elfnix_platform.cpp       | 434 ++++++++++++-----
 compiler-rt/lib/orc/elfnix_platform.h         |  56 +--
 .../llvm/ExecutionEngine/Orc/ELFNixPlatform.h | 126 ++---
 .../ExecutionEngine/Orc/ELFNixPlatform.cpp    | 435 ++++++++++--------
 5 files changed, 692 insertions(+), 456 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 bd76e3ed0ece7c..d8c5c5eb31d4f8 100644
--- a/compiler-rt/lib/orc/elfnix_platform.cpp
+++ b/compiler-rt/lib/orc/elfnix_platform.cpp
@@ -30,6 +30,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)
 
@@ -57,21 +58,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;
@@ -92,10 +78,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:
@@ -113,6 +107,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();
@@ -122,6 +122,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>>
@@ -132,18 +134,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;
 
@@ -207,6 +214,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)
@@ -218,28 +309,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;
 }
 
@@ -265,15 +354,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();
@@ -298,8 +388,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;
 }
 
@@ -316,24 +411,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,74 +437,140 @@ 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(JITDispatch(&__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(JITDispatch(&__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);
@@ -473,11 +616,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();
 }
@@ -488,6 +632,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 3efac4b2327f32..5ecbdf0cb9c86f 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>>;
 
 } // namespace orc_rt
 
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h b/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h
index 2c5b1104926960..b4ec36aa28ca84 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 {
@@ -155,7 +139,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);
@@ -163,7 +148,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);
 
@@ -171,11 +157,8 @@ class ELFNixPlatform : public Platform {
     ELFNixPlatform &MP;
   };
 
-  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>)>;
 
@@ -189,18 +172,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);
@@ -208,22 +183,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;
@@ -231,10 +218,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;
 };
 
@@ -266,63 +253,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 67c920a40ea2e5..6bb45d2553a036 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);
-      }
-    }
-  }
+void ELFNixPlatform::pushInitializersLoop(
+    PushInitializersSendResultFn SendResult, JITDylibSP JD) {
+  DenseMap<JITDylib *, SymbolLookupSet> NewInitSymbols;
+  DenseMap<JITDylib *, SmallVector<JITDylib *>> JDDepMap;
+  SmallVector<JITDylib *, 16> Worklist({JD.get()});
 
-  SendResult(std::move(FullInitSeq));
-}
+  ES.runSessionLocked([&]() {
+    while (!Worklist.empty()) {
+      // FIXME: Check for defunct dylibs.
 
-void ELFNixPlatform::getInitializersLookupPhase(
-    SendInitializerSequenceFn SendResult, JITDylib &JD) {
+      auto DepJD = Worklist.back();
+      Worklist.pop_back();
 
-  auto DFSLinkOrder = JD.getDFSLinkOrder();
-  if (!DFSLinkOrder) {
-    SendResult(DFSLinkOrder.takeError());
-    return;
-  }
+      // If we've already visited this JITDylib on this iteration then continue.
+      if (JDDepMap.count(DepJD))
+        continue;
 
-  DenseMap<JITDylib *, SymbolLookupSet> NewInitSymbols;
-  ES.runSessionLocked([&]() {
-    for (auto &InitJD : *DFSLinkOrder) {
-      auto RISItr = RegisteredInitSymbols.find(InitJD.get());
+      // 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,43 +597,54 @@ 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);
-}
 
-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(),
+                                    IsBootstrapping](jitlink::LinkGraph &G) {
+    return registerInitSections(G, JD, IsBootstrapping);
   });
+}
 
-  Config.PostFixupPasses.push_back(
-      [this, &JD = MR.getTargetJITDylib()](jitlink::LinkGraph &G) {
-        return registerInitSections(G, JD);
-      });
+ObjectLinkingLayer::Plugin::SyntheticSymbolDependenciesMap
+ELFNixPlatform::ELFNixPlatformPlugin::getSyntheticSymbolDependencies(
+    MaterializationResponsibility &MR) {
+  std::lock_guard<std::mutex> Lock(PluginMutex);
+  auto I = InitSymbolDeps.find(&MR);
+  if (I != InitSymbolDeps.end()) {
+    SyntheticSymbolDependenciesMap Result;
+    Result[MR.getInitializerSymbol()] = std::move(I->second);
+    InitSymbolDeps.erase(&MR);
+    return Result;
+  }
+  return SyntheticSymbolDependenciesMap();
 }
 
 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;
@@ -645,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();
   });
@@ -757,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");
@@ -766,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());
     }
   }
 
@@ -777,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 91625ef8fa2a59ea30565b9c12c7125c5eb15b19 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/9] 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 d8c5c5eb31d4f8..58699dc960ba95 100644
--- a/compiler-rt/lib/orc/elfnix_platform.cpp
+++ b/compiler-rt/lib/orc/elfnix_platform.cpp
@@ -214,6 +214,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);
 
@@ -298,14 +306,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 b4ec36aa28ca84..0fa50ed05a52dd 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h
@@ -183,7 +183,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 6bb45d2553a036..81da9c6938496e 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 a939bd1fb14ec2018b3bbaf5a47d961e72bcf256 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/9] Update and fix bootstrap process logic

---
 compiler-rt/lib/orc/common.h                  | 160 +++++++++---------
 compiler-rt/lib/orc/elfnix_platform.cpp       | 132 +++++++--------
 .../llvm/ExecutionEngine/Orc/ELFNixPlatform.h |   5 +-
 .../ExecutionEngine/Orc/ELFNixPlatform.cpp    |  92 +++++-----
 4 files changed, 183 insertions(+), 206 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 58699dc960ba95..e253107d124de4 100644
--- a/compiler-rt/lib/orc/elfnix_platform.cpp
+++ b/compiler-rt/lib/orc/elfnix_platform.cpp
@@ -29,9 +29,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
@@ -109,10 +107,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();
@@ -140,17 +138,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;
 
@@ -222,7 +220,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)) {
@@ -259,16 +258,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;
@@ -289,7 +289,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;
@@ -309,10 +310,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 {
@@ -323,10 +320,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;
@@ -362,7 +357,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()) {
@@ -388,12 +384,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,7 +466,8 @@ 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>(
@@ -501,8 +495,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())
@@ -518,9 +512,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;
@@ -560,7 +553,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)
@@ -621,7 +615,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();
 }
@@ -634,46 +628,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 0fa50ed05a52dd..239fc67bbc43c9 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h
@@ -140,7 +140,7 @@ class ELFNixPlatform : public Platform {
 
     void addDSOHandleSupportPasses(MaterializationResponsibility &MR,
                                    jitlink::PassConfiguration &Config,
-                                   bool IsBootstraping);
+                                   bool IsBootstrapping);
 
     void addEHAndTLVSupportPasses(MaterializationResponsibility &MR,
                                   jitlink::PassConfiguration &Config);
@@ -149,7 +149,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);
 
@@ -201,7 +201,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 81da9c6938496e..eb07aa126074d1 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>(

>From c869442ab48dbd83c20438bf47b382c06118a85c Mon Sep 17 00:00:00 2001
From: SahilPatidar <patidarsahil2001 at gmail.com>
Date: Mon, 12 Aug 2024 10:52:27 +0530
Subject: [PATCH 4/9] Fix Minor issues

---
 compiler-rt/lib/orc/common.h                        |  2 +-
 .../llvm/ExecutionEngine/Orc/ELFNixPlatform.h       | 10 +++++-----
 llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp     | 13 +++++++------
 3 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/compiler-rt/lib/orc/common.h b/compiler-rt/lib/orc/common.h
index 0db540c4dfde8f..e659f1ae13e954 100644
--- a/compiler-rt/lib/orc/common.h
+++ b/compiler-rt/lib/orc/common.h
@@ -17,8 +17,8 @@
 #include "error.h"
 #include "executor_address.h"
 #include "orc_rt/c_api.h"
-#include <type_traits>
 #include <algorithm>
+#include <type_traits>
 #include <vector>
 
 /// This macro should be used to define tags that will be associated with
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h b/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h
index 239fc67bbc43c9..660be6f1b49a6a 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h
@@ -172,7 +172,8 @@ class ELFNixPlatform : public Platform {
   // Associate ELFNixPlatform JIT-side runtime support functions with handlers.
   Error associateRuntimeSupportFunctions(JITDylib &PlatformJD);
 
-  void pushInitializersLoop(PushInitializersSendResultFn SendResult, JITDylibSP JD);
+  void pushInitializersLoop(PushInitializersSendResultFn SendResult,
+                            JITDylibSP JD);
 
   void rt_recordInitializers(PushInitializersSendResultFn SendResult,
                              ExecutorAddr JDHeader);
@@ -183,8 +184,8 @@ class ELFNixPlatform : public Platform {
   // Records the addresses of runtime symbols used by the platform.
   Error bootstrapELFNixRuntime(JITDylib &PlatformJD);
 
-  Error registerPerObjectSections(
-      jitlink::LinkGraph &G,const ELFPerObjectSectionsToRegister &POSR);
+  Error registerPerObjectSections(jitlink::LinkGraph &G,
+                                  const ELFPerObjectSectionsToRegister &POSR);
 
   Expected<uint64_t> createPThreadKey();
 
@@ -260,5 +261,4 @@ using SPSELFNixJITDylibDepInfoMap =
 } // end namespace orc
 } // end namespace llvm
 
-#endif // LLVM_EXECUTIONENGINE_ORC_ELFNIXPLATFORM_H
-
+#endif // LLVM_EXECUTIONENGINE_ORC_ELFNIXPLATFORM_H
\ No newline at end of file
diff --git a/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp b/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
index eb07aa126074d1..dbf2aa63607f31 100644
--- a/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
@@ -1,4 +1,5 @@
-//===------ ELFNixPlatform.cpp - Utilities for executing ELFNix 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.
@@ -812,11 +813,11 @@ Error ELFNixPlatform::ELFNixPlatformPlugin::registerInitSections(
     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>(

>From 8c6a82121a5c8bf8bd38d6fa15c87124ab866c05 Mon Sep 17 00:00:00 2001
From: SahilPatidar <patidarsahil2001 at gmail.com>
Date: Thu, 22 Aug 2024 11:52:30 +0530
Subject: [PATCH 5/9] Refactor Bootstrap Logic to Utilize Mach-O Specific Logic

---
 compiler-rt/lib/orc/elfnix_platform.cpp       |   9 +-
 .../llvm/ExecutionEngine/Orc/ELFNixPlatform.h | 107 +++--
 .../ExecutionEngine/Orc/ELFNixPlatform.cpp    | 435 ++++++++++++------
 3 files changed, 379 insertions(+), 172 deletions(-)

diff --git a/compiler-rt/lib/orc/elfnix_platform.cpp b/compiler-rt/lib/orc/elfnix_platform.cpp
index e253107d124de4..f0467f9a3ade51 100644
--- a/compiler-rt/lib/orc/elfnix_platform.cpp
+++ b/compiler-rt/lib/orc/elfnix_platform.cpp
@@ -622,8 +622,13 @@ __orc_rt_elfnix_platform_bootstrap(char *ArgData, size_t ArgSize) {
 
 ORC_RT_INTERFACE orc_rt_CWrapperFunctionResult
 __orc_rt_elfnix_platform_shutdown(char *ArgData, size_t ArgSize) {
-  ELFNixPlatformRuntimeState::destroy();
-  return WrapperFunctionResult().release();
+  return WrapperFunction<SPSError()>::handle(
+             ArgData, ArgSize,
+             []() {
+               ELFNixPlatformRuntimeState::destroy();
+               return Error::success();
+             })
+      .release();
 }
 
 ORC_RT_INTERFACE orc_rt_CWrapperFunctionResult
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h b/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h
index 660be6f1b49a6a..22500e8f6fad44 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h
@@ -21,6 +21,7 @@
 
 #include <future>
 #include <thread>
+#include <unordered_map>
 #include <vector>
 
 namespace llvm {
@@ -35,6 +36,34 @@ using ELFNixJITDylibDepInfo = std::vector<ExecutorAddr>;
 using ELFNixJITDylibDepInfoMap =
     std::vector<std::pair<ExecutorAddr, ELFNixJITDylibDepInfo>>;
 
+struct RuntimeFunction {
+  RuntimeFunction(SymbolStringPtr Name) : Name(std::move(Name)) {}
+  SymbolStringPtr Name;
+  ExecutorAddr Addr;
+};
+
+struct FunctionPairKeyHash {
+  std::size_t
+  operator()(const std::pair<RuntimeFunction *, RuntimeFunction *> &key) const {
+    return std::hash<void *>()(key.first->Addr.toPtr<void *>()) ^
+           std::hash<void *>()(key.second->Addr.toPtr<void *>());
+  }
+};
+
+struct FunctionPairKeyEqual {
+  std::size_t
+  operator()(const std::pair<RuntimeFunction *, RuntimeFunction *> &lhs,
+             const std::pair<RuntimeFunction *, RuntimeFunction *> &rhs) const {
+    return lhs.first == rhs.first && lhs.second == rhs.second;
+  }
+};
+
+using DeferredRuntimeFnMap = std::unordered_map<
+    std::pair<RuntimeFunction *, RuntimeFunction *>,
+    SmallVector<std::pair<shared::WrapperFunctionCall::ArgDataBufferType,
+                          shared::WrapperFunctionCall::ArgDataBufferType>>,
+    FunctionPairKeyHash, FunctionPairKeyEqual>;
+
 /// Mediates between ELFNix initialization and ExecutionSession state.
 class ELFNixPlatform : public Platform {
 public:
@@ -110,6 +139,23 @@ class ELFNixPlatform : public Platform {
   standardRuntimeUtilityAliases();
 
 private:
+  // Data needed for bootstrap only.
+  struct BootstrapInfo {
+    std::mutex Mutex;
+    std::condition_variable CV;
+    size_t ActiveGraphs = 0;
+    ExecutorAddr ELFNixHeaderAddr;
+    DeferredRuntimeFnMap DeferredRTFnMap;
+
+    void addArgumentsToRTFnMap(
+        RuntimeFunction *func1, RuntimeFunction *func2,
+        const shared::WrapperFunctionCall::ArgDataBufferType &arg1,
+        const shared::WrapperFunctionCall::ArgDataBufferType &arg2) {
+      auto &argList = DeferredRTFnMap[std::make_pair(func1, func2)];
+      argList.emplace_back(arg1, arg2);
+    }
+  };
+
   // The ELFNixPlatformPlugin scans/modifies LinkGraphs to support ELF
   // platform features including initializers, exceptions, TLV, and language
   // runtime registration.
@@ -135,15 +181,22 @@ class ELFNixPlatform : public Platform {
                                      ResourceKey SrcKey) override {}
 
   private:
+    using InitSymbolDepMap =
+        DenseMap<MaterializationResponsibility *, JITLinkSymbolSet>;
+
+    Error bootstrapPipelineStart(jitlink::LinkGraph &G);
+    Error bootstrapPipelineRecordRuntimeFunctions(jitlink::LinkGraph &G);
+    Error bootstrapPipelineEnd(jitlink::LinkGraph &G);
+
     void addInitializerSupportPasses(MaterializationResponsibility &MR,
                                      jitlink::PassConfiguration &Config);
 
     void addDSOHandleSupportPasses(MaterializationResponsibility &MR,
-                                   jitlink::PassConfiguration &Config,
-                                   bool IsBootstrapping);
+                                   jitlink::PassConfiguration &Config);
 
     void addEHAndTLVSupportPasses(MaterializationResponsibility &MR,
-                                  jitlink::PassConfiguration &Config);
+                                  jitlink::PassConfiguration &Config,
+                                  bool IsBootstrapping);
 
     Error preserveInitSections(jitlink::LinkGraph &G,
                                MaterializationResponsibility &MR);
@@ -181,38 +234,36 @@ class ELFNixPlatform : public Platform {
   void rt_lookupSymbol(SendSymbolAddressFn SendResult, ExecutorAddr Handle,
                        StringRef SymbolName);
 
-  // Records the addresses of runtime symbols used by the platform.
-  Error bootstrapELFNixRuntime(JITDylib &PlatformJD);
-
   Error registerPerObjectSections(jitlink::LinkGraph &G,
-                                  const ELFPerObjectSectionsToRegister &POSR);
+                                  const ELFPerObjectSectionsToRegister &POSR,
+                                  bool IsBootstrapping);
 
   Expected<uint64_t> createPThreadKey();
 
-  struct JDBootstrapState {
-    JITDylib *JD = nullptr;
-    std::string JDName;
-    ExecutorAddr HeaderAddr;
-    SmallVector<ExecutorAddrRange> Initializers;
-  };
-
-  std::map<JITDylib *, JDBootstrapState> JDBootstrapStates;
-
   ExecutionSession &ES;
+  JITDylib &PlatformJD;
   ObjectLinkingLayer &ObjLinkingLayer;
 
   SymbolStringPtr DSOHandleSymbol;
-  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;
+
+  RuntimeFunction PlatformBootstrap{
+      ES.intern("__orc_rt_elfnix_platform_bootstrap")};
+  RuntimeFunction PlatformShutdown{
+      ES.intern("__orc_rt_elfnix_platform_shutdown")};
+  RuntimeFunction RegisterJITDylib{
+      ES.intern("__orc_rt_elfnix_register_jitdylib")};
+  RuntimeFunction DeregisterJITDylib{
+      ES.intern("__orc_rt_elfnix_deregister_jitdylib")};
+  RuntimeFunction RegisterObjectSections{
+      ES.intern("__orc_rt_elfnix_register_object_sections")};
+  RuntimeFunction DeregisterObjectSections{
+      ES.intern("__orc_rt_elfnix_deregister_object_sections")};
+  RuntimeFunction RegisterInitSections{
+      ES.intern("__orc_rt_elfnix_register_init_sections")};
+  RuntimeFunction DeregisterInitSections{
+      ES.intern("__orc_rt_elfnix_deregister_init_sections")};
+  RuntimeFunction CreatePThreadKey{
+      ES.intern("__orc_rt_elfnix_create_pthread_key")};
 
   DenseMap<JITDylib *, SymbolLookupSet> RegisteredInitSymbols;
 
@@ -224,6 +275,8 @@ class ELFNixPlatform : public Platform {
   DenseMap<ExecutorAddr, JITDylib *> HandleAddrToJITDylib;
   DenseMap<JITDylib *, ExecutorAddr> JITDylibToHandleAddr;
   DenseMap<JITDylib *, uint64_t> JITDylibToPThreadKey;
+
+  std::atomic<BootstrapInfo *> Bootstrap;
 };
 
 namespace shared {
diff --git a/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp b/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
index dbf2aa63607f31..839a10c3723082 100644
--- a/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
@@ -30,6 +30,125 @@ using namespace llvm::orc::shared;
 
 namespace {
 
+template <typename SPSSerializer, typename... ArgTs>
+shared::WrapperFunctionCall::ArgDataBufferType
+getArgDataBufferType(const ArgTs &...Args) {
+  shared::WrapperFunctionCall::ArgDataBufferType ArgData;
+  ArgData.resize(SPSSerializer::size(Args...));
+  SPSOutputBuffer OB(ArgData.empty() ? nullptr : ArgData.data(),
+                     ArgData.size());
+  if (SPSSerializer::serialize(OB, Args...))
+    return ArgData;
+  return {};
+}
+
+std::unique_ptr<jitlink::LinkGraph> createPlatformGraph(ELFNixPlatform &MOP,
+                                                        std::string Name) {
+  unsigned PointerSize;
+  llvm::endianness Endianness;
+  const auto &TT = MOP.getExecutionSession().getTargetTriple();
+
+  switch (TT.getArch()) {
+  case Triple::x86_64:
+    PointerSize = 8;
+    Endianness = llvm::endianness::little;
+    break;
+  case Triple::aarch64:
+    PointerSize = 8;
+    Endianness = llvm::endianness::little;
+    break;
+  case Triple::ppc64:
+    PointerSize = 8;
+    Endianness = llvm::endianness::big;
+    break;
+  case Triple::ppc64le:
+    PointerSize = 8;
+    Endianness = llvm::endianness::little;
+    break;
+  default:
+    llvm_unreachable("Unrecognized architecture");
+  }
+
+  return std::make_unique<jitlink::LinkGraph>(std::move(Name), TT, PointerSize,
+                                              Endianness,
+                                              jitlink::getGenericEdgeKindName);
+}
+
+// Creates a Bootstrap-Complete LinkGraph to run deferred actions.
+class ELFNixPlatformCompleteBootstrapMaterializationUnit
+    : public MaterializationUnit {
+public:
+  ELFNixPlatformCompleteBootstrapMaterializationUnit(
+      ELFNixPlatform &MOP, StringRef PlatformJDName,
+      SymbolStringPtr CompleteBootstrapSymbol, DeferredRuntimeFnMap DeferredAAs,
+      ExecutorAddr ELFNixHeaderAddr, ExecutorAddr PlatformBootstrap,
+      ExecutorAddr PlatformShutdown, ExecutorAddr RegisterJITDylib,
+      ExecutorAddr DeregisterJITDylib)
+      : MaterializationUnit(
+            {{{CompleteBootstrapSymbol, JITSymbolFlags::None}}, nullptr}),
+        MOP(MOP), PlatformJDName(PlatformJDName),
+        CompleteBootstrapSymbol(std::move(CompleteBootstrapSymbol)),
+        DeferredAAsMap(std::move(DeferredAAs)),
+        ELFNixHeaderAddr(ELFNixHeaderAddr),
+        PlatformBootstrap(PlatformBootstrap),
+        PlatformShutdown(PlatformShutdown), RegisterJITDylib(RegisterJITDylib),
+        DeregisterJITDylib(DeregisterJITDylib) {}
+
+  StringRef getName() const override {
+    return "ELFNixPlatformCompleteBootstrap";
+  }
+
+  void materialize(std::unique_ptr<MaterializationResponsibility> R) override {
+    using namespace jitlink;
+    auto G = createPlatformGraph(MOP, "<OrcRTCompleteBootstrap>");
+    auto &PlaceholderSection =
+        G->createSection("__orc_rt_cplt_bs", MemProt::Read);
+    auto &PlaceholderBlock =
+        G->createZeroFillBlock(PlaceholderSection, 1, ExecutorAddr(), 1, 0);
+    G->addDefinedSymbol(PlaceholderBlock, 0, *CompleteBootstrapSymbol, 1,
+                        Linkage::Strong, Scope::Hidden, false, true);
+
+    // 1. Bootstrap the platform support code.
+    G->allocActions().push_back(
+        {cantFail(WrapperFunctionCall::Create<SPSArgList<SPSExecutorAddr>>(
+             PlatformBootstrap, ELFNixHeaderAddr)),
+         cantFail(
+             WrapperFunctionCall::Create<SPSArgList<>>(PlatformShutdown))});
+
+    // 2. Register the platform JITDylib.
+    G->allocActions().push_back(
+        {cantFail(WrapperFunctionCall::Create<
+                  SPSArgList<SPSString, SPSExecutorAddr>>(
+             RegisterJITDylib, PlatformJDName, ELFNixHeaderAddr)),
+         cantFail(WrapperFunctionCall::Create<SPSArgList<SPSExecutorAddr>>(
+             DeregisterJITDylib, ELFNixHeaderAddr))});
+
+    // 4. Add the deferred actions to the graph.
+    for (auto &[Fn, CallDatas] : DeferredAAsMap) {
+      for (auto &CallData : CallDatas) {
+        G->allocActions().push_back(
+            {WrapperFunctionCall(Fn.first->Addr, std::move(CallData.first)),
+             WrapperFunctionCall(Fn.second->Addr, std::move(CallData.second))});
+      }
+    }
+
+    MOP.getObjectLinkingLayer().emit(std::move(R), std::move(G));
+  }
+
+  void discard(const JITDylib &JD, const SymbolStringPtr &Sym) override {}
+
+private:
+  ELFNixPlatform &MOP;
+  StringRef PlatformJDName;
+  SymbolStringPtr CompleteBootstrapSymbol;
+  DeferredRuntimeFnMap DeferredAAsMap;
+  ExecutorAddr ELFNixHeaderAddr;
+  ExecutorAddr PlatformBootstrap;
+  ExecutorAddr PlatformShutdown;
+  ExecutorAddr RegisterJITDylib;
+  ExecutorAddr DeregisterJITDylib;
+};
+
 class DSOHandleMaterializationUnit : public MaterializationUnit {
 public:
   DSOHandleMaterializationUnit(ELFNixPlatform &ENP,
@@ -276,13 +395,16 @@ ELFNixPlatform::ELFNixPlatform(
     ExecutionSession &ES, ObjectLinkingLayer &ObjLinkingLayer,
     JITDylib &PlatformJD,
     std::unique_ptr<DefinitionGenerator> OrcRuntimeGenerator, Error &Err)
-    : ES(ES), ObjLinkingLayer(ObjLinkingLayer),
+    : ES(ES), PlatformJD(PlatformJD), ObjLinkingLayer(ObjLinkingLayer),
       DSOHandleSymbol(ES.intern("__dso_handle")) {
   ErrorAsOutParameter _(&Err);
   ObjLinkingLayer.addPlugin(std::make_unique<ELFNixPlatformPlugin>(*this));
 
   PlatformJD.addGenerator(std::move(OrcRuntimeGenerator));
 
+  BootstrapInfo BI;
+  Bootstrap = &BI;
+
   // 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)) {
@@ -290,16 +412,44 @@ ELFNixPlatform::ELFNixPlatform(
     return;
   }
 
-  // Associate wrapper function tags with JIT-side function implementations.
-  if (auto E2 = associateRuntimeSupportFunctions(PlatformJD)) {
-    Err = std::move(E2);
+  // Step (2) Request runtime registration functions to trigger
+  // materialization..
+  if ((Err = ES.lookup(
+                   makeJITDylibSearchOrder(&PlatformJD),
+                   SymbolLookupSet(
+                       {PlatformBootstrap.Name, PlatformShutdown.Name,
+                        RegisterJITDylib.Name, DeregisterJITDylib.Name,
+                        RegisterInitSections.Name, DeregisterInitSections.Name,
+                        RegisterObjectSections.Name,
+                        DeregisterObjectSections.Name, CreatePThreadKey.Name}))
+                 .takeError()))
     return;
+
+  // Step (3) Wait for any incidental linker work to complete.
+  {
+    std::unique_lock<std::mutex> Lock(BI.Mutex);
+    BI.CV.wait(Lock, [&]() { return BI.ActiveGraphs == 0; });
+    Bootstrap = nullptr;
   }
 
-  // Lookup addresses of runtime functions callable by the platform,
-  // call the platform bootstrap function to initialize the platform-state
-  // object in the executor.
-  if (auto E2 = bootstrapELFNixRuntime(PlatformJD)) {
+  // Step (4) Add complete-bootstrap materialization unit and request.
+  auto BootstrapCompleteSymbol =
+      ES.intern("__orc_rt_elfnix_complete_bootstrap");
+  if ((Err = PlatformJD.define(
+           std::make_unique<ELFNixPlatformCompleteBootstrapMaterializationUnit>(
+               *this, PlatformJD.getName(), BootstrapCompleteSymbol,
+               std::move(BI.DeferredRTFnMap), BI.ELFNixHeaderAddr,
+               PlatformBootstrap.Addr, PlatformShutdown.Addr,
+               RegisterJITDylib.Addr, DeregisterJITDylib.Addr))))
+    return;
+  if ((Err = ES.lookup(makeJITDylibSearchOrder(
+                           &PlatformJD, JITDylibLookupFlags::MatchAllSymbols),
+                       std::move(BootstrapCompleteSymbol))
+                 .takeError()))
+    return;
+
+  // Associate wrapper function tags with JIT-side function implementations.
+  if (auto E2 = associateRuntimeSupportFunctions(PlatformJD)) {
     Err = std::move(E2);
     return;
   }
@@ -487,94 +637,98 @@ void ELFNixPlatform::rt_lookupSymbol(SendSymbolAddressFn SendResult,
       RtLookupNotifyComplete(std::move(SendResult)), NoDependenciesToRegister);
 }
 
-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"),
-               &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;
-
-  if (auto Err = ES.callSPSWrapper<void(SPSExecutorAddr)>(
-          orc_rt_elfnix_platform_bootstrap, DSOHandleAddr))
-    return Err;
+Error ELFNixPlatform::ELFNixPlatformPlugin::bootstrapPipelineStart(
+    jitlink::LinkGraph &G) {
+  // Increment the active graphs count in BootstrapInfo.
+  std::lock_guard<std::mutex> Lock(MP.Bootstrap.load()->Mutex);
+  ++MP.Bootstrap.load()->ActiveGraphs;
+  return Error::success();
+}
 
-  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;
+Error ELFNixPlatform::ELFNixPlatformPlugin::
+    bootstrapPipelineRecordRuntimeFunctions(jitlink::LinkGraph &G) {
+  // Record bootstrap function names.
+  std::pair<StringRef, ExecutorAddr *> RuntimeSymbols[] = {
+      {*MP.DSOHandleSymbol, &MP.Bootstrap.load()->ELFNixHeaderAddr},
+      {*MP.PlatformBootstrap.Name, &MP.PlatformBootstrap.Addr},
+      {*MP.PlatformShutdown.Name, &MP.PlatformShutdown.Addr},
+      {*MP.RegisterJITDylib.Name, &MP.RegisterJITDylib.Addr},
+      {*MP.DeregisterJITDylib.Name, &MP.DeregisterJITDylib.Addr},
+      {*MP.RegisterObjectSections.Name, &MP.RegisterObjectSections.Addr},
+      {*MP.DeregisterObjectSections.Name, &MP.DeregisterObjectSections.Addr},
+      {*MP.RegisterInitSections.Name, &MP.RegisterInitSections.Addr},
+      {*MP.DeregisterInitSections.Name, &MP.DeregisterInitSections.Addr},
+      {*MP.CreatePThreadKey.Name, &MP.CreatePThreadKey.Addr}};
+
+  bool RegisterELFNixHeader = false;
+
+  for (auto *Sym : G.defined_symbols()) {
+    for (auto &RTSym : RuntimeSymbols) {
+      if (Sym->hasName() && Sym->getName() == RTSym.first) {
+        if (*RTSym.second)
+          return make_error<StringError>(
+              "Duplicate " + RTSym.first +
+                  " detected during ELFNixPlatform bootstrap",
+              inconvertibleErrorCode());
+
+        if (Sym->getName() == *MP.DSOHandleSymbol)
+          RegisterELFNixHeader = true;
+
+        *RTSym.second = Sym->getAddress();
+      }
+    }
   }
 
-  // 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.
-  RuntimeBootstrapped = true;
-  std::vector<ELFPerObjectSectionsToRegister> DeferredPOSRs;
-  {
-    std::lock_guard<std::mutex> Lock(PlatformMutex);
-    DeferredPOSRs = std::move(BootstrapPOSRs);
+  if (RegisterELFNixHeader) {
+    // If this graph defines the elfnix header symbol then create the internal
+    // mapping between it and PlatformJD.
+    std::lock_guard<std::mutex> Lock(MP.PlatformMutex);
+    MP.JITDylibToHandleAddr[&MP.PlatformJD] =
+        MP.Bootstrap.load()->ELFNixHeaderAddr;
+    MP.HandleAddrToJITDylib[MP.Bootstrap.load()->ELFNixHeaderAddr] =
+        &MP.PlatformJD;
   }
 
-  for (auto &D : DeferredPOSRs)
-    if (auto Err = ES.callSPSWrapper<void(SPSELFPerObjectSectionsToRegister)>(
-            orc_rt_elfnix_register_object_sections, D))
-      return Err;
-
-  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();
+}
 
+Error ELFNixPlatform::ELFNixPlatformPlugin::bootstrapPipelineEnd(
+    jitlink::LinkGraph &G) {
+  std::lock_guard<std::mutex> Lock(MP.Bootstrap.load()->Mutex);
+  assert(MP.Bootstrap && "DeferredAAs reset before bootstrap completed");
+  --MP.Bootstrap.load()->ActiveGraphs;
+  // Notify Bootstrap->CV while holding the mutex because the mutex is
+  // also keeping Bootstrap->CV alive.
+  if (MP.Bootstrap.load()->ActiveGraphs == 0)
+    MP.Bootstrap.load()->CV.notify_all();
   return Error::success();
 }
 
 Error ELFNixPlatform::registerPerObjectSections(
-    jitlink::LinkGraph &G, const ELFPerObjectSectionsToRegister &POSR) {
+    jitlink::LinkGraph &G, const ELFPerObjectSectionsToRegister &POSR,
+    bool IsBootstrapping) {
+  using SPSRegisterPerObjSectionsArgs =
+      SPSArgList<SPSELFPerObjectSectionsToRegister>;
 
-  if (!orc_rt_elfnix_register_object_sections)
-    return make_error<StringError>("Attempting to register per-object "
-                                   "sections, but runtime support has not "
-                                   "been loaded yet",
-                                   inconvertibleErrorCode());
+  if (LLVM_UNLIKELY(IsBootstrapping)) {
+    Bootstrap.load()->addArgumentsToRTFnMap(
+        &RegisterObjectSections, &DeregisterObjectSections,
+        getArgDataBufferType<SPSRegisterPerObjSectionsArgs>(POSR),
+        getArgDataBufferType<SPSRegisterPerObjSectionsArgs>(POSR));
+    return Error::success();
+  }
 
-  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<SPSRegisterPerObjSectionsArgs>(
+           RegisterObjectSections.Addr, POSR)),
+       cantFail(WrapperFunctionCall::Create<SPSRegisterPerObjSectionsArgs>(
+           DeregisterObjectSections.Addr, POSR))});
 
   return Error::success();
 }
 
 Expected<uint64_t> ELFNixPlatform::createPThreadKey() {
-  if (!orc_rt_elfnix_create_pthread_key)
+  if (!CreatePThreadKey.Addr)
     return make_error<StringError>(
         "Attempting to create pthread key in target, but runtime support has "
         "not been loaded yet",
@@ -582,7 +736,7 @@ Expected<uint64_t> ELFNixPlatform::createPThreadKey() {
 
   Expected<uint64_t> Result(0);
   if (auto Err = ES.callSPSWrapper<SPSExpected<uint64_t>(void)>(
-          orc_rt_elfnix_create_pthread_key, Result))
+          CreatePThreadKey.Addr, Result))
     return std::move(Err);
   return Result;
 }
@@ -590,18 +744,31 @@ Expected<uint64_t> ELFNixPlatform::createPThreadKey() {
 void ELFNixPlatform::ELFNixPlatformPlugin::modifyPassConfig(
     MaterializationResponsibility &MR, jitlink::LinkGraph &LG,
     jitlink::PassConfiguration &Config) {
+  using namespace jitlink;
+
+  bool InBootstrapPhase =
+      &MR.getTargetJITDylib() == &MP.PlatformJD && MP.Bootstrap;
 
-  bool IsBootstrapping = !MP.RuntimeBootstrapped.load();
+  // If we're in the bootstrap phase then increment the active graphs.
+  if (InBootstrapPhase) {
+    Config.PrePrunePasses.push_back(
+        [this](LinkGraph &G) { return bootstrapPipelineStart(G); });
+    Config.PostAllocationPasses.push_back([this](LinkGraph &G) {
+      return bootstrapPipelineRecordRuntimeFunctions(G);
+    });
+  }
 
-  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);
+  // If the initializer symbol is the __dso_handle symbol then just add
+  // the DSO handle support passes.
+  if (auto InitSymbol = MR.getInitializerSymbol()) {
+    if (InitSymbol == MP.DSOHandleSymbol && !InBootstrapPhase) {
+      addDSOHandleSupportPasses(MR, Config);
       // The DSOHandle materialization unit doesn't require any other
       // support, so we can bail out early.
       return;
     }
+
+    /// Preserve init sections.
     Config.PrePrunePasses.push_back(
         [this, &MR](jitlink::LinkGraph &G) -> Error {
           if (auto Err = preserveInitSections(G, MR))
@@ -611,12 +778,19 @@ void ELFNixPlatform::ELFNixPlatformPlugin::modifyPassConfig(
   }
 
   // Add passes for eh-frame and TLV support.
-  addEHAndTLVSupportPasses(MR, Config);
+  addEHAndTLVSupportPasses(MR, Config, InBootstrapPhase);
 
+  // If the object contains initializers then add passes to record them.
   Config.PostFixupPasses.push_back([this, &JD = MR.getTargetJITDylib(),
-                                    IsBootstrapping](jitlink::LinkGraph &G) {
-    return registerInitSections(G, JD, IsBootstrapping);
+                                    InBootstrapPhase](jitlink::LinkGraph &G) {
+    return registerInitSections(G, JD, InBootstrapPhase);
   });
+
+  // If we're in the bootstrap phase then steal allocation actions and then
+  // decrement the active graphs.
+  if (InBootstrapPhase)
+    Config.PostFixupPasses.push_back(
+        [this](LinkGraph &G) { return bootstrapPipelineEnd(G); });
 }
 
 ObjectLinkingLayer::Plugin::SyntheticSymbolDependenciesMap
@@ -634,11 +808,9 @@ ELFNixPlatform::ELFNixPlatformPlugin::getSyntheticSymbolDependencies(
 }
 
 void ELFNixPlatform::ELFNixPlatformPlugin::addDSOHandleSupportPasses(
-    MaterializationResponsibility &MR, jitlink::PassConfiguration &Config,
-    bool IsBootstrapping) {
+    MaterializationResponsibility &MR, jitlink::PassConfiguration &Config) {
 
-  Config.PostAllocationPasses.push_back([this, &JD = MR.getTargetJITDylib(),
-                                         IsBootstrapping](
+  Config.PostAllocationPasses.push_back([this, &JD = MR.getTargetJITDylib()](
                                             jitlink::LinkGraph &G) -> Error {
     auto I = llvm::find_if(G.defined_symbols(), [this](jitlink::Symbol *Sym) {
       return Sym->getName() == *MP.DSOHandleSymbol;
@@ -650,31 +822,20 @@ void ELFNixPlatform::ELFNixPlatformPlugin::addDSOHandleSupportPasses(
       MP.HandleAddrToJITDylib[HandleAddr] = &JD;
       MP.JITDylibToHandleAddr[&JD] = HandleAddr;
 
-      if (!IsBootstrapping) {
-        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);
-      }
+      G.allocActions().push_back(
+          {cantFail(WrapperFunctionCall::Create<
+                    SPSArgList<SPSString, SPSExecutorAddr>>(
+               MP.RegisterJITDylib.Addr, JD.getName(), HandleAddr)),
+           cantFail(WrapperFunctionCall::Create<SPSArgList<SPSExecutorAddr>>(
+               MP.DeregisterJITDylib.Addr, HandleAddr))});
     }
     return Error::success();
   });
 }
 
 void ELFNixPlatform::ELFNixPlatformPlugin::addEHAndTLVSupportPasses(
-    MaterializationResponsibility &MR, jitlink::PassConfiguration &Config) {
+    MaterializationResponsibility &MR, jitlink::PassConfiguration &Config,
+    bool IsBootstrapping) {
 
   // Insert TLV lowering at the start of the PostPrunePasses, since we want
   // it to run before GOT/PLT lowering.
@@ -688,7 +849,8 @@ void ELFNixPlatform::ELFNixPlatformPlugin::addEHAndTLVSupportPasses(
 
   // Add a pass to register the final addresses of the eh-frame and TLV sections
   // with the runtime.
-  Config.PostFixupPasses.push_back([this](jitlink::LinkGraph &G) -> Error {
+  Config.PostFixupPasses.push_back([this, IsBootstrapping](
+                                       jitlink::LinkGraph &G) -> Error {
     ELFPerObjectSectionsToRegister POSR;
 
     if (auto *EHFrameSection = G.findSectionByName(ELFEHFrameSectionName)) {
@@ -722,17 +884,7 @@ void ELFNixPlatform::ELFNixPlatformPlugin::addEHAndTLVSupportPasses(
     }
 
     if (POSR.EHFrameSection.Start || POSR.ThreadDataSection.Start) {
-
-      // If we're still bootstrapping the runtime then just record this
-      // frame for now.
-      if (!MP.RuntimeBootstrapped) {
-        std::lock_guard<std::mutex> Lock(MP.PlatformMutex);
-        MP.BootstrapPOSRs.push_back(POSR);
-        return Error::success();
-      }
-
-      // Otherwise register it immediately.
-      if (auto Err = MP.registerPerObjectSections(G, POSR))
+      if (auto Err = MP.registerPerObjectSections(G, POSR, IsBootstrapping))
         return Err;
     }
 
@@ -779,7 +931,6 @@ Error ELFNixPlatform::ELFNixPlatformPlugin::preserveInitSections(
 Error ELFNixPlatform::ELFNixPlatformPlugin::registerInitSections(
     jitlink::LinkGraph &G, JITDylib &JD, bool IsBootstrapping) {
   SmallVector<ExecutorAddrRange> ELFNixPlatformSecs;
-
   LLVM_DEBUG(dbgs() << "ELFNixPlatform::registerInitSections\n");
 
   for (auto &Sec : G.sections()) {
@@ -797,8 +948,6 @@ Error ELFNixPlatform::ELFNixPlatformPlugin::registerInitSections(
       dbgs() << "  " << Sec.getName() << ": " << R.getRange() << "\n";
     }
   });
-  using SPSRegisterInitSectionsArgs =
-      SPSArgList<SPSExecutorAddr, SPSSequence<SPSExecutorAddrRange>>;
 
   ExecutorAddr HeaderAddr;
   {
@@ -809,25 +958,25 @@ Error ELFNixPlatform::ELFNixPlatformPlugin::registerInitSections(
     HeaderAddr = I->second;
   }
 
-  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))});*/
-  } 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))});
+  using SPSRegisterInitSectionsArgs =
+      SPSArgList<SPSExecutorAddr, SPSSequence<SPSExecutorAddrRange>>;
+
+  if (LLVM_UNLIKELY(IsBootstrapping)) {
+    MP.Bootstrap.load()->addArgumentsToRTFnMap(
+        &MP.RegisterInitSections, &MP.DeregisterInitSections,
+        getArgDataBufferType<SPSRegisterInitSectionsArgs>(HeaderAddr,
+                                                          ELFNixPlatformSecs),
+        getArgDataBufferType<SPSRegisterInitSectionsArgs>(HeaderAddr,
+                                                          ELFNixPlatformSecs));
+    return Error::success();
   }
 
+  G.allocActions().push_back(
+      {cantFail(WrapperFunctionCall::Create<SPSRegisterInitSectionsArgs>(
+           MP.RegisterInitSections.Addr, HeaderAddr, ELFNixPlatformSecs)),
+       cantFail(WrapperFunctionCall::Create<SPSRegisterInitSectionsArgs>(
+           MP.DeregisterInitSections.Addr, HeaderAddr, ELFNixPlatformSecs))});
+
   return Error::success();
 }
 

>From 714d6a71d09b79992a5957d765dc09c089f8465d Mon Sep 17 00:00:00 2001
From: SahilPatidar <patidarsahil2001 at gmail.com>
Date: Tue, 10 Sep 2024 12:02:19 +0530
Subject: [PATCH 6/9] Address minor issue caused by recent main branch updates

---
 compiler-rt/lib/orc/common.h | 24 ++++++++++++------------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/compiler-rt/lib/orc/common.h b/compiler-rt/lib/orc/common.h
index e659f1ae13e954..3518adc3dd2c67 100644
--- a/compiler-rt/lib/orc/common.h
+++ b/compiler-rt/lib/orc/common.h
@@ -54,7 +54,7 @@ __orc_rt_jit_dispatch(__orc_rt_Opaque *DispatchCtx, const void *FnTag,
 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)); }
+  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(); }
@@ -65,7 +65,7 @@ template <typename RecordElement> class RecordSectionsTracker {
   /// Process all new sections.
   template <typename ProcessSectionFunc>
   std::enable_if_t<std::is_void_v<
-      std::invoke_result_t<ProcessSectionFunc, __orc_rt::span<RecordElement>>>>
+      std::invoke_result_t<ProcessSectionFunc, orc_rt::span<RecordElement>>>>
   processNewSections(ProcessSectionFunc &&ProcessSection) {
     for (auto &Sec : New)
       ProcessSection(Sec);
@@ -78,10 +78,10 @@ template <typename RecordElement> class RecordSectionsTracker {
   /// list.
   template <typename ProcessSectionFunc>
   std::enable_if_t<
-      std::is_same_v<__orc_rt::Error,
+      std::is_same_v<orc_rt::Error,
                      std::invoke_result_t<ProcessSectionFunc,
-                                          __orc_rt::span<RecordElement>>>,
-      __orc_rt::Error>
+                                          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])) {
@@ -92,7 +92,7 @@ template <typename RecordElement> class RecordSectionsTracker {
       }
     }
     moveNewToProcessed();
-    return __orc_rt::Error::success();
+    return orc_rt::Error::success();
   }
 
   /// Move all sections back to New for reprocessing.
@@ -102,7 +102,7 @@ template <typename RecordElement> class RecordSectionsTracker {
   }
 
   /// Remove the section with the given range.
-  bool removeIfPresent(__orc_rt::ExecutorAddrRange R) {
+  bool removeIfPresent(orc_rt::ExecutorAddrRange R) {
     if (removeIfPresent(New, R))
       return true;
     return removeIfPresent(Processed, R);
@@ -119,11 +119,11 @@ template <typename RecordElement> class RecordSectionsTracker {
     }
   }
 
-  bool removeIfPresent(std::vector<__orc_rt::span<RecordElement>> &V,
-                       __orc_rt::ExecutorAddrRange R) {
+  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) {
+                               const orc_rt::span<RecordElement> &E) {
                              return E.data() == RS.data();
                            });
     if (RI != V.rend()) {
@@ -133,7 +133,7 @@ template <typename RecordElement> class RecordSectionsTracker {
     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

>From a01380289da5bebfbfab4ccc13a41534cd7fd6b2 Mon Sep 17 00:00:00 2001
From: SahilPatidar <patidarsahil2001 at gmail.com>
Date: Wed, 11 Sep 2024 21:28:12 +0530
Subject: [PATCH 7/9] Remove `RecordSectionsTracker` from `macho_platform.cpp`

---
 compiler-rt/lib/orc/macho_platform.cpp | 87 --------------------------
 1 file changed, 87 deletions(-)

diff --git a/compiler-rt/lib/orc/macho_platform.cpp b/compiler-rt/lib/orc/macho_platform.cpp
index 8cc3594b5d0cf7..94b341ce45f003 100644
--- a/compiler-rt/lib/orc/macho_platform.cpp
+++ b/compiler-rt/lib/orc/macho_platform.cpp
@@ -168,93 +168,6 @@ class MachOPlatformRuntimeState {
 
   using AtExitsVector = std::vector<AtExitEntry>;
 
-  /// 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(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, 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<Error, std::invoke_result_t<ProcessSectionFunc,
-                                                   span<RecordElement>>>,
-        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 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(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<span<RecordElement>> &V,
-                         ExecutorAddrRange R) {
-      auto RI = std::find_if(
-          V.rbegin(), V.rend(),
-          [RS = R.toSpan<RecordElement>()](const span<RecordElement> &E) {
-            return E.data() == RS.data();
-          });
-      if (RI != V.rend()) {
-        V.erase(std::next(RI).base());
-        return true;
-      }
-      return false;
-    }
-
-    std::vector<span<RecordElement>> Processed;
-    std::vector<span<RecordElement>> New;
-  };
-
   struct UnwindSections {
     UnwindSections(const UnwindSectionInfo &USI)
         : DwarfSection(USI.DwarfSection.toSpan<char>()),

>From e046c0cb066b1c0fc9b14c24f353373cc4c9fecb Mon Sep 17 00:00:00 2001
From: SahilPatidar <patidarsahil2001 at gmail.com>
Date: Wed, 11 Sep 2024 21:45:58 +0530
Subject: [PATCH 8/9] Fix coding style

---
 compiler-rt/lib/orc/common.h | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/compiler-rt/lib/orc/common.h b/compiler-rt/lib/orc/common.h
index 3518adc3dd2c67..683eec977df6ad 100644
--- a/compiler-rt/lib/orc/common.h
+++ b/compiler-rt/lib/orc/common.h
@@ -121,11 +121,11 @@ template <typename RecordElement> class RecordSectionsTracker {
 
   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();
-                           });
+    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;

>From 8686784e3e0544c55e7d3feab5daaf69c7cfa4d6 Mon Sep 17 00:00:00 2001
From: SahilPatidar <patidarsahil2001 at gmail.com>
Date: Mon, 23 Sep 2024 11:10:21 +0530
Subject: [PATCH 9/9] Update and resolve conflict

---
 .../llvm/ExecutionEngine/Orc/ELFNixPlatform.h      |  5 -----
 llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp    | 14 --------------
 2 files changed, 19 deletions(-)

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h b/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h
index 22500e8f6fad44..3d1c81f5c3efc1 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h
@@ -181,16 +181,11 @@ class ELFNixPlatform : public Platform {
                                      ResourceKey SrcKey) override {}
 
   private:
-    using InitSymbolDepMap =
-        DenseMap<MaterializationResponsibility *, JITLinkSymbolSet>;
 
     Error bootstrapPipelineStart(jitlink::LinkGraph &G);
     Error bootstrapPipelineRecordRuntimeFunctions(jitlink::LinkGraph &G);
     Error bootstrapPipelineEnd(jitlink::LinkGraph &G);
 
-    void addInitializerSupportPasses(MaterializationResponsibility &MR,
-                                     jitlink::PassConfiguration &Config);
-
     void addDSOHandleSupportPasses(MaterializationResponsibility &MR,
                                    jitlink::PassConfiguration &Config);
 
diff --git a/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp b/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
index 839a10c3723082..d92077dbcbd034 100644
--- a/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp
@@ -793,20 +793,6 @@ void ELFNixPlatform::ELFNixPlatformPlugin::modifyPassConfig(
         [this](LinkGraph &G) { return bootstrapPipelineEnd(G); });
 }
 
-ObjectLinkingLayer::Plugin::SyntheticSymbolDependenciesMap
-ELFNixPlatform::ELFNixPlatformPlugin::getSyntheticSymbolDependencies(
-    MaterializationResponsibility &MR) {
-  std::lock_guard<std::mutex> Lock(PluginMutex);
-  auto I = InitSymbolDeps.find(&MR);
-  if (I != InitSymbolDeps.end()) {
-    SyntheticSymbolDependenciesMap Result;
-    Result[MR.getInitializerSymbol()] = std::move(I->second);
-    InitSymbolDeps.erase(&MR);
-    return Result;
-  }
-  return SyntheticSymbolDependenciesMap();
-}
-
 void ELFNixPlatform::ELFNixPlatformPlugin::addDSOHandleSupportPasses(
     MaterializationResponsibility &MR, jitlink::PassConfiguration &Config) {
 



More information about the llvm-commits mailing list