tvm
memory_manager.h
Go to the documentation of this file.
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 
24 #ifndef TVM_RUNTIME_MEMORY_MEMORY_MANAGER_H_
25 #define TVM_RUNTIME_MEMORY_MEMORY_MANAGER_H_
26 
27 #include <tvm/runtime/base.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 <string>
35 #include <unordered_map>
36 #include <vector>
37 
38 namespace tvm {
39 namespace runtime {
40 namespace memory {
41 
43  kNaive = 1,
45 };
46 
47 struct Buffer {
49  void* data{nullptr};
51  size_t size{0};
56 };
57 
58 class Allocator {
59  public:
60  explicit Allocator(AllocatorType type) : type_(type) {}
61  virtual ~Allocator() = default;
69  TVM_DLL NDArray Empty(ffi::Shape shape, DLDataType dtype, Device dev,
70  Optional<String> mem_scope = std::nullopt);
72  inline AllocatorType type() const { return type_; }
80  TVM_DLL virtual Buffer Alloc(Device dev, size_t nbytes, size_t alignment,
81  DLDataType type_hint) = 0;
89  TVM_DLL virtual Buffer Alloc(Device dev, ffi::Shape shape, DLDataType type_hint,
90  const std::string& mem_scope = "");
91 
99  TVM_DLL virtual void* CreateView(const Buffer& buffer, ffi::Shape shape, DLDataType type_hint,
100  const std::string& mem_scope = "global") {
101  return buffer.data;
102  }
103 
108  TVM_DLL virtual void FreeView(Device dev, void* data) {}
109 
113  TVM_DLL virtual void Free(const Buffer& buffer) = 0;
115  TVM_DLL virtual void Clear();
119  TVM_DLL virtual size_t UsedMemory() const = 0;
120 
121  protected:
123  TVM_DLL virtual bool AllowMemoryScope(const std::string& mem_scope) const;
124 
125  private:
126  AllocatorType type_;
127 };
128 
130  public:
131  TVM_DLL static MemoryManager* Global();
145  TVM_DLL static Allocator* GetAllocator(Device dev, AllocatorType type);
147  static void Clear();
148 
149  private:
150  MemoryManager() {}
151 
152  protected:
153  std::mutex mu_;
154  std::unordered_map<Device, std::unordered_map<AllocatorType, std::unique_ptr<Allocator>>>
156 };
157 
159 class StorageObj : public Object {
160  public:
164  Allocator* allocator = nullptr;
165 
167  TVM_DLL NDArray AllocNDArray(int64_t offset, ffi::Shape shape, DLDataType dtype);
168 
170  TVM_DLL NDArray AllocNDArrayScoped(int64_t offset, ffi::Shape shape, DLDataType dtype,
171  String scope = "global");
172 
174  if (allocator) {
176  }
177  }
178 
179  static constexpr const char* _type_key = "vm.Storage";
181 };
182 
184 class Storage : public ObjectRef {
185  public:
186  TVM_DLL explicit Storage(Buffer buffer, Allocator* allocator);
187 
189 };
190 
191 } // namespace memory
192 
193 using memory::Allocator;
196 using memory::StorageObj;
197 
198 } // namespace runtime
199 } // namespace tvm
200 
201 #endif // TVM_RUNTIME_MEMORY_MEMORY_MANAGER_H_
Managed NDArray. The array is backed by reference counted blocks.
Definition: ndarray.h:53
Definition: memory_manager.h:58
virtual void Free(const Buffer &buffer)=0
Free a buffer allocated by the allocator.
virtual Buffer Alloc(Device dev, ffi::Shape shape, DLDataType type_hint, const std::string &mem_scope="")
Allocate a buffer given a shape and type.
virtual size_t UsedMemory() const =0
The amount of memory currently allocated.
NDArray Empty(ffi::Shape shape, DLDataType dtype, Device dev, Optional< String > mem_scope=std::nullopt)
Allocate an empty NDArray using from the allocator.
AllocatorType type() const
Return the allocator type.
Definition: memory_manager.h:72
virtual bool AllowMemoryScope(const std::string &mem_scope) const
Check if the given memory scope is allowed to allocate by the allocator.
Allocator(AllocatorType type)
Definition: memory_manager.h:60
virtual Buffer Alloc(Device dev, size_t nbytes, size_t alignment, DLDataType type_hint)=0
Allocate a buffer given a size, alignment and type.
virtual void Clear()
Clear the allocated memory.
virtual void FreeView(Device dev, void *data)
Release the view .
Definition: memory_manager.h:108
virtual void * CreateView(const Buffer &buffer, ffi::Shape shape, DLDataType type_hint, const std::string &mem_scope="global")
Create a view for the buffer given a shape, type and scope.
Definition: memory_manager.h:99
Definition: memory_manager.h:129
static void Clear()
Clear the allocators.
static Allocator * GetOrCreateAllocator(Device dev, AllocatorType type)
Get or create an allocator given the context and allocator type.
std::mutex mu_
Definition: memory_manager.h:153
static Allocator * GetAllocator(Device dev, AllocatorType type)
Get an allocator given the context.
static MemoryManager * Global()
std::unordered_map< Device, std::unordered_map< AllocatorType, std::unique_ptr< Allocator > > > allocators_
Definition: memory_manager.h:155
An object representing a storage allocation.
Definition: memory_manager.h:159
Buffer buffer
The index into the VM function table.
Definition: memory_manager.h:162
Allocator * allocator
The allocator where the storage buffer is allocated from.
Definition: memory_manager.h:164
~StorageObj()
Definition: memory_manager.h:173
NDArray AllocNDArray(int64_t offset, ffi::Shape shape, DLDataType dtype)
Allocate an NDArray from a given piece of storage.
static constexpr const char * _type_key
Definition: memory_manager.h:179
NDArray AllocNDArrayScoped(int64_t offset, ffi::Shape shape, DLDataType dtype, String scope="global")
Allocate an NDArray with memory scope from a given piece of storage.
TVM_DECLARE_FINAL_OBJECT_INFO(StorageObj, Object)
reference to storage.
Definition: memory_manager.h:184
TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(Storage, ObjectRef, StorageObj)
Storage(Buffer buffer, Allocator *allocator)
AllocatorType
Definition: memory_manager.h:42
@ kNaive
Definition: memory_manager.h:43
@ kPooled
Definition: memory_manager.h:44
Tensor shape(const Tensor &src, DataType dtype, const std::string name="T_shape", const std::string tag=kInjective)
Get the shape of input tensor.
Definition: transform.h:1945
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
DLDevice Device
Definition: device_api.h:42
A device-independent managed NDArray abstraction.
A managed object in the TVM runtime.
Definition: memory_manager.h:47
Device device
The context of the allocated buffers.
Definition: memory_manager.h:53
AllocatorType alloc_type
The allocator that created this buffer.
Definition: memory_manager.h:55
void * data
The pointer to the allocated block of memory.
Definition: memory_manager.h:49
size_t size
The size of the block.
Definition: memory_manager.h:51