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/target/metadata.h |
22 | * \brief Extends Metadata for use in the compiler. |
23 | */ |
24 | #ifndef TVM_TARGET_METADATA_H_ |
25 | #define TVM_TARGET_METADATA_H_ |
26 | #include <tvm/ir/memory_pools.h> |
27 | #include <tvm/runtime/metadata.h> |
28 | |
29 | #include <memory> |
30 | #include <string> |
31 | #include <vector> |
32 | |
33 | namespace tvm { |
34 | namespace target { |
35 | namespace metadata { |
36 | |
37 | /*! |
38 | * \brief Subclass of MetadataNode that implements the VisitAttrs reflection method. |
39 | * |
40 | * This implementation (and other such Visitable subclasses) is compiled into libtvm.so, but not |
41 | * libtvm_runtime.so, because reflection is not supported in libtvm_runtime.so over code size |
42 | * concerns. It is used during compilation by the generic metadata code-generators. |
43 | */ |
44 | class VisitableMetadataNode : public ::tvm::runtime::metadata::MetadataNode { |
45 | public: |
46 | explicit VisitableMetadataNode(const struct ::TVMMetadata* data) : MetadataNode{data} {} |
47 | VisitableMetadataNode() : MetadataNode{nullptr} {} |
48 | |
49 | void VisitAttrs(AttrVisitor* v) { |
50 | int64_t version_cpp{version()}; |
51 | v->Visit("version" , &version_cpp); |
52 | auto inputs_array = Array<ObjectRef>(); |
53 | auto inputs_accessor = inputs(); |
54 | inputs_array.reserve(num_inputs()); |
55 | for (int64_t i = 0; i < num_inputs(); ++i) { |
56 | inputs_array.push_back(::tvm::runtime::metadata::TensorInfo{inputs_accessor[i]}); |
57 | } |
58 | ::tvm::runtime::metadata::MetadataArray inputs_metadata_array{ |
59 | inputs_array, ::tvm::runtime::metadata::MetadataKind::kMetadata, |
60 | ::tvm::runtime::metadata::TensorInfoNode::_type_key}; |
61 | v->Visit("inputs" , &inputs_metadata_array); |
62 | int64_t num_inputs_cpp = num_inputs(); |
63 | v->Visit("num_inputs" , &num_inputs_cpp); |
64 | auto outputs_array = Array<ObjectRef>(); |
65 | auto outputs_accessor = outputs(); |
66 | outputs_array.reserve(num_outputs()); |
67 | for (int64_t i = 0; i < num_outputs(); ++i) { |
68 | outputs_array.push_back(::tvm::runtime::metadata::TensorInfo{outputs_accessor[i]}); |
69 | } |
70 | ::tvm::runtime::metadata::MetadataArray outputs_metadata_array{ |
71 | outputs_array, ::tvm::runtime::metadata::MetadataKind::kMetadata, |
72 | ::tvm::runtime::metadata::TensorInfoNode::_type_key}; |
73 | v->Visit("outputs" , &outputs_metadata_array); |
74 | int64_t num_outputs_cpp = num_outputs(); |
75 | v->Visit("num_outputs" , &num_outputs_cpp); |
76 | auto pools_array = Array<ObjectRef>(); |
77 | auto pools_accessor = workspace_pools(); |
78 | pools_array.reserve(num_workspace_pools()); |
79 | for (int64_t i = 0; i < num_workspace_pools(); ++i) { |
80 | pools_array.push_back(::tvm::runtime::metadata::TensorInfo{pools_accessor[i]}); |
81 | } |
82 | ::tvm::runtime::metadata::MetadataArray workspace_pools_metadata_array{ |
83 | pools_array, ::tvm::runtime::metadata::MetadataKind::kMetadata, |
84 | ::tvm::runtime::metadata::TensorInfoNode::_type_key}; |
85 | v->Visit("workspace_pools" , &workspace_pools_metadata_array); |
86 | int64_t num_workspace_pools_cpp = num_workspace_pools(); |
87 | v->Visit("num_workspace_pools" , &num_workspace_pools_cpp); |
88 | |
89 | auto consts_array = Array<ObjectRef>(); |
90 | auto consts_accessor = constant_pools(); |
91 | consts_array.reserve(num_constant_pools()); |
92 | for (int64_t i = 0; i < num_constant_pools(); ++i) { |
93 | consts_array.push_back(::tvm::runtime::metadata::ConstantInfoMetadata{consts_accessor[i]}); |
94 | } |
95 | |
96 | int64_t num_const_pools_cpp = num_constant_pools(); |
97 | ::tvm::runtime::metadata::MetadataArray constant_pools_metadata_array{ |
98 | consts_array, ::tvm::runtime::metadata::MetadataKind::kMetadata, |
99 | ::tvm::runtime::metadata::ConstantInfoMetadataNode::_type_key}; |
100 | v->Visit("constant_pools" , &constant_pools_metadata_array); |
101 | v->Visit("num_constant_pools" , &num_const_pools_cpp); |
102 | ::std::string mod_name_cpp{data()->mod_name}; |
103 | v->Visit("mod_name" , &mod_name_cpp); |
104 | } |
105 | }; |
106 | |
107 | /*! |
108 | * \brief Subclass of MetadataNode which also owns the backing C structures. |
109 | * |
110 | * This class (and other InMemory subclasses) are used during compilation to instantiate Metadata |
111 | * instances whose storage lives outside of .rodata. This class exists because the Module returned |
112 | * from tvm.relay.build must also be ready to run inference. |
113 | */ |
114 | class InMemoryMetadataNode : public ::tvm::target::metadata::VisitableMetadataNode { |
115 | public: |
116 | InMemoryMetadataNode() |
117 | : InMemoryMetadataNode(0 /* version */, {} /* inputs */, {} /* outputs */, |
118 | {} /* workspace_pools */, {} /* constant_pools */, "" /* mod_name */) { |
119 | } |
120 | InMemoryMetadataNode(int64_t version, |
121 | const ::std::vector<::tvm::runtime::metadata::TensorInfo>& inputs, |
122 | const ::std::vector<::tvm::runtime::metadata::TensorInfo>& outputs, |
123 | const ::std::vector<::tvm::runtime::metadata::TensorInfo>& workspace_pools, |
124 | const ::std::vector<::tvm::ConstantInfo>& constant_pools, |
125 | const ::tvm::runtime::String mod_name) |
126 | : VisitableMetadataNode{&storage_}, |
127 | inputs_{new struct TVMTensorInfo[inputs.size()]}, |
128 | inputs_objs_{inputs}, |
129 | outputs_{new struct TVMTensorInfo[outputs.size()]}, |
130 | outputs_objs_{outputs}, |
131 | workspace_pools_{new struct TVMTensorInfo[workspace_pools.size()]}, |
132 | workspace_pools_objs_{workspace_pools}, |
133 | constant_pools_{new struct TVMConstantInfo[constant_pools.size()]}, |
134 | constant_pools_objs_{constant_pools}, |
135 | mod_name_{mod_name}, |
136 | storage_{version, nullptr, 0ull, nullptr, 0ull, |
137 | nullptr, 0ull, nullptr, 0ull, mod_name_.c_str()} { |
138 | storage_.inputs = inputs_.get(); |
139 | storage_.num_inputs = inputs.size(); |
140 | for (unsigned int i = 0; i < inputs.size(); ++i) { |
141 | inputs_.get()[i] = *inputs[i]->data(); |
142 | } |
143 | storage_.outputs = outputs_.get(); |
144 | storage_.num_outputs = outputs.size(); |
145 | for (unsigned int i = 0; i < outputs.size(); ++i) { |
146 | outputs_.get()[i] = *outputs[i]->data(); |
147 | } |
148 | storage_.workspace_pools = workspace_pools_.get(); |
149 | storage_.num_workspace_pools = workspace_pools.size(); |
150 | for (unsigned int i = 0; i < workspace_pools.size(); ++i) { |
151 | workspace_pools_.get()[i] = *workspace_pools[i]->data(); |
152 | } |
153 | storage_.constant_pools = constant_pools_.get(); |
154 | storage_.num_constant_pools = constant_pools.size(); |
155 | for (size_t i = 0; i < constant_pools.size(); ++i) { |
156 | constant_pools_.get()[i].name_hint = constant_pools[i]->name_hint.c_str(); |
157 | constant_pools_.get()[i].byte_offset = constant_pools[i]->byte_offset.IntValue(); |
158 | |
159 | std::string bytes; |
160 | dmlc::MemoryStringStream stream(&bytes); |
161 | auto data = constant_pools[i]->data; |
162 | data.Save(&stream); |
163 | // Allocated mem freed in destructor |
164 | constant_pools_.get()[i].data_len = bytes.size(); |
165 | char* a = reinterpret_cast<char*>(malloc(bytes.size())); |
166 | constant_pools_.get()[i].data_bytes = a; |
167 | memcpy(a, bytes.c_str(), bytes.size()); |
168 | } |
169 | } |
170 | |
171 | ~InMemoryMetadataNode() { |
172 | // frees allocated mem for const_objs_ |
173 | for (int i = 0; i < storage_.num_constant_pools; ++i) { |
174 | free(const_cast<void*>(constant_pools_.get()[i].data_bytes)); |
175 | } |
176 | } |
177 | |
178 | private: |
179 | ::std::unique_ptr<struct TVMTensorInfo[]> inputs_; |
180 | std::vector<::tvm::runtime::metadata::TensorInfo> inputs_objs_; |
181 | ::std::unique_ptr<struct TVMTensorInfo[]> outputs_; |
182 | std::vector<::tvm::runtime::metadata::TensorInfo> outputs_objs_; |
183 | ::std::unique_ptr<struct TVMTensorInfo[]> workspace_pools_; |
184 | std::vector<::tvm::runtime::metadata::TensorInfo> workspace_pools_objs_; |
185 | ::std::unique_ptr<struct TVMConstantInfo[]> constant_pools_; |
186 | std::vector<::tvm::ConstantInfo> constant_pools_objs_; |
187 | ::std::string mod_name_; |
188 | struct ::TVMMetadata storage_; |
189 | }; |
190 | |
191 | class VisitableTensorInfoNode : public ::tvm::runtime::metadata::TensorInfoNode { |
192 | public: |
193 | explicit VisitableTensorInfoNode(const struct ::TVMTensorInfo* data) : TensorInfoNode{data} {} |
194 | VisitableTensorInfoNode() : TensorInfoNode{nullptr} {} |
195 | |
196 | void VisitAttrs(AttrVisitor* v) { |
197 | ::std::string name_cpp{data()->name}; |
198 | v->Visit("name" , &name_cpp); |
199 | auto shape_array = Array<ObjectRef>(); |
200 | auto shape_accessor = shape(); |
201 | shape_array.reserve(num_shape()); |
202 | for (int64_t i = 0; i < num_shape(); ++i) { |
203 | shape_array.push_back(::tvm::Integer{static_cast<int>(shape_accessor[i])}); |
204 | } |
205 | ::tvm::runtime::metadata::MetadataArray shape_metadata_array{ |
206 | shape_array, ::tvm::runtime::metadata::MetadataKind::kInt64, nullptr}; |
207 | v->Visit("shape" , &shape_metadata_array); |
208 | int64_t num_shape_cpp = num_shape(); |
209 | v->Visit("num_shape" , &num_shape_cpp); |
210 | ::tvm::runtime::DataType dtype_cpp{dtype()}; |
211 | v->Visit("dtype" , &dtype_cpp); |
212 | } |
213 | }; |
214 | |
215 | class InMemoryTensorInfoNode : public ::tvm::target::metadata::VisitableTensorInfoNode { |
216 | public: |
217 | InMemoryTensorInfoNode() : InMemoryTensorInfoNode("" , {}, ::tvm::runtime::DataType(0, 0, 0)) {} |
218 | InMemoryTensorInfoNode(const ::tvm::runtime::String& name, const ::std::vector<int64_t>& shape, |
219 | ::tvm::runtime::DataType dtype) |
220 | : VisitableTensorInfoNode{&storage_}, |
221 | name_{name}, |
222 | shape_{new int64_t[shape.size()]()}, |
223 | storage_{name_.c_str(), nullptr, 0, dtype} { |
224 | storage_.shape = shape_.get(); |
225 | storage_.num_shape = shape.size(); |
226 | for (unsigned int i = 0; i < shape.size(); ++i) { |
227 | shape_.get()[i] = shape[i]; |
228 | } |
229 | } |
230 | |
231 | private: |
232 | ::std::string name_; |
233 | ::std::unique_ptr<int64_t[]> shape_; |
234 | struct ::TVMTensorInfo storage_; |
235 | }; |
236 | |
237 | class VisitableConstantInfoMetadataNode |
238 | : public ::tvm::runtime::metadata::ConstantInfoMetadataNode { |
239 | public: |
240 | explicit VisitableConstantInfoMetadataNode(const struct ::TVMConstantInfo* data) |
241 | : ConstantInfoMetadataNode{data} {} |
242 | VisitableConstantInfoMetadataNode() : ConstantInfoMetadataNode{nullptr} {} |
243 | |
244 | void VisitAttrs(AttrVisitor* v) { |
245 | ::std::string name_cpp{name_hint()}; |
246 | v->Visit("name_hint" , &name_cpp); |
247 | |
248 | uint64_t byte_offset_cpp{byte_offset()}; |
249 | v->Visit("byte_offset" , &byte_offset_cpp); |
250 | |
251 | ::tvm::runtime::NDArray data_cpp = data(); |
252 | v->Visit("data" , &data_cpp); |
253 | } |
254 | }; |
255 | |
256 | } // namespace metadata |
257 | } // namespace target |
258 | } // namespace tvm |
259 | |
260 | #endif // TVM_TARGET_METADATA_H_ |
261 | |