Automerge: [ORC][Runtime] Add dlupdate for coff (#115448)

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/coff_platform.cpp

@@ -110,6 +110,7 @@ class COFFPlatformRuntimeState {
 
   const char *dlerror();
   void *dlopen(std::string_view Name, int Mode);
+  int dlupdate(void *DSOHandle);
   int dlclose(void *Header);
   void *dlsym(void *Header, std::string_view Symbol);
 
@@ -141,6 +142,10 @@ class COFFPlatformRuntimeState {
   Error dlopenFull(JITDylibState &JDS);
   Error dlopenInitialize(JITDylibState &JDS, COFFJITDylibDepInfoMap &DepInfo);
 
+  Error dlupdateImpl(void *DSOHandle);
+  Error dlupdateFull(JITDylibState &JDS);
+  Error dlupdateInitialize(JITDylibState &JDS);
+
   Error dlcloseImpl(void *DSOHandle);
   Error dlcloseDeinitialize(JITDylibState &JDS);
 
@@ -265,6 +270,20 @@ void *COFFPlatformRuntimeState::dlopen(std::string_view Path, int Mode) {
   }
 }
 
+int COFFPlatformRuntimeState::dlupdate(void *DSOHandle) {
+  ORC_RT_DEBUG({
+    std::string S;
+    printdbg("COFFPlatform::dlupdate(%p) (%s)\n", DSOHandle, S.c_str());
+  });
+  std::lock_guard<std::recursive_mutex> Lock(JDStatesMutex);
+  if (auto Err = dlupdateImpl(DSOHandle)) {
+    // FIXME: Make dlerror thread safe.
+    DLFcnError = toString(std::move(Err));
+    return -1;
+  }
+  return 0;
+}
+
 int COFFPlatformRuntimeState::dlclose(void *DSOHandle) {
   ORC_RT_DEBUG({
     auto *JDS = getJITDylibStateByHeader(DSOHandle);
@@ -390,6 +409,55 @@ Error COFFPlatformRuntimeState::dlopenInitialize(
   return Error::success();
 }
 
+Error COFFPlatformRuntimeState::dlupdateImpl(void *DSOHandle) {
+  // Try to find JITDylib state by header.
+  auto *JDS = getJITDylibStateByHeader(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(*JDS))
+    return Err;
+
+  return Error::success();
+}
+
+Error COFFPlatformRuntimeState::dlupdateFull(JITDylibState &JDS) {
+  // Call back to the JIT to push the initializers.
+  Expected<COFFJITDylibDepInfoMap> DepInfoMap((COFFJITDylibDepInfoMap()));
+  if (auto Err = WrapperFunction<SPSExpected<SPSCOFFJITDylibDepInfoMap>(
+          SPSExecutorAddr)>::
+          call(JITDispatch(&__orc_rt_coff_push_initializers_tag), DepInfoMap,
+               ExecutorAddr::fromPtr(JDS.Header)))
+    return Err;
+  if (!DepInfoMap)
+    return DepInfoMap.takeError();
+
+  if (auto Err = dlupdateInitialize(JDS))
+    return Err;
+
+  return Error::success();
+}
+
+Error COFFPlatformRuntimeState::dlupdateInitialize(JITDylibState &JDS) {
+  ORC_RT_DEBUG({
+    printdbg("COFFPlatformRuntimeState::dlupdateInitialize(\"%s\")\n",
+             JDS.Name.c_str());
+  });
+
+  // Run static initializers.
+  JDS.CInitSection.RunAllNewAndFlush();
+  JDS.CXXInitSection.RunAllNewAndFlush();
+
+  return Error::success();
+}
+
 Error COFFPlatformRuntimeState::dlcloseImpl(void *DSOHandle) {
   // Try to find JITDylib state by header.
   auto *JDS = getJITDylibStateByHeader(DSOHandle);
@@ -667,6 +735,10 @@ void *__orc_rt_coff_jit_dlopen(const char *path, int mode) {
   return COFFPlatformRuntimeState::get().dlopen(path, mode);
 }
 
+int __orc_rt_coff_jit_dlupdate(void *dso_handle) {
+  return COFFPlatformRuntimeState::get().dlupdate(dso_handle);
+}
+
 int __orc_rt_coff_jit_dlclose(void *header) {
   return COFFPlatformRuntimeState::get().dlclose(header);
 }

compiler-rt/lib/orc/coff_platform.h

@@ -19,6 +19,7 @@
 // dlfcn functions.
 ORC_RT_INTERFACE const char *__orc_rt_coff_jit_dlerror();
 ORC_RT_INTERFACE void *__orc_rt_coff_jit_dlopen(const char *path, int mode);
+ORC_RT_INTERFACE int __orc_rt_coff_jit_dlupdate(void *dso_handle);
 ORC_RT_INTERFACE int __orc_rt_coff_jit_dlclose(void *header);
 ORC_RT_INTERFACE void *__orc_rt_coff_jit_dlsym(void *header,
                                                const char *symbol);

compiler-rt/lib/orc/dlfcn_wrapper.cpp

@@ -42,7 +42,6 @@ __orc_rt_jit_dlopen_wrapper(const char *ArgData, size_t ArgSize) {
       .release();
 }
 
-#ifndef _WIN32
 ORC_RT_INTERFACE orc_rt_WrapperFunctionResult
 __orc_rt_jit_dlupdate_wrapper(const char *ArgData, size_t ArgSize) {
   return WrapperFunction<int32_t(SPSExecutorAddr)>::handle(
@@ -52,7 +51,6 @@ __orc_rt_jit_dlupdate_wrapper(const char *ArgData, size_t ArgSize) {
              })
       .release();
 }
-#endif
 
 ORC_RT_INTERFACE orc_rt_WrapperFunctionResult
 __orc_rt_jit_dlclose_wrapper(const char *ArgData, size_t ArgSize) {

llvm/lib/ExecutionEngine/Orc/COFFPlatform.cpp

@@ -361,6 +361,7 @@ COFFPlatform::standardRuntimeUtilityAliases() {
           {"__orc_rt_run_program", "__orc_rt_coff_run_program"},
           {"__orc_rt_jit_dlerror", "__orc_rt_coff_jit_dlerror"},
           {"__orc_rt_jit_dlopen", "__orc_rt_coff_jit_dlopen"},
+          {"__orc_rt_jit_dlupdate", "__orc_rt_coff_jit_dlupdate"},
           {"__orc_rt_jit_dlclose", "__orc_rt_coff_jit_dlclose"},
           {"__orc_rt_jit_dlsym", "__orc_rt_coff_jit_dlsym"},
           {"__orc_rt_log_error", "__orc_rt_log_error_to_stderr"}};

llvm/lib/ExecutionEngine/Orc/LLJIT.cpp

@@ -617,14 +617,11 @@ Error ORCPlatformSupport::initialize(orc::JITDylib &JD) {
       [](const JITDylibSearchOrder &SO) { return SO; });
   StringRef WrapperToCall = "__orc_rt_jit_dlopen_wrapper";
   bool dlupdate = false;
-  const Triple &TT = ES.getTargetTriple();
-  if (TT.isOSBinFormatMachO() || TT.isOSBinFormatELF()) {
-    if (InitializedDylib.contains(&JD)) {
-      WrapperToCall = "__orc_rt_jit_dlupdate_wrapper";
-      dlupdate = true;
-    } else
-      InitializedDylib.insert(&JD);
-  }
+  if (InitializedDylib.contains(&JD)) {
+    WrapperToCall = "__orc_rt_jit_dlupdate_wrapper";
+    dlupdate = true;
+  } else
+    InitializedDylib.insert(&JD);
 
   if (auto WrapperAddr =
           ES.lookup(MainSearchOrder, J.mangleAndIntern(WrapperToCall))) {