From a2f684ea992f46e368858ffbc2c12efd7a08a421 Mon Sep 17 00:00:00 2001
From: liuyu <liuyu195@huawei.com>
Date: Tue, 16 Mar 2021 18:46:23 +0800
Subject: [PATCH] remove while and if

---
 mindspore/lite/src/ops/ops_def.cc             |  8 +++++
 mindspore/lite/src/ops/ops_def.h              | 29 +++++++++++++++++++
 mindspore/lite/src/ops/ops_utils.cc           |  2 +-
 .../tools/optimizer/graph/infershape_pass.cc  |  3 +-
 4 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/mindspore/lite/src/ops/ops_def.cc b/mindspore/lite/src/ops/ops_def.cc
index ded8368eb4..4cf237c54b 100644
--- a/mindspore/lite/src/ops/ops_def.cc
+++ b/mindspore/lite/src/ops/ops_def.cc
@@ -1006,6 +1006,11 @@ OP_SCHEMA_DEF(Unstack)
 OP_ATTR_WITH_VALUE(axis, long, 0)
 OP_SCHEMA_DEF_END(Unstack)
 
+OP_SCHEMA_DEF_ONLY(While)
+OP_ATTR_ONLY(cond_subgraph_index, long)
+OP_ATTR_ONLY(body_subgraph_index, long)
+OP_SCHEMA_DEF_ONLY_END(While)
+
 OP_SCHEMA_DEF(Where)
 OP_SCHEMA_DEF_END(Where)
 
@@ -1015,6 +1020,9 @@ OP_SCHEMA_DEF_END(ZerosLike)
 OP_SCHEMA_DEF(Select)
 OP_SCHEMA_DEF_END(Select)
 
+OP_SCHEMA_DEF_ONLY(If)
+OP_SCHEMA_DEF_ONLY_END(If)
+
 OP_SCHEMA_DEF(GRU)
 OP_ATTR_WITH_VALUE(bidirectional, bool, false)
 OP_SCHEMA_DEF_END(GRU)
diff --git a/mindspore/lite/src/ops/ops_def.h b/mindspore/lite/src/ops/ops_def.h
index 9f7c1f875b..4ceb61408e 100644
--- a/mindspore/lite/src/ops/ops_def.h
+++ b/mindspore/lite/src/ops/ops_def.h
@@ -154,4 +154,33 @@
 #else
 #define OP_SCHEMA_DEF_END(OP)
 #endif
+
+#ifdef GEN_SCHEMA_DEF
+#define OP_SCHEMA_DEF_ONLY(OP)     \
+  namespace mindspore::lite::ops { \
+  std::string Gen##OP##Def() {     \
+    std::string op_def = "table "; \
+    op_def.append(#OP);            \
+    op_def.append(" {\n");
+#else
+#define OP_SCHEMA_DEF_ONLY(OP)
+#endif
+
+#ifdef GEN_SCHEMA_DEF
+#define OP_ATTR_ONLY(key, type) op_def.append("    ").append(#key).append(": ").append(#type).append(";\n");
+#else
+#define OP_ATTR_ONLY(key, type)
+#endif
+
+#ifdef GEN_SCHEMA_DEF
+#define OP_SCHEMA_DEF_ONLY_END(OP)                 \
+  op_def.append("}\n\n");                          \
+  return op_def;                                   \
+  }                                                \
+  SchemaOpRegister g_schema_op_##OP(Gen##OP##Def); \
+  }  // namespace mindspore::lite::ops
+#else
+#define OP_SCHEMA_DEF_ONLY_END(OP)
+#endif
+
 #endif  // MINDSPORE_LITE_SRC_OPS_OPS_DEF_H_
diff --git a/mindspore/lite/src/ops/ops_utils.cc b/mindspore/lite/src/ops/ops_utils.cc
index 429a64b13c..62992afeec 100644
--- a/mindspore/lite/src/ops/ops_utils.cc
+++ b/mindspore/lite/src/ops/ops_utils.cc
@@ -39,7 +39,7 @@ schema::PrimitiveT *GetPrimitiveT(const AnfNodePtr &node) {
   if (creator != nullptr) {
     return creator(node);
   } else {
-    MS_LOG(ERROR) << "can not find MSOpsRegistry for op: " << prim->name();
+    MS_LOG(WARNING) << "can not find MSOpsRegistry for op: " << prim->name();
     return nullptr;
   }
 }
diff --git a/mindspore/lite/tools/optimizer/graph/infershape_pass.cc b/mindspore/lite/tools/optimizer/graph/infershape_pass.cc
index 7aa8b95772..a4367f2283 100644
--- a/mindspore/lite/tools/optimizer/graph/infershape_pass.cc
+++ b/mindspore/lite/tools/optimizer/graph/infershape_pass.cc
@@ -42,7 +42,8 @@ ParamValueLitePtr NewParamValueLitePtr(lite::Tensor *tensor) {
 bool IsSpecialType(const CNodePtr &cnode) {
   if (CheckPrimitiveType(cnode, prim::kPrimTupleGetItem) || CheckPrimitiveType(cnode, prim::kPrimDepend) ||
       CheckPrimitiveType(cnode, prim::kPrimControlDepend) || CheckPrimitiveType(cnode, kPrimMakeTuple) ||
-      CheckPrimitiveType(cnode, kPrimReturn)) {
+      CheckPrimitiveType(cnode, kPrimReturn) || CheckPrimitiveType(cnode, std::make_shared<Primitive>("While")) ||
+      CheckPrimitiveType(cnode, std::make_shared<Primitive>("If"))) {
     return true;
   }
   return false;