From 83e15327f6aab0b8b0e3d3d7278b22f613f535a5 Mon Sep 17 00:00:00 2001
From: lianghao <lianghao24@hisilicon.com>
Date: Mon, 25 Jan 2021 15:22:28 +0800
Subject: [PATCH] memory assign

---
 ge/graph/build/memory/graph_mem_assigner.cc  | 21 ++++++++++++++++++++
 ge/graph/build/memory/graph_mem_assigner.h   |  4 ++++
 ge/graph/build/memory/memory_assigner.cc     |  5 +++++
 ge/graph/build/memory/var_mem_assign_util.cc |  5 +----
 4 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/ge/graph/build/memory/graph_mem_assigner.cc b/ge/graph/build/memory/graph_mem_assigner.cc
index a868fdcd..8c5d8940 100755
--- a/ge/graph/build/memory/graph_mem_assigner.cc
+++ b/ge/graph/build/memory/graph_mem_assigner.cc
@@ -88,6 +88,14 @@ Status VariableMemoryAssigner::AssignVarAttr2Nodes() {
   return ge::SUCCESS;
 }
 
+Status VariableMemoryAssigner::AssignMemory2HasRefAttrNode() {
+  Status result = ge::VarMemAssignUtil::AssignMemory2HasRefAttrNode(compute_graph_);
+  if (result != ge::SUCCESS) {
+    return result;
+  }
+  return ge::SUCCESS;
+}
+
 Status GraphMemoryAssigner::AssignMemory() {
   ge::HybridMemAssignerPtr mem_assigner(new(std::nothrow) HybridMemAssigner(compute_graph_));
   if (mem_assigner->Assign() != ge::SUCCESS) {
@@ -135,6 +143,19 @@ ge::Status GraphMemoryAssigner::AssignVarAttr2Nodes() {
   return ge::SUCCESS;
 }
 
+ge::Status GraphMemoryAssigner::AssignMemory2HasRefAttrNode() {
+  auto variable_assigner =
+      std::unique_ptr<ge::VariableMemoryAssigner>(new(std::nothrow) ge::VariableMemoryAssigner(compute_graph_));
+  if (variable_assigner == nullptr) {
+    GELOGE(ge::FAILED, "Alloc VariableMemoryAssigner failed.");
+    return ge::FAILED;
+  }
+  if (variable_assigner->AssignMemory2HasRefAttrNode() != ge::SUCCESS) {
+    return ge::FAILED;
+  }
+  return ge::SUCCESS;
+}
+
 ge::Status CalculateTensorRealSizeAndOutSize(const ge::ConstGeTensorDescPtr &output_desc,
                                                                   int64_t dim_index, int64_t &output_mem_size,
                                                                   int64_t &batch_dim_num, int64_t &out_size) {
diff --git a/ge/graph/build/memory/graph_mem_assigner.h b/ge/graph/build/memory/graph_mem_assigner.h
index a380e594..be6c47b0 100755
--- a/ge/graph/build/memory/graph_mem_assigner.h
+++ b/ge/graph/build/memory/graph_mem_assigner.h
@@ -63,6 +63,8 @@ class VariableMemoryAssigner {
   ///
   ge::Status AssignVarAttr2Nodes();
 
+  ge::Status AssignMemory2HasRefAttrNode();
+
  private:
   ge::ComputeGraphPtr compute_graph_;
 };
@@ -99,6 +101,8 @@ class GraphMemoryAssigner {
   ///
   ge::Status AssignVarAttr2Nodes();
 
+  ge::Status AssignMemory2HasRefAttrNode();
+
   ge::Status ReAssignMemory(bool is_loop_graph, map<int64_t, size_t> &mem_type_to_offset);
 
   ge::Status AssignZeroCopyMemory(map<int64_t, size_t> &mem_offset, size_t &zero_mem_copy_size);
diff --git a/ge/graph/build/memory/memory_assigner.cc b/ge/graph/build/memory/memory_assigner.cc
index 055103a9..0f58a040 100755
--- a/ge/graph/build/memory/memory_assigner.cc
+++ b/ge/graph/build/memory/memory_assigner.cc
@@ -40,6 +40,11 @@ Status MemoryAssigner::AssignMemory(bool is_loop_graph, map<int64_t, size_t> &me
     return ge::FAILED;
   }
 
+  if (graph_mem_assigner.AssignMemory2HasRefAttrNode() != ge::SUCCESS) {
+    GELOGE(ge::FAILED, "Assign memory to node which has ref attr failed!");
+    return ge::FAILED;
+  }
+
   // Assign memory for reference
   if (graph_mem_assigner.AssignReferenceMemory() != ge::SUCCESS) {
     GELOGE(ge::FAILED, "Assign reference memory failed!");
diff --git a/ge/graph/build/memory/var_mem_assign_util.cc b/ge/graph/build/memory/var_mem_assign_util.cc
index dfc633af..f910d2e2 100755
--- a/ge/graph/build/memory/var_mem_assign_util.cc
+++ b/ge/graph/build/memory/var_mem_assign_util.cc
@@ -33,10 +33,7 @@ using std::vector;
 
 namespace ge {
 Status VarMemAssignUtil::AssignVarMemory(ge::ComputeGraphPtr &compute_graph) {
-  GE_CHK_STATUS_RET(AssignMemory2VariableNode(compute_graph));
-  GE_CHK_STATUS_RET(AssignMemory2HasRefAttrNode(compute_graph));
-
-  return SUCCESS;
+  return AssignMemory2VariableNode(compute_graph);
 }
 
 Status VarMemAssignUtil::AssignConstantOpMemory(ge::ComputeGraphPtr &compute_graph) {