[ORC] Add JITLink-compatible remote memory-manager and LLJITWithChildProcess example

This adds RemoteJITLinkMemoryManager is a new subclass of OrcRemoteTargetClient. It implements jitlink::JITLinkMemoryManager and targets the OrcRemoteTargetRPCAPI.

Behavior should be very similar to RemoteRTDyldMemoryManager. The essential differnce with JITLink is that allocations work in isolation from its memory manager. Thus, the RemoteJITLinkMemoryManager might be seen as "JITLink allocation factory".

RPCMMAlloc is another subclass of OrcRemoteTargetClient and implements the actual functionality. It allocates working memory on the host and target memory on the remote target. Upon finalization working memory is copied over to the tagrte address space. Finalization can be asynchronous for JITLink allocations, but I don't see that it makes a difference here.

Differential Revision: https://reviews.llvm.org/D85919

examples/OrcV2Examples/CMakeLists.txt

@@ -1,4 +1,5 @@
 add_subdirectory(LLJITDumpObjects)
+add_subdirectory(LLJITWithChildProcess)
 add_subdirectory(LLJITWithCustomObjectLinkingLayer)
 add_subdirectory(LLJITWithGDBRegistrationListener)
 add_subdirectory(LLJITWithInitializers)

examples/OrcV2Examples/LLJITWithChildProcess/CMakeLists.txt

@@ -0,0 +1,12 @@
+set(LLVM_LINK_COMPONENTS
+  Core
+  ExecutionEngine
+  IRReader
+  OrcJIT
+  Support
+  nativecodegen
+  )
+
+add_llvm_example(LLJITInChildProcess
+  LLJITWithChildProcess.cpp
+  )

examples/OrcV2Examples/LLJITWithChildProcess/LLJITWithChildProcess.cpp

