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
37namespace tvm {
38namespace runtime {
39namespace vm {
40
41struct 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
50enum AllocatorType {
51 kNaive = 1,
52 kPooled,
53};
54
55class 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
88class 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. */
114class 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. */
136class 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