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
39namespace tvm {
40namespace codegen {
41
42using runtime::PackedFunc;
43using namespace tvm::relay::backend;
44
45class 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 EmitUpperHeaderGuard(std::stringstream& code_stream) {
141 std::string header_guard_name = 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 EmitLowerHeaderGuard(std::stringstream& code_stream) {
151 std::string header_guard_name = 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
306runtime::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
317TVM_REGISTER_GLOBAL("runtime.InterfaceCCreate").set_body_typed(InterfaceCCreate);
318
319} // namespace codegen
320} // namespace tvm
321