[llvm] [Orc] Move TaskDispatcher ownership from EPC to derived classes (PR #85071)

[Stefan Gränitz via llvm-commits]

https://github.com/weliveindetail created https://github.com/llvm/llvm-project/pull/85071

We use this approach for MemoryManager and MemoryAccess already. This patch applies it for TaskDispatcher too.

>From 21ba6039b2deb158211cafb4e851a7ea2672fb9b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stefan=20Gr=C3=A4nitz?= <stefan.graenitz at gmail.com>
Date: Wed, 13 Mar 2024 13:22:42 +0100
Subject: [PATCH] [Orc] Move TaskDispatcher ownership from EPC to derived
 classes

We use the same approach for MemoryManager and MemoryAccess. This patch applies it for TaskDispatcher too.
---
 .../Orc/ExecutorProcessControl.h              | 13 ++++++------
 .../ExecutionEngine/Orc/SimpleRemoteEPC.h     |  5 ++++-
 .../Orc/ExecutorProcessControl.cpp            | 20 ++++++++++---------
 .../Orc/ObjectLinkingLayerTest.cpp            |  2 +-
 4 files changed, 22 insertions(+), 18 deletions(-)

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/ExecutorProcessControl.h b/llvm/include/llvm/ExecutionEngine/Orc/ExecutorProcessControl.h
index 6468f2dfc11ad0..2aacc1ea2acc5c 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/ExecutorProcessControl.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/ExecutorProcessControl.h
@@ -187,9 +187,8 @@ class ExecutorProcessControl {
     ExecutorAddr JITDispatchContext;
   };
 
-  ExecutorProcessControl(std::shared_ptr<SymbolStringPool> SSP,
-                         std::unique_ptr<TaskDispatcher> D)
-    : SSP(std::move(SSP)), D(std::move(D)) {}
+  ExecutorProcessControl(std::shared_ptr<SymbolStringPool> SSP)
+    : SSP(std::move(SSP)) {}
 
   virtual ~ExecutorProcessControl();
 
@@ -421,11 +420,11 @@ class ExecutorProcessControl {
 protected:
 
   std::shared_ptr<SymbolStringPool> SSP;
-  std::unique_ptr<TaskDispatcher> D;
   ExecutionSession *ES = nullptr;
   Triple TargetTriple;
   unsigned PageSize = 0;
   JITDispatchInfo JDI;
+  TaskDispatcher *D;
   MemoryAccess *MemAccess = nullptr;
   jitlink::JITLinkMemoryManager *MemMgr = nullptr;
   StringMap<std::vector<char>> BootstrapMap;
@@ -465,11 +464,10 @@ class UnsupportedExecutorProcessControl : public ExecutorProcessControl,
 public:
   UnsupportedExecutorProcessControl(
       std::shared_ptr<SymbolStringPool> SSP = nullptr,
-      std::unique_ptr<TaskDispatcher> D = nullptr, const std::string &TT = "",
+      const std::string &TT = "",
       unsigned PageSize = 0)
       : ExecutorProcessControl(
-            SSP ? std::move(SSP) : std::make_shared<SymbolStringPool>(),
-            D ? std::move(D) : std::make_unique<InPlaceTaskDispatcher>()),
+            SSP ? std::move(SSP) : std::make_shared<SymbolStringPool>()),
         InProcessMemoryAccess(Triple(TT).isArch64Bit()) {
     this->TargetTriple = Triple(TT);
     this->PageSize = PageSize;
@@ -549,6 +547,7 @@ class SelfExecutorProcessControl : public ExecutorProcessControl,
   jitDispatchViaWrapperFunctionManager(void *Ctx, const void *FnTag,
                                        const char *Data, size_t Size);
 
+  std::unique_ptr<TaskDispatcher> OwnedTaskDispatcher;
   std::unique_ptr<jitlink::JITLinkMemoryManager> OwnedMemMgr;
   char GlobalManglingPrefix = 0;
 };
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/SimpleRemoteEPC.h b/llvm/include/llvm/ExecutionEngine/Orc/SimpleRemoteEPC.h
index c10b8df01cc0a4..0d7267dc4dbee0 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/SimpleRemoteEPC.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/SimpleRemoteEPC.h
@@ -96,7 +96,9 @@ class SimpleRemoteEPC : public ExecutorProcessControl,
 private:
   SimpleRemoteEPC(std::shared_ptr<SymbolStringPool> SSP,
                   std::unique_ptr<TaskDispatcher> D)
-    : ExecutorProcessControl(std::move(SSP), std::move(D)) {}
+      : ExecutorProcessControl(std::move(SSP)), OwnedTaskDispatcher(std::move(D)) {
+    this->D = OwnedTaskDispatcher.get();
+  }
 
   static Expected<std::unique_ptr<jitlink::JITLinkMemoryManager>>
   createDefaultMemoryManager(SimpleRemoteEPC &SREPC);
@@ -130,6 +132,7 @@ class SimpleRemoteEPC : public ExecutorProcessControl,
   std::unique_ptr<SimpleRemoteEPCTransport> T;
   std::unique_ptr<jitlink::JITLinkMemoryManager> OwnedMemMgr;
   std::unique_ptr<MemoryAccess> OwnedMemAccess;
+  std::unique_ptr<TaskDispatcher> OwnedTaskDispatcher;
 
   std::unique_ptr<EPCGenericDylibManager> DylibMgr;
   ExecutorAddr RunAsMainAddr;
diff --git a/llvm/lib/ExecutionEngine/Orc/ExecutorProcessControl.cpp b/llvm/lib/ExecutionEngine/Orc/ExecutorProcessControl.cpp
index efafca949e61ef..759d3254d3a3fc 100644
--- a/llvm/lib/ExecutionEngine/Orc/ExecutorProcessControl.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ExecutorProcessControl.cpp
@@ -29,7 +29,7 @@ SelfExecutorProcessControl::SelfExecutorProcessControl(
     std::shared_ptr<SymbolStringPool> SSP, std::unique_ptr<TaskDispatcher> D,
     Triple TargetTriple, unsigned PageSize,
     std::unique_ptr<jitlink::JITLinkMemoryManager> MemMgr)
-    : ExecutorProcessControl(std::move(SSP), std::move(D)),
+    : ExecutorProcessControl(std::move(SSP)),
       InProcessMemoryAccess(TargetTriple.isArch64Bit()) {
 
   OwnedMemMgr = std::move(MemMgr);
@@ -37,10 +37,20 @@ SelfExecutorProcessControl::SelfExecutorProcessControl(
     OwnedMemMgr = std::make_unique<jitlink::InProcessMemoryManager>(
         sys::Process::getPageSizeEstimate());
 
+  OwnedTaskDispatcher = std::move(D);
+  if (!OwnedTaskDispatcher) {
+#if LLVM_ENABLE_THREADS
+    OwnedTaskDispatcher = std::make_unique<DynamicThreadPoolTaskDispatcher>();
+#else
+    OwnedTaskDispatcher = std::make_unique<InPlaceTaskDispatcher>();
+#endif
+  }
+
   this->TargetTriple = std::move(TargetTriple);
   this->PageSize = PageSize;
   this->MemMgr = OwnedMemMgr.get();
   this->MemAccess = this;
+  this->D = OwnedTaskDispatcher.get();
   this->JDI = {ExecutorAddr::fromPtr(jitDispatchViaWrapperFunctionManager),
                ExecutorAddr::fromPtr(this)};
   if (this->TargetTriple.isOSBinFormatMachO())
@@ -61,14 +71,6 @@ SelfExecutorProcessControl::Create(
   if (!SSP)
     SSP = std::make_shared<SymbolStringPool>();
 
-  if (!D) {
-#if LLVM_ENABLE_THREADS
-    D = std::make_unique<DynamicThreadPoolTaskDispatcher>();
-#else
-    D = std::make_unique<InPlaceTaskDispatcher>();
-#endif
-  }
-
   auto PageSize = sys::Process::getPageSize();
   if (!PageSize)
     return PageSize.takeError();
diff --git a/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp b/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp
index 7ab3e40df7459d..6643ecab749c15 100644
--- a/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp
@@ -182,7 +182,7 @@ TEST(ObjectLinkingLayerSearchGeneratorTest, AbsoluteSymbolsObjectLayer) {
   class TestEPC : public UnsupportedExecutorProcessControl {
   public:
     TestEPC()
-        : UnsupportedExecutorProcessControl(nullptr, nullptr,
+        : UnsupportedExecutorProcessControl(nullptr,
                                             "x86_64-apple-darwin") {}
 
     Expected<tpctypes::DylibHandle> loadDylib(const char *DylibPath) override {