1 /**
2 * Copyright 2021 Huawei Technologies Co., Ltd
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 #include "coder/generator/component/weight_component.h"
18 #include <memory>
19 #include <utility>
20 #include "coder/generator/component/const_blocks/license.h"
21 #include "coder/utils/coder_utils.h"
22 #include "coder/opcoders/parallel.h"
23
24 namespace mindspore::lite::micro {
CodeWeightFileHeader(std::ofstream & ofs,const std::unique_ptr<CoderContext> & ctx)25 void CodeWeightFileHeader(std::ofstream &ofs, const std::unique_ptr<CoderContext> &ctx) {
26 ofs << g_hwLicense;
27 // include all operator header
28 for (const auto &h_file : ctx->h_files()) {
29 ofs << "#include \"" << h_file << "\"\n";
30 }
31 ofs << "#include <stdlib.h>\n"
32 << "#include <string.h>\n"
33 << "extern unsigned char *" << ctx->buffer_name() << ";\n";
34 ofs << "enum STATUS {\n"
35 " RET_OK = 0,\n"
36 " RET_ERROR = 1,\n"
37 "};\n\n";
38 // set a global var for thread_pool
39 ofs << "extern int " << gThreadNum << ";\n";
40 }
41
CodeModelParamsState(std::ofstream & ofs,const std::map<std::string,Tensor * > & weights)42 void CodeModelParamsState(std::ofstream &ofs, const std::map<std::string, Tensor *> &weights) {
43 for (auto &item : weights) {
44 std::string name = item.first;
45 Tensor *tensor = item.second;
46 if (CheckConstantTensor(tensor)) {
47 if (tensor->data() == nullptr) {
48 continue;
49 }
50 ofs << "extern const " << GetTensorDataType(tensor->data_type()) << name << "[];\n";
51 }
52 }
53 }
54
CodeModelParamsData(std::ofstream & ofs,const std::map<std::string,Tensor * > & weights)55 void CodeModelParamsData(std::ofstream &ofs, const std::map<std::string, Tensor *> &weights) {
56 for (auto &item : weights) {
57 std::string name = item.first;
58 Tensor *tensor = item.second;
59 if (CheckConstantTensor(tensor)) {
60 if (tensor->data() == nullptr) {
61 continue;
62 }
63 ofs << "const " << GetTensorDataType(tensor->data_type()) << name << "[] = ";
64 PrintTensorData(tensor, ofs);
65 }
66 }
67 }
68
CodeModelParamsForNet(std::ofstream & hofs,std::ofstream & cofs,const std::unique_ptr<CoderContext> & ctx)69 void CodeModelParamsForNet(std::ofstream &hofs, std::ofstream &cofs, const std::unique_ptr<CoderContext> &ctx) {
70 // reverse key and value of tensors_map
71 std::map<std::string, Tensor *> address_map;
72 for (const auto &item : ctx->tensors_map()) {
73 address_map.insert(std::make_pair(item.second, item.first));
74 }
75 for (auto &item : address_map) {
76 std::string name = item.first;
77 Tensor *tensor = item.second;
78 if (tensor->data() == nullptr) {
79 continue;
80 }
81 if (CheckConstantTensor(tensor)) {
82 hofs << "extern " << GetTensorDataType(tensor->data_type()) << name << "[];\n";
83 cofs << GetTensorDataType(tensor->data_type()) << name << "[" << tensor->ElementsNum() << "];\n";
84 } else if (tensor->category() == Tensor::Category::VAR) {
85 hofs << "extern " << GetTensorDataType(tensor->data_type()) << "*" << name << ";\n";
86 cofs << GetTensorDataType(tensor->data_type()) << "*" << name << " = NULL;\n";
87 }
88 }
89 cofs << "\n";
90 }
91
CodeInitWeightState(std::ofstream & ofs)92 void CodeInitWeightState(std::ofstream &ofs) {
93 ofs << "/**\n"
94 << " * @param weight_buffer, the address of the weight binary file\n"
95 << " * @param weight_size, the size of the model file in bytes\n"
96 << " **/\n"
97 << "int Init(void *weight_buffer, int weight_size);\n\n";
98 }
99
CodeWeightInitFunc(std::ofstream & ofs,const std::unique_ptr<CoderContext> & ctx)100 void CodeWeightInitFunc(std::ofstream &ofs, const std::unique_ptr<CoderContext> &ctx) {
101 ofs << "int Init(void *weight_buffer, int weight_size) {\n"
102 << " if (weight_buffer == NULL) {\n"
103 << " return RET_ERROR;\n"
104 << " }\n";
105 ofs << " struct ModelParameter {\n"
106 << " void *addr;\n"
107 << " size_t size;\n"
108 << " size_t offset;\n"
109 << " };\n";
110 size_t params_num = 0;
111 size_t offset = 0;
112 std::string params;
113 std::string origins;
114 for (const auto &item : ctx->saved_weights()) {
115 std::string name = item.first;
116 Tensor *tensor = item.second;
117 if (!CheckConstantTensor(tensor)) {
118 continue;
119 }
120 std::map<Tensor *, std::string> ctx_tensor_map = ctx->tensors_map();
121 auto iter = ctx_tensor_map.find(tensor);
122 if (iter != ctx_tensor_map.end()) {
123 origins += " {" + name + ", " + std::to_string(tensor->Size()) + ", " + std::to_string(offset) + "},\n";
124 params_num++;
125 } else {
126 TypeId data_type = tensor->data_type();
127 params +=
128 " " + GetTensorDataType(data_type) + "*" + name + " = (weight_buffer + " + std::to_string(offset) + ");\n";
129 }
130 offset += tensor->Size();
131 }
132 ofs << params << "\n";
133 ofs << " struct ModelParameter model_params[] = {\n" << origins << " };\n";
134
135 ofs << "\n";
136 ofs << " for(int i = 0; i < " << params_num << "; ++i) {\n"
137 << " if (model_params[i].offset + model_params[i].size > weight_size) {\n"
138 " return RET_ERROR;\n"
139 " }\n"
140 << " memcpy(model_params[i].addr, (weight_buffer + model_params[i].offset), model_params[i].size);\n"
141 << " }\n";
142 for (const auto &block : ctx->init_contents()) {
143 ofs << "{\n" << block << "}\n";
144 }
145 ofs << " return RET_OK;\n";
146 ofs << "}\n\n";
147 }
148
SaveDataToNet(const std::map<std::string,Tensor * > & saved_weights,const std::string & net_file)149 void SaveDataToNet(const std::map<std::string, Tensor *> &saved_weights, const std::string &net_file) {
150 std::ofstream net(net_file, std::ios::out | std::ios::trunc | std::ios::binary);
151 MS_CHECK_TRUE_WITHOUT_RET(net.is_open(), "net file open failed!");
152 for (auto &item : saved_weights) {
153 std::string name = item.first;
154 Tensor *tensor = item.second;
155 if ((CheckConstantTensor(tensor)) && tensor->data() != nullptr) {
156 net.write(reinterpret_cast<const char *>(tensor->data()), tensor->Size());
157 }
158 }
159 net.close();
160 }
161 } // namespace mindspore::lite::micro
162