Fix custom EHFrame registration on llvm21

Author: yuyichao

src/jitlayers.cpp

@@ -16,6 +16,7 @@
 #include <llvm/ExecutionEngine/Orc/DebugObjectManagerPlugin.h>
 #if JL_LLVM_VERSION >= 210000
 #  include <llvm/ExecutionEngine/Orc/SelfExecutorProcessControl.h>
+#  include <llvm/ExecutionEngine/Orc/Shared/OrcRTBridge.h>
 #endif
 #include <llvm/ExecutionEngine/Orc/TargetProcess/JITLoaderGDB.h>
 #if JL_LLVM_VERSION >= 200000
@@ -1226,6 +1227,8 @@ std::unique_ptr<jitlink::JITLinkMemoryManager> createJITLinkMemoryManager() JL_N
 #endif
 }
 
+#if defined(JL_USE_JITLINK) && defined(LLVM_SHLIB)
+#  if JL_LLVM_VERSION < 210000
 class JLEHFrameRegistrar final : public jitlink::EHFrameRegistrar {
 public:
     Error registerEHFrames(orc::ExecutorAddrRange EHFrameSection) override {
@@ -1238,6 +1241,72 @@ class JLEHFrameRegistrar final : public jitlink::EHFrameRegistrar {
         return Error::success();
     }
 };
+#  else
+class JLEHFrameRegistrationPlugin final : public LinkGraphLinkingLayer::Plugin {
+    static Error registerEHFrameWrapper(orc::ExecutorAddrRange EHFrame) {
+        register_eh_frames(EHFrame.Start.toPtr<uint8_t *>(), static_cast<size_t>(EHFrame.size()));
+        return Error::success();
+    }
+
+    static Error deregisterEHFrameWrapper(orc::ExecutorAddrRange EHFrame) {
+        deregister_eh_frames(EHFrame.Start.toPtr<uint8_t *>(), static_cast<size_t>(EHFrame.size()));
+        return Error::success();
+    }
+
+    static orc::shared::CWrapperFunctionResult
+    registerEHFrameSectionAllocAction(const char *ArgData, size_t ArgSize) {
+        using namespace llvm::orc::shared;
+        return WrapperFunction<SPSError(SPSExecutorAddrRange)>::handle(
+            ArgData, ArgSize, registerEHFrameWrapper)
+            .release();
+    }
+
+    static orc::shared::CWrapperFunctionResult
+    deregisterEHFrameSectionAllocAction(const char *ArgData, size_t ArgSize) {
+        using namespace llvm::orc::shared;
+        return WrapperFunction<SPSError(SPSExecutorAddrRange)>::handle(
+            ArgData, ArgSize, deregisterEHFrameWrapper)
+            .release();
+    }
+
+    static Error postFixup(jitlink::LinkGraph &G)
+    {
+        using namespace llvm::orc::shared;
+        auto registerFrame = ExecutorAddr::fromPtr(registerEHFrameSectionAllocAction);
+        auto deregisterFrame = ExecutorAddr::fromPtr(deregisterEHFrameSectionAllocAction);
+        if (auto *EHFrame = jitlink::getEHFrameSection(G)) {
+            auto R = jitlink::SectionRange(*EHFrame).getRange();
+            G.allocActions().push_back(
+                {cantFail(
+                        WrapperFunctionCall::Create<SPSArgList<SPSExecutorAddrRange>>(
+                            registerFrame, R)),
+                 cantFail(
+                     WrapperFunctionCall::Create<SPSArgList<SPSExecutorAddrRange>>(
+                         deregisterFrame, R))});
+        }
+        return Error::success();
+    }
+
+public:
+    JLEHFrameRegistrationPlugin() {}
+
+    void modifyPassConfig(MaterializationResponsibility&,
+                          jitlink::LinkGraph&,
+                          jitlink::PassConfiguration &PassConfig) override
+    {
+        PassConfig.PostFixupPasses.push_back(postFixup);
+    }
+    Error notifyFailed(MaterializationResponsibility&) override {
+        return Error::success();
+    }
+    Error notifyRemovingResources(JITDylib&, ResourceKey) override {
+        return Error::success();
+    }
+    void notifyTransferringResources(JITDylib&, ResourceKey,
+                                     ResourceKey) override {}
+};
+#  endif
+#endif
 
 RTDyldMemoryManager *createRTDyldMemoryManager(void) JL_NOTSAFEPOINT;
 
@@ -1950,15 +2019,26 @@ JuliaOJIT::JuliaOJIT()
     OptSelLayer(ES, OptimizeLayer, static_cast<orc::ThreadSafeModule (*)(orc::ThreadSafeModule, orc::MaterializationResponsibility&)>(selectOptLevel))
 {
 #ifdef JL_USE_JITLINK
-# if defined(LLVM_SHLIB)
+# if JL_LLVM_VERSION < 210000
+#  if defined(LLVM_SHLIB)
     // When dynamically linking against LLVM, use our custom EH frame registration code
     // also used with RTDyld to inform both our and the libc copy of libunwind.
     auto ehRegistrar = std::make_unique<JLEHFrameRegistrar>();
-# else
+#  else
     auto ehRegistrar = std::make_unique<jitlink::InProcessEHFrameRegistrar>();
-# endif
+#  endif
     ObjectLayer.addPlugin(std::make_unique<EHFrameRegistrationPlugin>(
         ES, std::move(ehRegistrar)));
+#else
+    // llvm's EHFrameRegistrationPlugin does not seem to have any customization
+    // hooks in 21+. Do our own registration with a separate plugin instead.
+#  if defined(LLVM_SHLIB)
+    // When dynamically linking against LLVM, use our custom EH frame registration code
+    // also used with RTDyld to inform both our and the libc copy of libunwind.
+    ObjectLayer.addPlugin(std::make_unique<JLEHFrameRegistrationPlugin>());
+#  endif
+    ObjectLayer.addPlugin(std::move(EHFrameRegistrationPlugin::Create(ES).get()));
+#endif
 
     ObjectLayer.addPlugin(std::make_unique<JLDebuginfoPlugin>());
     ObjectLayer.addPlugin(std::make_unique<JLMemoryUsagePlugin>(&jit_bytes_size));