xref: /llvm-project/llvm/lib/ExecutionEngine/Orc/EPCGenericJITLinkMemoryManager.cpp (revision da7f993a8d615c7306d3576f9669eb4d0a4b65f4)
1 //===---- EPCGenericJITLinkMemoryManager.cpp -- Mem management via EPC ----===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "llvm/ExecutionEngine/Orc/EPCGenericJITLinkMemoryManager.h"
10 #include "llvm/ExecutionEngine/Orc/LookupAndRecordAddrs.h"
11 #include "llvm/ExecutionEngine/Orc/Shared/OrcRTBridge.h"
12 
13 #include <limits>
14 
15 namespace llvm {
16 namespace orc {
17 
18 class EPCGenericJITLinkMemoryManager::Alloc
19     : public jitlink::JITLinkMemoryManager::Allocation {
20 public:
21   struct SegInfo {
22     char *WorkingMem = nullptr;
23     ExecutorAddr TargetAddr;
24     uint64_t ContentSize = 0;
25     uint64_t ZeroFillSize = 0;
26   };
27   using SegInfoMap = DenseMap<unsigned, SegInfo>;
28 
29   Alloc(EPCGenericJITLinkMemoryManager &Parent, ExecutorAddr TargetAddr,
30         std::unique_ptr<char[]> WorkingBuffer, SegInfoMap Segs)
31       : Parent(Parent), TargetAddr(TargetAddr),
32         WorkingBuffer(std::move(WorkingBuffer)), Segs(std::move(Segs)) {}
33 
34   MutableArrayRef<char> getWorkingMemory(ProtectionFlags Seg) override {
35     auto I = Segs.find(Seg);
36     assert(I != Segs.end() && "No allocation for seg");
37     assert(I->second.ContentSize <= std::numeric_limits<size_t>::max());
38     return {I->second.WorkingMem, static_cast<size_t>(I->second.ContentSize)};
39   }
40 
41   JITTargetAddress getTargetMemory(ProtectionFlags Seg) override {
42     auto I = Segs.find(Seg);
43     assert(I != Segs.end() && "No allocation for seg");
44     return I->second.TargetAddr.getValue();
45   }
46 
47   void finalizeAsync(FinalizeContinuation OnFinalize) override {
48     char *WorkingMem = WorkingBuffer.get();
49     tpctypes::FinalizeRequest FR;
50     for (auto &KV : Segs) {
51       assert(KV.second.ContentSize <= std::numeric_limits<size_t>::max());
52       FR.Segments.push_back(tpctypes::SegFinalizeRequest{
53           tpctypes::toWireProtectionFlags(
54               static_cast<sys::Memory::ProtectionFlags>(KV.first)),
55           KV.second.TargetAddr,
56           alignTo(KV.second.ContentSize + KV.second.ZeroFillSize,
57                   Parent.EPC.getPageSize()),
58           {WorkingMem, static_cast<size_t>(KV.second.ContentSize)}});
59       WorkingMem += KV.second.ContentSize;
60     }
61     Parent.EPC.callSPSWrapperAsync<
62         rt::SPSSimpleExecutorMemoryManagerFinalizeSignature>(
63         Parent.SAs.Finalize,
64         [OnFinalize = std::move(OnFinalize)](Error SerializationErr,
65                                              Error FinalizeErr) {
66           if (SerializationErr)
67             OnFinalize(std::move(SerializationErr));
68           else
69             OnFinalize(std::move(FinalizeErr));
70         },
71         Parent.SAs.Allocator, std::move(FR));
72   }
73 
74   Error deallocate() override {
75     Error Err = Error::success();
76     if (auto E2 = Parent.EPC.callSPSWrapper<
77                   rt::SPSSimpleExecutorMemoryManagerDeallocateSignature>(
78             Parent.SAs.Deallocate, Err, Parent.SAs.Allocator,
79             ArrayRef<ExecutorAddr>(TargetAddr)))
80       return E2;
81     return Err;
82   }
83 
84 private:
85   EPCGenericJITLinkMemoryManager &Parent;
86   ExecutorAddr TargetAddr;
87   std::unique_ptr<char[]> WorkingBuffer;
88   SegInfoMap Segs;
89 };
90 
91 Expected<std::unique_ptr<jitlink::JITLinkMemoryManager::Allocation>>
92 EPCGenericJITLinkMemoryManager::allocate(const jitlink::JITLinkDylib *JD,
93                                          const SegmentsRequestMap &Request) {
94   Alloc::SegInfoMap Segs;
95   uint64_t AllocSize = 0;
96   size_t WorkingSize = 0;
97   for (auto &KV : Request) {
98     if (!isPowerOf2_64(KV.second.getAlignment()))
99       return make_error<StringError>("Alignment is not a power of two",
100                                      inconvertibleErrorCode());
101     if (KV.second.getAlignment() > EPC.getPageSize())
102       return make_error<StringError>("Alignment exceeds page size",
103                                      inconvertibleErrorCode());
104 
105     auto &Seg = Segs[KV.first];
106     Seg.ContentSize = KV.second.getContentSize();
107     Seg.ZeroFillSize = KV.second.getZeroFillSize();
108     AllocSize += alignTo(Seg.ContentSize + Seg.ZeroFillSize, EPC.getPageSize());
109     WorkingSize += Seg.ContentSize;
110   }
111 
112   std::unique_ptr<char[]> WorkingBuffer;
113   if (WorkingSize > 0)
114     WorkingBuffer = std::make_unique<char[]>(WorkingSize);
115   Expected<ExecutorAddr> TargetAllocAddr((ExecutorAddr()));
116   if (auto Err = EPC.callSPSWrapper<
117                  rt::SPSSimpleExecutorMemoryManagerReserveSignature>(
118           SAs.Reserve, TargetAllocAddr, SAs.Allocator, AllocSize))
119     return std::move(Err);
120   if (!TargetAllocAddr)
121     return TargetAllocAddr.takeError();
122 
123   char *WorkingMem = WorkingBuffer.get();
124   JITTargetAddress SegAddr = TargetAllocAddr->getValue();
125   for (auto &KV : Segs) {
126     auto &Seg = KV.second;
127     Seg.TargetAddr.setValue(SegAddr);
128     SegAddr += alignTo(Seg.ContentSize + Seg.ZeroFillSize, EPC.getPageSize());
129     Seg.WorkingMem = WorkingMem;
130     WorkingMem += Seg.ContentSize;
131   }
132 
133   return std::make_unique<Alloc>(*this, *TargetAllocAddr,
134                                  std::move(WorkingBuffer), std::move(Segs));
135 }
136 
137 } // end namespace orc
138 } // end namespace llvm
139