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
30struct TVMConstantInfo;
31namespace tvm {
32
33/*!
34 * \brief Describes a pool of memory accessible by one or more targets.
35 */
36struct 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 */
111static 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 */
118static constexpr const char* kTargetPoolReadOnlyAccess = "ro";
119
120/*! \brief The PoolSize is unrestricted for the memory planner */
121static const int kUnrestrictedPoolSizeHint = -1;
122
123/*! \brief The clock frequency is not known */
124static const int kUnknownClockFrequency = -1;
125
126/*! \brief The read bandwidth is not known */
127static const int kUnknownReadBandwidth = -1;
128
129/*! \brief The write bandwidth is not known */
130static const int kUnknownWriteBandwidth = -1;
131
132/*! \brief Base class for WorkspacePoolInfo and ConstantPoolInfo */
133class 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 */
149struct 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
210class 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 */
223struct 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
236class 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 */
248struct 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
276class 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 */
285struct 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
307class 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 */
316struct 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
331class 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 */
338struct 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
353class 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