From acafc1b57875b38ed661665fd77ccfe3c7d5ed62 Mon Sep 17 00:00:00 2001 From: lianghao Date: Mon, 28 Dec 2020 11:47:33 +0800 Subject: [PATCH] AssignMemory2HasRefAttrNode --- 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 | 1 - 4 files changed, 30 insertions(+), 1 deletion(-) diff --git a/ge/graph/build/memory/graph_mem_assigner.cc b/ge/graph/build/memory/graph_mem_assigner.cc index 1cdb2efa..55f9ecae 100644 --- a/ge/graph/build/memory/graph_mem_assigner.cc +++ b/ge/graph/build/memory/graph_mem_assigner.cc @@ -89,6 +89,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) { @@ -131,6 +139,19 @@ ge::Status GraphMemoryAssigner::AssignVarAttr2Nodes() { return ge::SUCCESS; } +ge::Status GraphMemoryAssigner::AssignMemory2HasRefAttrNode() { + auto variable_assigner = + std::unique_ptr(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 GraphMemoryAssigner::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 201e6d01..74e45856 100644 --- 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_; }; @@ -97,6 +99,8 @@ class GraphMemoryAssigner { /// ge::Status AssignVarAttr2Nodes(); + ge::Status AssignMemory2HasRefAttrNode(); + ge::Status ReAssignMemory(bool is_loop_graph, size_t &mem_offset); ge::Status AssignZeroCopyMemory(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 e36f082e..27bb7ad7 100644 --- 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, size_t &mem_offset, size return ge::FAILED; } + if (graph_mem_assigner.AssignMemory2HasRefAttrNode() != ge::SUCCESS) { + GELOGE(ge::FAILED, "Assign memory 2 HasRef AttrNode 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 a352cf65..605f918d 100644 --- a/ge/graph/build/memory/var_mem_assign_util.cc +++ b/ge/graph/build/memory/var_mem_assign_util.cc @@ -34,7 +34,6 @@ 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; }