blob: 6435637f3e5a7381da704d4814a314524d1e7302 [file] [log] [blame]
//===----------- MemoryManager.h - Target independent memory manager ------===//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// Declarations for target independent memory manager.
#include <cassert>
#include <functional>
#include <list>
#include <mutex>
#include <set>
#include <unordered_map>
#include <vector>
// Forward declaration
struct DeviceTy;
class MemoryManagerTy {
/// A structure stores the meta data of a target pointer
struct NodeTy {
/// Memory size
const size_t Size;
/// Target pointer
void *Ptr;
/// Constructor
NodeTy(size_t Size, void *Ptr) : Size(Size), Ptr(Ptr) {}
/// To make \p NodePtrTy ordered when they're put into \p std::multiset.
struct NodeCmpTy {
bool operator()(const NodeTy &LHS, const NodeTy &RHS) const {
return LHS.Size < RHS.Size;
/// A \p FreeList is a set of Nodes. We're using \p std::multiset here to make
/// the look up procedure more efficient.
using FreeListTy = std::multiset<std::reference_wrapper<NodeTy>, NodeCmpTy>;
/// A list of \p FreeListTy entries, each of which is a \p std::multiset of
/// Nodes whose size is less or equal to a specific bucket size.
std::vector<FreeListTy> FreeLists;
/// A list of mutex for each \p FreeListTy entry
std::vector<std::mutex> FreeListLocks;
/// A table to map from a target pointer to its node
std::unordered_map<void *, NodeTy> PtrToNodeTable;
/// The mutex for the table \p PtrToNodeTable
std::mutex MapTableLock;
/// A reference to its corresponding \p DeviceTy object
DeviceTy &Device;
/// Request memory from target device
void *allocateOnDevice(size_t Size, void *HstPtr) const;
/// Deallocate data on device
int deleteOnDevice(void *Ptr) const;
/// This function is called when it tries to allocate memory on device but the
/// device returns out of memory. It will first free all memory in the
/// FreeList and try to allocate again.
void *freeAndAllocate(size_t Size, void *HstPtr);
/// The goal is to allocate memory on the device. It first tries to allocate
/// directly on the device. If a \p nullptr is returned, it might be because
/// the device is OOM. In that case, it will free all unused memory and then
/// try again.
void *allocateOrFreeAndAllocateOnDevice(size_t Size, void *HstPtr);
/// Constructor. If \p Threshold is non-zero, then the default threshold will
/// be overwritten by \p Threshold.
MemoryManagerTy(DeviceTy &Dev, size_t Threshold = 0);
/// Destructor
/// Allocate memory of size \p Size from target device. \p HstPtr is used to
/// assist the allocation.
void *allocate(size_t Size, void *HstPtr);
/// Deallocate memory pointed by \p TgtPtr
int free(void *TgtPtr);