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 relay/backend/token_allocator.h |
22 | * \brief Token allocation classes for backend |
23 | */ |
24 | #ifndef TVM_RELAY_BACKEND_TOKEN_ALLOCATOR_H_ |
25 | #define TVM_RELAY_BACKEND_TOKEN_ALLOCATOR_H_ |
26 | |
27 | #include <tvm/relay/type.h> |
28 | #include <tvm/target/virtual_device.h> |
29 | |
30 | #include <map> |
31 | #include <string> |
32 | #include <unordered_map> |
33 | #include <unordered_set> |
34 | #include <vector> |
35 | |
36 | #include "../../runtime/texture.h" |
37 | |
38 | namespace tvm { |
39 | namespace relay { |
40 | |
41 | /*! A representation of a block of memory required at runtime on some device. */ |
42 | struct StorageToken { |
43 | /*! \brief Reference counter */ |
44 | int ref_counter{0}; |
45 | /*! \brief number of bytes */ |
46 | size_t max_bytes{0}; |
47 | /*! \brief The corresponding tensor type. */ |
48 | TensorType ttype{nullptr}; |
49 | /*! \brief VirtualDevice on which the memory will reside. */ |
50 | VirtualDevice virtual_device = VirtualDevice::FullyUnconstrained(); |
51 | /*! \brief The storage id */ |
52 | int64_t storage_id{-1}; |
53 | |
54 | bool is_valid() const { return !virtual_device->IsFullyUnconstrained(); } |
55 | |
56 | bool is_compatible(const StorageToken& that) const { |
57 | return virtual_device == that.virtual_device; |
58 | } |
59 | |
60 | std::string ToString() const { |
61 | std::ostringstream os; |
62 | os << "{storage_id: " << storage_id << ", max_bytes: " << max_bytes |
63 | << ", ttype: " << PrettyPrint(ttype) << ", virtual_device: " << virtual_device << "}" ; |
64 | return os.str(); |
65 | } |
66 | }; |
67 | |
68 | /** |
69 | * @brief Memory manager for flattened 1d memory (buffers) |
70 | */ |
71 | class TokenAllocator1D { |
72 | public: |
73 | /*! |
74 | * \brief ceil(size/word_size) to get number of words. |
75 | * \param size The original size. |
76 | * \param word_size The element size. |
77 | */ |
78 | static size_t DivRoundUp(size_t size, size_t word_size) { |
79 | return (size + word_size - 1) / word_size; |
80 | } |
81 | |
82 | /*! |
83 | * \brief Get the memory requirement. |
84 | * \param prototype The prototype token. |
85 | * \return The required memory size. |
86 | * |
87 | * TODO(mbs): Gf GetMemorySizeBytes in aot_executor_codegen.cc, |
88 | * CalculateRelayExprSizeBytes in utils.cc |
89 | */ |
90 | size_t GetMemorySize(StorageToken* prototype); |
91 | /*! |
92 | * \brief Request a storage token for a given prototype. |
93 | * \param prototype. The prototype storage token. |
94 | * \return The result token. |
95 | */ |
96 | StorageToken* Request(StorageToken* prototype); |
97 | /*! |
98 | * \brief Alloacte a storage token by consuming prototype |
99 | * \param prototype The prototype token. |
100 | * \param size The size of memory being requested. |
101 | */ |
102 | StorageToken* Alloc(StorageToken* prototype, int64_t storage_id); |
103 | /*! |
104 | * \brief Check if we can release token. |
105 | * \param tok The token to be released. |
106 | */ |
107 | void CheckForRelease(StorageToken* tok); |
108 | |
109 | private: |
110 | // scale used for rough match |
111 | const size_t match_range_{16}; |
112 | // free list of storage entry |
113 | std::multimap<size_t, StorageToken*> free_; |
114 | // all the storage resources available |
115 | std::vector<StorageToken*> data_; |
116 | }; |
117 | |
118 | /** |
119 | * @brief Memory manager for 2d memory (textures) |
120 | */ |
121 | class TokenAllocator2D { |
122 | public: |
123 | /*! |
124 | * \brief Request a storage token for a given prototype. |
125 | * \param prototype. The prototype storage token. |
126 | * \return The result token. |
127 | */ |
128 | StorageToken* Request(StorageToken* prototype); |
129 | /*! |
130 | * \brief Alloacte a storage token by consuming prototype |
131 | * \param prototype The prototype token. |
132 | * \param size The size of memory being requested. |
133 | */ |
134 | StorageToken* Alloc(StorageToken* prototype, int64_t storage_id); |
135 | /*! |
136 | * \brief Check if we can release token. |
137 | * \param tok The token to be released. |
138 | */ |
139 | void CheckForRelease(StorageToken* tok); |
140 | /*! |
141 | * \brief Get the texture 2d size requirement |
142 | * \param prototype The prototype token. |
143 | * \return The required texture 2d memory size in (width, height, channel). |
144 | */ |
145 | runtime::Texture2DShape<int64_t> GetSize2D(StorageToken* prototype); |
146 | |
147 | protected: |
148 | struct MemBlock { |
149 | StorageToken* token_; |
150 | int64_t x_; |
151 | int64_t y_; |
152 | }; |
153 | |
154 | std::unordered_map<int64_t, MemBlock> blocks_; |
155 | std::unordered_set<int64_t> free_list_; |
156 | }; |
157 | |
158 | } // namespace relay |
159 | } // namespace tvm |
160 | |
161 | #endif // TVM_RELAY_BACKEND_TOKEN_ALLOCATOR_H_ |
162 | |