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/ir/memory_pools.h |
22 | * \brief The object definition for relay.build argument type of memory pools |
23 | */ |
24 | #ifndef TVM_IR_MEMORY_POOLS_H_ |
25 | #define TVM_IR_MEMORY_POOLS_H_ |
26 | |
27 | #include <tvm/runtime/registry.h> |
28 | #include <tvm/target/target.h> |
29 | |
30 | struct TVMConstantInfo; |
31 | namespace tvm { |
32 | |
33 | /*! |
34 | * \brief Describes a pool of memory accessible by one or more targets. |
35 | */ |
36 | struct PoolInfoNode : public Object { |
37 | public: |
38 | /*! \brief The name of the memory pool */ |
39 | String pool_name; |
40 | /*! \brief The expected size hint to be used by the allocator. |
41 | * The size_hint_bytes is set to kUnrestrictedPoolSizeHint |
42 | * to indicate the pool is not size restricted. |
43 | */ |
44 | Integer size_hint_bytes; |
45 | /*! \brief The clock frequency of the memory in Hz */ |
46 | Integer clock_frequency_hz; |
47 | /*! \brief The read bandwidth in bytes/cycle */ |
48 | Integer read_bandwidth_bytes_per_cycle; |
49 | /*! \brief The write bandwidth in bytes/cycle */ |
50 | Integer write_bandwidth_bytes_per_cycle; |
51 | /*! \brief The read latency in cycles */ |
52 | Integer read_latency_cycles; |
53 | /*! \brief The write latency in cycles */ |
54 | Integer write_latency_cycles; |
55 | /*! \brief The burst length in bytes for each Target */ |
56 | Map<Target, Integer> target_burst_bytes; |
57 | /*! \brief Whether pool is internally generated. |
58 | * The internal pools will be generated as part of |
59 | * the entry point code generation of the executor |
60 | */ |
61 | bool is_internal = false; |
62 | |
63 | /*! \brief The targets linked to the pool */ |
64 | Array<Target> targets; |
65 | |
66 | void VisitAttrs(tvm::AttrVisitor* v) { |
67 | v->Visit("pool_name" , &pool_name); |
68 | v->Visit("targets" , &targets); |
69 | v->Visit("size_hint_bytes" , &size_hint_bytes); |
70 | v->Visit("clock_frequency_hz" , &clock_frequency_hz); |
71 | v->Visit("read_bandwidth_bytes_per_cycle" , &read_bandwidth_bytes_per_cycle); |
72 | v->Visit("write_bandwidth_bytes_per_cycle" , &write_bandwidth_bytes_per_cycle); |
73 | v->Visit("read_latency_cycles" , &read_latency_cycles); |
74 | v->Visit("write_latency_cycles" , &write_latency_cycles); |
75 | v->Visit("target_burst_bytes" , &target_burst_bytes); |
76 | v->Visit("is_internal" , &is_internal); |
77 | } |
78 | |
79 | bool SEqualReduce(const PoolInfoNode* other, SEqualReducer equal) const { |
80 | return equal(pool_name, other->pool_name) && equal(size_hint_bytes, other->size_hint_bytes) && |
81 | equal(clock_frequency_hz, other->clock_frequency_hz) && |
82 | equal(read_bandwidth_bytes_per_cycle, other->read_bandwidth_bytes_per_cycle) && |
83 | equal(write_bandwidth_bytes_per_cycle, other->write_bandwidth_bytes_per_cycle) && |
84 | equal(read_latency_cycles, other->read_latency_cycles) && |
85 | equal(write_latency_cycles, other->write_latency_cycles) && |
86 | equal(target_burst_bytes, other->target_burst_bytes) && |
87 | equal(is_internal, other->is_internal); |
88 | } |
89 | |
90 | void SHashReduce(SHashReducer hash_reduce) const { |
91 | hash_reduce(pool_name); |
92 | hash_reduce(size_hint_bytes); |
93 | hash_reduce(clock_frequency_hz); |
94 | hash_reduce(read_bandwidth_bytes_per_cycle); |
95 | hash_reduce(write_bandwidth_bytes_per_cycle); |
96 | hash_reduce(read_latency_cycles); |
97 | hash_reduce(write_latency_cycles); |
98 | hash_reduce(target_burst_bytes); |
99 | hash_reduce(is_internal); |
100 | } |
101 | |
102 | static constexpr const char* _type_key = "ir.PoolInfo" ; |
103 | TVM_DECLARE_BASE_OBJECT_INFO(PoolInfoNode, Object); |
104 | }; |
105 | |
106 | /*! |
107 | * \brief The string parameter to indicate read and write access to a pool |
108 | * This needs to be kept in sync with PoolInfo.READ_WRITE_ACCESS in |
109 | * python/tvm/ir/memory_pools.py |
110 | */ |
111 | static constexpr const char* kTargetPoolReadWriteAccess = "rw" ; |
112 | |
113 | /*! |
114 | * \brief The string parameter to indicate read only access to a pool |
115 | * This needs to be kept in sync with PoolInfo.READ_ONLY_ACCESS in |
116 | * python/tvm/ir/memory_pools.py |
117 | */ |
118 | static constexpr const char* kTargetPoolReadOnlyAccess = "ro" ; |
119 | |
120 | /*! \brief The PoolSize is unrestricted for the memory planner */ |
121 | static const int kUnrestrictedPoolSizeHint = -1; |
122 | |
123 | /*! \brief The clock frequency is not known */ |
124 | static const int kUnknownClockFrequency = -1; |
125 | |
126 | /*! \brief The read bandwidth is not known */ |
127 | static const int kUnknownReadBandwidth = -1; |
128 | |
129 | /*! \brief The write bandwidth is not known */ |
130 | static const int kUnknownWriteBandwidth = -1; |
131 | |
132 | /*! \brief Base class for WorkspacePoolInfo and ConstantPoolInfo */ |
133 | class PoolInfo : public ObjectRef { |
134 | protected: |
135 | TVM_DLL PoolInfo(String pool_name, Integer size_hint_bytes = kUnrestrictedPoolSizeHint, |
136 | Integer clock_frequency_hz = kUnknownClockFrequency, |
137 | Integer read_bandwidth_bytes_per_cycle = kUnknownReadBandwidth, |
138 | Integer write_bandwidth_bytes_per_cycle = kUnknownWriteBandwidth, |
139 | Integer read_latency_cycles = 0, Integer write_latency_cycles = 0, |
140 | Map<Target, Integer> target_burst_bytes = {}, Bool is_internal = Bool(false)); |
141 | |
142 | public: |
143 | TVM_DEFINE_OBJECT_REF_METHODS(PoolInfo, ObjectRef, PoolInfoNode); |
144 | }; |
145 | |
146 | /*! |
147 | * \brief Describes a pool of memory properties |
148 | */ |
149 | struct PoolInfoPropertiesNode : public Object { |
150 | /*! \brief The expected size hint to be used by the allocator. |
151 | * The size_hint_bytes is set to kUnrestrictedPoolSizeHint |
152 | * to indicate the pool is not size restricted. |
153 | */ |
154 | Integer size_hint_bytes = kUnrestrictedPoolSizeHint; |
155 | /*! \brief The clock frequency of the memory in Hz */ |
156 | Integer clock_frequency_hz = kUnknownClockFrequency; |
157 | /*! \brief The read bandwidth in bytes/cycle */ |
158 | Integer read_bandwidth_bytes_per_cycle = kUnknownReadBandwidth; |
159 | /*! \brief The write bandwidth in bytes/cycle */ |
160 | Integer write_bandwidth_bytes_per_cycle = kUnknownWriteBandwidth; |
161 | /*! \brief The read latency in cycles */ |
162 | Integer read_latency_cycles = 0; |
163 | /*! \brief The write latency in cycles */ |
164 | Integer write_latency_cycles = 0; |
165 | /*! \brief The burst length in bytes for each Target */ |
166 | Map<Target, Integer> target_burst_bytes{}; |
167 | /*! \brief Whether pool is internally generated. |
168 | * The internal pools will be generated as part of |
169 | * the entry point code generation of the executor |
170 | */ |
171 | bool is_internal = false; |
172 | |
173 | void VisitAttrs(tvm::AttrVisitor* v) { |
174 | v->Visit("size_hint_bytes" , &size_hint_bytes); |
175 | v->Visit("clock_frequency_hz" , &clock_frequency_hz); |
176 | v->Visit("read_bandwidth_bytes_per_cycle" , &read_bandwidth_bytes_per_cycle); |
177 | v->Visit("write_bandwidth_bytes_per_cycle" , &write_bandwidth_bytes_per_cycle); |
178 | v->Visit("read_latency_cycles" , &read_latency_cycles); |
179 | v->Visit("write_latency_cycles" , &write_latency_cycles); |
180 | v->Visit("target_burst_bytes" , &target_burst_bytes); |
181 | v->Visit("is_internal" , &is_internal); |
182 | } |
183 | |
184 | bool SEqualReduce(const PoolInfoPropertiesNode* other, SEqualReducer equal) const { |
185 | return equal(size_hint_bytes, other->size_hint_bytes) && |
186 | equal(clock_frequency_hz, other->clock_frequency_hz) && |
187 | equal(read_bandwidth_bytes_per_cycle, other->read_bandwidth_bytes_per_cycle) && |
188 | equal(write_bandwidth_bytes_per_cycle, other->write_bandwidth_bytes_per_cycle) && |
189 | equal(read_latency_cycles, other->read_latency_cycles) && |
190 | equal(write_latency_cycles, other->write_latency_cycles) && |
191 | equal(target_burst_bytes, other->target_burst_bytes) && |
192 | equal(is_internal, other->is_internal); |
193 | } |
194 | |
195 | void SHashReduce(SHashReducer hash_reduce) const { |
196 | hash_reduce(size_hint_bytes); |
197 | hash_reduce(clock_frequency_hz); |
198 | hash_reduce(read_bandwidth_bytes_per_cycle); |
199 | hash_reduce(write_bandwidth_bytes_per_cycle); |
200 | hash_reduce(read_latency_cycles); |
201 | hash_reduce(write_latency_cycles); |
202 | hash_reduce(target_burst_bytes); |
203 | hash_reduce(is_internal); |
204 | } |
205 | |
206 | static constexpr const char* _type_key = "ir.PoolInfoProperties" ; |
207 | TVM_DECLARE_FINAL_OBJECT_INFO(PoolInfoPropertiesNode, Object); |
208 | }; |
209 | |
210 | class PoolInfoProperties : public ObjectRef { |
211 | public: |
212 | TVM_DLL PoolInfoProperties(Integer size_hint_bytes, |
213 | Integer clock_frequency_hz = kUnknownClockFrequency, |
214 | Integer read_bandwidth_bytes_per_cycle = kUnknownReadBandwidth, |
215 | Integer write_bandwidth_bytes_per_cycle = kUnknownWriteBandwidth, |
216 | Integer read_latency_cycles = 0, Integer write_latency_cycles = 0, |
217 | Map<Target, Integer> target_burst_bytes = {}, |
218 | Bool is_internal = Bool(false)); |
219 | TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(PoolInfoProperties, ObjectRef, PoolInfoPropertiesNode); |
220 | }; |
221 | |
222 | /* \brief Represents RW memory area */ |
223 | struct WorkspacePoolInfoNode : public PoolInfoNode { |
224 | void VisitAttrs(tvm::AttrVisitor* v) { PoolInfoNode::VisitAttrs(v); } |
225 | |
226 | bool SEqualReduce(const WorkspacePoolInfoNode* other, SEqualReducer equal) const { |
227 | return PoolInfoNode::SEqualReduce(other, equal); |
228 | } |
229 | |
230 | void SHashReduce(SHashReducer hash_reduce) const { PoolInfoNode::SHashReduce(hash_reduce); } |
231 | |
232 | static constexpr const char* _type_key = "ir.WorkspacePoolInfo" ; |
233 | TVM_DECLARE_FINAL_OBJECT_INFO(WorkspacePoolInfoNode, PoolInfoNode); |
234 | }; |
235 | |
236 | class WorkspacePoolInfo : public PoolInfo { |
237 | public: |
238 | TVM_DLL WorkspacePoolInfo( |
239 | String pool_name, Array<Target> targets, |
240 | PoolInfoProperties properties = PoolInfoProperties(kUnrestrictedPoolSizeHint)); |
241 | TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(WorkspacePoolInfo, PoolInfo, WorkspacePoolInfoNode); |
242 | }; |
243 | |
244 | /* |
245 | * \brief The ConstantInfoNode contains numeric literal in RO pool |
246 | * Used to initialise RO memory in ConstantPoolInfo |
247 | */ |
248 | struct ConstantInfoNode : public Object { |
249 | String name_hint; |
250 | Integer byte_offset; |
251 | runtime::NDArray data; |
252 | |
253 | void VisitAttrs(tvm::AttrVisitor* v) { |
254 | v->Visit("name_hint" , &name_hint); |
255 | v->Visit("byte_offset" , &byte_offset); |
256 | v->Visit("data" , &data); |
257 | } |
258 | |
259 | bool SEqualReduce(const ConstantInfoNode* other, SEqualReducer equal) const { |
260 | return equal(name_hint, other->name_hint) && equal(byte_offset, other->byte_offset) && |
261 | equal(data, other->data); |
262 | } |
263 | |
264 | void SHashReduce(SHashReducer hash_reduce) const { |
265 | hash_reduce(name_hint); |
266 | hash_reduce(byte_offset); |
267 | hash_reduce(data); |
268 | } |
269 | |
270 | static constexpr const char* _type_key = "ir.ConstantInfo" ; |
271 | static constexpr bool _type_has_method_sequal_reduce = true; |
272 | static constexpr bool _type_has_method_shash_reduce = true; |
273 | TVM_DECLARE_FINAL_OBJECT_INFO(ConstantInfoNode, Object); |
274 | }; |
275 | |
276 | class ConstantInfo : public ObjectRef { |
277 | public: |
278 | TVM_DLL ConstantInfo(const struct ::TVMConstantInfo* data); |
279 | ConstantInfo(String name, Integer byte_offset, runtime::NDArray data); |
280 | TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(ConstantInfo, ObjectRef, ConstantInfoNode); |
281 | }; |
282 | |
283 | /* \brief ConstantPoolInfoNode represents an RO memory area initialized with |
284 | * data from constant_info_array */ |
285 | struct ConstantPoolInfoNode : public PoolInfoNode { |
286 | Array<ConstantInfo> constant_info_array; |
287 | |
288 | void VisitAttrs(tvm::AttrVisitor* v) { |
289 | PoolInfoNode::VisitAttrs(v); |
290 | v->Visit("constant_info_array" , &constant_info_array); |
291 | } |
292 | |
293 | bool SEqualReduce(const ConstantPoolInfoNode* other, SEqualReducer equal) const { |
294 | return PoolInfoNode::SEqualReduce(other, equal) && |
295 | equal(constant_info_array, other->constant_info_array); |
296 | } |
297 | |
298 | void SHashReduce(SHashReducer hash_reduce) const { |
299 | PoolInfoNode::SHashReduce(hash_reduce); |
300 | hash_reduce(constant_info_array); |
301 | } |
302 | |
303 | static constexpr const char* _type_key = "ir.ConstantPoolInfo" ; |
304 | TVM_DECLARE_FINAL_OBJECT_INFO(ConstantPoolInfoNode, PoolInfoNode); |
305 | }; |
306 | |
307 | class ConstantPoolInfo : public PoolInfo { |
308 | public: |
309 | TVM_DLL ConstantPoolInfo( |
310 | String pool_name, Array<Target> targets, Array<ConstantInfo> constant_info_array, |
311 | PoolInfoProperties properties = PoolInfoProperties(kUnrestrictedPoolSizeHint)); |
312 | TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(ConstantPoolInfo, PoolInfo, ConstantPoolInfoNode); |
313 | }; |
314 | |
315 | /* \brief A container for WorkspacePoolInfo objects */ |
316 | struct WorkspaceMemoryPoolsNode : public Object { |
317 | Array<PoolInfo> pools; |
318 | |
319 | void VisitAttrs(tvm::AttrVisitor* v) { v->Visit("pools" , &pools); } |
320 | |
321 | bool SEqualReduce(const WorkspaceMemoryPoolsNode* other, SEqualReducer equal) const { |
322 | return equal(pools, other->pools); |
323 | } |
324 | |
325 | void SHashReduce(SHashReducer hash_reduce) const { hash_reduce(pools); } |
326 | |
327 | static constexpr const char* _type_key = "ir.WorkspaceMemoryPools" ; |
328 | TVM_DECLARE_FINAL_OBJECT_INFO(WorkspaceMemoryPoolsNode, Object); |
329 | }; |
330 | |
331 | class WorkspaceMemoryPools : public ObjectRef { |
332 | public: |
333 | TVM_DLL WorkspaceMemoryPools(Array<PoolInfo> pools); |
334 | TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(WorkspaceMemoryPools, ObjectRef, WorkspaceMemoryPoolsNode); |
335 | }; |
336 | |
337 | /* \brief A container for ConstantPoolInfo objects */ |
338 | struct ConstantMemoryPoolsNode : public Object { |
339 | Array<ConstantPoolInfo> pools; |
340 | |
341 | void VisitAttrs(tvm::AttrVisitor* v) { v->Visit("pools" , &pools); } |
342 | |
343 | bool SEqualReduce(const ConstantMemoryPoolsNode* other, SEqualReducer equal) const { |
344 | return equal(pools, other->pools); |
345 | } |
346 | |
347 | void SHashReduce(SHashReducer hash_reduce) const { hash_reduce(pools); } |
348 | |
349 | static constexpr const char* _type_key = "ir.ConstantMemoryPools" ; |
350 | TVM_DECLARE_FINAL_OBJECT_INFO(ConstantMemoryPoolsNode, Object); |
351 | }; |
352 | |
353 | class ConstantMemoryPools : public ObjectRef { |
354 | public: |
355 | TVM_DLL ConstantMemoryPools(Array<ConstantPoolInfo> pools); |
356 | TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(ConstantMemoryPools, ObjectRef, ConstantMemoryPoolsNode); |
357 | }; |
358 | |
359 | } // namespace tvm |
360 | |
361 | #endif // TVM_IR_MEMORY_POOLS_H_ |
362 | |