From 464f3e7ba29c3e58d848bd6a73b5d305946c72dc Mon Sep 17 00:00:00 2001
From: l00444296 <lixiwen1@huawei.com>
Date: Thu, 17 Dec 2020 15:31:22 +0800
Subject: [PATCH] Feature: reset shape of dynamic single op

---
 ge/generator/ge_generator.cc                  |  5 +-
 .../dynamic_single_op_reset_shape_pass.cc     | 59 +++++++++++++++----
 .../dynamic_single_op_reset_shape_pass.h      |  2 +
 3 files changed, 54 insertions(+), 12 deletions(-)

diff --git a/ge/generator/ge_generator.cc b/ge/generator/ge_generator.cc
index 2baa2c5b..7c083d2b 100644
--- a/ge/generator/ge_generator.cc
+++ b/ge/generator/ge_generator.cc
@@ -275,7 +275,10 @@ static void ResetTensorVecShape(const vector<GeTensor> &inputs, vector<GeTensor>
 
     ge::GeTensor inputTensor;
     ge::GeTensorDesc desc(input_desc);
-    if (shape_ori.GetDims().size() > 0) {
+
+    bool is_const = false;
+    (void)AttrUtils::GetBool(input_desc, CONST_ATTR_NAME_INPUT, is_const);
+    if (!is_const && shape_ori.GetDims().size() > 0) {
       desc.SetShape(dynamic_shape);
     }
 
diff --git a/ge/graph/passes/dynamic_single_op_reset_shape_pass.cc b/ge/graph/passes/dynamic_single_op_reset_shape_pass.cc
index e543a9a5..519b0574 100644
--- a/ge/graph/passes/dynamic_single_op_reset_shape_pass.cc
+++ b/ge/graph/passes/dynamic_single_op_reset_shape_pass.cc
@@ -38,6 +38,18 @@ Status DynamicSingleOpResetShapePass::Run(ComputeGraphPtr graph) {
     return ge::GE_CLI_GE_NOT_INITIALIZED;
   }
 
+  // pass if graph has not aicpu node.
+  bool is_not_aicpu = false;
+  if (CheckAllAicpuNodes(graph, is_not_aicpu) != SUCCESS) {
+    GELOGE(ge::GE_CLI_GE_NOT_INITIALIZED, "Check if graph has not aicpu node failed.");
+    return ge::GE_CLI_GE_NOT_INITIALIZED;
+  }
+  if (is_not_aicpu) {
+    GELOGI("The graph [%s] has not aicpu node, whose aicpu nodes would not be reset dynamic shape",
+           graph->GetName().c_str());
+    return SUCCESS;
+  }
+
   for (const auto &node : graph->GetDirectNode()) {
     GE_CHECK_NOTNULL(node->GetOpDesc());
     // pass input and output node
@@ -53,18 +65,8 @@ Status DynamicSingleOpResetShapePass::Run(ComputeGraphPtr graph) {
       continue;
     }
 
-    // pass not aicpu node.
-    auto op_desc = node->GetOpDesc();
-    string engine_name = op_desc->GetOpEngineName();
-    if (engine_name.empty()) {
-      GELOGE(GRAPH_FAILED, "Get engine failed of node[%s].", node->GetName().c_str());
-      return GRAPH_FAILED;
-    }
-    if (engine_name != kEngineNameAiCpu && engine_name != kEngineNameAiCpuTf) {
-      continue;
-    }
-
     // reset aicpu shape to unknown shape
+    auto op_desc = node->GetOpDesc();
     if (ResetOpShape(op_desc) != SUCCESS) {
       GELOGE(ge::GE_CLI_GE_NOT_INITIALIZED, "Reset node[%s] dynamic shapr failed.", node->GetName().c_str());
       return ge::GE_CLI_GE_NOT_INITIALIZED;
@@ -76,6 +78,37 @@ Status DynamicSingleOpResetShapePass::Run(ComputeGraphPtr graph) {
   return SUCCESS;
 }
 
+Status DynamicSingleOpResetShapePass::CheckAllAicpuNodes(const ComputeGraphPtr &graph, bool &is_not_aicpu) {
+  is_not_aicpu = false;
+  for (const auto &node : graph->GetDirectNode()) {
+    GE_CHECK_NOTNULL(node->GetOpDesc());
+    // pass input and output node
+    if (node->GetType() == DATA || node->GetType() == CONSTANT || node->GetType() == CONSTANTOP ||
+        node->GetType() == NETOUTPUT) {
+      continue;
+    }
+
+    // find if there are aicpu nodes.
+    auto op_desc = node->GetOpDesc();
+    string engine_name = op_desc->GetOpEngineName();
+    if (engine_name.empty()) {
+      GELOGE(GRAPH_FAILED, "Get engine failed of node[%s].", node->GetName().c_str());
+      return GRAPH_FAILED;
+    }
+    if (engine_name != kEngineNameAiCpu && engine_name != kEngineNameAiCpuTf) {
+      is_not_aicpu = true;
+      return SUCCESS;
+    }
+  }
+  return SUCCESS;
+}
+
+bool DynamicSingleOpResetShapePass::CheckIfConstInput(const GeTensorDescPtr &input_tensor_desc) {
+  bool is_const = false;
+  (void)AttrUtils::GetBool(input_desc, CONST_ATTR_NAME_INPUT, is_const);
+  return is_const;
+}
+
 Status DynamicSingleOpResetShapePass::ResetOpShape(OpDescPtr &op_desc) {
   GE_CHECK_NOTNULL(op_desc);
   std::vector<int64_t> dynamic_shape_dims = {kDynamicShapeDim};
@@ -88,6 +121,10 @@ Status DynamicSingleOpResetShapePass::ResetOpShape(OpDescPtr &op_desc) {
     if (dims_ori.size() == 0) {
       continue;
     }
+    // pass const input
+    if (CheckIfConstInput(input_desc)) {
+      continue;
+    }
     input_desc->SetShape(dynamic_shape);
   }
   for (size_t i = 0; i < op_desc->GetAllOutputsDesc().size(); i++) {
diff --git a/ge/graph/passes/dynamic_single_op_reset_shape_pass.h b/ge/graph/passes/dynamic_single_op_reset_shape_pass.h
index 9dcc7853..50ec13ba 100644
--- a/ge/graph/passes/dynamic_single_op_reset_shape_pass.h
+++ b/ge/graph/passes/dynamic_single_op_reset_shape_pass.h
@@ -27,6 +27,8 @@ class DynamicSingleOpResetShapePass : public GraphPass {
 
  private:
   Status ResetOpShape(OpDescPtr &op_desc);
+  Status CheckAllAicpuNodes(const ComputeGraphPtr &graph, bool &is_not_aicpu);
+  bool CheckIfConstInput();
 };
 }  // namespace ge
 #endif  // GE_GRAPH_PASSES_DYNAMIC_SINGLE_OP_RESET_SHAPE_PASS_H_