[ORC] Add LinkGraph& argument to MemoryMapper::prepare. (#163121)

This gives MemoryMapper implementations a chance to allocate working
memory using the LinkGraph's allocator.

Author: lhames

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

@@ -51,7 +51,11 @@ class LLVM_ABI MemoryMapper {
   virtual void reserve(size_t NumBytes, OnReservedFunction OnReserved) = 0;
 
   /// Provides working memory
-  virtual char *prepare(ExecutorAddr Addr, size_t ContentSize) = 0;
+  /// The LinkGraph parameter is included to allow implementations to allocate
+  /// working memory from the LinkGraph's allocator, in which case it will be
+  /// deallocated when the LinkGraph is destroyed.
+  virtual char *prepare(jitlink::LinkGraph &G, ExecutorAddr Addr,
+                        size_t ContentSize) = 0;
 
   using OnInitializedFunction = unique_function<void(Expected<ExecutorAddr>)>;
 
@@ -92,7 +96,8 @@ class LLVM_ABI InProcessMemoryMapper : public MemoryMapper {
 
   void initialize(AllocInfo &AI, OnInitializedFunction OnInitialized) override;
 
-  char *prepare(ExecutorAddr Addr, size_t ContentSize) override;
+  char *prepare(jitlink::LinkGraph &G, ExecutorAddr Addr,
+                size_t ContentSize) override;
 
   void deinitialize(ArrayRef<ExecutorAddr> Allocations,
                     OnDeinitializedFunction OnDeInitialized) override;
@@ -142,7 +147,8 @@ class LLVM_ABI SharedMemoryMapper final : public MemoryMapper {
 
   void reserve(size_t NumBytes, OnReservedFunction OnReserved) override;
 
-  char *prepare(ExecutorAddr Addr, size_t ContentSize) override;
+  char *prepare(jitlink::LinkGraph &G, ExecutorAddr Addr,
+                size_t ContentSize) override;
 
   void initialize(AllocInfo &AI, OnInitializedFunction OnInitialized) override;
 

llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp

@@ -90,7 +90,7 @@ void MapperJITLinkMemoryManager::allocate(const JITLinkDylib *JD, LinkGraph &G,
       auto TotalSize = Seg.ContentSize + Seg.ZeroFillSize;
 
       Seg.Addr = NextSegAddr;
-      Seg.WorkingMem = Mapper->prepare(NextSegAddr, TotalSize);
+      Seg.WorkingMem = Mapper->prepare(G, NextSegAddr, TotalSize);
 
       NextSegAddr += alignTo(TotalSize, Mapper->getPageSize());
 

llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp

@@ -58,7 +58,8 @@ void InProcessMemoryMapper::reserve(size_t NumBytes,
       ExecutorAddrRange(ExecutorAddr::fromPtr(MB.base()), MB.allocatedSize()));
 }
 
-char *InProcessMemoryMapper::prepare(ExecutorAddr Addr, size_t ContentSize) {
+char *InProcessMemoryMapper::prepare(jitlink::LinkGraph &G, ExecutorAddr Addr,
+                                     size_t ContentSize) {
   return Addr.toPtr<char *>();
 }
 
@@ -324,7 +325,8 @@ void SharedMemoryMapper::reserve(size_t NumBytes,
 #endif
 }
 
-char *SharedMemoryMapper::prepare(ExecutorAddr Addr, size_t ContentSize) {
+char *SharedMemoryMapper::prepare(jitlink::LinkGraph &G, ExecutorAddr Addr,
+                                  size_t ContentSize) {
   auto R = Reservations.upper_bound(Addr);
   assert(R != Reservations.begin() && "Attempt to prepare unreserved range");
   R--;

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

@@ -623,8 +623,9 @@ class InProcessDeltaMapper final : public InProcessMemoryMapper {
         });
   }
 
-  char *prepare(ExecutorAddr Addr, size_t ContentSize) override {
-    return InProcessMemoryMapper::prepare(Addr - DeltaAddr, ContentSize);
+  char *prepare(jitlink::LinkGraph &G, ExecutorAddr Addr,
+                size_t ContentSize) override {
+    return InProcessMemoryMapper::prepare(G, Addr - DeltaAddr, ContentSize);
   }
 
   void initialize(AllocInfo &AI, OnInitializedFunction OnInitialized) override {

llvm/unittests/ExecutionEngine/Orc/MapperJITLinkMemoryManagerTest.cpp

@@ -39,8 +39,8 @@ class CounterMapper final : public MemoryMapper {
     return Mapper->initialize(AI, std::move(OnInitialized));
   }
 
-  char *prepare(ExecutorAddr Addr, size_t ContentSize) override {
-    return Mapper->prepare(Addr, ContentSize);
+  char *prepare(LinkGraph &G, ExecutorAddr Addr, size_t ContentSize) override {
+    return Mapper->prepare(G, Addr, ContentSize);
   }
 
   void deinitialize(ArrayRef<ExecutorAddr> Allocations,

llvm/unittests/ExecutionEngine/Orc/MemoryMapperTest.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ExecutionEngine/Orc/MemoryMapper.h"
+#include "llvm/ExecutionEngine/JITLink/JITLink.h"
 #include "llvm/Support/Process.h"
 #include "llvm/Testing/Support/Error.h"
 #include "gtest/gtest.h"
@@ -66,6 +67,9 @@ TEST(MemoryMapperTest, InitializeDeinitialize) {
   {
     std::unique_ptr<MemoryMapper> Mapper =
         cantFail(InProcessMemoryMapper::Create());
+    jitlink::LinkGraph G("G", std::make_shared<SymbolStringPool>(),
+                         Triple("x86_64-apple-darwin"), SubtargetFeatures(),
+                         jitlink::getGenericEdgeKindName);
 
     // We will do two separate allocations
     auto PageSize = Mapper->getPageSize();
@@ -80,7 +84,7 @@ TEST(MemoryMapperTest, InitializeDeinitialize) {
 
     {
       // Provide working memory
-      char *WA1 = Mapper->prepare(Mem1->Start, HW.size() + 1);
+      char *WA1 = Mapper->prepare(G, Mem1->Start, HW.size() + 1);
       std::strcpy(WA1, HW.c_str());
     }
 
@@ -105,7 +109,7 @@ TEST(MemoryMapperTest, InitializeDeinitialize) {
     }
 
     {
-      char *WA2 = Mapper->prepare(Mem1->Start + PageSize, HW.size() + 1);
+      char *WA2 = Mapper->prepare(G, Mem1->Start + PageSize, HW.size() + 1);
       std::strcpy(WA2, HW.c_str());
     }
 
@@ -158,7 +162,7 @@ TEST(MemoryMapperTest, InitializeDeinitialize) {
       auto Mem2 = reserve(*Mapper, PageSize);
       EXPECT_THAT_ERROR(Mem2.takeError(), Succeeded());
 
-      char *WA = Mapper->prepare(Mem2->Start, HW.size() + 1);
+      char *WA = Mapper->prepare(G, Mem2->Start, HW.size() + 1);
       std::strcpy(WA, HW.c_str());
 
       MemoryMapper::AllocInfo Alloc3;

llvm/unittests/ExecutionEngine/Orc/SharedMemoryMapperTest.cpp

@@ -8,6 +8,7 @@
 
 #include "OrcTestCommon.h"
 #include "llvm/Config/llvm-config.h" // for LLVM_ON_UNIX
+#include "llvm/ExecutionEngine/JITLink/JITLink.h"
 #include "llvm/ExecutionEngine/Orc/MemoryMapper.h"
 #include "llvm/ExecutionEngine/Orc/SelfExecutorProcessControl.h"
 #include "llvm/ExecutionEngine/Orc/Shared/OrcRTBridge.h"
@@ -67,12 +68,16 @@ TEST(SharedMemoryMapperTest, MemReserveInitializeDeinitializeRelease) {
 
     auto PageSize = Mapper->getPageSize();
     size_t ReqSize = PageSize;
+    jitlink::LinkGraph G("G", std::make_shared<SymbolStringPool>(),
+                         Triple("x86_64-apple-darwin"), SubtargetFeatures(),
+                         jitlink::getGenericEdgeKindName);
 
     Mapper->reserve(ReqSize, [&](Expected<ExecutorAddrRange> Result) {
       EXPECT_THAT_ERROR(Result.takeError(), Succeeded());
       auto Reservation = std::move(*Result);
       {
-        char *Addr = Mapper->prepare(Reservation.Start, TestString.size() + 1);
+        char *Addr =
+            Mapper->prepare(G, Reservation.Start, TestString.size() + 1);
         std::strcpy(Addr, TestString.c_str());
       }
       MemoryMapper::AllocInfo AI;