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 interface_c.cc |
22 | * \brief Generates a C interface header for a given modules inputs and outputs |
23 | */ |
24 | |
25 | #include <tvm/runtime/container/array.h> |
26 | #include <tvm/runtime/container/string.h> |
27 | #include <tvm/runtime/module.h> |
28 | #include <tvm/runtime/name_transforms.h> |
29 | #include <tvm/runtime/packed_func.h> |
30 | #include <tvm/runtime/registry.h> |
31 | #include <tvm/tir/usmp/utils.h> |
32 | |
33 | #include <numeric> |
34 | #include <string> |
35 | |
36 | #include "../../relay/backend/name_transforms.h" |
37 | #include "codegen_params.h" |
38 | |
39 | namespace tvm { |
40 | namespace codegen { |
41 | |
42 | using runtime::PackedFunc; |
43 | using namespace tvm::relay::backend; |
44 | |
45 | class InterfaceCNode : public runtime::ModuleNode { |
46 | public: |
47 | InterfaceCNode(std::string module_name, Array<String> inputs, Array<String> outputs, |
48 | Array<tir::usmp::AllocatedPoolInfo> pools, |
49 | Map<String, tir::usmp::PoolAllocation> io_pool_allocations, Array<String> devices, |
50 | int workspace_size, Map<String, IntImm> input_sizes, |
51 | Map<String, IntImm> output_sizes) |
52 | : module_name_(module_name), |
53 | inputs_(inputs), |
54 | outputs_(outputs), |
55 | devices_(devices), |
56 | pools_(FilterExternalPools(pools)), |
57 | io_pool_allocations_(io_pool_allocations), |
58 | workspace_size_(workspace_size), |
59 | input_sizes_(input_sizes), |
60 | output_sizes_(output_sizes) {} |
61 | const char* type_key() const final { return "h" ; } |
62 | |
63 | std::string GetSource(const std::string& format) final { |
64 | std::stringstream code; |
65 | |
66 | EmitUpperHeaderGuard(code); |
67 | |
68 | // Emit macros for input sizes |
69 | for (auto const& it : input_sizes_) { |
70 | std::string input_name = SanitizeName(it.first); |
71 | std::string input_macro_name = input_name + "_size" ; |
72 | int input_size = it.second->value; |
73 | EmitIntegerValueMacro(code, "Input tensor " + input_name + " size (in bytes)" , |
74 | input_macro_name, input_size); |
75 | } |
76 | |
77 | // Emit macros for output sizes |
78 | for (auto const& it : output_sizes_) { |
79 | std::string output_name = SanitizeName(it.first); |
80 | std::string output_macro_name = output_name + "_size" ; |
81 | int output_size = it.second->value; |
82 | EmitIntegerValueMacro(code, "Output tensor " + output_name + " size (in bytes)" , |
83 | output_macro_name, output_size); |
84 | } |
85 | |
86 | EmitBrief(code, "Input tensor pointers" ); |
87 | EmitStruct(code, "inputs" , inputs_); |
88 | EmitBrief(code, "Output tensor pointers" ); |
89 | EmitStruct(code, "outputs" , outputs_); |
90 | |
91 | if (!devices_.empty()) { |
92 | EmitBrief(code, "Device context pointers" ); |
93 | EmitStruct(code, "devices" , devices_); |
94 | } |
95 | if (!pools_.empty()) { |
96 | EmitBrief(code, "Workspace pool pointers" ); |
97 | Array<String> pool_names; |
98 | for (const tir::usmp::AllocatedPoolInfo pool : pools_) { |
99 | pool_names.push_back(pool->pool_info->pool_name); |
100 | } |
101 | EmitStruct(code, "workspace_pools" , pool_names); |
102 | } |
103 | |
104 | if (!io_pool_allocations_.empty()) { |
105 | std::string inputs_struct = ToCVariableStyle(PrefixGeneratedName({module_name_, "inputs" })); |
106 | EmitMapIOToPoolsFunction(code, inputs_struct, "map_inputs" , inputs_); |
107 | std::string outputs_struct = ToCVariableStyle(PrefixGeneratedName({module_name_, "outputs" })); |
108 | EmitMapIOToPoolsFunction(code, outputs_struct, "map_outputs" , outputs_); |
109 | } |
110 | |
111 | EmitRunFunction(code); |
112 | // Emit workspace |
113 | EmitIntegerValueMacro(code, "Workspace size" , "WORKSPACE_SIZE" , workspace_size_); |
114 | // Emit memory pool sizes |
115 | for (const tir::usmp::AllocatedPoolInfo pool : pools_) { |
116 | String pool_name = pool->pool_info->pool_name; |
117 | Integer pool_size = pool->allocated_size; |
118 | if (const auto* pool_info = pool->pool_info.as<ConstantPoolInfoNode>()) { |
119 | EmitConstantPool(code, SanitizeName(pool_name) + " initialization data" , pool_info); |
120 | } else { |
121 | EmitIntegerValueMacro(code, SanitizeName(pool_name) + " size" , |
122 | SanitizeName(pool_name) + _macro_workspace_pool_size_postfix, |
123 | pool_size->value); |
124 | } |
125 | } |
126 | EmitLowerHeaderGuard(code); |
127 | |
128 | return code.str(); |
129 | } |
130 | |
131 | PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) final { |
132 | return PackedFunc(); |
133 | } |
134 | |
135 | private: |
136 | constexpr static const char* _macro_workspace_pool_size_postfix = "_WORKSPACE_POOL_SIZE" ; |
137 | constexpr static const char* _macro_constant_pool_size_postfix = "_CONSTANT_POOL_SIZE" ; |
138 | constexpr static const char* _macro_constant_pool_data_postfix = "_CONSTANT_POOL_DATA" ; |
139 | |
140 | void (std::stringstream& code_stream) { |
141 | std::string = ToCConstantStyle(PrefixGeneratedName({module_name_, "H" })); |
142 | code_stream << "#ifndef " << header_guard_name << "_\n" |
143 | << "#define " << header_guard_name << "_\n" |
144 | << "#include <stdint.h>\n\n" |
145 | << "#ifdef __cplusplus\n" |
146 | << "extern \"C\" {\n" |
147 | << "#endif\n\n" ; |
148 | } |
149 | |
150 | void (std::stringstream& code_stream) { |
151 | std::string = ToCConstantStyle(PrefixGeneratedName({module_name_, "H" })); |
152 | code_stream << "\n#ifdef __cplusplus\n" |
153 | << "}\n" |
154 | << "#endif\n\n" |
155 | << "#endif // " << header_guard_name << "_\n" ; |
156 | } |
157 | |
158 | void EmitBrief(std::stringstream& code_stream, const std::string& description) { |
159 | code_stream << "/*!\n" |
160 | << " * \\brief " << description << " for TVM module \"" << module_name_ << "\" \n" |
161 | << " */\n" ; |
162 | } |
163 | |
164 | void EmitStruct(std::stringstream& code_stream, const std::string& suffix, |
165 | Array<String> properties) { |
166 | std::string struct_name = ToCVariableStyle(PrefixGeneratedName({module_name_, suffix})); |
167 | code_stream << "struct " << struct_name << " {\n" ; |
168 | |
169 | std::vector<std::string> sanitized_properties; |
170 | for (const String& property : properties) { |
171 | std::string sanitized_property = SanitizeName(property); |
172 | ICHECK(std::find(sanitized_properties.begin(), sanitized_properties.end(), |
173 | sanitized_property) == sanitized_properties.end()) |
174 | << "Sanitized input tensor name clash" << sanitized_property; |
175 | code_stream << " void* " << sanitized_property << ";\n" ; |
176 | sanitized_properties.push_back(sanitized_property); |
177 | } |
178 | code_stream << "};\n\n" ; |
179 | } |
180 | |
181 | void EmitIntegerValueMacro(std::stringstream& code_stream, const std::string& brief_description, |
182 | const std::string& macro_name, int macro_value) { |
183 | EmitBrief(code_stream, brief_description); |
184 | std::string macro_name_prefixed = |
185 | ToCConstantStyle(PrefixGeneratedName({module_name_, macro_name})); |
186 | code_stream << "#define " << macro_name_prefixed << " " << macro_value << "\n" ; |
187 | } |
188 | |
189 | void EmitConstantPool(std::stringstream& code_, const std::string& brief_description, |
190 | const ConstantPoolInfoNode* pool_info) { |
191 | EmitBrief(code_, brief_description); |
192 | std::string name_prefixed = |
193 | ToCConstantStyle(PrefixGeneratedName({module_name_, SanitizeName(pool_info->pool_name)})); |
194 | |
195 | if (pool_info->constant_info_array.size() > 0) { |
196 | std::vector<ConstantInfo> const_info_vec(pool_info->constant_info_array.begin(), |
197 | pool_info->constant_info_array.end()); |
198 | std::sort(const_info_vec.begin(), const_info_vec.end(), |
199 | [](const ConstantInfo& a, const ConstantInfo& b) { |
200 | return a->byte_offset->value < b->byte_offset->value; |
201 | }); |
202 | int64_t accumulated_pool_len = |
203 | const_info_vec.back()->byte_offset.IntValue() + |
204 | runtime::GetDataSize(*const_info_vec.back()->data.operator->()); |
205 | const auto& accumulated_pool = runtime::NDArray::Empty( |
206 | {accumulated_pool_len}, DataType::UInt(8), const_info_vec.back()->data->device); |
207 | for (const auto& const_info : const_info_vec) { |
208 | const auto& data = const_info->data; |
209 | const auto& offs = const_info->byte_offset; |
210 | data.CopyToBytes(static_cast<uint8_t*>(accumulated_pool->data) + offs.IntValue(), |
211 | runtime::GetDataSize(*data.operator->())); |
212 | } |
213 | |
214 | code_ << "#define " << name_prefixed << _macro_constant_pool_size_postfix << " " |
215 | << accumulated_pool_len << "\n" ; |
216 | code_ << "#define " << name_prefixed << _macro_constant_pool_data_postfix << " \\\n" ; |
217 | codegen::NDArrayDataToC(accumulated_pool, 4, code_, "\\\n" ); |
218 | code_ << '\n'; |
219 | |
220 | } else { |
221 | LOG(FATAL) << "No constant data in constant pool found " << GetRef<ObjectRef>(pool_info); |
222 | } |
223 | } |
224 | |
225 | void EmitRunFunction(std::stringstream& code_stream) { |
226 | std::string run_function = ToCVariableStyle(PrefixGeneratedName({module_name_, "run" })); |
227 | std::string inputs_struct = ToCVariableStyle(PrefixGeneratedName({module_name_, "inputs" })); |
228 | std::string outputs_struct = ToCVariableStyle(PrefixGeneratedName({module_name_, "outputs" })); |
229 | std::string devices_struct = ToCVariableStyle(PrefixGeneratedName({module_name_, "devices" })); |
230 | std::string pools_struct = |
231 | ToCVariableStyle(PrefixGeneratedName({module_name_, "workspace_pools" })); |
232 | |
233 | code_stream << "/*!\n" |
234 | << " * \\brief entrypoint function for TVM module \"" << module_name_ << "\"\n" ; |
235 | if (io_pool_allocations_.empty()) { |
236 | code_stream << " * \\param inputs Input tensors for the module \n" ; |
237 | code_stream << " * \\param outputs Output tensors for the module \n" ; |
238 | } |
239 | |
240 | if (!pools_.empty()) { |
241 | code_stream << " * \\param workspace_pools Workspace memory pool pointers for the module \n" ; |
242 | } |
243 | if (!devices_.empty()) { |
244 | code_stream << " * \\param devices Device context pointers for the module \n" ; |
245 | } |
246 | |
247 | code_stream << " */\n" |
248 | << "int32_t " << run_function << "(\n" ; |
249 | |
250 | std::stringstream call_args_ss; |
251 | if (io_pool_allocations_.empty()) { |
252 | call_args_ss << " struct " << inputs_struct << "* inputs,\n" ; |
253 | call_args_ss << " struct " << outputs_struct << "* outputs,\n" ; |
254 | } |
255 | if (!pools_.empty()) { |
256 | call_args_ss << " struct " << pools_struct << "* workspace_pools,\n" ; |
257 | } |
258 | if (!devices_.empty()) { |
259 | call_args_ss << " struct " << devices_struct << "* devices,\n" ; |
260 | } |
261 | std::string call_args_str = call_args_ss.str(); |
262 | call_args_str.pop_back(); |
263 | call_args_str.pop_back(); |
264 | code_stream << call_args_str << "\n);\n" ; |
265 | } |
266 | |
267 | void EmitMapIOToPoolsFunction(std::stringstream& code_stream, const std::string& struct_type, |
268 | const std::string& function_name, |
269 | const Array<String>& tensor_names) { |
270 | code_stream << "/*!\n" |
271 | << " * \\brief Maps I/O inside the workspace pools for TVM module \"" |
272 | << module_name_ << "\"\n" |
273 | << " * \\param workspace_pools Workspace memory pool struct for the module \n" |
274 | << " * \\return I/O tensor struct for the module \n" ; |
275 | std::string map_function = ToCVariableStyle(PrefixGeneratedName({module_name_, function_name})); |
276 | code_stream << " */\n" |
277 | << "struct " << struct_type << " " << map_function << "(\n" ; |
278 | std::string pools_struct = |
279 | ToCVariableStyle(PrefixGeneratedName({module_name_, "workspace_pools" })); |
280 | code_stream << " struct " << pools_struct << "* workspace_pools\n" ; |
281 | code_stream << ");\n\n" ; |
282 | } |
283 | |
284 | Array<tir::usmp::AllocatedPoolInfo> FilterExternalPools( |
285 | const Array<tir::usmp::AllocatedPoolInfo>& pools) { |
286 | Array<tir::usmp::AllocatedPoolInfo> external_pools; |
287 | for (tir::usmp::AllocatedPoolInfo pool : pools) { |
288 | if (!pool->pool_info->is_internal) { |
289 | external_pools.push_back(pool); |
290 | } |
291 | } |
292 | return external_pools; |
293 | } |
294 | |
295 | std::string module_name_; |
296 | Array<String> inputs_; |
297 | Array<String> outputs_; |
298 | Array<String> devices_; |
299 | Array<tir::usmp::AllocatedPoolInfo> pools_; |
300 | Map<String, tir::usmp::PoolAllocation> io_pool_allocations_; |
301 | int workspace_size_; |
302 | Map<String, IntImm> input_sizes_; |
303 | Map<String, IntImm> output_sizes_; |
304 | }; |
305 | |
306 | runtime::Module InterfaceCCreate(std::string module_name, Array<String> inputs, |
307 | Array<String> outputs, Array<tir::usmp::AllocatedPoolInfo> pools, |
308 | Map<String, tir::usmp::PoolAllocation> io_pool_allocations, |
309 | Array<String> devices, int workspace_size, |
310 | Map<String, IntImm> input_sizes, |
311 | Map<String, IntImm> output_sizes) { |
312 | auto n = make_object<InterfaceCNode>(module_name, inputs, outputs, pools, io_pool_allocations, |
313 | devices, workspace_size, input_sizes, output_sizes); |
314 | return runtime::Module(n); |
315 | } |
316 | |
317 | TVM_REGISTER_GLOBAL("runtime.InterfaceCCreate" ).set_body_typed(InterfaceCCreate); |
318 | |
319 | } // namespace codegen |
320 | } // namespace tvm |
321 | |