You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.

dump_proto.cc 23 kB

6 years ago
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609
  1. /**
  2. * Copyright 2019-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. #include "debug/dump_proto.h"
  17. #include <algorithm>
  18. #include <fstream>
  19. #include <map>
  20. #include <memory>
  21. #include <utility>
  22. #include <vector>
  23. #include "debug/anf_ir_utils.h"
  24. #include "proto/anf_ir.pb.h"
  25. #include "ir/graph_utils.h"
  26. #include "utils/ms_context.h"
  27. #include "utils/symbolic.h"
  28. #include "utils/utils.h"
  29. #include "pipeline/jit/base.h"
  30. namespace mindspore {
  31. class ProtoExporter {
  32. public:
  33. ProtoExporter() {}
  34. ~ProtoExporter() {}
  35. std::string GetFuncGraphProtoString(const FuncGraphPtr &func_graph);
  36. void ExportFuncGraph(const FuncGraphPtr &func_graph, irpb::GraphProto *graph_proto);
  37. private:
  38. void InitModelInfo();
  39. void GetOpNodeTypeAndAttrs(const FuncGraphPtr &func_graph, const AnfNodePtr &node, irpb::NodeProto *node_proto);
  40. std::string GetOpNodeInputId(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
  41. const std::map<AnfNodePtr, size_t> &apply_map,
  42. std::map<AnfNodePtr, size_t> *const_map_ptr);
  43. void SetValueToProto(const ValuePtr &attr_value, irpb::ValueProto *value_proto);
  44. void SetScalarToProto(const ScalarPtr &val, irpb::ValueProto *value_proto);
  45. void SetSequenceToProto(const ValueSequeuePtr &val, irpb::ValueProto *value_proto);
  46. void SetDictionaryToProto(const ValueDictionaryPtr &val, irpb::ValueProto *value_proto);
  47. void SetNodeOutputType(const AnfNodePtr &node, irpb::TypeProto *type_proto);
  48. void SetNodeOutputType(const TypePtr &node, const BaseShapePtr &shape, irpb::TypeProto *type_proto);
  49. void ExportParameters(const FuncGraphPtr &func_graph, irpb::GraphProto *graph_proto);
  50. void ExportCNodes(const FuncGraphPtr &func_graph, irpb::GraphProto *graph_proto,
  51. std::map<AnfNodePtr, size_t> *const_map_ptr);
  52. void ExportCNode(const FuncGraphPtr &func_graph, const CNodePtr &node, std::map<AnfNodePtr, size_t> *apply_map_ptr,
  53. std::map<AnfNodePtr, size_t> *const_map_ptr, irpb::GraphProto *graph_proto);
  54. void ExportFuncGraphOutput(const FuncGraphPtr &func_graph, const CNodePtr &ret_node,
  55. const std::map<AnfNodePtr, size_t> &apply_map, std::map<AnfNodePtr, size_t> *const_map_ptr,
  56. irpb::GraphProto *graph_proto);
  57. void ExportValueNodes(const std::map<AnfNodePtr, size_t> &const_map, irpb::GraphProto *graph_proto);
  58. static std::string GetConstNodeId(size_t idx) { return std::string("cst") + std::to_string(idx); }
  59. irpb::ModelProto model_;
  60. };
  61. static irpb::DataType GetNumberDataType(const TypePtr &type) {
  62. switch (type->type_id()) {
  63. case kNumberTypeBool:
  64. return irpb::DT_BOOL;
  65. case kNumberTypeInt8:
  66. return irpb::DT_INT8;
  67. case kNumberTypeInt16:
  68. return irpb::DT_INT16;
  69. case kNumberTypeInt32:
  70. return irpb::DT_INT32;
  71. case kNumberTypeInt64:
  72. return irpb::DT_INT64;
  73. case kNumberTypeUInt8:
  74. return irpb::DT_UINT8;
  75. case kNumberTypeUInt16:
  76. return irpb::DT_UINT16;
  77. case kNumberTypeUInt32:
  78. return irpb::DT_UINT32;
  79. case kNumberTypeUInt64:
  80. return irpb::DT_UINT64;
  81. case kNumberTypeFloat16:
  82. return irpb::DT_FLOAT16;
  83. case kNumberTypeFloat32:
  84. return irpb::DT_FLOAT32;
  85. case kNumberTypeFloat64:
  86. return irpb::DT_FLOAT64;
  87. case kNumberTypeInt:
  88. return irpb::DT_BASE_INT;
  89. case kNumberTypeUInt:
  90. return irpb::DT_BASE_UINT;
  91. case kNumberTypeFloat:
  92. return irpb::DT_BASE_FLOAT;
  93. case kNumberTypeComplex64:
  94. return irpb::DT_COMPLEX64;
  95. case kNumberTypeComplex128:
  96. return irpb::DT_COMPLEX128;
  97. default:
  98. MS_LOG(EXCEPTION) << "Unexpected type " << type->type_name();
  99. }
  100. }
  101. void CheckIfValidType(const TypePtr &type) {
  102. MS_EXCEPTION_IF_NULL(type);
  103. if (type->isa<Problem>()) {
  104. MS_LOG(WARNING) << "The type: " << type->type_name();
  105. }
  106. if (!(type->isa<Number>() || type->isa<TensorType>() || type->isa<Tuple>() || type->isa<TypeType>() ||
  107. type->isa<List>() || type->isa<TypeAnything>() || type->isa<RefKeyType>() || type->isa<RefType>() ||
  108. type->isa<Function>() || type->isa<TypeNone>() || type->isa<Problem>() || type->isa<String>() ||
  109. type->isa<RowTensorType>() || type->isa<UndeterminedType>() || type->isa<SparseTensorType>() ||
  110. type->isa<SymbolicKeyType>() || type->isa<MonadType>())) {
  111. MS_LOG(EXCEPTION) << "Unknown type: " << type->type_name();
  112. }
  113. }
  114. void ProtoExporter::SetNodeOutputType(const TypePtr &type, const BaseShapePtr &shape, irpb::TypeProto *type_proto) {
  115. if (type_proto == nullptr) {
  116. return;
  117. }
  118. if (type != nullptr) {
  119. CheckIfValidType(type);
  120. }
  121. if (type == nullptr) {
  122. type_proto->set_data_type(irpb::DT_UNDEFINED);
  123. } else if (type->isa<Number>()) {
  124. type_proto->set_data_type(GetNumberDataType(type));
  125. } else if (type->isa<TensorType>()) {
  126. TypePtr elem_type = dyn_cast<TensorType>(type)->element();
  127. type_proto->mutable_tensor_type()->set_elem_type(GetNumberDataType(elem_type));
  128. type_proto->set_data_type(irpb::DT_TENSOR);
  129. if (shape != nullptr && shape->isa<abstract::Shape>()) {
  130. abstract::ShapePtr shape_info = dyn_cast<abstract::Shape>(shape);
  131. for (const auto &elem : shape_info->shape()) {
  132. type_proto->mutable_tensor_type()->mutable_shape()->add_dim()->set_size(elem);
  133. }
  134. }
  135. } else if (type->isa<Tuple>()) {
  136. TuplePtr tuple_type = dyn_cast<Tuple>(type);
  137. type_proto->set_data_type(irpb::DT_TUPLE);
  138. for (const auto &elem_type : tuple_type->elements()) {
  139. SetNodeOutputType(elem_type, nullptr, type_proto->mutable_sequence_type()->add_elem_types());
  140. }
  141. } else if (type->isa<TypeType>()) {
  142. type_proto->set_data_type(irpb::DT_TYPE);
  143. } else if (type->isa<List>()) {
  144. ListPtr list_type = dyn_cast<List>(type);
  145. type_proto->set_data_type(irpb::DT_LIST);
  146. for (const auto &elem_type : list_type->elements()) {
  147. SetNodeOutputType(elem_type, nullptr, type_proto->mutable_sequence_type()->add_elem_types());
  148. }
  149. } else if (type->isa<TypeAnything>()) {
  150. type_proto->set_data_type(irpb::DT_ANYTHING);
  151. } else if (type->isa<RefKeyType>()) {
  152. type_proto->set_data_type(irpb::DT_REFKEY);
  153. } else if (type->isa<RefType>()) {
  154. type_proto->set_data_type(irpb::DT_REF);
  155. } else if (type->isa<Function>()) {
  156. type_proto->set_data_type(irpb::DT_GRAPH);
  157. } else if (type->isa<TypeNone>()) {
  158. type_proto->set_data_type(irpb::DT_NONE);
  159. } else if (type->isa<String>()) {
  160. type_proto->set_data_type(irpb::DT_STRING);
  161. }
  162. }
  163. void ProtoExporter::SetNodeOutputType(const AnfNodePtr &node, irpb::TypeProto *type_proto) {
  164. if (node == nullptr || type_proto == nullptr) {
  165. return;
  166. }
  167. SetNodeOutputType(node->Type(), node->Shape(), type_proto);
  168. }
  169. void ProtoExporter::SetValueToProto(const ValuePtr &val, irpb::ValueProto *value_proto) {
  170. if (val == nullptr || value_proto == nullptr) {
  171. return;
  172. }
  173. if (val->isa<StringImm>()) {
  174. const StringImmPtr &value = dyn_cast<StringImm>(val);
  175. value_proto->set_dtype(irpb::DT_STRING);
  176. value_proto->set_str_val(value->value());
  177. } else if (val->isa<Scalar>()) {
  178. SetScalarToProto(dyn_cast<Scalar>(val), value_proto);
  179. } else if (val->isa<Bool>()) {
  180. value_proto->set_dtype(irpb::DT_TYPE);
  181. value_proto->mutable_type_val()->set_data_type(irpb::DT_BOOL);
  182. } else if (val->isa<Int>()) {
  183. value_proto->set_dtype(irpb::DT_TYPE);
  184. value_proto->mutable_type_val()->set_data_type(irpb::DT_BASE_INT);
  185. } else if (val->isa<UInt>()) {
  186. value_proto->set_dtype(irpb::DT_TYPE);
  187. value_proto->mutable_type_val()->set_data_type(irpb::DT_BASE_UINT);
  188. } else if (val->isa<Float>()) {
  189. value_proto->set_dtype(irpb::DT_TYPE);
  190. value_proto->mutable_type_val()->set_data_type(irpb::DT_BASE_FLOAT);
  191. } else if (val->isa<ValueSequeue>()) {
  192. SetSequenceToProto(dyn_cast<ValueSequeue>(val), value_proto);
  193. } else if (val->isa<None>()) {
  194. value_proto->set_dtype(irpb::DT_NONE);
  195. value_proto->set_str_val("None");
  196. } else if (val->isa<SymbolicKeyInstance>()) {
  197. SymbolicKeyInstancePtr sym_inst = dyn_cast<SymbolicKeyInstance>(val);
  198. ParameterPtr sym_node = dyn_cast<Parameter>(sym_inst->node());
  199. value_proto->set_dtype(irpb::DT_SYM_INST);
  200. value_proto->set_str_val(sym_node == nullptr ? std::string("nullptr") : sym_node->ToString());
  201. } else if (val->isa<ValueDictionary>()) {
  202. SetDictionaryToProto(dyn_cast<ValueDictionary>(val), value_proto);
  203. } else if (val->isa<tensor::Tensor>()) {
  204. tensor::TensorPtr tensor_ptr = dyn_cast<tensor::Tensor>(val);
  205. value_proto->set_dtype(irpb::DT_TENSOR);
  206. irpb::TensorProto *tensor_proto = value_proto->mutable_tensor_val();
  207. tensor_proto->set_data_type(GetNumberDataType(tensor_ptr->Dtype()));
  208. for (auto &elem : tensor_ptr->shape()) {
  209. tensor_proto->add_dims(elem);
  210. }
  211. } else if (val->isa<TensorType>()) {
  212. value_proto->set_dtype(irpb::DT_TYPE);
  213. irpb::TypeProto *type_proto = value_proto->mutable_type_val();
  214. type_proto->set_data_type(irpb::DT_TENSOR);
  215. TypePtr elem_type = dyn_cast<TensorType>(val)->element();
  216. type_proto->mutable_tensor_type()->set_elem_type(GetNumberDataType(elem_type));
  217. } else if (val->isa<Monad>() || val->isa<MonadType>()) {
  218. value_proto->set_str_val(val->ToString());
  219. } else if (val->isa<Complex>()) {
  220. value_proto->set_dtype(irpb::DT_TYPE);
  221. value_proto->mutable_type_val()->set_data_type(irpb::DT_BASE_COMPLEX);
  222. } else {
  223. MS_LOG(DEBUG) << "Unsupported type " << val->type_name();
  224. }
  225. }
  226. void ProtoExporter::SetScalarToProto(const ScalarPtr &val, irpb::ValueProto *value_proto) {
  227. if (val == nullptr || value_proto == nullptr) {
  228. return;
  229. }
  230. if (val->isa<BoolImm>()) {
  231. const BoolImmPtr &value = dyn_cast<BoolImm>(val);
  232. value_proto->set_dtype(irpb::DT_BOOL);
  233. value_proto->set_bool_val(value->value());
  234. } else if (val->isa<Int8Imm>()) {
  235. const Int8ImmPtr &value = dyn_cast<Int8Imm>(val);
  236. value_proto->set_dtype(irpb::DT_INT8);
  237. value_proto->set_int_val(value->value());
  238. } else if (val->isa<Int16Imm>()) {
  239. const Int16ImmPtr &value = dyn_cast<Int16Imm>(val);
  240. value_proto->set_dtype(irpb::DT_INT16);
  241. value_proto->set_int_val(value->value());
  242. } else if (val->isa<Int32Imm>()) {
  243. const Int32ImmPtr &value = dyn_cast<Int32Imm>(val);
  244. value_proto->set_dtype(irpb::DT_INT32);
  245. value_proto->set_int_val(value->value());
  246. } else if (val->isa<Int64Imm>()) {
  247. const Int64ImmPtr &value = dyn_cast<Int64Imm>(val);
  248. value_proto->set_dtype(irpb::DT_INT64);
  249. value_proto->set_int_val(value->value());
  250. } else if (val->isa<UInt8Imm>()) {
  251. const UInt8ImmPtr &value = dyn_cast<UInt8Imm>(val);
  252. value_proto->set_dtype(irpb::DT_UINT8);
  253. value_proto->set_uint_val(value->value());
  254. } else if (val->isa<UInt16Imm>()) {
  255. const UInt16ImmPtr &value = dyn_cast<UInt16Imm>(val);
  256. value_proto->set_dtype(irpb::DT_UINT16);
  257. value_proto->set_uint_val(value->value());
  258. } else if (val->isa<UInt32Imm>()) {
  259. const UInt32ImmPtr &value = dyn_cast<UInt32Imm>(val);
  260. value_proto->set_dtype(irpb::DT_UINT32);
  261. value_proto->set_uint_val(value->value());
  262. } else if (val->isa<UInt64Imm>()) {
  263. const UInt64ImmPtr &value = dyn_cast<UInt64Imm>(val);
  264. value_proto->set_dtype(irpb::DT_UINT64);
  265. value_proto->set_uint_val(value->value());
  266. } else if (val->isa<FP32Imm>()) {
  267. const FP32ImmPtr &value = dyn_cast<FP32Imm>(val);
  268. value_proto->set_dtype(irpb::DT_FLOAT32);
  269. value_proto->set_float_val(value->value());
  270. } else if (val->isa<FP64Imm>()) {
  271. const FP64ImmPtr &value = dyn_cast<FP64Imm>(val);
  272. value_proto->set_dtype(irpb::DT_FLOAT64);
  273. value_proto->set_double_val(value->value());
  274. } else {
  275. MS_LOG(EXCEPTION) << "Unknown scalar type " << val->ToString();
  276. }
  277. }
  278. void ProtoExporter::SetSequenceToProto(const ValueSequeuePtr &val, irpb::ValueProto *value_proto) {
  279. if (val == nullptr || value_proto == nullptr) {
  280. return;
  281. }
  282. if (val->isa<ValueTuple>()) {
  283. const ValueTuplePtr &value = dyn_cast<ValueTuple>(val);
  284. value_proto->set_dtype(irpb::DT_TUPLE);
  285. for (const auto &item : value->value()) {
  286. SetValueToProto(item, value_proto->add_values());
  287. }
  288. } else if (val->isa<ValueList>()) {
  289. const ValueListPtr &value = dyn_cast<ValueList>(val);
  290. value_proto->set_dtype(irpb::DT_LIST);
  291. for (const auto &item : value->value()) {
  292. SetValueToProto(item, value_proto->add_values());
  293. }
  294. }
  295. }
  296. void ProtoExporter::SetDictionaryToProto(const ValueDictionaryPtr &val, irpb::ValueProto *value_proto) {
  297. if (val == nullptr || value_proto == nullptr) {
  298. return;
  299. }
  300. value_proto->set_dtype(irpb::DT_DICT);
  301. for (const auto &item : val->value()) {
  302. irpb::NamedValueProto *named_val = value_proto->add_dict_val();
  303. named_val->set_key(item.first);
  304. SetValueToProto(item.second, named_val->mutable_value());
  305. }
  306. }
  307. void ProtoExporter::GetOpNodeTypeAndAttrs(const FuncGraphPtr &, const AnfNodePtr &node, irpb::NodeProto *node_proto) {
  308. if (node == nullptr || node_proto == nullptr) {
  309. return;
  310. }
  311. if (node->isa<CNode>() || node->isa<Parameter>() || IsValueNode<FuncGraph>(node)) {
  312. MS_LOG(EXCEPTION) << "Op node can not be CNode, Parameter or ValueNode Graph. But got " << node->ToString();
  313. }
  314. if (!IsValueNode<Primitive>(node)) {
  315. MS_LOG(EXCEPTION) << "Op node is not primitive: " << node->ToString();
  316. }
  317. const PrimitivePtr &prim = GetValueNode<PrimitivePtr>(node);
  318. // set node parallel info
  319. auto operator_info = node->user_data<parallel::OperatorInfo>();
  320. if (operator_info != nullptr) {
  321. auto strategy = operator_info->strategy();
  322. if (strategy != nullptr) {
  323. ValuePtr strategy_value = MakeValue(strategy->GetInputDim());
  324. // display the strategy generated by batch parallel
  325. auto attrs = prim->attrs();
  326. attrs[mindspore::parallel::IN_STRATEGY] = strategy_value;
  327. (void)prim->SetAttrs(attrs);
  328. }
  329. }
  330. node_proto->set_op_type(prim->name());
  331. for (const auto &attr : prim->attrs()) {
  332. irpb::AttributeProto *attr_proto = node_proto->add_attribute();
  333. attr_proto->set_name(attr.first);
  334. SetValueToProto(attr.second, attr_proto->mutable_value());
  335. }
  336. node_proto->set_scope(node->scope()->name());
  337. }
  338. std::string ProtoExporter::GetOpNodeInputId(const FuncGraphPtr &, const AnfNodePtr &node,
  339. const std::map<AnfNodePtr, size_t> &apply_map,
  340. std::map<AnfNodePtr, size_t> *const_map_ptr) {
  341. if (node == nullptr || const_map_ptr == nullptr) {
  342. return "";
  343. }
  344. if (node->isa<CNode>()) {
  345. auto iter = apply_map.find(node);
  346. if (iter == apply_map.end()) {
  347. MS_LOG(EXCEPTION) << "Can not find node '" << node->ToString() << "' in apply_map";
  348. }
  349. return std::to_string(iter->second);
  350. }
  351. if (node->isa<Parameter>()) {
  352. return node->ToString();
  353. }
  354. if (node->isa<ValueNode>()) {
  355. auto iter = const_map_ptr->find(node);
  356. if (iter == const_map_ptr->end()) {
  357. // Start index number from 1
  358. auto const_idx = const_map_ptr->size() + 1;
  359. (*const_map_ptr)[node] = const_idx;
  360. }
  361. return GetConstNodeId((*const_map_ptr)[node]);
  362. }
  363. MS_LOG(EXCEPTION) << "Unknown node type. node is '" << node->ToString() << "'";
  364. }
  365. std::string ProtoExporter::GetFuncGraphProtoString(const FuncGraphPtr &func_graph) {
  366. if (func_graph == nullptr) {
  367. return "";
  368. }
  369. InitModelInfo();
  370. irpb::GraphProto *graph_proto = model_.mutable_graph();
  371. ExportFuncGraph(func_graph, graph_proto);
  372. return model_.SerializeAsString();
  373. }
  374. void ProtoExporter::ExportFuncGraph(const FuncGraphPtr &func_graph, irpb::GraphProto *graph_proto) {
  375. if (func_graph == nullptr || graph_proto == nullptr) {
  376. return;
  377. }
  378. // map for store ValueNodes of this graph
  379. std::map<AnfNodePtr, size_t> const_map;
  380. // set graph name
  381. graph_proto->set_name(func_graph->ToString());
  382. ExportParameters(func_graph, graph_proto);
  383. ExportCNodes(func_graph, graph_proto, &const_map);
  384. ExportValueNodes(const_map, graph_proto);
  385. }
  386. void ProtoExporter::ExportParameters(const FuncGraphPtr &func_graph, irpb::GraphProto *graph_proto) {
  387. if (func_graph == nullptr || graph_proto == nullptr) {
  388. return;
  389. }
  390. std::vector<AnfNodePtr> parameters = func_graph->parameters();
  391. for (auto &param : parameters) {
  392. irpb::ParameterProto *param_proto = graph_proto->add_parameters();
  393. param_proto->set_name(param->ToString());
  394. SetNodeOutputType(param, param_proto->mutable_type());
  395. const ParameterPtr param_ptr = dyn_cast<Parameter>(param);
  396. if (param_ptr == nullptr) {
  397. MS_LOG(EXCEPTION) << "Parameter '" << param->ToString() << "' could not cast to parameter.";
  398. }
  399. }
  400. }
  401. void ProtoExporter::ExportCNodes(const FuncGraphPtr &func_graph, irpb::GraphProto *graph_proto,
  402. std::map<AnfNodePtr, size_t> *const_map_ptr) {
  403. if (func_graph == nullptr || graph_proto == nullptr || const_map_ptr == nullptr) {
  404. return;
  405. }
  406. // topo sort nodes
  407. std::vector<AnfNodePtr> nodes = TopoSort(func_graph->get_return(), SuccIncoming, AlwaysInclude);
  408. std::map<AnfNodePtr, size_t> apply_map;
  409. for (const AnfNodePtr &node : nodes) {
  410. MS_EXCEPTION_IF_NULL(node);
  411. if (!node->isa<CNode>()) {
  412. continue;
  413. }
  414. auto cnode = node->cast<CNodePtr>();
  415. if (cnode != func_graph->get_return()) {
  416. ExportCNode(func_graph, cnode, &apply_map, const_map_ptr, graph_proto);
  417. } else {
  418. ExportFuncGraphOutput(func_graph, cnode, apply_map, const_map_ptr, graph_proto);
  419. }
  420. }
  421. }
  422. void ProtoExporter::ExportCNode(const FuncGraphPtr &func_graph, const CNodePtr &node,
  423. std::map<AnfNodePtr, size_t> *apply_map_ptr,
  424. std::map<AnfNodePtr, size_t> *const_map_ptr, irpb::GraphProto *graph_proto) {
  425. if (func_graph == nullptr || node == nullptr || apply_map_ptr == nullptr || const_map_ptr == nullptr ||
  426. graph_proto == nullptr) {
  427. return;
  428. }
  429. auto apply_idx = apply_map_ptr->size() + 1;
  430. (*apply_map_ptr)[node] = apply_idx;
  431. auto &inputs = node->inputs();
  432. if (inputs.size() < 1) {
  433. MS_LOG(EXCEPTION) << "Inputs of apply node is empty";
  434. }
  435. AnfNodePtr op = inputs[0];
  436. irpb::NodeProto *node_proto = graph_proto->add_node();
  437. // CNode/ConstGraph/Const/Parameter
  438. if (op->isa<CNode>() || IsValueNode<FuncGraph>(op) || op->isa<Parameter>()) {
  439. MS_LOG(DEBUG) << "Operator must be a primitive";
  440. } else {
  441. GetOpNodeTypeAndAttrs(func_graph, op, node_proto);
  442. node_proto->set_name(std::to_string(apply_idx));
  443. node_proto->set_scope(node->scope()->name());
  444. node_proto->set_full_name(GetKernelNodeName(node));
  445. // process OP inputs
  446. for (size_t i = 1; i < inputs.size(); ++i) {
  447. irpb::InputProto *input_proto = node_proto->add_input();
  448. input_proto->set_type(irpb::InputProto_EdgeType_DATA_EDGE);
  449. std::string id = GetOpNodeInputId(func_graph, inputs[i], *apply_map_ptr, const_map_ptr);
  450. input_proto->set_name(id);
  451. }
  452. // set node output type
  453. SetNodeOutputType(node, node_proto->mutable_output_type());
  454. if (IsValueNode<Primitive>(op)) {
  455. PrimitivePtr primitive = GetValueNode<PrimitivePtr>(op);
  456. if (!primitive->instance_name().empty()) {
  457. node_proto->set_instance_name(primitive->instance_name());
  458. }
  459. }
  460. }
  461. }
  462. void ProtoExporter::ExportFuncGraphOutput(const FuncGraphPtr &func_graph, const CNodePtr &ret_node,
  463. const std::map<AnfNodePtr, size_t> &apply_map,
  464. std::map<AnfNodePtr, size_t> *const_map_ptr, irpb::GraphProto *graph_proto) {
  465. if (ret_node == nullptr || !ret_node->isa<CNode>()) {
  466. MS_LOG(EXCEPTION) << "Graph return node is illegal";
  467. }
  468. // ret node has two input 1 ret op + 1 value
  469. const size_t ret_input_size = 2;
  470. if (ret_node->inputs().size() != ret_input_size) {
  471. return;
  472. }
  473. AnfNodePtr arg = ret_node->input(1);
  474. if (graph_proto == nullptr) {
  475. MS_LOG(EXCEPTION) << "graph_proto is nullptr";
  476. }
  477. irpb::OutputProto *output_proto = graph_proto->add_outputs();
  478. if (output_proto == nullptr) {
  479. MS_LOG(EXCEPTION) << "output_proto is nullptr";
  480. }
  481. std::string id = GetOpNodeInputId(func_graph, arg, apply_map, const_map_ptr);
  482. output_proto->set_name(id);
  483. SetNodeOutputType(arg, output_proto->mutable_type());
  484. }
  485. static bool CompareValue(const std::pair<AnfNodePtr, size_t> &x, const std::pair<AnfNodePtr, size_t> &y) {
  486. return x.second < y.second;
  487. }
  488. void ProtoExporter::ExportValueNodes(const std::map<AnfNodePtr, size_t> &const_map, irpb::GraphProto *graph_proto) {
  489. std::vector<std::pair<AnfNodePtr, size_t>> nodes;
  490. (void)std::transform(const_map.cbegin(), const_map.cend(), std::back_inserter(nodes),
  491. [](const std::pair<AnfNodePtr, size_t> &item) { return item; });
  492. sort(nodes.begin(), nodes.end(), CompareValue);
  493. for (auto &item : nodes) {
  494. if (graph_proto == nullptr) {
  495. MS_LOG(EXCEPTION) << "graph_proto is nullptr";
  496. }
  497. irpb::NamedValueProto *named_value = graph_proto->add_const_vals();
  498. MS_EXCEPTION_IF_NULL(named_value);
  499. named_value->set_key(GetConstNodeId(item.second));
  500. SetValueToProto(GetValueNode(item.first), named_value->mutable_value());
  501. }
  502. }
  503. void ProtoExporter::InitModelInfo() { model_.set_ir_version(irpb::IR_VERSION); }
  504. std::string GetFuncGraphProtoString(const FuncGraphPtr &func_graph) {
  505. ProtoExporter exporter;
  506. return exporter.GetFuncGraphProtoString(func_graph);
  507. }
  508. void GetFuncGraphProto(const FuncGraphPtr &func_graph, irpb::GraphProto *graph_proto) {
  509. ProtoExporter exporter;
  510. exporter.ExportFuncGraph(func_graph, graph_proto);
  511. }
  512. #ifdef ENABLE_DUMP_IR
  513. void DumpIRProto(const FuncGraphPtr &func_graph, const std::string &suffix) {
  514. if (func_graph == nullptr) {
  515. MS_LOG(ERROR) << "Func graph is nullptr";
  516. return;
  517. }
  518. std::string file_path = GetSaveGraphsPathName("ms_output_" + suffix + ".pb");
  519. auto realpath = Common::CreatePrefixPath(file_path);
  520. if (!realpath.has_value()) {
  521. MS_LOG(ERROR) << "Get real path failed, path=" << file_path;
  522. return;
  523. }
  524. ChangeFileMode(realpath.value(), S_IWUSR);
  525. // write to pb file
  526. std::ofstream ofs(file_path);
  527. if (!ofs.is_open()) {
  528. MS_LOG(ERROR) << "Open file '" << file_path << "' failed!" << ErrnoToString(errno);
  529. return;
  530. }
  531. ofs << GetFuncGraphProtoString(func_graph);
  532. ofs.close();
  533. // set file mode to read only by user
  534. ChangeFileMode(file_path, S_IRUSR);
  535. }
  536. #else
  537. void DumpIRProto(const FuncGraphPtr &, const std::string &) {
  538. static bool already_printed = false;
  539. if (already_printed) {
  540. return;
  541. }
  542. already_printed = true;
  543. MS_LOG(WARNING) << "The functionality of dumping function graph IR in protobuf format is disabled, "
  544. << "please recompile source to enable it. See help of building script.";
  545. }
  546. #endif
  547. } // namespace mindspore