vendor/llvm-project/llvmorg-15-init-17485-ga3e38b4a206b

1 files changed, 135 insertions, 0 deletions

diff --git a/llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp b/llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp
new file mode 100644
index 000000000000..c2e7baabb994
--- /dev/null
+++ b/llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp
@@ -0,0 +1,135 @@
+//=== MapperJITLinkMemoryManager.cpp - Memory management with MemoryMapper ===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ExecutionEngine/Orc/MapperJITLinkMemoryManager.h"
+
+#include "llvm/ExecutionEngine/JITLink/JITLink.h"
+#include "llvm/Support/Process.h"
+
+#include <limits>
+
+using namespace llvm::jitlink;
+
+namespace llvm {
+namespace orc {
+
+class MapperJITLinkMemoryManager::InFlightAlloc
+    : public JITLinkMemoryManager::InFlightAlloc {
+public:
+  InFlightAlloc(MapperJITLinkMemoryManager &Parent, LinkGraph &G,
+                ExecutorAddr AllocAddr,
+                std::vector<MemoryMapper::AllocInfo::SegInfo> Segs)
+      : Parent(Parent), G(G), AllocAddr(AllocAddr), Segs(std::move(Segs)) {}
+
+  void finalize(OnFinalizedFunction OnFinalize) override {
+    MemoryMapper::AllocInfo AI;
+    AI.MappingBase = AllocAddr;
+
+    std::swap(AI.Segments, Segs);
+    std::swap(AI.Actions, G.allocActions());
+
+    Parent.Mapper->initialize(AI, [&](Expected<ExecutorAddr> Result) {
+      if (!Result) {
+        OnFinalize(Result.takeError());
+        return;
+      }
+
+      OnFinalize(FinalizedAlloc(*Result));
+    });
+  }
+
+  void abandon(OnAbandonedFunction OnFinalize) override {
+    Parent.Mapper->release({AllocAddr}, std::move(OnFinalize));
+  }
+
+private:
+  MapperJITLinkMemoryManager &Parent;
+  LinkGraph &G;
+  ExecutorAddr AllocAddr;
+  std::vector<MemoryMapper::AllocInfo::SegInfo> Segs;
+};
+
+MapperJITLinkMemoryManager::MapperJITLinkMemoryManager(
+    std::unique_ptr<MemoryMapper> Mapper)
+    : Mapper(std::move(Mapper)) {}
+
+void MapperJITLinkMemoryManager::allocate(const JITLinkDylib *JD, LinkGraph &G,
+                                          OnAllocatedFunction OnAllocated) {
+  BasicLayout BL(G);
+
+  // find required address space
+  auto SegsSizes = BL.getContiguousPageBasedLayoutSizes(Mapper->getPageSize());
+  if (!SegsSizes) {
+    OnAllocated(SegsSizes.takeError());
+    return;
+  }
+
+  // Check if total size fits in address space
+  if (SegsSizes->total() > std::numeric_limits<size_t>::max()) {
+    OnAllocated(make_error<JITLinkError>(
+        formatv("Total requested size {:x} for graph {} exceeds address space",
+                SegsSizes->total(), G.getName())));
+    return;
+  }
+
+  Mapper->reserve(
+      SegsSizes->total(),
+      [this, &G, BL = std::move(BL), OnAllocated = std::move(OnAllocated)](
+          Expected<ExecutorAddrRange> Result) mutable {
+        if (!Result) {
+          return OnAllocated(Result.takeError());
+        }
+
+        auto NextSegAddr = Result->Start;
+
+        std::vector<MemoryMapper::AllocInfo::SegInfo> SegInfos;
+
+        for (auto &KV : BL.segments()) {
+          auto &AG = KV.first;
+          auto &Seg = KV.second;
+
+          auto TotalSize = Seg.ContentSize + Seg.ZeroFillSize;
+
+          Seg.Addr = NextSegAddr;
+          Seg.WorkingMem = Mapper->prepare(NextSegAddr, TotalSize);
+
+          NextSegAddr += alignTo(TotalSize, Mapper->getPageSize());
+
+          MemoryMapper::AllocInfo::SegInfo SI;
+          SI.Offset = Seg.Addr - Result->Start;
+          SI.ContentSize = Seg.ContentSize;
+          SI.ZeroFillSize = Seg.ZeroFillSize;
+          SI.Prot = (toSysMemoryProtectionFlags(AG.getMemProt()));
+          SI.WorkingMem = Seg.WorkingMem;
+
+          SegInfos.push_back(SI);
+        }
+
+        if (auto Err = BL.apply()) {
+          OnAllocated(std::move(Err));
+          return;
+        }
+
+        OnAllocated(std::make_unique<InFlightAlloc>(*this, G, Result->Start,
+                                                    std::move(SegInfos)));
+      });
+}
+
+void MapperJITLinkMemoryManager::deallocate(
+    std::vector<FinalizedAlloc> Allocs, OnDeallocatedFunction OnDeallocated) {
+  std::vector<ExecutorAddr> Bases;
+  Bases.reserve(Allocs.size());
+  for (auto &FA : Allocs) {
+    Bases.push_back(FA.getAddress());
+    FA.release();
+  }
+  Mapper->release(Bases, std::move(OnDeallocated));
+}
+
+} // end namespace orc
+} // end namespace llvm