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-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // Declarations for target independent memory manager.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #ifndef LLVM_OPENMP_LIBOMPTARGET_SRC_MEMORYMANAGER_H
14 #define LLVM_OPENMP_LIBOMPTARGET_SRC_MEMORYMANAGER_H
15 
16 #include <cassert>
17 #include <functional>
18 #include <list>
19 #include <mutex>
20 #include <set>
21 #include <unordered_map>
22 #include <vector>
23 
24 // Forward declaration
25 struct DeviceTy;
26 
27 class MemoryManagerTy {
28   /// A structure stores the meta data of a target pointer
29   struct NodeTy {
30     /// Memory size
31     const size_t Size;
32     /// Target pointer
33     void *Ptr;
34 
35     /// Constructor
NodeTyNodeTy36     NodeTy(size_t Size, void *Ptr) : Size(Size), Ptr(Ptr) {}
37   };
38 
39   /// To make \p NodePtrTy ordered when they're put into \p std::multiset.
40   struct NodeCmpTy {
operatorNodeCmpTy41     bool operator()(const NodeTy &LHS, const NodeTy &RHS) const {
42       return LHS.Size < RHS.Size;
43     }
44   };
45 
46   /// A \p FreeList is a set of Nodes. We're using \p std::multiset here to make
47   /// the look up procedure more efficient.
48   using FreeListTy = std::multiset<std::reference_wrapper<NodeTy>, NodeCmpTy>;
49 
50   /// A list of \p FreeListTy entries, each of which is a \p std::multiset of
51   /// Nodes whose size is less or equal to a specific bucket size.
52   std::vector<FreeListTy> FreeLists;
53   /// A list of mutex for each \p FreeListTy entry
54   std::vector<std::mutex> FreeListLocks;
55   /// A table to map from a target pointer to its node
56   std::unordered_map<void *, NodeTy> PtrToNodeTable;
57   /// The mutex for the table \p PtrToNodeTable
58   std::mutex MapTableLock;
59   /// A reference to its corresponding \p DeviceTy object
60   DeviceTy &Device;
61 
62   /// Request memory from target device
63   void *allocateOnDevice(size_t Size, void *HstPtr) const;
64 
65   /// Deallocate data on device
66   int deleteOnDevice(void *Ptr) const;
67 
68   /// This function is called when it tries to allocate memory on device but the
69   /// device returns out of memory. It will first free all memory in the
70   /// FreeList and try to allocate again.
71   void *freeAndAllocate(size_t Size, void *HstPtr);
72 
73   /// The goal is to allocate memory on the device. It first tries to allocate
74   /// directly on the device. If a \p nullptr is returned, it might be because
75   /// the device is OOM. In that case, it will free all unused memory and then
76   /// try again.
77   void *allocateOrFreeAndAllocateOnDevice(size_t Size, void *HstPtr);
78 
79 public:
80   /// Constructor. If \p Threshold is non-zero, then the default threshold will
81   /// be overwritten by \p Threshold.
82   MemoryManagerTy(DeviceTy &Dev, size_t Threshold = 0);
83 
84   /// Destructor
85   ~MemoryManagerTy();
86 
87   /// Allocate memory of size \p Size from target device. \p HstPtr is used to
88   /// assist the allocation.
89   void *allocate(size_t Size, void *HstPtr);
90 
91   /// Deallocate memory pointed by \p TgtPtr
92   int free(void *TgtPtr);
93 };
94 
95 #endif // LLVM_OPENMP_LIBOMPTARGET_SRC_MEMORYMANAGER_H
96