From 2ecf2a7d0c3bb13c4dc6bb30163111b85ba07d61 Mon Sep 17 00:00:00 2001
From: wangzhengjun <wangzhengjun3@huawei.com>
Date: Mon, 21 Jun 2021 19:32:53 +0800
Subject: [PATCH] skip control flow op when replace node with empty tensor

---
 .../passes/replace_with_empty_const_pass.cc   | 20 +++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/ge/graph/passes/replace_with_empty_const_pass.cc b/ge/graph/passes/replace_with_empty_const_pass.cc
index 9459c852..3176d1ee 100644
--- a/ge/graph/passes/replace_with_empty_const_pass.cc
+++ b/ge/graph/passes/replace_with_empty_const_pass.cc
@@ -21,7 +21,23 @@
 #include "framework/common/debug/ge_log.h"
 #include "framework/common/ge_inner_error_codes.h"
 #include "graph/utils/graph_utils.h"
+#include "graph/utils/node_utils.h"
 
+namespace {
+const std::unordered_set<std::string> kControlFlowOps = {
+  ge::SWITCH,
+  ge::REFSWITCH,
+  ge::MERGE,
+  ge::REFMERGE,
+  ge::ENTER,
+  ge::REFENTER,
+  ge::NEXTITERATION,
+  ge::REFNEXTITERATION,
+  ge::EXIT,
+  ge::REFEXIT,
+  ge::LOOPCOND
+};
+}
 namespace ge {
 Status ReplaceWithEmptyConstPass::Run(NodePtr &node) {
   GELOGD("ReplaceWithEmptyConstPass in.");
@@ -39,6 +55,10 @@ Status ReplaceWithEmptyConstPass::Run(NodePtr &node) {
     GELOGI("Node %s is const. Ignore current pass.", node->GetName().c_str());
     return SUCCESS;
   }
+  if (kControlFlowOps.count(NodeUtils::GetNodeType(node)) != 0) {
+    GELOGI("Node %s is control flow op. Ignore current pass.", node->GetName().c_str());
+    return SUCCESS;
+  }
   // Node like no op, it has no output
   if (node->GetOpDesc()->GetAllOutputsDescPtr().empty()) {
     GELOGI("Node %s has no output desc. Ignore current pass.", node->GetName().c_str());