410 lines · cpp
1//===- SimpleNativeMemoryMap.cpp ------------------------------------------===//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-exception6//7//===----------------------------------------------------------------------===//8//9// SimpleNativeMemoryMap and related APIs.10//11// TODO: We don't reset / uncommit pages on deinitialize, or on failure during12// initialize. We should do that to reduce memory pressure.13//14//===----------------------------------------------------------------------===//15 16#include "orc-rt/SimpleNativeMemoryMap.h"17#include "orc-rt/SPSAllocAction.h"18#include "orc-rt/SPSMemoryFlags.h"19#include <sstream>20 21#if defined(__APPLE__) || defined(__linux__)22#include "Unix/NativeMemoryAPIs.inc"23#else24#error "Target OS memory APIs unsupported"25#endif26 27namespace orc_rt {28 29struct SPSSimpleNativeMemoryMapSegment;30 31template <>32class SPSSerializationTraits<33 SPSSimpleNativeMemoryMapSegment,34 SimpleNativeMemoryMap::InitializeRequest::Segment> {35 using SPSType =36 SPSTuple<SPSAllocGroup, SPSExecutorAddr, uint64_t, SPSSequence<char>>;37 38public:39 static bool40 deserialize(SPSInputBuffer &IB,41 SimpleNativeMemoryMap::InitializeRequest::Segment &S) {42 AllocGroup AG;43 ExecutorAddr Address;44 uint64_t Size;45 span<const char> Content;46 if (!SPSType::AsArgList::deserialize(IB, AG, Address, Size, Content))47 return false;48 if (Size > std::numeric_limits<size_t>::max())49 return false;50 S = {AG, Address.toPtr<char *>(), static_cast<size_t>(Size), Content};51 return true;52 }53};54 55struct SPSSimpleNativeMemoryMapInitializeRequest;56 57template <>58class SPSSerializationTraits<SPSSimpleNativeMemoryMapInitializeRequest,59 SimpleNativeMemoryMap::InitializeRequest> {60 using SPSType = SPSTuple<SPSSequence<SPSSimpleNativeMemoryMapSegment>,61 SPSSequence<SPSAllocActionPair>>;62 63public:64 static bool deserialize(SPSInputBuffer &IB,65 SimpleNativeMemoryMap::InitializeRequest &FR) {66 return SPSType::AsArgList::deserialize(IB, FR.Segments, FR.AAPs);67 }68};69 70void SimpleNativeMemoryMap::reserve(OnReserveCompleteFn &&OnComplete,71 size_t Size) {72 // FIXME: Get page size from session object.73 if (Size % (64 * 1024)) {74 return OnComplete(make_error<StringError>(75 (std::ostringstream()76 << "SimpleNativeMemoryMap error: reserved size " << std::hex << Size77 << " is not a page-size multiple")78 .str()));79 }80 81 auto Addr = hostOSMemoryReserve(Size);82 if (!Addr)83 return OnComplete(Addr.takeError());84 85 {86 std::scoped_lock<std::mutex> Lock(M);87 assert(!Slabs.count(*Addr) &&88 "hostOSMemoryReserve returned duplicate addresses");89 Slabs.emplace(std::make_pair(*Addr, SlabInfo(Size)));90 }91 92 OnComplete(*Addr);93}94 95void SimpleNativeMemoryMap::release(OnReleaseCompleteFn &&OnComplete,96 void *Addr) {97 std::optional<SlabInfo> SI;98 {99 std::scoped_lock<std::mutex> Lock(M);100 auto I = Slabs.find(Addr);101 if (I != Slabs.end()) {102 SI = std::move(I->second);103 Slabs.erase(I);104 }105 }106 107 if (!SI) {108 std::ostringstream ErrMsg;109 ErrMsg << "SimpleNativeMemoryMap error: release called on unrecognized "110 "address "111 << Addr;112 return OnComplete(make_error<StringError>(ErrMsg.str()));113 }114 115 for (auto &[Addr, DAAs] : SI->DeallocActions)116 runDeallocActions(std::move(DAAs));117 118 OnComplete(hostOSMemoryRelease(Addr, SI->Size));119}120 121void SimpleNativeMemoryMap::releaseMultiple(OnReleaseCompleteFn &&OnComplete,122 std::vector<void *> Addrs) {123 releaseNext(std::move(OnComplete), std::move(Addrs), false, Error::success());124}125 126void SimpleNativeMemoryMap::initialize(OnInitializeCompleteFn &&OnComplete,127 InitializeRequest IR) {128 129 void *Base = nullptr;130 131 // TODO: Record initialize segments for release.132 // std::vector<std::pair<void*, size_t>> InitializeSegments;133 134 // Check segment validity before proceeding.135 for (auto &S : IR.Segments) {136 137 if (S.Content.size() > S.Size) {138 return OnComplete(make_error<StringError>(139 (std::ostringstream()140 << "For segment [" << (void *)S.Address << ".."141 << (void *)(S.Address + S.Size) << "), "142 << " content size (" << std::hex << S.Content.size()143 << ") exceeds segment size (" << S.Size << ")")144 .str()));145 }146 147 // Copy any requested content.148 if (!S.Content.empty())149 memcpy(S.Address, S.Content.data(), S.Content.size());150 151 // Zero-fill the rest of the section.152 if (size_t ZeroFillSize = S.Size - S.Content.size())153 memset(S.Address + S.Content.size(), 0, ZeroFillSize);154 155 if (auto Err = hostOSMemoryProtect(S.Address, S.Size, S.AG.getMemProt()))156 return OnComplete(std::move(Err));157 158 switch (S.AG.getMemLifetime()) {159 case MemLifetime::Standard:160 if (!Base || S.Address < Base)161 Base = S.Address;162 break;163 case MemLifetime::Finalize:164 // TODO: Record finalize segment for release.165 // FinalizeSegments.push_back({S.Address, S.Size});166 break;167 }168 }169 170 if (!Base)171 return OnComplete(172 make_error<StringError>("SimpleNativeMemoryMap initialize error: "173 "finalization requires at least "174 "one standard-lifetime segment"));175 176 auto DeallocActions = runFinalizeActions(std::move(IR.AAPs));177 if (!DeallocActions)178 return OnComplete(DeallocActions.takeError());179 180 if (auto Err = recordDeallocActions(Base, std::move(*DeallocActions))) {181 runDeallocActions(std::move(*DeallocActions));182 return OnComplete(std::move(Err));183 }184 185 OnComplete(Base);186}187 188void SimpleNativeMemoryMap::deinitialize(OnDeinitializeCompleteFn &&OnComplete,189 void *Base) {190 std::vector<AllocAction> DAAs;191 192 {193 std::unique_lock<std::mutex> Lock(M);194 auto *SI = findSlabInfoFor(Base);195 if (!SI) {196 Lock.unlock();197 return OnComplete(makeBadSlabError(Base, "deinitialize"));198 }199 200 auto I = SI->DeallocActions.find(Base);201 if (I == SI->DeallocActions.end()) {202 Lock.unlock();203 std::ostringstream ErrMsg;204 ErrMsg205 << "SimpleNativeMemoryMap deinitialize error: no deallocate actions "206 "registered for segment base address "207 << Base;208 return OnComplete(make_error<StringError>(ErrMsg.str()));209 }210 211 DAAs = std::move(I->second);212 SI->DeallocActions.erase(I);213 }214 215 runDeallocActions(std::move(DAAs));216 OnComplete(Error::success());217}218 219void SimpleNativeMemoryMap::deinitializeMultiple(220 OnDeinitializeCompleteFn &&OnComplete, std::vector<void *> Bases) {221 deinitializeNext(std::move(OnComplete), std::move(Bases), false,222 Error::success());223}224 225void SimpleNativeMemoryMap::detach(ResourceManager::OnCompleteFn OnComplete) {226 // Detach is a noop for now: we just retain all actions to run at shutdown227 // time.228 OnComplete(Error::success());229}230 231void SimpleNativeMemoryMap::shutdown(ResourceManager::OnCompleteFn OnComplete) {232 // TODO: Establish a clear order to run dealloca actions across slabs,233 // object boundaries.234 235 // Collect slab base addresses for removal.236 std::vector<void *> Bases;237 {238 std::scoped_lock<std::mutex> Lock(M);239 for (auto &[Base, _] : Slabs)240 Bases.push_back(Base);241 }242 243 shutdownNext(std::move(OnComplete), std::move(Bases));244}245 246void SimpleNativeMemoryMap::releaseNext(OnReleaseCompleteFn &&OnComplete,247 std::vector<void *> Addrs,248 bool AnyError, Error LastErr) {249 // TODO: Log error?250 if (LastErr) {251 consumeError(std::move(LastErr));252 AnyError |= true;253 }254 255 if (Addrs.empty()) {256 if (!AnyError)257 return OnComplete(Error::success());258 259 return OnComplete(260 make_error<StringError>("Failed to release some addresses"));261 }262 263 void *NextAddr = Addrs.back();264 Addrs.pop_back();265 266 release(267 [this, OnComplete = std::move(OnComplete), AnyError = AnyError,268 Addrs = std::move(Addrs)](Error Err) mutable {269 releaseNext(std::move(OnComplete), std::move(Addrs), AnyError,270 std::move(Err));271 },272 NextAddr);273}274 275void SimpleNativeMemoryMap::deinitializeNext(276 OnDeinitializeCompleteFn &&OnComplete, std::vector<void *> Addrs,277 bool AnyError, Error LastErr) {278 // TODO: Log error?279 if (LastErr) {280 consumeError(std::move(LastErr));281 AnyError |= true;282 }283 284 if (Addrs.empty()) {285 if (!AnyError)286 return OnComplete(Error::success());287 288 return OnComplete(289 make_error<StringError>("Failed to deinitialize some addresses"));290 }291 292 void *NextAddr = Addrs.back();293 Addrs.pop_back();294 295 deinitialize(296 [this, OnComplete = std::move(OnComplete), AnyError = AnyError,297 Addrs = std::move(Addrs)](Error Err) mutable {298 deinitializeNext(std::move(OnComplete), std::move(Addrs), AnyError,299 std::move(Err));300 },301 NextAddr);302}303 304void SimpleNativeMemoryMap::shutdownNext(305 ResourceManager::OnCompleteFn OnComplete, std::vector<void *> Bases) {306 if (Bases.empty())307 return OnComplete(Error::success());308 309 auto *Base = Bases.back();310 Bases.pop_back();311 312 release(313 [this, Bases = std::move(Bases),314 OnComplete = std::move(OnComplete)](Error Err) mutable {315 if (Err) {316 // TODO: Log release error?317 consumeError(std::move(Err));318 }319 shutdownNext(std::move(OnComplete), std::move(Bases));320 },321 Base);322}323 324Error SimpleNativeMemoryMap::makeBadSlabError(void *Base, const char *Op) {325 std::ostringstream ErrMsg;326 ErrMsg << "SimpleNativeMemoryMap " << Op << " error: segment base address "327 << Base << " does not fall within an allocated slab";328 return make_error<StringError>(ErrMsg.str());329}330 331SimpleNativeMemoryMap::SlabInfo *332SimpleNativeMemoryMap::findSlabInfoFor(void *Base) {333 // NOTE: We assume that the caller is holding a lock for M.334 auto I = Slabs.upper_bound(Base);335 if (I == Slabs.begin())336 return nullptr;337 338 --I;339 if (reinterpret_cast<char *>(I->first) + I->second.Size <=340 reinterpret_cast<char *>(Base))341 return nullptr;342 343 return &I->second;344}345 346Error SimpleNativeMemoryMap::recordDeallocActions(347 void *Base, std::vector<AllocAction> DeallocActions) {348 349 std::unique_lock<std::mutex> Lock(M);350 auto *SI = findSlabInfoFor(Base);351 if (!SI) {352 Lock.unlock();353 return makeBadSlabError(Base, "deinitialize");354 }355 356 auto I = SI->DeallocActions.find(Base);357 if (I != SI->DeallocActions.end()) {358 Lock.unlock();359 std::ostringstream ErrMsg;360 ErrMsg << "SimpleNativeMemoryMap initialize error: segment base address "361 "reused in subsequent initialize call";362 return make_error<StringError>(ErrMsg.str());363 }364 365 SI->DeallocActions[Base] = std::move(DeallocActions);366 return Error::success();367}368 369ORC_RT_SPS_INTERFACE void orc_rt_SimpleNativeMemoryMap_reserve_sps_wrapper(370 orc_rt_SessionRef S, uint64_t CallId, orc_rt_WrapperFunctionReturn Return,371 orc_rt_WrapperFunctionBuffer ArgBytes) {372 using Sig = SPSExpected<SPSExecutorAddr>(SPSExecutorAddr, SPSSize);373 SPSWrapperFunction<Sig>::handle(374 S, CallId, Return, ArgBytes,375 WrapperFunction::handleWithAsyncMethod(&SimpleNativeMemoryMap::reserve));376}377 378ORC_RT_SPS_INTERFACE void379orc_rt_SimpleNativeMemoryMap_releaseMultiple_sps_wrapper(380 orc_rt_SessionRef S, uint64_t CallId, orc_rt_WrapperFunctionReturn Return,381 orc_rt_WrapperFunctionBuffer ArgBytes) {382 using Sig = SPSError(SPSExecutorAddr, SPSSequence<SPSExecutorAddr>);383 SPSWrapperFunction<Sig>::handle(S, CallId, Return, ArgBytes,384 WrapperFunction::handleWithAsyncMethod(385 &SimpleNativeMemoryMap::releaseMultiple));386}387 388ORC_RT_SPS_INTERFACE void orc_rt_SimpleNativeMemoryMap_initialize_sps_wrapper(389 orc_rt_SessionRef S, uint64_t CallId, orc_rt_WrapperFunctionReturn Return,390 orc_rt_WrapperFunctionBuffer ArgBytes) {391 using Sig = SPSExpected<SPSExecutorAddr>(392 SPSExecutorAddr, SPSSimpleNativeMemoryMapInitializeRequest);393 SPSWrapperFunction<Sig>::handle(S, CallId, Return, ArgBytes,394 WrapperFunction::handleWithAsyncMethod(395 &SimpleNativeMemoryMap::initialize));396}397 398ORC_RT_SPS_INTERFACE void399orc_rt_SimpleNativeMemoryMap_deinitializeMultiple_sps_wrapper(400 orc_rt_SessionRef S, uint64_t CallId, orc_rt_WrapperFunctionReturn Return,401 orc_rt_WrapperFunctionBuffer ArgBytes) {402 using Sig = SPSError(SPSExecutorAddr, SPSSequence<SPSExecutorAddr>);403 SPSWrapperFunction<Sig>::handle(404 S, CallId, Return, ArgBytes,405 WrapperFunction::handleWithAsyncMethod(406 &SimpleNativeMemoryMap::deinitializeMultiple));407}408 409} // namespace orc_rt410