[ORC] Use a Setup object for SimpleRemoteEPC construction.

SimpleRemoteEPC notionally allowed subclasses to override the
createMemoryManager and createMemoryAccess methods to use custom objects, but
could not actually be subclassed in practice (The construction process in
SimpleRemoteEPC::Create could not be re-used).

Instead of subclassing, this commit adds a SimpleRemoteEPC::Setup class that
can be used by clients to set up the memory manager and memory access members.
A default-constructed Setup object results in no change from previous behavior
(EPCGeneric* memory manager and memory access objects used by default).

Author: lhames

llvm/include/llvm/ExecutionEngine/Orc/SimpleRemoteEPC.h

@@ -31,10 +31,24 @@ namespace orc {
 class SimpleRemoteEPC : public ExecutorProcessControl,
                         public SimpleRemoteEPCTransportClient {
 public:
+  /// A setup object containing callbacks to construct a memory manager and
+  /// memory access object. Both are optional. If not specified,
+  /// EPCGenericJITLinkMemoryManager and EPCGenericMemoryAccess will be used.
+  struct Setup {
+    using CreateMemoryManagerFn =
+        Expected<std::unique_ptr<jitlink::JITLinkMemoryManager>>(
+            SimpleRemoteEPC &);
+    using CreateMemoryAccessFn =
+        Expected<std::unique_ptr<MemoryAccess>>(SimpleRemoteEPC &);
+
+    unique_function<CreateMemoryManagerFn> CreateMemoryManager;
+    unique_function<CreateMemoryAccessFn> CreateMemoryAccess;
+  };
+
   /// Create a SimpleRemoteEPC using the given transport type and args.
   template <typename TransportT, typename... TransportTCtorArgTs>
   static Expected<std::unique_ptr<SimpleRemoteEPC>>
-  Create(std::unique_ptr<TaskDispatcher> D,
+  Create(std::unique_ptr<TaskDispatcher> D, Setup S,
          TransportTCtorArgTs &&...TransportTCtorArgs) {
     std::unique_ptr<SimpleRemoteEPC> SREPC(
         new SimpleRemoteEPC(std::make_shared<SymbolStringPool>(),
@@ -44,7 +58,7 @@ class SimpleRemoteEPC : public ExecutorProcessControl,
     if (!T)
       return T.takeError();
     SREPC->T = std::move(*T);
-    if (auto Err = SREPC->setup())
+    if (auto Err = SREPC->setup(std::move(S)))
       return joinErrors(std::move(Err), SREPC->disconnect());
     return std::move(SREPC);
   }
@@ -75,22 +89,22 @@ class SimpleRemoteEPC : public ExecutorProcessControl,
 
   void handleDisconnect(Error Err) override;
 
-protected:
-  virtual Expected<std::unique_ptr<jitlink::JITLinkMemoryManager>>
-  createMemoryManager();
-  virtual Expected<std::unique_ptr<MemoryAccess>> createMemoryAccess();
-
 private:
   SimpleRemoteEPC(std::shared_ptr<SymbolStringPool> SSP,
                   std::unique_ptr<TaskDispatcher> D)
     : ExecutorProcessControl(std::move(SSP), std::move(D)) {}
 
+  static Expected<std::unique_ptr<jitlink::JITLinkMemoryManager>>
+  createDefaultMemoryManager(SimpleRemoteEPC &SREPC);
+  static Expected<std::unique_ptr<MemoryAccess>>
+  createDefaultMemoryAccess(SimpleRemoteEPC &SREPC);
+
   Error sendMessage(SimpleRemoteEPCOpcode OpC, uint64_t SeqNo,
                     ExecutorAddr TagAddr, ArrayRef<char> ArgBytes);
 
   Error handleSetup(uint64_t SeqNo, ExecutorAddr TagAddr,
                     SimpleRemoteEPCArgBytesVector ArgBytes);
-  Error setup();
+  Error setup(Setup S);
 
   Error handleResult(uint64_t SeqNo, ExecutorAddr TagAddr,
                      SimpleRemoteEPCArgBytesVector ArgBytes);

llvm/lib/ExecutionEngine/Orc/SimpleRemoteEPC.cpp

@@ -180,22 +180,21 @@ void SimpleRemoteEPC::handleDisconnect(Error Err) {
 }
 
 Expected<std::unique_ptr<jitlink::JITLinkMemoryManager>>
-SimpleRemoteEPC::createMemoryManager() {
+SimpleRemoteEPC::createDefaultMemoryManager(SimpleRemoteEPC &SREPC) {
   EPCGenericJITLinkMemoryManager::SymbolAddrs SAs;
-  if (auto Err = getBootstrapSymbols(
+  if (auto Err = SREPC.getBootstrapSymbols(
           {{SAs.Allocator, rt::SimpleExecutorMemoryManagerInstanceName},
            {SAs.Reserve, rt::SimpleExecutorMemoryManagerReserveWrapperName},
            {SAs.Finalize, rt::SimpleExecutorMemoryManagerFinalizeWrapperName},
            {SAs.Deallocate,
             rt::SimpleExecutorMemoryManagerDeallocateWrapperName}}))
     return std::move(Err);
 
-  return std::make_unique<EPCGenericJITLinkMemoryManager>(*this, SAs);
+  return std::make_unique<EPCGenericJITLinkMemoryManager>(SREPC, SAs);
 }
 
 Expected<std::unique_ptr<ExecutorProcessControl::MemoryAccess>>
-SimpleRemoteEPC::createMemoryAccess() {
-
+SimpleRemoteEPC::createDefaultMemoryAccess(SimpleRemoteEPC &SREPC) {
   return nullptr;
 }
 
@@ -260,7 +259,7 @@ Error SimpleRemoteEPC::handleSetup(uint64_t SeqNo, ExecutorAddr TagAddr,
   return Error::success();
 }
 
-Error SimpleRemoteEPC::setup() {
+Error SimpleRemoteEPC::setup(Setup S) {
   using namespace SimpleRemoteEPCDefaultBootstrapSymbolNames;
 
   std::promise<MSVCPExpected<SimpleRemoteEPCExecutorInfo>> EIP;
@@ -322,13 +321,21 @@ Error SimpleRemoteEPC::setup() {
   else
     return DM.takeError();
 
-  if (auto MemMgr = createMemoryManager()) {
+  // Set a default CreateMemoryManager if none is specified.
+  if (!S.CreateMemoryManager)
+    S.CreateMemoryManager = createDefaultMemoryManager;
+
+  if (auto MemMgr = S.CreateMemoryManager(*this)) {
     OwnedMemMgr = std::move(*MemMgr);
     this->MemMgr = OwnedMemMgr.get();
   } else
     return MemMgr.takeError();
 
-  if (auto MemAccess = createMemoryAccess()) {
+  // Set a default CreateMemoryAccess if none is specified.
+  if (!S.CreateMemoryAccess)
+    S.CreateMemoryAccess = createDefaultMemoryAccess;
+
+  if (auto MemAccess = S.CreateMemoryAccess(*this)) {
     OwnedMemAccess = std::move(*MemAccess);
     this->MemAccess = OwnedMemAccess.get();
   } else

llvm/tools/lli/lli.cpp

@@ -1151,6 +1151,6 @@ Expected<std::unique_ptr<orc::ExecutorProcessControl>> launchRemote() {
   // Return a SimpleRemoteEPC instance connected to our end of the pipes.
   return orc::SimpleRemoteEPC::Create<orc::FDSimpleRemoteEPCTransport>(
       std::make_unique<llvm::orc::InPlaceTaskDispatcher>(),
-      PipeFD[1][0], PipeFD[0][1]);
+      llvm::orc::SimpleRemoteEPC::Setup(), PipeFD[1][0], PipeFD[0][1]);
 #endif
 }

llvm/tools/llvm-jitlink/llvm-jitlink.cpp

@@ -798,7 +798,7 @@ static Expected<std::unique_ptr<ExecutorProcessControl>> launchExecutor() {
 
   return SimpleRemoteEPC::Create<FDSimpleRemoteEPCTransport>(
       std::make_unique<DynamicThreadPoolTaskDispatcher>(),
-      FromExecutor[ReadEnd], ToExecutor[WriteEnd]);
+      SimpleRemoteEPC::Setup(), FromExecutor[ReadEnd], ToExecutor[WriteEnd]);
 #endif
 }
 
@@ -883,7 +883,8 @@ static Expected<std::unique_ptr<ExecutorProcessControl>> connectToExecutor() {
     return SockFD.takeError();
 
   return SimpleRemoteEPC::Create<FDSimpleRemoteEPCTransport>(
-      std::make_unique<DynamicThreadPoolTaskDispatcher>(), *SockFD, *SockFD);
+      std::make_unique<DynamicThreadPoolTaskDispatcher>(),
+      SimpleRemoteEPC::Setup(), *SockFD, *SockFD);
 #endif
 }