372 lines · c
1//===----------- MemoryManager.h - Target independent memory manager ------===//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// Target independent memory manager.10//11//===----------------------------------------------------------------------===//12 13#ifndef LLVM_OPENMP_LIBOMPTARGET_PLUGINS_COMMON_MEMORYMANAGER_H14#define LLVM_OPENMP_LIBOMPTARGET_PLUGINS_COMMON_MEMORYMANAGER_H15 16#include <cassert>17#include <functional>18#include <list>19#include <mutex>20#include <set>21#include <unordered_map>22#include <vector>23 24#include "Shared/Debug.h"25#include "Shared/Utils.h"26#include "omptarget.h"27 28#include "llvm/Support/Error.h"29 30namespace llvm {31 32/// Base class of per-device allocator.33class DeviceAllocatorTy {34public:35 virtual ~DeviceAllocatorTy() = default;36 37 /// Allocate a memory of size \p Size . \p HstPtr is used to assist the38 /// allocation.39 virtual Expected<void *>40 allocate(size_t Size, void *HstPtr,41 TargetAllocTy Kind = TARGET_ALLOC_DEFAULT) = 0;42 43 /// Delete the pointer \p TgtPtr on the device44 virtual Error free(void *TgtPtr,45 TargetAllocTy Kind = TARGET_ALLOC_DEFAULT) = 0;46};47 48/// Class of memory manager. The memory manager is per-device by using49/// per-device allocator. Therefore, each plugin using memory manager should50/// have an allocator for each device.51class MemoryManagerTy {52 static constexpr const size_t BucketSize[] = {53 0, 1U << 2, 1U << 3, 1U << 4, 1U << 5, 1U << 6, 1U << 7,54 1U << 8, 1U << 9, 1U << 10, 1U << 11, 1U << 12, 1U << 13};55 56 static constexpr const int NumBuckets =57 sizeof(BucketSize) / sizeof(BucketSize[0]);58 59 /// Find the previous number that is power of 2 given a number that is not60 /// power of 2.61 static size_t floorToPowerOfTwo(size_t Num) {62 Num |= Num >> 1;63 Num |= Num >> 2;64 Num |= Num >> 4;65 Num |= Num >> 8;66 Num |= Num >> 16;67#if INTPTR_MAX == INT64_MAX68 Num |= Num >> 32;69#elif INTPTR_MAX == INT32_MAX70 // Do nothing with 32-bit71#else72#error Unsupported architecture73#endif74 Num += 1;75 return Num >> 1;76 }77 78 /// Find a suitable bucket79 static int findBucket(size_t Size) {80 const size_t F = floorToPowerOfTwo(Size);81 82 DP("findBucket: Size %zu is floored to %zu.\n", Size, F);83 84 int L = 0, H = NumBuckets - 1;85 while (H - L > 1) {86 int M = (L + H) >> 1;87 if (BucketSize[M] == F)88 return M;89 if (BucketSize[M] > F)90 H = M - 1;91 else92 L = M;93 }94 95 assert(L >= 0 && L < NumBuckets && "L is out of range");96 97 DP("findBucket: Size %zu goes to bucket %d\n", Size, L);98 99 return L;100 }101 102 /// A structure stores the meta data of a target pointer103 struct NodeTy {104 /// Memory size105 const size_t Size;106 /// Target pointer107 void *Ptr;108 109 /// Constructor110 NodeTy(size_t Size, void *Ptr) : Size(Size), Ptr(Ptr) {}111 };112 113 /// To make \p NodePtrTy ordered when they're put into \p std::multiset.114 struct NodeCmpTy {115 bool operator()(const NodeTy &LHS, const NodeTy &RHS) const {116 return LHS.Size < RHS.Size;117 }118 };119 120 /// A \p FreeList is a set of Nodes. We're using \p std::multiset here to make121 /// the look up procedure more efficient.122 using FreeListTy = std::multiset<std::reference_wrapper<NodeTy>, NodeCmpTy>;123 124 /// A list of \p FreeListTy entries, each of which is a \p std::multiset of125 /// Nodes whose size is less or equal to a specific bucket size.126 std::vector<FreeListTy> FreeLists;127 /// A list of mutex for each \p FreeListTy entry128 std::vector<std::mutex> FreeListLocks;129 /// A table to map from a target pointer to its node130 std::unordered_map<void *, NodeTy> PtrToNodeTable;131 /// The mutex for the table \p PtrToNodeTable132 std::mutex MapTableLock;133 134 /// The reference to a device allocator135 DeviceAllocatorTy &DeviceAllocator;136 137 /// The threshold to manage memory using memory manager. If the request size138 /// is larger than \p SizeThreshold, the allocation will not be managed by the139 /// memory manager.140 size_t SizeThreshold = 1U << 13;141 142 /// Request memory from target device143 Expected<void *> allocateOnDevice(size_t Size, void *HstPtr) const {144 return DeviceAllocator.allocate(Size, HstPtr, TARGET_ALLOC_DEVICE);145 }146 147 /// Deallocate data on device148 Error deleteOnDevice(void *Ptr) const { return DeviceAllocator.free(Ptr); }149 150 /// This function is called when it tries to allocate memory on device but the151 /// device returns out of memory. It will first free all memory in the152 /// FreeList and try to allocate again.153 Expected<void *> freeAndAllocate(size_t Size, void *HstPtr) {154 std::vector<void *> RemoveList;155 156 // Deallocate all memory in FreeList157 for (int I = 0; I < NumBuckets; ++I) {158 FreeListTy &List = FreeLists[I];159 std::lock_guard<std::mutex> Lock(FreeListLocks[I]);160 if (List.empty())161 continue;162 for (const NodeTy &N : List) {163 if (auto Err = deleteOnDevice(N.Ptr))164 return Err;165 RemoveList.push_back(N.Ptr);166 }167 FreeLists[I].clear();168 }169 170 // Remove all nodes in the map table which have been released171 if (!RemoveList.empty()) {172 std::lock_guard<std::mutex> LG(MapTableLock);173 for (void *P : RemoveList)174 PtrToNodeTable.erase(P);175 }176 177 // Try allocate memory again178 return allocateOnDevice(Size, HstPtr);179 }180 181 /// The goal is to allocate memory on the device. It first tries to182 /// allocate directly on the device. If a \p nullptr is returned, it might183 /// be because the device is OOM. In that case, it will free all unused184 /// memory and then try again.185 Expected<void *> allocateOrFreeAndAllocateOnDevice(size_t Size,186 void *HstPtr) {187 auto TgtPtrOrErr = allocateOnDevice(Size, HstPtr);188 if (!TgtPtrOrErr)189 return TgtPtrOrErr.takeError();190 191 void *TgtPtr = *TgtPtrOrErr;192 // We cannot get memory from the device. It might be due to OOM. Let's193 // free all memory in FreeLists and try again.194 if (TgtPtr == nullptr) {195 DP("Failed to get memory on device. Free all memory in FreeLists and "196 "try again.\n");197 TgtPtrOrErr = freeAndAllocate(Size, HstPtr);198 if (!TgtPtrOrErr)199 return TgtPtrOrErr.takeError();200 TgtPtr = *TgtPtrOrErr;201 }202 203 if (TgtPtr == nullptr)204 DP("Still cannot get memory on device probably because the device is "205 "OOM.\n");206 207 return TgtPtr;208 }209 210public:211 /// Constructor. If \p Threshold is non-zero, then the default threshold will212 /// be overwritten by \p Threshold.213 MemoryManagerTy(DeviceAllocatorTy &DeviceAllocator, size_t Threshold = 0)214 : FreeLists(NumBuckets), FreeListLocks(NumBuckets),215 DeviceAllocator(DeviceAllocator) {216 if (Threshold)217 SizeThreshold = Threshold;218 }219 220 /// Destructor221 ~MemoryManagerTy() {222 for (auto &PtrToNode : PtrToNodeTable) {223 assert(PtrToNode.second.Ptr && "nullptr in map table");224 if (auto Err = deleteOnDevice(PtrToNode.second.Ptr))225 REPORT("Failure to delete memory: %s\n",226 toString(std::move(Err)).data());227 }228 }229 230 /// Allocate memory of size \p Size from target device. \p HstPtr is used to231 /// assist the allocation.232 Expected<void *> allocate(size_t Size, void *HstPtr) {233 // If the size is zero, we will not bother the target device. Just return234 // nullptr directly.235 if (Size == 0)236 return nullptr;237 238 DP("MemoryManagerTy::allocate: size %zu with host pointer " DPxMOD ".\n",239 Size, DPxPTR(HstPtr));240 241 // If the size is greater than the threshold, allocate it directly from242 // device.243 if (Size > SizeThreshold) {244 DP("%zu is greater than the threshold %zu. Allocate it directly from "245 "device\n",246 Size, SizeThreshold);247 auto TgtPtrOrErr = allocateOrFreeAndAllocateOnDevice(Size, HstPtr);248 if (!TgtPtrOrErr)249 return TgtPtrOrErr.takeError();250 251 DP("Got target pointer " DPxMOD ". Return directly.\n",252 DPxPTR(*TgtPtrOrErr));253 254 return *TgtPtrOrErr;255 }256 257 NodeTy *NodePtr = nullptr;258 259 // Try to get a node from FreeList260 {261 const int B = findBucket(Size);262 FreeListTy &List = FreeLists[B];263 264 NodeTy TempNode(Size, nullptr);265 std::lock_guard<std::mutex> LG(FreeListLocks[B]);266 const auto Itr = List.find(TempNode);267 268 if (Itr != List.end()) {269 NodePtr = &Itr->get();270 List.erase(Itr);271 }272 }273 274 if (NodePtr != nullptr)275 DP("Find one node " DPxMOD " in the bucket.\n", DPxPTR(NodePtr));276 277 // We cannot find a valid node in FreeLists. Let's allocate on device and278 // create a node for it.279 if (NodePtr == nullptr) {280 DP("Cannot find a node in the FreeLists. Allocate on device.\n");281 // Allocate one on device282 auto TgtPtrOrErr = allocateOrFreeAndAllocateOnDevice(Size, HstPtr);283 if (!TgtPtrOrErr)284 return TgtPtrOrErr.takeError();285 286 void *TgtPtr = *TgtPtrOrErr;287 if (TgtPtr == nullptr)288 return nullptr;289 290 // Create a new node and add it into the map table291 {292 std::lock_guard<std::mutex> Guard(MapTableLock);293 auto Itr = PtrToNodeTable.emplace(TgtPtr, NodeTy(Size, TgtPtr));294 NodePtr = &Itr.first->second;295 }296 297 DP("Node address " DPxMOD ", target pointer " DPxMOD ", size %zu\n",298 DPxPTR(NodePtr), DPxPTR(TgtPtr), Size);299 }300 301 assert(NodePtr && "NodePtr should not be nullptr at this point");302 303 return NodePtr->Ptr;304 }305 306 /// Deallocate memory pointed by \p TgtPtr307 Error free(void *TgtPtr) {308 DP("MemoryManagerTy::free: target memory " DPxMOD ".\n", DPxPTR(TgtPtr));309 310 NodeTy *P = nullptr;311 312 // Look it up into the table313 {314 std::lock_guard<std::mutex> G(MapTableLock);315 auto Itr = PtrToNodeTable.find(TgtPtr);316 317 // We don't remove the node from the map table because the map does not318 // change.319 if (Itr != PtrToNodeTable.end())320 P = &Itr->second;321 }322 323 // The memory is not managed by the manager324 if (P == nullptr) {325 DP("Cannot find its node. Delete it on device directly.\n");326 return deleteOnDevice(TgtPtr);327 }328 329 // Insert the node to the free list330 const int B = findBucket(P->Size);331 332 DP("Found its node " DPxMOD ". Insert it to bucket %d.\n", DPxPTR(P), B);333 334 {335 std::lock_guard<std::mutex> G(FreeListLocks[B]);336 FreeLists[B].insert(*P);337 }338 339 return Error::success();340 }341 342 /// Get the size threshold from the environment variable343 /// \p LIBOMPTARGET_MEMORY_MANAGER_THRESHOLD . Returns a <tt>344 /// std::pair<size_t, bool> </tt> where the first element represents the345 /// threshold and the second element represents whether user disables memory346 /// manager explicitly by setting the var to 0. If user doesn't specify347 /// anything, returns <0, true>.348 static std::pair<size_t, bool> getSizeThresholdFromEnv() {349 static UInt64Envar MemoryManagerThreshold(350 "LIBOMPTARGET_MEMORY_MANAGER_THRESHOLD", 0);351 352 size_t Threshold = MemoryManagerThreshold.get();353 354 if (MemoryManagerThreshold.isPresent() && Threshold == 0) {355 DP("Disabled memory manager as user set "356 "LIBOMPTARGET_MEMORY_MANAGER_THRESHOLD=0.\n");357 return std::make_pair(0, false);358 }359 360 return std::make_pair(Threshold, true);361 }362};363 364// GCC still cannot handle the static data member like Clang so we still need365// this part.366constexpr const size_t MemoryManagerTy::BucketSize[];367constexpr const int MemoryManagerTy::NumBuckets;368 369} // namespace llvm370 371#endif // LLVM_OPENMP_LIBOMPTARGET_PLUGINS_COMMON_MEMORYMANAGER_H372