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.

flow_ctrl_pass.cc 21 kB

5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
5 years ago
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486
  1. /**
  2. * Copyright 2019-2020 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 "graph/passes/flow_ctrl_pass.h"
  17. #include <memory>
  18. #include <string>
  19. #include <vector>
  20. #include "framework/common/debug/ge_log.h"
  21. #include "graph/debug/ge_attr_define.h"
  22. #include "graph/common/omg_util.h"
  23. #include "common/ge/ge_util.h"
  24. #include "graph/manager/graph_var_manager.h"
  25. #include "graph/passes/pass_utils.h"
  26. namespace ge {
  27. // when namespace change to ge, please delete the using code.
  28. Status FlowCtrlPass::Run(ComputeGraphPtr compute_graph) {
  29. GE_CHECK_NOTNULL(compute_graph);
  30. if (!PassUtils::IsNeedTrainIteFlowCtrl(compute_graph)) {
  31. GELOGI("No need FlowCtrl for graph %u", compute_graph->GetGraphID());
  32. return NOT_CHANGED;
  33. }
  34. GELOGI("FlowCtrl pass begin");
  35. bool graph_change = false;
  36. // 1. Add FP/BP flow ctrl (big cycle)
  37. for (auto &node : compute_graph->GetDirectNode()) {
  38. if (node == nullptr) {
  39. continue;
  40. }
  41. GE_IF_BOOL_EXEC(node->GetOpDesc() == nullptr, continue);
  42. uint32_t true_stream_id = 0;
  43. bool is_found = AttrUtils::GetInt(node->GetOpDesc(), ATTR_NAME_TRUE_BRANCH_STREAM, true_stream_id);
  44. // FP/BP cycle flag is true_stream_id == 0
  45. if (is_found && (true_stream_id == TRUE_STREAM_ID)) {
  46. // Add big cycle
  47. Status ret = AddFpBpIteratorCtrl(compute_graph, node);
  48. if (ret != SUCCESS) {
  49. GELOGE(ret, "AddFpBpIteratorCtrl fail, node: %s.", node->GetName().c_str());
  50. return ret;
  51. }
  52. graph_change = true;
  53. // only one big cycle, so break.
  54. break;
  55. }
  56. }
  57. // 2. Add special node flow ctrl. eg, IteratorGetNext. (small cycle)
  58. // NOTE: Small cycle share the variables with big cycle.
  59. for (auto &node : compute_graph->GetDirectNode()) {
  60. if (node == nullptr) {
  61. continue;
  62. }
  63. GE_IF_BOOL_EXEC(node->GetOpDesc() == nullptr, continue);
  64. bool need_cycle_flag = false;
  65. bool is_found = AttrUtils::GetBool(node->GetOpDesc(), ATTR_NAME_STREAM_CYCLE_EVENT_FLAG, need_cycle_flag);
  66. // small cycle flag is need_stream_cycle_event == true
  67. if (is_found && need_cycle_flag) {
  68. Status ret = AddSpecialNodeIteratorCtrl(compute_graph, node);
  69. if (ret != SUCCESS) {
  70. GELOGE(ret, "AddSpecialNodeIteratorCtrl fail, node: %s.", node->GetName().c_str());
  71. return ret;
  72. }
  73. graph_change = true;
  74. }
  75. }
  76. GELOGI("FlowCtrl pass end, graph is %s.", graph_change ? "changed" : "not changed");
  77. return graph_change ? SUCCESS : NOT_CHANGED;
  78. }
  79. NodePtr FlowCtrlPass::InsertOp(ComputeGraphPtr &compute_graph, const string &node_type, const string &node_name,
  80. const std::vector<GeTensorDesc> &input_list,
  81. const std::vector<GeTensorDesc> &output_list) {
  82. OpDescPtr op_desc = MakeShared<OpDesc>(node_name, node_type);
  83. if (op_desc == nullptr) {
  84. GELOGE(FAILED, "Make OpDesc failed, name:%s, type:%s.", node_name.c_str(), node_type.c_str());
  85. return nullptr;
  86. }
  87. for (auto &input_desc : input_list) {
  88. graphStatus graph_status = op_desc->AddInputDesc(input_desc);
  89. if (graph_status != GRAPH_SUCCESS) {
  90. GELOGE(FAILED, "Add node:%s intput desc failed, error=%u.", node_name.c_str(), graph_status);
  91. return nullptr;
  92. }
  93. }
  94. for (auto &output_desc : output_list) {
  95. graphStatus graph_status = op_desc->AddOutputDesc(output_desc);
  96. if (graph_status != GRAPH_SUCCESS) {
  97. GELOGE(FAILED, "Add node:%s output desc failed, error=%u.", node_name.c_str(), graph_status);
  98. return nullptr;
  99. }
  100. }
  101. GE_IF_BOOL_EXEC(compute_graph == nullptr, DOMI_LOGE("compute_graph is nullptr"); return nullptr);
  102. NodePtr node = compute_graph->AddNode(op_desc);
  103. if (node == nullptr) {
  104. GELOGE(FAILED, "add node failed, name:%s, type:%s.", node_name.c_str(), node_type.c_str());
  105. return nullptr;
  106. }
  107. GELOGI("Insert op success, name:%s, type:%s.", node_name.c_str(), node_type.c_str());
  108. return node;
  109. }
  110. NodePtr FlowCtrlPass::InsertStreamSwitchOp(ComputeGraphPtr &compute_graph, const string &switch_name,
  111. const NodePtr &loop_cond, const NodePtr &iter_per_loop) {
  112. GE_IF_BOOL_EXEC(loop_cond == nullptr || loop_cond->GetOpDesc() == nullptr,
  113. GELOGE(FAILED, "loop_cond is null"); return nullptr);
  114. GE_IF_BOOL_EXEC(iter_per_loop == nullptr || iter_per_loop->GetOpDesc() == nullptr,
  115. GELOGE(FAILED, "iter_per_loop is nullptr"); return nullptr);
  116. std::vector<GeTensorDesc> input_desc_list = {loop_cond->GetOpDesc()->GetOutputDesc(0),
  117. iter_per_loop->GetOpDesc()->GetOutputDesc(0)};
  118. std::vector<GeTensorDesc> output_desc_list;
  119. NodePtr stream_switch = InsertOp(compute_graph, STREAMSWITCH, switch_name, input_desc_list, output_desc_list);
  120. if (stream_switch == nullptr) {
  121. GELOGE(FAILED, "InsertStreamSwitchOp failed, name:%s.", switch_name.c_str());
  122. return nullptr;
  123. }
  124. // set input 0
  125. graphStatus add_ret = GraphUtils::AddEdge(loop_cond->GetOutDataAnchor(0), stream_switch->GetInDataAnchor(0));
  126. if (add_ret != GRAPH_SUCCESS) {
  127. GELOGE(FAILED, "Add loop_cond_node to switch_node:%s edge failed, ret = %u.", switch_name.c_str(), add_ret);
  128. return nullptr;
  129. }
  130. // set input 1
  131. add_ret = GraphUtils::AddEdge(iter_per_loop->GetOutDataAnchor(0), stream_switch->GetInDataAnchor(1));
  132. if (add_ret != GRAPH_SUCCESS) {
  133. GELOGE(FAILED, "Add iter_per_loop_node to switch_node:%s edge failed, ret = %u.", switch_name.c_str(), add_ret);
  134. return nullptr;
  135. }
  136. // stream switch op need switch cond by attr.
  137. GE_IF_BOOL_EXEC(!AttrUtils::SetInt(stream_switch->GetOpDesc(), ATTR_NAME_STREAM_SWITCH_COND,
  138. static_cast<int64_t>(RT_LESS)),
  139. DOMI_LOGE("set ATTR_NAME_STREAM_SWITCH_COND failed"); return nullptr);
  140. return stream_switch;
  141. }
  142. NodePtr FlowCtrlPass::AddVariableNode(ComputeGraphPtr &compute_graph, const string &name) {
  143. GE_IF_BOOL_EXEC(compute_graph == nullptr, DOMI_LOGE("compute_graph is nullptr"); return nullptr);
  144. NodePtr exist_node = compute_graph->FindNode(name);
  145. if (exist_node != nullptr) {
  146. GELOGD("Node %s already exist, no need add.", name.c_str());
  147. return exist_node;
  148. }
  149. // fetch and set tensor desc
  150. GeTensorDesc tensor_desc;
  151. if (ge::VarManager::Instance(compute_graph->GetSessionID()) == nullptr) {
  152. return nullptr;
  153. }
  154. Status ret = ge::VarManager::Instance(compute_graph->GetSessionID())->GetCurVarDesc(name, tensor_desc);
  155. if (ret != SUCCESS) {
  156. GELOGE(FAILED, "Get var desc fail, name:%s", name.c_str());
  157. return nullptr;
  158. }
  159. std::vector<GeTensorDesc> input_desc_list;
  160. std::vector<GeTensorDesc> output_desc_list = {tensor_desc};
  161. // insert node
  162. return InsertOp(compute_graph, VARIABLE, name, input_desc_list, output_desc_list);
  163. }
  164. Status FlowCtrlPass::AddGlobalStepVariableNode(ComputeGraphPtr &compute_graph) {
  165. NodePtr output_node = compute_graph->FindFirstNodeMatchType(NETOUTPUT);
  166. if (output_node == nullptr) {
  167. GELOGD("Node type %s can't be found in graph %u", NETOUTPUT, compute_graph->GetGraphID());
  168. return SUCCESS;
  169. }
  170. // Global step just add to main graph's netoutput node.And the main graph must be known shape
  171. if ((compute_graph->GetParentGraph() != nullptr) ||
  172. ((compute_graph->GetParentGraph() == nullptr) && (GraphUtils::IsUnknownShapeGraph(compute_graph)))) {
  173. GELOGD("Subgraph %s no need global step variable.", compute_graph->GetName().c_str());
  174. return SUCCESS;
  175. }
  176. NodePtr exist_node = compute_graph->FindNode(NODE_NAME_GLOBAL_STEP);
  177. if (exist_node != nullptr) {
  178. GELOGD("Node %s already exist, no need add.", NODE_NAME_GLOBAL_STEP.c_str());
  179. return SUCCESS;
  180. }
  181. // set global step tensor desc
  182. GeTensorDesc tensor_desc(GeShape({1}), FORMAT_ND, DT_UINT64);
  183. std::vector<GeTensorDesc> input_desc_list = {};
  184. std::vector<GeTensorDesc> output_desc_list = {tensor_desc};
  185. NodePtr global_step = InsertOp(compute_graph, VARIABLE, NODE_NAME_GLOBAL_STEP,
  186. input_desc_list, output_desc_list);
  187. if (global_step == nullptr) {
  188. GELOGE(FAILED, "Add global_step node failed, global_step is null.");
  189. return FAILED;
  190. }
  191. // add ctrl edges
  192. graphStatus add_ret = GraphUtils::AddEdge(global_step->GetOutControlAnchor(), output_node->GetInControlAnchor());
  193. if (add_ret != GRAPH_SUCCESS) {
  194. GELOGE(FAILED, "Add global_step to netoutput edge failed, add_ret=%u.", add_ret);
  195. return FAILED;
  196. }
  197. GELOGD("Add global_step to netoutput edge in graph %u success", compute_graph->GetGraphID());
  198. return SUCCESS;
  199. }
  200. NodePtr FlowCtrlPass::InsertAssignOp(ge::ComputeGraphPtr &compute_graph, const string &node_type,
  201. const string &node_name, const NodePtr &ref_node, const NodePtr &value_node) {
  202. GE_IF_BOOL_EXEC(ref_node == nullptr || value_node == nullptr ||
  203. ref_node->GetOpDesc() == nullptr || value_node->GetOpDesc() == nullptr,
  204. GELOGE(FAILED, "ref node or value node is null");
  205. return nullptr);
  206. GeTensorDesc ref_tensor_desc = ref_node->GetOpDesc()->GetOutputDesc(0);
  207. GeTensorDesc val_tensor_desc = value_node->GetOpDesc()->GetOutputDesc(0);
  208. std::vector<GeTensorDesc> input_desc_list = {ref_tensor_desc, val_tensor_desc};
  209. std::vector<GeTensorDesc> output_desc_list = {ref_tensor_desc};
  210. NodePtr assign_node = InsertOp(compute_graph, node_type, node_name, input_desc_list, output_desc_list);
  211. if (assign_node == nullptr) {
  212. GELOGE(FAILED, "Insert node %s(%s) failed.", node_name.c_str(), node_type.c_str());
  213. return nullptr;
  214. }
  215. // assign node input 0 = ref_node
  216. graphStatus add_ret = GraphUtils::AddEdge(ref_node->GetOutDataAnchor(0), assign_node->GetInDataAnchor(0));
  217. if (add_ret != GRAPH_SUCCESS) {
  218. GELOGE(FAILED, "Add ref_node to %s edge failed, add_ret=%u.", node_name.c_str(), add_ret);
  219. return nullptr;
  220. }
  221. // assign input 1 = value_node
  222. add_ret = GraphUtils::AddEdge(value_node->GetOutDataAnchor(0), assign_node->GetInDataAnchor(1));
  223. if (add_ret != GRAPH_SUCCESS) {
  224. GELOGE(FAILED, "Add value_node to %s edge failed, add_ret=%u.", node_name.c_str(), add_ret);
  225. return nullptr;
  226. }
  227. (void)ge::AttrUtils::SetBool(assign_node->GetOpDesc(), ATTR_NEED_COMPILE, true);
  228. return assign_node;
  229. }
  230. Status FlowCtrlPass::CreateIterCtrlTrueBranch(ComputeGraphPtr &compute_graph, const NodePtr &loop_cond_node,
  231. const NodePtr &loop_inc_node, NodePtr &switch_node) {
  232. /*
  233. * loopCond
  234. * |
  235. * v
  236. * switch --> AssignAdd --> active
  237. * ^
  238. * |
  239. * loopIncrement
  240. */
  241. // Insert AssignAdd node
  242. NodePtr assign_add_node =
  243. InsertAssignOp(compute_graph, ASSIGNADD, NODE_NAME_FLOWCTRL_LOOP_ASSIGNADD, loop_cond_node, loop_inc_node);
  244. if (assign_add_node == nullptr || switch_node == nullptr) {
  245. GELOGE(PARAM_INVALID, "assign add node or switch node is null");
  246. return FAILED;
  247. }
  248. string active_name = switch_node->GetName() + "_StreamActive";
  249. // add attr for stream assign model to break branch.
  250. GE_CHK_STATUS_RET(SetStreamLabel(assign_add_node, active_name), "set stream label failed");
  251. // used for stream assign to find true branch
  252. GE_CHK_STATUS_RET(SetActiveLabelList(switch_node, { active_name }), "set active label list failed");
  253. // 2. Insert active node
  254. NodePtr active_node = InsertOp(compute_graph, STREAMACTIVE, active_name, {}, {});
  255. if (active_node == nullptr) {
  256. GELOGE(FAILED, "Insert stream active node:%s for IterCtrlTrueStream failed.", active_name.c_str());
  257. return FAILED;
  258. }
  259. GE_CHK_STATUS_RET(SetStreamLabel(active_node, active_name), "set stream label failed");
  260. GE_IF_BOOL_EXEC(!AttrUtils::SetBool(active_node->GetOpDesc(), ATTR_NAME_IS_LOOP_ACTIVE, true),
  261. DOMI_LOGE("set ATTR_NAME_IS_LOOP_ACTIVE failed"); return FAILED);
  262. // add ctrl edges
  263. graphStatus add_ret = GraphUtils::AddEdge(switch_node->GetOutControlAnchor(), assign_add_node->GetInControlAnchor());
  264. if (add_ret != GRAPH_SUCCESS) {
  265. GELOGE(FAILED, "Add switch_node to assign_add_node ctrl edge failed, add_ret=%u.", add_ret);
  266. return FAILED;
  267. }
  268. add_ret = GraphUtils::AddEdge(assign_add_node->GetOutControlAnchor(), active_node->GetInControlAnchor());
  269. if (add_ret != GRAPH_SUCCESS) {
  270. GELOGE(FAILED, "Add assign_add_node to active_node ctrl edge failed, add_ret=%u.", add_ret);
  271. return FAILED;
  272. }
  273. GELOGI("CreateIterCtrlTrueBranch success. StreamActive op:%s.", active_node->GetName().c_str());
  274. return SUCCESS;
  275. }
  276. Status FlowCtrlPass::CreateIterCtrlFalseBranch(ComputeGraphPtr &compute_graph, const NodePtr &loop_cond_node,
  277. const NodePtr &loop_reset_node, NodePtr &switch_node) {
  278. /*
  279. * loopCond
  280. * |
  281. * v
  282. * switch --> Assign
  283. * ^
  284. * |
  285. * loopReset
  286. */
  287. // Insert Assign node
  288. NodePtr assign_node =
  289. InsertAssignOp(compute_graph, ASSIGN, NODE_NAME_FLOWCTRL_LOOP_ASSIGN, loop_cond_node, loop_reset_node);
  290. if (assign_node == nullptr || switch_node == nullptr) {
  291. GELOGE(PARAM_INVALID, "assign_node or switch node is null");
  292. return FAILED;
  293. }
  294. GE_CHK_STATUS_RET(SetStreamLabel(assign_node, switch_node->GetName()), "set stream label failed");
  295. // 3. Insert ctrl edges
  296. graphStatus add_ret = GraphUtils::AddEdge(switch_node->GetOutControlAnchor(), assign_node->GetInControlAnchor());
  297. if (add_ret != GRAPH_SUCCESS) {
  298. GELOGE(FAILED, "Add switch_node to assign_node ctrl edge failed, add_ret=%u.", add_ret);
  299. return FAILED;
  300. }
  301. GELOGI("CreateIterCtrlFalseBranch success.");
  302. return SUCCESS;
  303. }
  304. Status FlowCtrlPass::AddFpBpIteratorCtrl(ComputeGraphPtr &compute_graph, NodePtr &pre_node) {
  305. GE_IF_BOOL_EXEC(pre_node == nullptr, DOMI_LOGE("pre_node is nullptr"); return FAILED);
  306. string pre_node_name = pre_node->GetName();
  307. GELOGI("Add FpBp Iterator ctrl, pre node:%s.", pre_node_name.c_str());
  308. // 1. Get or add variables
  309. NodePtr loop_cond_node = AddVariableNode(compute_graph, NODE_NAME_FLOWCTRL_LOOP_COND);
  310. if (loop_cond_node == nullptr) {
  311. GELOGE(FAILED, "Add variable:%s failed.", NODE_NAME_FLOWCTRL_LOOP_COND.c_str());
  312. return FAILED;
  313. }
  314. NodePtr loop_inc_node = AddVariableNode(compute_graph, NODE_NAME_FLOWCTRL_LOOP_INCREMENT);
  315. if (loop_inc_node == nullptr) {
  316. GELOGE(FAILED, "Add variable:%s failed.", NODE_NAME_FLOWCTRL_LOOP_INCREMENT.c_str());
  317. return FAILED;
  318. }
  319. NodePtr loop_reset_node = AddVariableNode(compute_graph, NODE_NAME_FLOWCTRL_LOOP_RESETVALUE);
  320. if (loop_reset_node == nullptr) {
  321. GELOGE(FAILED, "Add variable:%s failed.", NODE_NAME_FLOWCTRL_LOOP_RESETVALUE.c_str());
  322. return FAILED;
  323. }
  324. NodePtr iter_per_loop_node = AddVariableNode(compute_graph, NODE_NAME_FLOWCTRL_LOOP_PER_ITER);
  325. if (iter_per_loop_node == nullptr) {
  326. GELOGE(FAILED, "Add variable:%s failed.", NODE_NAME_FLOWCTRL_LOOP_PER_ITER.c_str());
  327. return FAILED;
  328. }
  329. // 2. Add StreamSwitch
  330. string switch_name = pre_node_name + "_" + NODE_NAME_STREAM_SWITCH;
  331. NodePtr switch_node = InsertStreamSwitchOp(compute_graph, switch_name, loop_cond_node, iter_per_loop_node);
  332. if (switch_node == nullptr) {
  333. GELOGE(FAILED, "InsertStreamSwitchOp:%s failed.", switch_name.c_str());
  334. return FAILED;
  335. }
  336. GE_CHK_STATUS_RET(SetStreamLabel(switch_node, switch_name), "set stream label failed");
  337. graphStatus add_ret = GraphUtils::AddEdge(pre_node->GetOutControlAnchor(), switch_node->GetInControlAnchor());
  338. if (add_ret != GRAPH_SUCCESS) {
  339. GELOGE(FAILED, "Add pre node:%s to switch_node:%s ctrl edge failed, ret = %u.", pre_node_name.c_str(),
  340. switch_name.c_str(), add_ret);
  341. return FAILED;
  342. }
  343. // 3. Create switch false branch: return results and reset the loopCond
  344. Status ret = CreateIterCtrlFalseBranch(compute_graph, loop_cond_node, loop_reset_node, switch_node);
  345. if (ret != SUCCESS) {
  346. GELOGE(ret, "CreateIterCtrlFalseBranch fail, pre node:%s.", pre_node_name.c_str());
  347. return ret;
  348. }
  349. // 4. Create switch true branch:
  350. // active train streams and increase the loopCond
  351. ret = CreateIterCtrlTrueBranch(compute_graph, loop_cond_node, loop_inc_node, switch_node);
  352. if (ret != SUCCESS) {
  353. GELOGE(ret, "CreateIterCtrlTrueBranch fail, pre node:%s.", pre_node_name.c_str());
  354. return ret;
  355. }
  356. return SUCCESS;
  357. }
  358. Status FlowCtrlPass::AddSpecialNodeIteratorCtrl(ComputeGraphPtr &compute_graph, NodePtr &loop_after_node) {
  359. /*
  360. * before add:
  361. * iterator
  362. * |
  363. * v
  364. * MemcpyAsync
  365. *
  366. * after add:
  367. * iterator ----------┐
  368. * | ┆c
  369. * v c v c
  370. * MemcpyAsync-----> switch -----> active
  371. * ^
  372. * / \
  373. * itersPerLoop loopCond
  374. */
  375. GE_IF_BOOL_EXEC(loop_after_node == nullptr || compute_graph == nullptr,
  376. DOMI_LOGE("loop after node or compute graph is null"); return FAILED);
  377. InDataAnchorPtr in_anchor = loop_after_node->GetInDataAnchor(0);
  378. if (in_anchor == nullptr || in_anchor->GetPeerOutAnchor() == nullptr) {
  379. GELOGE(FAILED, "Find %s in data anchor failed.", loop_after_node->GetName().c_str());
  380. return FAILED;
  381. }
  382. NodePtr loop_pre_node = in_anchor->GetPeerOutAnchor()->GetOwnerNode();
  383. // 1. Get variables
  384. NodePtr loop_cond_node = compute_graph->FindNode(NODE_NAME_FLOWCTRL_LOOP_COND);
  385. if (loop_cond_node == nullptr) {
  386. GELOGE(FAILED, "Find node :%s failed.", NODE_NAME_FLOWCTRL_LOOP_COND.c_str());
  387. return FAILED;
  388. }
  389. NodePtr iter_per_loop_node = compute_graph->FindNode(NODE_NAME_FLOWCTRL_LOOP_PER_ITER);
  390. if (iter_per_loop_node == nullptr) {
  391. GELOGE(FAILED, "Find node :%s failed.", NODE_NAME_FLOWCTRL_LOOP_PER_ITER.c_str());
  392. return FAILED;
  393. }
  394. // 2. Add StreamSwitch and edges to switch_node.
  395. GE_IF_BOOL_EXEC(loop_pre_node == nullptr, DOMI_LOGE("loop pre node is null"); return FAILED);
  396. string switch_name = loop_pre_node->GetName() + "_" + NODE_NAME_STREAM_SWITCH;
  397. NodePtr switch_node = InsertStreamSwitchOp(compute_graph, switch_name, loop_cond_node, iter_per_loop_node);
  398. if (switch_node == nullptr) {
  399. GELOGE(FAILED, "InsertStreamSwitchOp:%s failed.", switch_name.c_str());
  400. return FAILED;
  401. }
  402. GE_CHK_STATUS_RET(SetStreamLabel(switch_node, switch_name), "set stream label failed");
  403. graphStatus add_ret = GraphUtils::AddEdge(loop_pre_node->GetOutControlAnchor(), switch_node->GetInControlAnchor());
  404. if (add_ret != GRAPH_SUCCESS) {
  405. GELOGE(FAILED, "Add loop_pre_node:%s to switch_node:%s ctrl edge failed, ret = %u.",
  406. loop_pre_node->GetName().c_str(), switch_name.c_str(), add_ret);
  407. return FAILED;
  408. }
  409. add_ret = GraphUtils::AddEdge(loop_after_node->GetOutControlAnchor(), switch_node->GetInControlAnchor());
  410. if (add_ret != GRAPH_SUCCESS) {
  411. GELOGE(FAILED, "Add node:%s to switch_node:%s ctrl edge failed, ret = %u.", loop_after_node->GetName().c_str(),
  412. switch_name.c_str(), add_ret);
  413. return FAILED;
  414. }
  415. // 3. Create switch true branch: only active
  416. string active_name = switch_name + "_StreamActive";
  417. NodePtr active_node = InsertOp(compute_graph, STREAMACTIVE, active_name, {}, {});
  418. if (active_node == nullptr) {
  419. GELOGE(FAILED, "Insert stream active node:%s for SpecialNodeIteratorCtrl failed.", active_name.c_str());
  420. return FAILED;
  421. }
  422. GE_CHK_STATUS_RET(SetStreamLabel(active_node, active_name), "set stream label failed");
  423. GE_IF_BOOL_EXEC(!AttrUtils::SetBool(active_node->GetOpDesc(), ATTR_NAME_IS_LOOP_ACTIVE, true),
  424. DOMI_LOGE("set ATTR_NAME_IS_LOOP_ACTIVE failed"); return FAILED);
  425. add_ret = GraphUtils::AddEdge(switch_node->GetOutControlAnchor(), active_node->GetInControlAnchor());
  426. if (add_ret != GRAPH_SUCCESS) {
  427. GELOGE(FAILED, "Add switch_node:%s to active_node:%s ctrl edge failed, ret = %u.", switch_name.c_str(),
  428. active_name.c_str(), add_ret);
  429. return FAILED;
  430. }
  431. // used for stream assign to find true branch
  432. GE_CHK_STATUS_RET(SetActiveLabelList(switch_node, { active_name }), "set active label list failed");
  433. // used for stream assign to find active stream
  434. GE_CHK_STATUS_RET(SetActiveLabelList(active_node, { loop_pre_node->GetName() }), "set active label list failed");
  435. return SUCCESS;
  436. }
  437. } // namespace ge

图引擎模块(GE)是MindSpore的一个子模块,其代码由C++实现,位于前端模块ME和底层硬件之间,起到承接作用。图引擎模块以ME下发的图作为输入,然后进行一系列的深度图优化操作,最后输出一张可以在底层硬件上高效运行的图。GE针对昇腾AI处理器的硬件结构特点,做了特定的优化工作,以此来充分发挥出昇腾AI处理器的强大算力。在进行模型训练/推理时,GE会被自动调用而用户并不感知。GE主要由GE API和GE Core两部分组成,详细的架构图如下所示