1 | /* |
2 | * Licensed to the Apache Software Foundation (ASF) under one |
3 | * or more contributor license agreements. See the NOTICE file |
4 | * distributed with this work for additional information |
5 | * regarding copyright ownership. The ASF licenses this file |
6 | * to you under the Apache License, Version 2.0 (the |
7 | * "License"); you may not use this file except in compliance |
8 | * with the License. You may obtain a copy of the License at |
9 | * |
10 | * http://www.apache.org/licenses/LICENSE-2.0 |
11 | * |
12 | * Unless required by applicable law or agreed to in writing, |
13 | * software distributed under the License is distributed on an |
14 | * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
15 | * KIND, either express or implied. See the License for the |
16 | * specific language governing permissions and limitations |
17 | * under the License. |
18 | */ |
19 | |
20 | /*! |
21 | * \file tvm/runtime/vm/memory_manager.h |
22 | * \brief Abstract device memory management API |
23 | */ |
24 | #ifndef TVM_RUNTIME_VM_MEMORY_MANAGER_H_ |
25 | #define TVM_RUNTIME_VM_MEMORY_MANAGER_H_ |
26 | |
27 | #include <tvm/runtime/c_runtime_api.h> |
28 | #include <tvm/runtime/ndarray.h> |
29 | #include <tvm/runtime/object.h> |
30 | |
31 | #include <functional> |
32 | #include <memory> |
33 | #include <mutex> |
34 | #include <unordered_map> |
35 | #include <vector> |
36 | |
37 | namespace tvm { |
38 | namespace runtime { |
39 | namespace vm { |
40 | |
41 | struct Buffer { |
42 | /*! \brief The pointer to the allocated block of memory. */ |
43 | void* data{nullptr}; |
44 | /*! \brief The size of the block. */ |
45 | size_t size{0}; |
46 | /*! \brief The context of the allocated buffers. */ |
47 | Device device; |
48 | }; |
49 | |
50 | enum AllocatorType { |
51 | kNaive = 1, |
52 | kPooled, |
53 | }; |
54 | |
55 | class Allocator { |
56 | public: |
57 | explicit Allocator(AllocatorType type) : type_(type) {} |
58 | virtual ~Allocator() = default; |
59 | /*! \brief Allocate an empty NDArray using from the allocator. |
60 | * \param shape The shape of the NDArray. |
61 | * \param dtype The datatype of the NDArray. |
62 | * \param dev The device where the array is allocated. |
63 | * \return The empty NDArray. |
64 | */ |
65 | NDArray Empty(std::vector<int64_t> shape, DLDataType dtype, Device dev); |
66 | /*! \brief Return the allocator type. */ |
67 | inline AllocatorType type() const { return type_; } |
68 | /*! \brief Allocate a buffer given a size, alignment and type. |
69 | * \param nbytes The size of the buffer. |
70 | * \param alignment The alignment of the buffer. |
71 | * \param type_hint A type hint to the allocator. |
72 | * \return A sized allocation in the form of a buffer. |
73 | */ |
74 | virtual Buffer Alloc(size_t nbytes, size_t alignment, DLDataType type_hint) = 0; |
75 | /*! \brief Free a buffer allocated by the allocator. |
76 | * \param buffer The buffer to free. |
77 | */ |
78 | virtual void Free(const Buffer& buffer) = 0; |
79 | /*! \brief The amount of memory currently allocated. |
80 | * \return The amount of memory currently allocated. |
81 | */ |
82 | virtual size_t UsedMemory() const = 0; |
83 | |
84 | private: |
85 | AllocatorType type_; |
86 | }; |
87 | |
88 | class MemoryManager { |
89 | public: |
90 | static MemoryManager* Global(); |
91 | /*! |
92 | * \brief Get or create an allocator given the context and allocator type. |
93 | * \param dev The TVM device |
94 | * \param type The allocator type |
95 | * \return The memory allocator. |
96 | */ |
97 | static Allocator* GetOrCreateAllocator(Device dev, AllocatorType type); |
98 | /*! |
99 | * \brief Get an allocator given the context. |
100 | * \param dev The TVM device |
101 | * \return The memory allocator. |
102 | */ |
103 | static Allocator* GetAllocator(Device dev); |
104 | |
105 | private: |
106 | MemoryManager() {} |
107 | |
108 | private: |
109 | std::mutex mu_; |
110 | std::unordered_map<Device, std::unique_ptr<Allocator>> allocators_; |
111 | }; |
112 | |
113 | /*! \brief An object representing a storage allocation. */ |
114 | class StorageObj : public Object { |
115 | public: |
116 | /*! \brief The index into the VM function table. */ |
117 | Buffer buffer; |
118 | |
119 | /*! \brief Allocate an NDArray from a given piece of storage. */ |
120 | NDArray AllocNDArray(size_t offset, std::vector<int64_t> shape, DLDataType dtype); |
121 | |
122 | /*! \brief The deleter for an NDArray when allocated from underlying storage. */ |
123 | static void Deleter(Object* ptr); |
124 | |
125 | ~StorageObj() { |
126 | auto alloc = MemoryManager::Global()->GetAllocator(buffer.device); |
127 | alloc->Free(buffer); |
128 | } |
129 | |
130 | static constexpr const uint32_t _type_index = TypeIndex::kDynamic; |
131 | static constexpr const char* _type_key = "vm.Storage" ; |
132 | TVM_DECLARE_FINAL_OBJECT_INFO(StorageObj, Object); |
133 | }; |
134 | |
135 | /*! \brief reference to storage. */ |
136 | class Storage : public ObjectRef { |
137 | public: |
138 | explicit Storage(Buffer buffer); |
139 | |
140 | TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(Storage, ObjectRef, StorageObj); |
141 | }; |
142 | |
143 | } // namespace vm |
144 | } // namespace runtime |
145 | } // namespace tvm |
146 | |
147 | #endif // TVM_RUNTIME_VM_MEMORY_MANAGER_H_ |
148 | |