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
33namespace tvm {
34namespace target {
35namespace 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 */
44class 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 */
114class 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
191class 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
215class 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
237class 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