@@ -0,0 +1,128 @@
+//===--- LLJITWithLazyReexports.cpp - LLJIT example with custom laziness --===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// In this example we will execute JITed code in a child process:
+//
+// 1. Launch a remote process.
+// 2. Create a JITLink-compatible remote memory manager.
+// 3. Use LLJITBuilder to create a (greedy) LLJIT instance.
+// 4. Add the Add1Example module and execute add1().
+// 5. Terminate the remote target session.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h"
+#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
+#include "llvm/ExecutionEngine/Orc/LLJIT.h"
+#include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
+#include "llvm/ExecutionEngine/Orc/OrcRemoteTargetClient.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/InitLLVM.h"
+#include "llvm/Support/TargetSelect.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include "../ExampleModules.h"
+#include "RemoteJITUtils.h"
+
+#include <memory>
+#include <string>
+
+#define DEBUG_TYPE "orc"
+
+using namespace llvm;
+using namespace llvm::orc;
+
+// Executable running in the child process for remote execution. It communicates
+// via stdin/stdout pipes.
+cl::opt<std::string>
+    ChildExecPath("remote-process", cl::Required,
+                  cl::desc("Specify the filename of the process to launch for "
+                           "remote JITing."),
+                  cl::value_desc("filename"));
+
+int main(int argc, char *argv[]) {
+  InitLLVM X(argc, argv);
+
+  InitializeNativeTarget();
+  InitializeNativeTargetAsmPrinter();
+
+  cl::ParseCommandLineOptions(argc, argv, "LLJITWithChildProcess");
+
+  ExitOnError ExitOnErr;
+  ExitOnErr.setBanner(std::string(argv[0]) + ": ");
+
+  if (!sys::fs::can_execute(ChildExecPath)) {
+    WithColor::error(errs(), argv[0])
+        << "Child executable invalid: '" << ChildExecPath << "'\n";
+    return -1;
+  }
+
+  ExecutionSession ES;
+  ES.setErrorReporter([&](Error Err) { ExitOnErr(std::move(Err)); });
+
+  // Launch the remote process and get a channel to it.
+  pid_t ChildPID;
+  std::unique_ptr<FDRawChannel> Ch = launchRemote(ChildExecPath, ChildPID);
+  if (!Ch) {
+    WithColor::error(errs(), argv[0]) << "Failed to launch remote JIT.\n";
+    exit(1);
+  }
+
+  LLVM_DEBUG({
+    dbgs()
+        << "Launched executable in subprocess " << ChildPID << ":\n"
+        << ChildExecPath << "\n\n"
+        << "You may want to attach a debugger now. Press enter to continue.\n";
+    fflush(stdin);
+    getchar();
+  });
+
+  std::unique_ptr<remote::OrcRemoteTargetClient> Client =
+      ExitOnErr(remote::OrcRemoteTargetClient::Create(*Ch, ES));
+
+  // Create a JITLink-compatible remote memory manager.
+  using MemManager = remote::OrcRemoteTargetClient::RemoteJITLinkMemoryManager;
+  std::unique_ptr<MemManager> RemoteMM =
+      ExitOnErr(Client->createRemoteJITLinkMemoryManager());
+
+  // Our remote target is running on the host system.
+  auto JTMB = ExitOnErr(JITTargetMachineBuilder::detectHost());
+  JTMB.setCodeModel(CodeModel::Small);
+
+  // Create an LLJIT instance with a JITLink ObjectLinkingLayer.
+  auto J = ExitOnErr(
+      LLJITBuilder()
+          .setJITTargetMachineBuilder(std::move(JTMB))
+          .setObjectLinkingLayerCreator(
+              [&](ExecutionSession &ES,
+                  const Triple &TT) -> std::unique_ptr<ObjectLayer> {
+                return std::make_unique<ObjectLinkingLayer>(ES, *RemoteMM);
+              })
+          .create());
+
+  auto M = ExitOnErr(parseExampleModule(Add1Example, "add1"));
+
+  ExitOnErr(J->addIRModule(std::move(M)));
+
+  // Look up the JIT'd function.
+  auto Add1Sym = ExitOnErr(J->lookup("add1"));
+
+  // Run in child target.
+  Expected<int> Result = Client->callIntInt(Add1Sym.getAddress(), 42);
+  if (Result)
+    outs() << "add1(42) = " << *Result << "\n";
+  else
+    ES.reportError(Result.takeError());
+
+  // Signal the remote target that we're done JITing.
+  ExitOnErr(Client->terminateSession());
+  LLVM_DEBUG(dbgs() << "Subprocess terminated\n");
+
+  return 0;
+}

examples/OrcV2Examples/LLJITWithChildProcess/RemoteJITUtils.h

