[Orc] Use MapperJITLinkMemoryManager with InProcessMapper in llvm-jitlink tool

[argentite]

Summary

MapperJITLinkMemoryManager has slab allocation. Combined with
InProcessMapper, it can replace InProcessMemoryManager.

It can also replace JITLinkSlabAllocator through the InProcessDeltaMapper
that adds an offset to the executor addresses for use in tests.

[lhames]

Otherwise LGTM!

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

458

I'd change this to "only used for -noexec testing" -- I don't think there's any other use case.

llvm/include/llvm/ExecutionEngine/Orc/MemoryMapper.h

public:
 
   /// Release address space acquired through reserve()
   virtual void release(ArrayRef<ExecutorAddr> Reservations,
                        OnReleasedFunction OnRelease) = 0;
 
   virtual ~MemoryMapper();
 };
 
-class InProcessMemoryMapper final : public MemoryMapper {
+class InProcessMemoryMapper : public MemoryMapper {
 public:
   InProcessMemoryMapper(size_t PageSize);
 
   static Expected<std::unique_ptr<InProcessMemoryMapper>> Create();
 
   unsigned int getPageSize() override { return PageSize; }
 
   void reserve(size_t NumBytes, OnReservedFunction OnReserved) override;

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

 #include "llvm/ExecutionEngine/Orc/DebuggerSupportPlugin.h"
 #include "llvm/ExecutionEngine/Orc/ELFNixPlatform.h"
 #include "llvm/ExecutionEngine/Orc/EPCDebugObjectRegistrar.h"
 #include "llvm/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.h"
 #include "llvm/ExecutionEngine/Orc/EPCEHFrameRegistrar.h"
 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
 #include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
 #include "llvm/ExecutionEngine/Orc/MachOPlatform.h"
+#include "llvm/ExecutionEngine/Orc/MapperJITLinkMemoryManager.h"
 #include "llvm/ExecutionEngine/Orc/ObjectFileInterface.h"
 #include "llvm/ExecutionEngine/Orc/TargetProcess/JITLoaderGDB.h"
 #include "llvm/ExecutionEngine/Orc/TargetProcess/RegisterEHFrames.h"
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCDisassembler/MCDisassembler.h"
 #include "llvm/MC/MCInstPrinter.h"
 #include "llvm/MC/MCInstrAnalysis.h"
for (auto *Sym : Syms) {
           OS << formatv(" {0:x-2}", SymData[NextAddr - SymStart]);
         ++NextAddr;
       }
     }
     OS << "\n";
   }
 }
 
