[ORC][Runtime] Add dlupdate for elf (#110406)

With the help of @lhames, This pull request introduces the dlupdate
function in the ORC runtime. dlupdate enables incremental execution of
new initializers introduced in the REPL environment. Unlike traditional
dlopen, which manages initializers, code mapping, and library reference
counts, dlupdate focuses exclusively on running new initializers.

Author: SahilPatidar

compiler-rt/lib/orc/dlfcn_wrapper.cpp

@@ -42,7 +42,7 @@ __orc_rt_jit_dlopen_wrapper(const char *ArgData, size_t ArgSize) {
       .release();
 }
 
-#ifdef __APPLE__
+#ifndef _WIN32
 ORC_RT_INTERFACE orc_rt_CWrapperFunctionResult
 __orc_rt_jit_dlupdate_wrapper(const char *ArgData, size_t ArgSize) {
   return WrapperFunction<int32_t(SPSExecutorAddr)>::handle(

compiler-rt/lib/orc/elfnix_platform.cpp

@@ -105,6 +105,7 @@ class ELFNixPlatformRuntimeState {
 
   const char *dlerror();
   void *dlopen(std::string_view Name, int Mode);
+  int dlupdate(void *DSOHandle);
   int dlclose(void *DSOHandle);
   void *dlsym(void *DSOHandle, std::string_view Symbol);
 
@@ -136,6 +137,10 @@ class ELFNixPlatformRuntimeState {
   Error dlopenInitialize(std::unique_lock<std::recursive_mutex> &JDStatesLock,
                          PerJITDylibState &JDS,
                          ELFNixJITDylibDepInfoMap &DepInfo);
+  Error dlupdateImpl(void *DSOHandle);
+  Error dlupdateFull(std::unique_lock<std::recursive_mutex> &JDStatesLock,
+                     PerJITDylibState &JDS);
+
   Error dlcloseImpl(void *DSOHandle);
   Error dlcloseInitialize(std::unique_lock<std::recursive_mutex> &JDStatesLock,
                           PerJITDylibState &JDS);
@@ -309,6 +314,15 @@ void *ELFNixPlatformRuntimeState::dlopen(std::string_view Path, int Mode) {
   }
 }
 
+int ELFNixPlatformRuntimeState::dlupdate(void *DSOHandle) {
+  if (auto Err = dlupdateImpl(DSOHandle)) {
+    // FIXME: Make dlerror thread safe.
+    DLFcnError = toString(std::move(Err));
+    return -1;
+  }
+  return 0;
+}
+
 int ELFNixPlatformRuntimeState::dlclose(void *DSOHandle) {
   if (auto Err = dlcloseImpl(DSOHandle)) {
     DLFcnError = toString(std::move(Err));
@@ -523,6 +537,50 @@ Error ELFNixPlatformRuntimeState::dlopenInitialize(
   return Error::success();
 }
 
+Error ELFNixPlatformRuntimeState::dlupdateImpl(void *DSOHandle) {
+  std::unique_lock<std::recursive_mutex> Lock(JDStatesMutex);
+
+  // Try to find JITDylib state by name.
+  auto *JDS = getJITDylibStateByHeaderAddr(DSOHandle);
+
+  if (!JDS) {
+    std::ostringstream ErrStream;
+    ErrStream << "No registered JITDylib for " << DSOHandle;
+    return make_error<StringError>(ErrStream.str());
+  }
+
+  if (!JDS->referenced())
+    return make_error<StringError>("dlupdate failed, JITDylib must be open.");
+
+  if (auto Err = dlupdateFull(Lock, *JDS))
+    return Err;
+
+  return Error::success();
+}
+
+Error ELFNixPlatformRuntimeState::dlupdateFull(
+    std::unique_lock<std::recursive_mutex> &JDStatesLock,
+    PerJITDylibState &JDS) {
+  // Call back to the JIT to push the initializers.
+  Expected<ELFNixJITDylibDepInfoMap> DepInfo((ELFNixJITDylibDepInfoMap()));
+  // Unlock so that we can accept the initializer update.
+  JDStatesLock.unlock();
+  if (auto Err = WrapperFunction<SPSExpected<SPSELFNixJITDylibDepInfoMap>(
+          SPSExecutorAddr)>::
+          call(JITDispatch(&__orc_rt_elfnix_push_initializers_tag), DepInfo,
+               ExecutorAddr::fromPtr(JDS.Header)))
+    return Err;
+  JDStatesLock.lock();
+
+  if (!DepInfo)
+    return DepInfo.takeError();
+
+  if (auto Err = runInits(JDStatesLock, JDS))
+    return Err;
+
+  return Error::success();
+}
+
 Error ELFNixPlatformRuntimeState::dlcloseImpl(void *DSOHandle) {
 
   std::unique_lock<std::recursive_mutex> Lock(JDStatesMutex);
@@ -765,6 +823,10 @@ void *__orc_rt_elfnix_jit_dlopen(const char *path, int mode) {
   return ELFNixPlatformRuntimeState::get().dlopen(path, mode);
 }
 
+int __orc_rt_elfnix_jit_dlupdate(void *dso_handle) {
+  return ELFNixPlatformRuntimeState::get().dlupdate(dso_handle);
+}
+
 int __orc_rt_elfnix_jit_dlclose(void *dso_handle) {
   return ELFNixPlatformRuntimeState::get().dlclose(dso_handle);
 }

compiler-rt/lib/orc/elfnix_platform.h

@@ -25,6 +25,7 @@ ORC_RT_INTERFACE void __orc_rt_elfnix_cxa_finalize(void *dso_handle);
 // dlfcn functions.
 ORC_RT_INTERFACE const char *__orc_rt_elfnix_jit_dlerror();
 ORC_RT_INTERFACE void *__orc_rt_elfnix_jit_dlopen(const char *path, int mode);
+ORC_RT_INTERFACE int __orc_rt_elfnix_jit_dlupdate(void *dso_handle);
 ORC_RT_INTERFACE int __orc_rt_elfnix_jit_dlclose(void *dso_handle);
 ORC_RT_INTERFACE void *__orc_rt_elfnix_jit_dlsym(void *dso_handle,
                                                  const char *symbol);

llvm/lib/ExecutionEngine/Orc/ELFNixPlatform.cpp

@@ -372,6 +372,7 @@ ELFNixPlatform::standardRuntimeUtilityAliases() {
           {"__orc_rt_run_program", "__orc_rt_elfnix_run_program"},
           {"__orc_rt_jit_dlerror", "__orc_rt_elfnix_jit_dlerror"},
           {"__orc_rt_jit_dlopen", "__orc_rt_elfnix_jit_dlopen"},
+          {"__orc_rt_jit_dlupdate", "__orc_rt_elfnix_jit_dlupdate"},
           {"__orc_rt_jit_dlclose", "__orc_rt_elfnix_jit_dlclose"},
           {"__orc_rt_jit_dlsym", "__orc_rt_elfnix_jit_dlsym"},
           {"__orc_rt_log_error", "__orc_rt_log_error_to_stderr"}};

llvm/lib/ExecutionEngine/Orc/LLJIT.cpp

@@ -625,7 +625,8 @@ Error ORCPlatformSupport::initialize(orc::JITDylib &JD) {
       [](const JITDylibSearchOrder &SO) { return SO; });
   StringRef WrapperToCall = "__orc_rt_jit_dlopen_wrapper";
   bool dlupdate = false;
-  if (ES.getTargetTriple().isOSBinFormatMachO()) {
+  const Triple &TT = ES.getTargetTriple();
+  if (TT.isOSBinFormatMachO() || TT.isOSBinFormatELF()) {
     if (InitializedDylib.contains(&JD)) {
       WrapperToCall = "__orc_rt_jit_dlupdate_wrapper";
       dlupdate = true;
@@ -639,12 +640,10 @@ Error ORCPlatformSupport::initialize(orc::JITDylib &JD) {
       int32_t result;
       auto E = ES.callSPSWrapper<SPSDLUpdateSig>(WrapperAddr->getAddress(),
                                                  result, DSOHandles[&JD]);
-      if (E)
-        return E;
-      else if (result)
+      if (result)
         return make_error<StringError>("dlupdate failed",
                                        inconvertibleErrorCode());
-      return Error::success();
+      return E;
     }
     return ES.callSPSWrapper<SPSDLOpenSig>(WrapperAddr->getAddress(),
                                            DSOHandles[&JD], JD.getName(),