@@ -0,0 +1,121 @@
+//===-- RemoteJITUtils.h - Utilities for remote-JITing ----------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Utilities for remote-JITing
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXAMPLES_ORCV2EXAMPLES_LLJITWITHCHILDPROCESS_REMOTEJITUTILS_H
+#define LLVM_EXAMPLES_ORCV2EXAMPLES_LLJITWITHCHILDPROCESS_REMOTEJITUTILS_H
+
+#include "llvm/ExecutionEngine/Orc/RPC/RawByteChannel.h"
+#include <mutex>
+
+#if !defined(_MSC_VER) && !defined(__MINGW32__)
+#include <unistd.h>
+#else
+#include <io.h>
+#endif
+
+/// RPC channel that reads from and writes from file descriptors.
+class FDRawChannel final : public llvm::orc::rpc::RawByteChannel {
+public:
+  FDRawChannel(int InFD, int OutFD) : InFD(InFD), OutFD(OutFD) {}
+
+  llvm::Error readBytes(char *Dst, unsigned Size) override {
+    assert(Dst && "Attempt to read into null.");
+    ssize_t Completed = 0;
+    while (Completed < static_cast<ssize_t>(Size)) {
+      ssize_t Read = ::read(InFD, Dst + Completed, Size - Completed);
+      if (Read <= 0) {
+        auto ErrNo = errno;
+        if (ErrNo == EAGAIN || ErrNo == EINTR)
+          continue;
+        else
+          return llvm::errorCodeToError(
+              std::error_code(errno, std::generic_category()));
+      }
+      Completed += Read;
+    }
+    return llvm::Error::success();
+  }
+
+  llvm::Error appendBytes(const char *Src, unsigned Size) override {
+    assert(Src && "Attempt to append from null.");
+    ssize_t Completed = 0;
+    while (Completed < static_cast<ssize_t>(Size)) {
+      ssize_t Written = ::write(OutFD, Src + Completed, Size - Completed);
+      if (Written < 0) {
+        auto ErrNo = errno;
+        if (ErrNo == EAGAIN || ErrNo == EINTR)
+          continue;
+        else
+          return llvm::errorCodeToError(
+              std::error_code(errno, std::generic_category()));
+      }
+      Completed += Written;
+    }
+    return llvm::Error::success();
+  }
+
+  llvm::Error send() override { return llvm::Error::success(); }
+
+private:
+  int InFD, OutFD;
+};
+
+// Launch child process and return a channel to it.
+std::unique_ptr<FDRawChannel> launchRemote(std::string ExecPath,
+                                           pid_t &ChildPID) {
+  // Create two pipes.
+  int PipeFD[2][2];
+  if (pipe(PipeFD[0]) != 0 || pipe(PipeFD[1]) != 0)
+    perror("Error creating pipe: ");
+
+  ChildPID = fork();
+
+  if (ChildPID == 0) {
+    // In the child...
+
+    // Close the parent ends of the pipes
+    close(PipeFD[0][1]);
+    close(PipeFD[1][0]);
+
+    // Execute the child process.
+    std::unique_ptr<char[]> ChildPath, ChildIn, ChildOut;
+    {
+      ChildPath.reset(new char[ExecPath.size() + 1]);
+      std::copy(ExecPath.begin(), ExecPath.end(), &ChildPath[0]);
+      ChildPath[ExecPath.size()] = '\0';
+      std::string ChildInStr = llvm::utostr(PipeFD[0][0]);
+      ChildIn.reset(new char[ChildInStr.size() + 1]);
+      std::copy(ChildInStr.begin(), ChildInStr.end(), &ChildIn[0]);
+      ChildIn[ChildInStr.size()] = '\0';
+      std::string ChildOutStr = llvm::utostr(PipeFD[1][1]);
+      ChildOut.reset(new char[ChildOutStr.size() + 1]);
+      std::copy(ChildOutStr.begin(), ChildOutStr.end(), &ChildOut[0]);
+      ChildOut[ChildOutStr.size()] = '\0';
+    }
+
+    char *const args[] = {&ChildPath[0], &ChildIn[0], &ChildOut[0], nullptr};
+    int rc = execv(ExecPath.c_str(), args);
+    if (rc != 0)
+      perror("Error executing child process: ");
+    llvm_unreachable("Error executing child process");
+  }
+  // else we're the parent...
+
+  // Close the child ends of the pipes
+  close(PipeFD[0][0]);
+  close(PipeFD[1][1]);
+
+  // Return an RPC channel connected to our end of the pipes.
+  return std::make_unique<FDRawChannel>(PipeFD[1][0], PipeFD[0][1]);
+}
+
+#endif

include/llvm/ExecutionEngine/Orc/OrcRemoteTargetClient.h

@@ -329,6 +329,218 @@ public:
     std::vector<EHFrame> RegisteredEHFrames;
   };
 
