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 src/ir/memory_pools.cc
22 * \brief The object definition for relay.build argument type of memory pools
23 */
24
25#include <tvm/ir/memory_pools.h>
26#include <tvm/relay/executor.h>
27
28namespace tvm {
29
30PoolInfo::PoolInfo(String pool_name, Integer size_hint_bytes, Integer clock_frequency_hz,
31 Integer read_bandwidth_bytes_per_cycle, Integer write_bandwidth_bytes_per_cycle,
32 Integer read_latency_cycles, Integer write_latency_cycles,
33 Map<Target, Integer> target_burst_bytes, Bool is_internal) {
34 auto poolinfo_node = make_object<PoolInfoNode>();
35 poolinfo_node->pool_name = pool_name;
36 poolinfo_node->size_hint_bytes = size_hint_bytes;
37 poolinfo_node->clock_frequency_hz = clock_frequency_hz;
38 poolinfo_node->read_bandwidth_bytes_per_cycle = read_bandwidth_bytes_per_cycle;
39 poolinfo_node->write_bandwidth_bytes_per_cycle = write_bandwidth_bytes_per_cycle;
40 poolinfo_node->read_latency_cycles = read_latency_cycles;
41 poolinfo_node->write_latency_cycles = write_latency_cycles;
42 poolinfo_node->target_burst_bytes = target_burst_bytes;
43 poolinfo_node->is_internal = is_internal;
44 data_ = std::move(poolinfo_node);
45}
46
47TVM_REGISTER_NODE_TYPE(PoolInfoNode);
48
49TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
50 .set_dispatch<PoolInfoNode>([](const ObjectRef& ref, ReprPrinter* p) {
51 auto* node = static_cast<const PoolInfoNode*>(ref.get());
52 p->stream << "PoolInfoNode(\n"
53 << " pool_name=" << node->pool_name
54 << ",\n size_hint_bytes=" << node->size_hint_bytes
55 << ",\n clock_frequency_hz=" << node->clock_frequency_hz
56 << ",\n read_bandwidth_bytes_per_cycle=" << node->read_bandwidth_bytes_per_cycle
57 << ",\n write_bandwidth_bytes_per_cycle=" << node->write_bandwidth_bytes_per_cycle
58 << ",\n read_latency_cycles=" << node->read_latency_cycles
59 << ",\n write_latency_cycles=" << node->write_latency_cycles
60 << ",\n target_burst_bytes=" << node->target_burst_bytes << ")";
61 });
62
63PoolInfoProperties::PoolInfoProperties(Integer size_hint_bytes, Integer clock_frequency_hz,
64 Integer read_bandwidth_bytes_per_cycle,
65 Integer write_bandwidth_bytes_per_cycle,
66 Integer read_latency_cycles, Integer write_latency_cycles,
67 Map<Target, Integer> target_burst_bytes, Bool is_internal) {
68 auto poolinfo_properties_node = make_object<PoolInfoPropertiesNode>();
69 poolinfo_properties_node->size_hint_bytes = size_hint_bytes;
70 poolinfo_properties_node->clock_frequency_hz = clock_frequency_hz;
71 poolinfo_properties_node->read_bandwidth_bytes_per_cycle = read_bandwidth_bytes_per_cycle;
72 poolinfo_properties_node->write_bandwidth_bytes_per_cycle = write_bandwidth_bytes_per_cycle;
73 poolinfo_properties_node->read_latency_cycles = read_latency_cycles;
74 poolinfo_properties_node->write_latency_cycles = write_latency_cycles;
75 poolinfo_properties_node->target_burst_bytes = target_burst_bytes;
76 poolinfo_properties_node->is_internal = is_internal;
77 data_ = std::move(poolinfo_properties_node);
78}
79
80TVM_REGISTER_NODE_TYPE(PoolInfoPropertiesNode);
81TVM_REGISTER_GLOBAL("ir.PoolInfoProperties")
82 .set_body_typed([](Integer size_hint_bytes, Integer clock_frequency_hz,
83 Integer read_bandwidth_bytes_per_cycle,
84 Integer write_bandwidth_bytes_per_cycle, Integer read_latency_cycles,
85 Integer write_latency_cycles, Map<Target, Integer> target_burst_bytes) {
86 return PoolInfoProperties(size_hint_bytes, clock_frequency_hz, read_bandwidth_bytes_per_cycle,
87 write_bandwidth_bytes_per_cycle, read_latency_cycles,
88 write_latency_cycles, target_burst_bytes, Bool(false));
89 });
90
91TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
92 .set_dispatch<PoolInfoPropertiesNode>([](const ObjectRef& ref, ReprPrinter* p) {
93 auto* node = static_cast<const PoolInfoPropertiesNode*>(ref.get());
94 p->stream << "PoolInfoPropertiesNode(\n"
95 << " size_hint_bytes=" << node->size_hint_bytes
96 << ",\n clock_frequency_hz=" << node->clock_frequency_hz
97 << ",\n read_bandwidth_bytes_per_cycle=" << node->read_bandwidth_bytes_per_cycle
98 << ",\n write_bandwidth_bytes_per_cycle=" << node->write_bandwidth_bytes_per_cycle
99 << ",\n read_latency_cycles=" << node->read_latency_cycles
100 << ",\n write_latency_cycles=" << node->write_latency_cycles
101 << ",\n target_burst_bytes=" << node->target_burst_bytes << ")";
102 });
103
104WorkspacePoolInfo::WorkspacePoolInfo(String pool_name, Array<Target> targets,
105 PoolInfoProperties properties) {
106 auto poolinfo_node = make_object<WorkspacePoolInfoNode>();
107 poolinfo_node->pool_name = pool_name;
108 poolinfo_node->size_hint_bytes = properties->size_hint_bytes;
109 poolinfo_node->targets = targets;
110 poolinfo_node->clock_frequency_hz = properties->clock_frequency_hz;
111 poolinfo_node->read_bandwidth_bytes_per_cycle = properties->read_bandwidth_bytes_per_cycle;
112 poolinfo_node->write_bandwidth_bytes_per_cycle = properties->write_bandwidth_bytes_per_cycle;
113 poolinfo_node->read_latency_cycles = properties->read_latency_cycles;
114 poolinfo_node->write_latency_cycles = properties->write_latency_cycles;
115 poolinfo_node->target_burst_bytes = properties->target_burst_bytes;
116 poolinfo_node->is_internal = properties->is_internal;
117 data_ = std::move(poolinfo_node);
118}
119
120TVM_REGISTER_NODE_TYPE(WorkspacePoolInfoNode);
121TVM_REGISTER_GLOBAL("ir.WorkspacePoolInfo")
122 .set_body_typed([](String pool_name, Array<Target> targets, PoolInfoProperties properties) {
123 return WorkspacePoolInfo(pool_name, targets, properties);
124 });
125
126TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
127 .set_dispatch<WorkspacePoolInfoNode>([](const ObjectRef& ref, ReprPrinter* p) {
128 auto* node = static_cast<const WorkspacePoolInfoNode*>(ref.get());
129 p->stream << "WorkspacePoolInfoNode(\n"
130 << " pool_name=" << node->pool_name << ",\n targets=" << node->targets
131 << ",\n size_hint_bytes=" << node->size_hint_bytes
132 << ",\n clock_frequency_hz=" << node->clock_frequency_hz
133 << ",\n read_bandwidth_bytes_per_cycle=" << node->read_bandwidth_bytes_per_cycle
134 << ",\n write_bandwidth_bytes_per_cycle=" << node->write_bandwidth_bytes_per_cycle
135 << ",\n read_latency_cycles=" << node->read_latency_cycles
136 << ",\n write_latency_cycles=" << node->write_latency_cycles
137 << ",\n target_burst_bytes=" << node->target_burst_bytes
138 << ",\n is_internal=" << node->is_internal << ")"
139 << "\n";
140 });
141
142ConstantInfo::ConstantInfo(String name_hint, Integer byte_offset, runtime::NDArray data) {
143 auto constant_info_node = make_object<ConstantInfoNode>();
144 constant_info_node->name_hint = name_hint;
145 constant_info_node->byte_offset = byte_offset;
146 constant_info_node->data = data;
147 data_ = std::move(constant_info_node);
148}
149
150TVM_REGISTER_NODE_TYPE(ConstantInfoNode);
151TVM_REGISTER_GLOBAL("ir.ConstantInfo")
152 .set_body_typed([](String name_hint, Integer byte_offset, runtime::NDArray data) {
153 return ConstantInfo(name_hint, byte_offset, data);
154 });
155TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
156 .set_dispatch<ConstantInfoNode>([](const ObjectRef& ref, ReprPrinter* p) {
157 auto* node = static_cast<const ConstantInfoNode*>(ref.get());
158 p->stream << "ConstantInfoNode(\n"
159 << "name_hint=" << node->name_hint << ",\n byte_offset=" << node->byte_offset
160 << ",\n data=" << node->data << ")";
161 });
162
163ConstantPoolInfo::ConstantPoolInfo(String pool_name, Array<Target> targets,
164 Array<ConstantInfo> constant_info_array,
165 PoolInfoProperties properties) {
166 auto constant_poolinfo_node = make_object<ConstantPoolInfoNode>();
167 constant_poolinfo_node->pool_name = pool_name;
168 constant_poolinfo_node->constant_info_array = constant_info_array;
169 constant_poolinfo_node->targets = targets;
170
171 constant_poolinfo_node->size_hint_bytes = properties->size_hint_bytes;
172 constant_poolinfo_node->clock_frequency_hz = properties->clock_frequency_hz;
173 constant_poolinfo_node->read_bandwidth_bytes_per_cycle =
174 properties->read_bandwidth_bytes_per_cycle;
175 constant_poolinfo_node->write_bandwidth_bytes_per_cycle =
176 properties->write_bandwidth_bytes_per_cycle;
177 constant_poolinfo_node->read_latency_cycles = properties->read_latency_cycles;
178 constant_poolinfo_node->write_latency_cycles = properties->write_latency_cycles;
179 constant_poolinfo_node->target_burst_bytes = properties->target_burst_bytes;
180 constant_poolinfo_node->is_internal = properties->is_internal;
181 data_ = std::move(constant_poolinfo_node);
182}
183
184TVM_REGISTER_NODE_TYPE(ConstantPoolInfoNode);
185TVM_REGISTER_GLOBAL("ir.ConstantPoolInfo")
186 .set_body_typed([](String pool_name, Array<Target> targets,
187 Array<ConstantInfo> constant_info_array, PoolInfoProperties properties) {
188 return ConstantPoolInfo(pool_name, targets, constant_info_array, properties);
189 });
190
191TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
192 .set_dispatch<ConstantPoolInfoNode>([](const ObjectRef& ref, ReprPrinter* p) {
193 auto* node = static_cast<const ConstantPoolInfoNode*>(ref.get());
194 p->stream << "ConstantPoolInfoNode(\n"
195 << " pool_name=" << node->pool_name << ",\n targets=" << node->targets
196 << ",\n constant_info_array=" << node->constant_info_array
197 << ",\n size_hint_bytes=" << node->size_hint_bytes
198 << ",\n clock_frequency_hz=" << node->clock_frequency_hz
199 << ",\n read_bandwidth_bytes_per_cycle=" << node->read_bandwidth_bytes_per_cycle
200 << ",\n write_bandwidth_bytes_per_cycle=" << node->write_bandwidth_bytes_per_cycle
201 << ",\n read_latency_cycles=" << node->read_latency_cycles
202 << ",\n write_latency_cycles=" << node->write_latency_cycles
203 << ",\n target_burst_bytes=" << node->target_burst_bytes
204 << ",\n is_internal=" << node->is_internal << ")"
205 << "\n";
206 });
207
208WorkspaceMemoryPools::WorkspaceMemoryPools(Array<PoolInfo> pools) {
209 auto workspace_memory_pools_node = make_object<WorkspaceMemoryPoolsNode>();
210 workspace_memory_pools_node->pools = pools;
211 data_ = std::move(workspace_memory_pools_node);
212}
213
214TVM_REGISTER_NODE_TYPE(WorkspaceMemoryPoolsNode);
215TVM_REGISTER_GLOBAL("ir.WorkspaceMemoryPools").set_body_typed([](Array<PoolInfo> pools) {
216 return WorkspaceMemoryPools(pools);
217});
218
219TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
220 .set_dispatch<WorkspaceMemoryPoolsNode>([](const ObjectRef& ref, ReprPrinter* p) {
221 auto* node = static_cast<const WorkspaceMemoryPoolsNode*>(ref.get());
222 p->stream << "WorkspaceMemoryPoolsNode(\n"
223 << "pools=" << node->pools << ")";
224 });
225
226ConstantMemoryPools::ConstantMemoryPools(Array<ConstantPoolInfo> pools) {
227 auto constant_memory_pools_node = make_object<ConstantMemoryPoolsNode>();
228 constant_memory_pools_node->pools = pools;
229 data_ = std::move(constant_memory_pools_node);
230}
231
232TVM_REGISTER_NODE_TYPE(ConstantMemoryPoolsNode);
233TVM_REGISTER_GLOBAL("ir.ConstantMemoryPools").set_body_typed([](Array<ConstantPoolInfo> pools) {
234 return ConstantMemoryPools(pools);
235});
236
237TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
238 .set_dispatch<ConstantMemoryPoolsNode>([](const ObjectRef& ref, ReprPrinter* p) {
239 auto* node = static_cast<const ConstantMemoryPoolsNode*>(ref.get());
240 p->stream << "ConstantMemoryPoolsNode(\n"
241 << "pools=" << node->pools << ")";
242 });
243} // namespace tvm
244