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/session.h"
18 #include <set>
19 #include <vector>
20 #include <utility>
21 #include "coder/context.h"
22 #include "coder/train.h"
23 #include "coder/allocator/allocator.h"
24 #include "coder/generator/generator.h"
25 #include "coder/generator/inference/inference_generator.h"
26 #include "coder/generator/train/train_generator.h"
27 #include "coder/opcoders/op_coder_builder.h"
28 #include "coder/opcoders/kernel_registry.h"
29 #include "coder/utils/coder_utils.h"
30 #include "coder/log.h"
31 #include "src/ops/populate/populate_register.h"
32 #include "src/common/version_manager.h"
33 #include "src/runtime/infer_manager.h"
34 #include "src/scheduler.h"
35 #include "src/lite_model.h"
36 #include "include/errorcode.h"
37 #include "include/model.h"
38 #include "src/common/file_utils.h"
39 #include "src/common/prim_util.h"
40 #include "coder/opcoders/nnacl/dequant/de_quant.h"
41
42 namespace mindspore::lite::micro {
CoderSession()43 CoderSession::CoderSession() { allocator_ = MemoryAllocator::GetInstance(); }
44
EndCode()45 void CoderSession::EndCode() {
46 context_->set_tensor_map(allocator_->tensors_map());
47 context_->set_saved_weights(allocator_->saved_weights());
48 size_t de_quant_max_workspace_size = nnacl::Dequant::GetInstance()->de_quant_max_workspace();
49 size_t final_total_size = allocator_->total_buffer_size() > de_quant_max_workspace_size
50 ? allocator_->total_buffer_size()
51 : de_quant_max_workspace_size;
52 context_->set_total_buffer_size(final_total_size);
53 context_->set_graph_inputs(coder_graph_->input_tensors());
54 context_->set_graph_outputs(coder_graph_->output_tensors());
55 Configurator *config = Configurator::GetInstance();
56 if (config->debug_mode()) {
57 std::vector<std::string> blocks;
58 blocks = AddDumpDataInfo(context_->code_blocks(), op_coders_);
59 context_->set_code_blocks(blocks);
60 }
61 if (config->code_mode() == Train) {
62 Train::TransformGraphForTrain(context_.get(), op_coders_, schema_version_);
63 }
64 }
65
Run()66 int CoderSession::Run() {
67 MS_LOG(INFO) << "start run opcoders";
68 // 1. assign memory
69 std::vector<lite::Tensor *> inputs = coder_graph_->input_tensors();
70 int ret = allocator_->Assign(inputs, op_coders_);
71 MS_CHECK_RET_CODE(ret, "assign memory failed");
72 // 2. prepare, init model parameters
73 for (const auto &op_coder : op_coders_) {
74 MS_CHECK_PTR(op_coder);
75 MS_LOG(DEBUG) << "prepare: " << op_coder->name();
76 ret = op_coder->Prepare(context_.get());
77 MS_CHECK_RET_CODE(ret, "prepare coder " << op_coder->name() << " failed");
78 allocator_->enable_is_next();
79 }
80 // 3. docode, write operator code
81 for (const auto &op_coder : op_coders_) {
82 MS_CHECK_PTR(op_coder);
83 MS_LOG(DEBUG) << "code: " << op_coder->name();
84 ret = op_coder->DoCode(this->context_.get());
85 MS_CHECK_RET_CODE(ret, "do coder " << op_coder->name() << " failed");
86 }
87
88 this->EndCode();
89 MS_LOG(INFO) << "run opcoders success";
90 return RET_OK;
91 }
92
GenerateCode()93 int CoderSession::GenerateCode() {
94 MS_LOG(INFO) << "CoderSession::GenerateCode start";
95 std::shared_ptr<Generator> generator;
96 Configurator *config = Configurator::GetInstance();
97 CodeMode code_mode = config->code_mode();
98 switch (code_mode) {
99 case Inference:
100 MS_LOG(INFO) << "generate code for Inference";
101 generator = std::make_shared<InferenceGenerator>(std::move(context_));
102 break;
103 case Train:
104 MS_LOG(INFO) << "generate code for Train";
105 generator = std::make_shared<TrainGenerator>(std::move(context_));
106 break;
107 default:
108 MS_LOG(ERROR) << "unsupported generator code mode, " << code_mode;
109 return RET_ERROR;
110 }
111 // when use file, coder context need to remove initial parameters from tensors info
112 // we use tmp_tensor_list to storage
113 MS_CHECK_PTR(generator);
114 int ret = generator->GenerateCode();
115 if (ret != RET_OK) {
116 MS_LOG(ERROR) << "generate code failed";
117 }
118 MS_LOG(INFO) << "CoderSession::GenerateCode done";
119 return ret;
120 }
121
Init(const std::string & model_path)122 int CoderSession::Init(const std::string &model_path) {
123 MS_LOG(INFO) << "CoderSession::Init start";
124 // Load graph
125 MS_LOG(DEBUG) << "start reading model file";
126 size_t size = 0;
127 char *graph_buf = ReadFile(model_path.c_str(), &size);
128 if (graph_buf == nullptr) {
129 MS_LOG(ERROR) << "read model file from path \"" << model_path << "\" failed.";
130 return RET_ERROR;
131 }
132 // new a context for session
133 if (size >= UINT_MAX) {
134 MS_LOG(ERROR) << "the size is invalid";
135 delete[] graph_buf;
136 return RET_ERROR;
137 }
138 Model *model = lite::Model::Import(graph_buf, size);
139 delete[] graph_buf;
140 MS_CHECK_PTR(model);
141 coder_graph_ = std::make_unique<CoderGraph>(model);
142 context_ = std::make_unique<CoderContext>();
143 MS_LOG(INFO) << "CoderSession::Init done";
144 return RET_OK;
145 }
146
Build()147 int CoderSession::Build() {
148 if (coder_graph_ == nullptr) {
149 return RET_ERROR;
150 }
151 int ret = this->CompileGraph();
152 if (ret != RET_OK) {
153 MS_LOG(ERROR) << "CompileGraph failed: " << ret;
154 return ret;
155 }
156 return RET_OK;
157 }
158
InitOpcodersInputsAndOutputs()159 int CoderSession::InitOpcodersInputsAndOutputs() {
160 std::map<Tensor *, OperatorCoder *> input_node_map;
161 std::map<Tensor *, OperatorCoder *> output_node_map;
162 for (const auto &op_coder : op_coders_) {
163 std::vector<Tensor *> inputs = op_coder->input_tensors();
164 std::for_each(inputs.begin(), inputs.end(),
165 [&](Tensor *t) { input_node_map.insert(std::make_pair(t, op_coder.get())); });
166 std::vector<Tensor *> outputs = op_coder->input_tensors();
167 std::for_each(outputs.begin(), outputs.end(),
168 [&](Tensor *t) { output_node_map.insert(std::make_pair(t, op_coder.get())); });
169 }
170 for (const auto &op_coder : op_coders_) {
171 std::vector<Tensor *> inputs = op_coder->input_tensors();
172 for (const auto &tensor : inputs) {
173 auto item = output_node_map.find(tensor);
174 if (item != output_node_map.end()) {
175 op_coder->AddInputOp(item->second);
176 }
177 }
178 std::vector<Tensor *> outputs = op_coder->output_tensors();
179 for (const auto &tensor : outputs) {
180 auto item = input_node_map.find(tensor);
181 if (item != input_node_map.end()) {
182 op_coder->AddOutputOp(item->second);
183 }
184 }
185 }
186 return RET_OK;
187 }
188
InitTensorsRef()189 int CoderSession::InitTensorsRef() {
190 auto all_tensors = coder_graph_->all_tensors();
191 for (auto &tensor : all_tensors) {
192 size_t refcount = 0;
193 for (const auto &node : this->op_coders_) {
194 auto inputs = node->input_tensors();
195 auto iter = std::find(inputs.begin(), inputs.end(), tensor);
196 if (iter != inputs.end()) {
197 refcount++;
198 }
199 }
200 tensor->set_ref_count(refcount);
201 }
202 return RET_OK;
203 }
204
GenParameterAndInfer(const Model::Node * node,const std::vector<lite::Tensor * > & inputs,std::vector<lite::Tensor * > * outputs) const205 OpParameter *CoderSession::GenParameterAndInfer(const Model::Node *node, const std::vector<lite::Tensor *> &inputs,
206 std::vector<lite::Tensor *> *outputs) const {
207 auto primitive = node->primitive_;
208 MS_CHECK_PTR_RET_NULL(primitive);
209 auto parame_gen =
210 PopulateRegistry::GetInstance()->GetParameterCreator(GetPrimitiveType(primitive, schema_version_), schema_version_);
211 MS_CHECK_PTR_RET_NULL(parame_gen);
212 auto parameter = parame_gen(primitive);
213 MS_CHECK_PTR_RET_NULL(parameter);
214 auto ret = KernelInferShape(inputs, *outputs, parameter);
215 if (ret == RET_INFER_INVALID) {
216 MS_LOG(INFO) << "InferShape shouldn't be done before runtime, name: " << node->name_
217 << ", type: " << GetPrimitiveTypeName(primitive, schema_version_) << "flag set to false.";
218 } else if (ret != RET_OK) {
219 MS_LOG(ERROR) << "InferShape failed, name: " << node->name_
220 << ", type: " << GetPrimitiveTypeName(primitive, schema_version_);
221 return nullptr;
222 }
223 return parameter;
224 }
225
CreateOpCoders()226 int CoderSession::CreateOpCoders() {
227 const Model *model = coder_graph_->model();
228 if (model == nullptr) {
229 MS_LOG(ERROR) << "Graph model is nullptr";
230 return RET_ERROR;
231 }
232 schema_version_ = reinterpret_cast<const lite::LiteModel *>(model)->GetSchemaVersion();
233 Configurator *config = Configurator::GetInstance();
234 Target code_target = config->target();
235 CodeMode code_mode = config->code_mode();
236 bool support_parallel = config->support_parallel();
237 uint32_t nodes_size = model->all_nodes_.size();
238 OpCoderBuilder builder;
239 for (uint32_t i = 0; i < nodes_size; ++i) {
240 const auto *node = model->all_nodes_.at(i);
241 if (node == nullptr) {
242 MS_LOG(ERROR) << "node is nullptr";
243 return RET_ERROR;
244 }
245 std::vector<lite::Tensor *> all_tensors = coder_graph_->all_tensors();
246 if (all_tensors.empty()) {
247 MS_LOG(ERROR) << "coder_graph has no any tensors";
248 return RET_ERROR;
249 }
250 // set op_coder's inputs && outputs info
251 std::vector<uint32_t> input_indices;
252 Uint32Vector node_input_indices = node->input_indices_;
253 input_indices.insert(input_indices.end(), node_input_indices.begin(), node_input_indices.end());
254 std::vector<uint32_t> output_indices;
255 Uint32Vector node_output_indices = node->output_indices_;
256 output_indices.insert(output_indices.end(), node_output_indices.begin(), node_output_indices.end());
257
258 std::vector<lite::Tensor *> inputs;
259 std::vector<lite::Tensor *> outputs;
260 for (auto in_index : input_indices) {
261 in_index = static_cast<size_t>(in_index);
262 if (in_index > all_tensors.size()) {
263 MS_LOG(ERROR) << "in_index is invalid";
264 return RET_ERROR;
265 }
266 inputs.push_back(all_tensors.at(in_index));
267 }
268 for (auto out_index : output_indices) {
269 out_index = static_cast<size_t>(out_index);
270 if (out_index > all_tensors.size()) {
271 MS_LOG(ERROR) << "out_index is invalid";
272 return RET_ERROR;
273 }
274 outputs.push_back(all_tensors.at(out_index));
275 }
276 if (inputs.empty()) {
277 MS_LOG(ERROR) << "node: " << node->name_ << "has no inputs tensor";
278 return RET_ERROR;
279 }
280 if (outputs.empty()) {
281 MS_LOG(ERROR) << "node: " << node->name_ << "has no outputs tensor";
282 return RET_ERROR;
283 }
284
285 OpParameter *parameter = nullptr;
286 if (IsCustomNode(node->primitive_, schema_version_)) {
287 KernelRegistry::GetInstance()->RegisterKernel(schema::PrimitiveType_Custom);
288 } else {
289 parameter = GenParameterAndInfer(node, inputs, &outputs); // built-in ops infer
290 MS_CHECK_PTR(parameter);
291 }
292
293 TypeId tensor_data_type = inputs.at(0)->data_type();
294 std::unique_ptr<OperatorCoder> op_coder = builder.inputs(inputs)
295 .outputs(outputs)
296 .node(node)
297 .parameter(parameter)
298 .target(code_target)
299 .support_parallel(support_parallel)
300 .data_type(tensor_data_type)
301 .mode(code_mode)
302 .input_indices(input_indices)
303 .output_indices(output_indices)
304 .build(schema_version_);
305 if (op_coder == nullptr) {
306 coder_graph_->DumpUnSupportLayer(code_target);
307 return RET_ERROR;
308 }
309 op_coders_.push_back(std::move(op_coder));
310 builder.Reset();
311 }
312 InitOpcodersInputsAndOutputs();
313 return RET_OK;
314 }
315
InitCodeGraph()316 int CoderSession::InitCodeGraph() {
317 MS_CHECK_RET_CODE(coder_graph_->ConvertTensors(), "convert tensors failed");
318 MS_CHECK_RET_CODE(coder_graph_->InitGraphInOutTensors(), "init graph inputs and outputs failed");
319 return RET_OK;
320 }
321
CompileGraph()322 int CoderSession::CompileGraph() {
323 MS_LOG(INFO) << "CompileGraph";
324 MS_CHECK_RET_CODE(InitCodeGraph(), "InitGraphInOutTensors failed");
325 MS_CHECK_RET_CODE(CreateOpCoders(), "CreateOpCoders failed!");
326 MS_CHECK_RET_CODE(InitTensorsRef(), "InitTensorsRefcount failed!");
327 return RET_OK;
328 }
329
CreateCoderSession()330 std::shared_ptr<CoderSession> CreateCoderSession() {
331 auto session = std::make_shared<CoderSession>();
332 return session;
333 }
334
~CoderSession()335 CoderSession::~CoderSession() { allocator_->Free(); }
336 } // namespace mindspore::lite::micro
337