+  class RPCMMAlloc : public jitlink::JITLinkMemoryManager::Allocation {
+    using AllocationMap = DenseMap<unsigned, sys::MemoryBlock>;
+    using FinalizeContinuation =
+        jitlink::JITLinkMemoryManager::Allocation::FinalizeContinuation;
+    using ProtectionFlags = sys::Memory::ProtectionFlags;
+    using SegmentsRequestMap =
+        DenseMap<unsigned, jitlink::JITLinkMemoryManager::SegmentRequest>;
+
+    RPCMMAlloc(OrcRemoteTargetClient &Client, ResourceIdMgr::ResourceId Id)
+        : Client(Client), Id(Id) {}
+
+  public:
+    static Expected<std::unique_ptr<RPCMMAlloc>>
+    Create(OrcRemoteTargetClient &Client, ResourceIdMgr::ResourceId Id,
+           const SegmentsRequestMap &Request) {
+      auto *MM = new RPCMMAlloc(Client, Id);
+
+      if (Error Err = MM->allocateHostBlocks(Request))
+        return std::move(Err);
+
+      if (Error Err = MM->allocateTargetBlocks())
+        return std::move(Err);
+
+      return std::unique_ptr<RPCMMAlloc>(MM);
+    }
+
+    MutableArrayRef<char> getWorkingMemory(ProtectionFlags Seg) override {
+      assert(HostSegBlocks.count(Seg) && "No allocation for segment");
+      return {static_cast<char *>(HostSegBlocks[Seg].base()),
+              HostSegBlocks[Seg].allocatedSize()};
+    }
+
+    JITTargetAddress getTargetMemory(ProtectionFlags Seg) override {
+      assert(TargetSegBlocks.count(Seg) && "No allocation for segment");
+      return pointerToJITTargetAddress(TargetSegBlocks[Seg].base());
+    }
+
+    void finalizeAsync(FinalizeContinuation OnFinalize) override {
+      // Host allocations (working memory) remain ReadWrite.
+      OnFinalize(copyAndProtect());
+    }
+
+    Error deallocate() override {
+      // TODO: Cannot release target allocation. RPCAPI has no function
+      // symmetric to reserveMem(). Add RPC call like freeMem()?
+      return errorCodeToError(sys::Memory::releaseMappedMemory(HostAllocation));
+    }
+
+  private:
+    OrcRemoteTargetClient &Client;
+    ResourceIdMgr::ResourceId Id;
+    AllocationMap HostSegBlocks;
+    AllocationMap TargetSegBlocks;
+    JITTargetAddress TargetSegmentAddr;
+    sys::MemoryBlock HostAllocation;
+
+    Error allocateHostBlocks(const SegmentsRequestMap &Request) {
+      unsigned TargetPageSize = Client.getPageSize();
+
+      if (!isPowerOf2_64(static_cast<uint64_t>(TargetPageSize)))
+        return make_error<StringError>("Host page size is not a power of 2",
+                                       inconvertibleErrorCode());
+
+      auto TotalSize = calcTotalAllocSize(Request, TargetPageSize);
+      if (!TotalSize)
+        return TotalSize.takeError();
+
+      // Allocate one slab to cover all the segments.
+      const sys::Memory::ProtectionFlags ReadWrite =
+          static_cast<sys::Memory::ProtectionFlags>(sys::Memory::MF_READ |
+                                                    sys::Memory::MF_WRITE);
+      std::error_code EC;
+      HostAllocation =
+          sys::Memory::allocateMappedMemory(*TotalSize, nullptr, ReadWrite, EC);
+      if (EC)
+        return errorCodeToError(EC);
+
+      char *SlabAddr = static_cast<char *>(HostAllocation.base());
+      char *SlabAddrEnd = SlabAddr + HostAllocation.allocatedSize();
+
+      // Allocate segment memory from the slab.
+      for (auto &KV : Request) {
+        const auto &Seg = KV.second;
+
+        uint64_t SegmentSize = Seg.getContentSize() + Seg.getZeroFillSize();
+        uint64_t AlignedSegmentSize = alignTo(SegmentSize, TargetPageSize);
+
+        // Zero out zero-fill memory.
+        char *ZeroFillBegin = SlabAddr + Seg.getContentSize();
+        memset(ZeroFillBegin, 0, Seg.getZeroFillSize());
+
+        // Record the block for this segment.
+        HostSegBlocks[KV.first] =
+            sys::MemoryBlock(SlabAddr, AlignedSegmentSize);
+
+        SlabAddr += AlignedSegmentSize;
+        assert(SlabAddr <= SlabAddrEnd && "Out of range");
+      }
+
+      return Error::success();
+    }
+
+    Error allocateTargetBlocks() {
+      // Reserve memory for all blocks on the target. We need as much space on
+      // the target as we allocated on the host.
+      TargetSegmentAddr = Client.reserveMem(Id, HostAllocation.allocatedSize(),
+                                            Client.getPageSize());
+      if (!TargetSegmentAddr)
+        return make_error<StringError>("Failed to reserve memory on the target",
+                                       inconvertibleErrorCode());
+
+      // Map memory blocks into the allocation, that match the host allocation.
+      JITTargetAddress TargetAllocAddr = TargetSegmentAddr;
+      for (const auto &KV : HostSegBlocks) {
+        size_t TargetAllocSize = KV.second.allocatedSize();
+
+        TargetSegBlocks[KV.first] =
+            sys::MemoryBlock(jitTargetAddressToPointer<void *>(TargetAllocAddr),
+                             TargetAllocSize);
+
+        TargetAllocAddr += TargetAllocSize;
+        assert(TargetAllocAddr - TargetSegmentAddr <=
+                   HostAllocation.allocatedSize() &&
+               "Out of range on target");
+      }
+
+      return Error::success();
+    }
+
+    Error copyAndProtect() {
+      unsigned Permissions = 0u;
+
+      // Copy segments one by one.
+      for (auto &KV : TargetSegBlocks) {
+        Permissions |= KV.first;
+
+        const sys::MemoryBlock &TargetBlock = KV.second;
+        const sys::MemoryBlock &HostBlock = HostSegBlocks.lookup(KV.first);
+
+        size_t TargetAllocSize = TargetBlock.allocatedSize();
+        auto TargetAllocAddr = pointerToJITTargetAddress(TargetBlock.base());
+        auto *HostAllocBegin = static_cast<const char *>(HostBlock.base());
+
+        bool CopyErr =
+            Client.writeMem(TargetAllocAddr, HostAllocBegin, TargetAllocSize);
+        if (CopyErr)
+          return createStringError(inconvertibleErrorCode(),
+                                   "Failed to copy %d segment to the target",
+                                   KV.first);
+      }
+
+      // Set permission flags for all segments at once.
+      bool ProtectErr =
+          Client.setProtections(Id, TargetSegmentAddr, Permissions);
+      if (ProtectErr)
+        return createStringError(inconvertibleErrorCode(),
+                                 "Failed to apply permissions for %d segment "
+                                 "on the target",
+                                 Permissions);
+      return Error::success();
+    }
+
+    static Expected<size_t>
+    calcTotalAllocSize(const SegmentsRequestMap &Request,
+                       unsigned TargetPageSize) {
+      size_t TotalSize = 0;
+      for (const auto &KV : Request) {
+        const auto &Seg = KV.second;
+
+        if (Seg.getAlignment() > TargetPageSize)
+          return make_error<StringError>("Cannot request alignment higher than "
+                                         "page alignment on target",
+                                         inconvertibleErrorCode());
+
+        TotalSize = alignTo(TotalSize, TargetPageSize);
+        TotalSize += Seg.getContentSize();
+        TotalSize += Seg.getZeroFillSize();
+      }
+
+      return TotalSize;
+    }
+  };
+
+  class RemoteJITLinkMemoryManager : public jitlink::JITLinkMemoryManager {
+  public:
+    RemoteJITLinkMemoryManager(OrcRemoteTargetClient &Client,
+                               ResourceIdMgr::ResourceId Id)
+        : Client(Client), Id(Id) {}
+
+    RemoteJITLinkMemoryManager(const RemoteJITLinkMemoryManager &) = delete;
+    RemoteJITLinkMemoryManager(RemoteJITLinkMemoryManager &&) = default;
+
+    RemoteJITLinkMemoryManager &
+    operator=(const RemoteJITLinkMemoryManager &) = delete;
+    RemoteJITLinkMemoryManager &
+    operator=(RemoteJITLinkMemoryManager &&) = delete;
+
+    ~RemoteJITLinkMemoryManager() {
+      Client.destroyRemoteAllocator(Id);
+      LLVM_DEBUG(dbgs() << "Destroyed remote allocator " << Id << "\n");
+    }
+
+    Expected<std::unique_ptr<Allocation>>
+    allocate(const SegmentsRequestMap &Request) override {
+      return RPCMMAlloc::Create(Client, Id, Request);
+    }
+
+  private:
+    OrcRemoteTargetClient &Client;
+    ResourceIdMgr::ResourceId Id;
+  };
+
   /// Remote indirect stubs manager.
   class RemoteIndirectStubsManager : public IndirectStubsManager {
   public:
@@ -504,6 +716,14 @@ public:
     return callB<exec::CallIntVoid>(Addr);
   }
 