+// A memory mapper with a fake offset applied only used for -noexec testing
+class InProcessDeltaMapper final : public InProcessMemoryMapper {
+public:
+  InProcessDeltaMapper(size_t PageSize, uint64_t TargetAddr)
+      : InProcessMemoryMapper(PageSize), TargetMapAddr(TargetAddr),
+        DeltaAddr(0) {}
+
+  static Expected<std::unique_ptr<InProcessDeltaMapper>> Create() {
+    auto PageSize = sys::Process::getPageSize();
+    if (!PageSize)
+      return PageSize.takeError();
+    return std::make_unique<InProcessDeltaMapper>(*PageSize, SlabAddress);
+  }
+
+  void reserve(size_t NumBytes, OnReservedFunction OnReserved) override {
+    InProcessMemoryMapper::reserve(
+        NumBytes, [this, OnReserved = std::move(OnReserved)](
+                      Expected<ExecutorAddrRange> Result) mutable {
+          if (!Result)
+            return OnReserved(Result.takeError());
+
+          assert(DeltaAddr == 0 && "Overwriting previous offset");
+          if (TargetMapAddr != ~0ULL)
+            DeltaAddr = TargetMapAddr - Result->Start.getValue();
+          auto OffsetRange = ExecutorAddrRange(Result->Start + DeltaAddr,
+                                               Result->End + DeltaAddr);
+
+          OnReserved(OffsetRange);
+        });
+  }
+
+  char *prepare(ExecutorAddr Addr, size_t ContentSize) override {
+    return InProcessMemoryMapper::prepare(Addr - DeltaAddr, ContentSize);
+  }
+
+  void initialize(AllocInfo &AI, OnInitializedFunction OnInitialized) override {
+    auto FixedAI = AI;
+    FixedAI.MappingBase -= DeltaAddr;
+    InProcessMemoryMapper::initialize(
+        FixedAI, [this, OnInitialized = std::move(OnInitialized)](
+                     Expected<ExecutorAddr> Result) mutable {
+          if (!Result)
+            return OnInitialized(Result.takeError());
+
+          OnInitialized(ExecutorAddr(Result->getValue() + DeltaAddr));
+        });
+  }
+
+  void deinitialize(ArrayRef<ExecutorAddr> Allocations,
+                    OnDeinitializedFunction OnDeInitialized) override {
+    std::vector<ExecutorAddr> Addrs(Allocations.size());
+    for (const auto Base : Allocations) {
+      Addrs.push_back(Base - DeltaAddr);
+    }
+
+    InProcessMemoryMapper::deinitialize(Addrs, std::move(OnDeInitialized));
+  }
+
+  void release(ArrayRef<ExecutorAddr> Reservations,
+               OnReleasedFunction OnRelease) override {
+    std::vector<ExecutorAddr> Addrs(Reservations.size());
+    for (const auto Base : Reservations) {
+      Addrs.push_back(Base - DeltaAddr);
+    }
+    InProcessMemoryMapper::release(Addrs, std::move(OnRelease));
+  }
+
+private:
+  uint64_t TargetMapAddr;
+  uint64_t DeltaAddr;
+};
+
 class JITLinkSlabAllocator final : public JITLinkMemoryManager {
 private:
   struct FinalizedAllocInfo {
     FinalizedAllocInfo(sys::MemoryBlock Mem,
                        std::vector<shared::WrapperFunctionCall> DeallocActions)
         : Mem(Mem), DeallocActions(std::move(DeallocActions)) {}
     sys::MemoryBlock Mem;
     std::vector<shared::WrapperFunctionCall> DeallocActions;
Expected<uint64_t> getSlabAllocSize(StringRef SizeString) {
   uint64_t SlabSize = 0;
   if (SizeString.getAsInteger(10, SlabSize))
     return make_error<StringError>("Invalid numeric format for slab size",
                                    inconvertibleErrorCode());
 
   return SlabSize * Units;
 }
 
-static std::unique_ptr<JITLinkMemoryManager> createMemoryManager() {
+static std::unique_ptr<JITLinkMemoryManager> createInProcessMemoryManager() {
   if (!SlabAllocateSizeString.empty()) {
     auto SlabSize = ExitOnErr(getSlabAllocSize(SlabAllocateSizeString));
-    return ExitOnErr(JITLinkSlabAllocator::Create(SlabSize));
+
+    return ExitOnErr(
+        MapperJITLinkMemoryManager::CreateWithMapper<InProcessDeltaMapper>(
+            SlabSize));
   }
-  return ExitOnErr(InProcessMemoryManager::Create());
+
+#ifdef _WIN32
+  return ExitOnErr(
+      MapperJITLinkMemoryManager::CreateWithMapper<InProcessMemoryMapper>(
+          1024 * 1024));
+#else
+  return ExitOnErr(
+      MapperJITLinkMemoryManager::CreateWithMapper<InProcessMemoryMapper>(
+          1024 * 1024 * 1024));
+#endif
 }
 
 static Expected<MaterializationUnit::Interface>
 getTestObjectFileInterface(Session &S, MemoryBufferRef O) {
 
   // Get the standard interface for this object, but ignore the symbols field.
   // We'll handle that manually to include promotion.
   auto I = getObjectFileInterface(S.ES, O);
Expected<std::unique_ptr<Session>> Session::Create(Triple TT) {
   } else {
     /// Otherwise use SelfExecutorProcessControl to target the current process.
     auto PageSize = sys::Process::getPageSize();
     if (!PageSize)
       return PageSize.takeError();
     EPC = std::make_unique<SelfExecutorProcessControl>(
         std::make_shared<SymbolStringPool>(),
         std::make_unique<InPlaceTaskDispatcher>(), std::move(TT), *PageSize,
-        createMemoryManager());
+        createInProcessMemoryManager());
   }
 
   Error Err = Error::success();
   std::unique_ptr<Session> S(new Session(std::move(EPC), Err));
   if (Err)
     return std::move(Err);
   return std::move(S);
 }