+  /// Call the int(int) function at the given address in the target and return
+  /// its result.
+  Expected<int> callIntInt(JITTargetAddress Addr, int Arg) {
+    LLVM_DEBUG(dbgs() << "Calling int(*)(int) " << format("0x%016" PRIx64, Addr)
+                      << "\n");
+    return callB<exec::CallIntInt>(Addr, Arg);
+  }
+
   /// Call the int(int, char*[]) function at the given address in the target and
   /// return its result.
   Expected<int> callMain(JITTargetAddress Addr,
@@ -532,6 +752,18 @@ public:
         new RemoteRTDyldMemoryManager(*this, Id));
   }
 
+  /// Create a JITLink-compatible memory manager which will allocate working
+  /// memory on the host and target memory on the remote target.
+  Expected<std::unique_ptr<RemoteJITLinkMemoryManager>>
+  createRemoteJITLinkMemoryManager() {
+    auto Id = AllocatorIds.getNext();
+    if (auto Err = callB<mem::CreateRemoteAllocator>(Id))
+      return std::move(Err);
+    LLVM_DEBUG(dbgs() << "Created remote allocator " << Id << "\n");
+    return std::unique_ptr<RemoteJITLinkMemoryManager>(
+        new RemoteJITLinkMemoryManager(*this, Id));
+  }
+
   /// Create an RCIndirectStubsManager that will allocate stubs on the remote
   /// target.
   Expected<std::unique_ptr<RemoteIndirectStubsManager>>

include/llvm/ExecutionEngine/Orc/OrcRemoteTargetRPCAPI.h

@@ -195,6 +195,14 @@ namespace exec {
     static const char *getName() { return "CallIntVoid"; }
   };
 
+  /// Call an 'int32_t(int32_t)'-type function on the remote, returns the called
+  /// function's return value.
+  class CallIntInt
+      : public rpc::Function<CallIntInt, int32_t(JITTargetAddress Addr, int)> {
+  public:
+    static const char *getName() { return "CallIntInt"; }
+  };
+
   /// Call an 'int32_t(int32_t, char**)'-type function on the remote, returns the
   /// called function's return value.
   class CallMain

include/llvm/ExecutionEngine/Orc/OrcRemoteTargetServer.h

@@ -63,6 +63,7 @@ public:
         EHFramesDeregister(std::move(EHFramesDeregister)) {
     using ThisT = std::remove_reference_t<decltype(*this)>;
     addHandler<exec::CallIntVoid>(*this, &ThisT::handleCallIntVoid);
+    addHandler<exec::CallIntInt>(*this, &ThisT::handleCallIntInt);
     addHandler<exec::CallMain>(*this, &ThisT::handleCallMain);
     addHandler<exec::CallVoidVoid>(*this, &ThisT::handleCallVoidVoid);
     addHandler<mem::CreateRemoteAllocator>(*this,
@@ -168,6 +169,19 @@ private:
     return Result;
   }
 
+  Expected<int32_t> handleCallIntInt(JITTargetAddress Addr, int Arg) {
+    using IntIntFnTy = int (*)(int);
+
+    IntIntFnTy Fn = reinterpret_cast<IntIntFnTy>(static_cast<uintptr_t>(Addr));
+
+    LLVM_DEBUG(dbgs() << "  Calling " << format("0x%016x", Addr)
+                      << " with argument " << Arg << "\n");
+    int Result = Fn(Arg);
+    LLVM_DEBUG(dbgs() << "  Result = " << Result << "\n");
+
+    return Result;
+  }
+
   Expected<int32_t> handleCallMain(JITTargetAddress Addr,
                                    std::vector<std::string> Args) {
     using MainFnTy = int (*)(int, const char *[]);