diff --git a/ge/graph/build/memory/graph_mem_assigner.cc b/ge/graph/build/memory/graph_mem_assigner.cc index 33fda096..f4df8268 100755 --- a/ge/graph/build/memory/graph_mem_assigner.cc +++ b/ge/graph/build/memory/graph_mem_assigner.cc @@ -430,7 +430,7 @@ Status GraphMemoryAssigner::ReAssignContinuousMemory(bool is_loop_graph) { GELOGE(FAILED, "node %s has no continuous type!", node->GetName().c_str()); return FAILED; } - GE_CHK_STATUS_RET(AssignContinuousInputMemoryWithAtomicProcess(node, iter->second), + GE_CHK_STATUS_RET(AssignContinuousInputMemoryWithAtomicProcess(node, iter->second, true), "Assign node %s continuous input memory failed.", node->GetName().c_str()) } for (auto pair : memory_offset_) { @@ -441,7 +441,7 @@ Status GraphMemoryAssigner::ReAssignContinuousMemory(bool is_loop_graph) { } Status GraphMemoryAssigner::AssignContinuousInputMemory(const ge::NodePtr &node, int64_t &continuous_mem_start, - int64_t &continuous_mem_size, int64_t memory_type, uint32_t continuous_type) { + int64_t &continuous_mem_size, int64_t memory_type, uint32_t continuous_type, bool reverse_refresh) { GELOGI("Current node %s needs continuous input.", node->GetName().c_str()); auto iter = memory_offset_.find(memory_type); if (iter == memory_offset_.end()) { @@ -508,12 +508,16 @@ Status GraphMemoryAssigner::AssignContinuousInputMemory(const ge::NodePtr &node, std::map out2ins; GE_CHK_STATUS_RET(GetAllRef(node, out2ins), "Node: %s get all ref failed", node->GetName().c_str()); // output is beginning offset, set offset for input; only support this case now - if (out2ins.size() == 1 && out2ins.begin()->second == 0) { + if ((out2ins.size() == 1) && (out2ins.begin()->second == 0) && (reverse_refresh)) { + auto peer_output_offset = output_list.at(peer_out_data_anchor->GetIdx()); output_list.at(peer_out_data_anchor->GetIdx()) = output_list_this.at(out2ins.begin()->first); peer_op_desc->SetOutputOffset(output_list); + GELOGI("Node %s out %d ref in %d input node %s, use output offset %ld update %ld", node->GetName().c_str(), + out2ins.begin()->first, out2ins.begin()->second, peer_op_desc->GetName().c_str(), + output_list_this.at(out2ins.begin()->first), peer_output_offset); } else { - GELOGW("Node %s out %d ref in %d with total ref numbers %zu", node->GetName().c_str(), out2ins.begin()->first, - out2ins.begin()->second, out2ins.size()); + GELOGD("Node %s out %d ref in %d input node %s with total ref numbers %zu", node->GetName().c_str(), + out2ins.begin()->first, out2ins.begin()->second, peer_op_desc->GetName().c_str(), out2ins.size()); } // first input is beginning offset mem_offset = output_list.at(peer_out_data_anchor->GetIdx()); @@ -1535,6 +1539,11 @@ ge::Status GraphMemoryAssigner::GetAllRef(const NodePtr &node, map &node_2_continuous_type) { for (const auto &in_node : input_continuous_node->GetInDataNodes()) { + if (in_node->GetType() == VARIABLE) { + GELOGI("node %s 's precursor node %s is variable, do not store.", input_continuous_node->GetName().c_str(), + in_node->GetName().c_str()); + return true; + } auto iter = node_2_continuous_type.find(in_node); // In node's topo order in the front, so function can not be exception auto continuous_type = iter->second; @@ -1560,13 +1569,15 @@ bool GraphMemoryAssigner::AssignContinuousInputMemoryWithAtomicProcessDirectly( } ge::Status GraphMemoryAssigner::AssignContinuousInputMemoryWithAtomicProcess(const NodePtr &input_continuous_node, - uint32_t continuous_type) { + uint32_t continuous_type, + bool reverse_refresh) { int64_t mem_clean_start = 0; int64_t mem_clean_size = 0; int64_t memory_type = RT_MEMORY_HBM; GE_CHK_STATUS_RET(GetNodeMemoryType(input_continuous_node, memory_type, "input"), "Get node memory type failed."); - auto ret = AssignContinuousInputMemory(input_continuous_node, mem_clean_start, mem_clean_size, memory_type, continuous_type); + auto ret = AssignContinuousInputMemory(input_continuous_node, mem_clean_start, mem_clean_size, memory_type, + continuous_type, reverse_refresh); if (ret != ge::SUCCESS) { GELOGE(ret, "Assign continuous input memory failed!"); return ret; diff --git a/ge/graph/build/memory/graph_mem_assigner.h b/ge/graph/build/memory/graph_mem_assigner.h index b56c3716..f4d1366d 100755 --- a/ge/graph/build/memory/graph_mem_assigner.h +++ b/ge/graph/build/memory/graph_mem_assigner.h @@ -131,13 +131,14 @@ class GraphMemoryAssigner { std::map &node_2_continuous_type); ge::Status AssignContinuousInputMemoryWithAtomicProcess(const NodePtr &input_continuous_node, - uint32_t continuous_type); + uint32_t continuous_type, bool reverse_refresh=false); ge::Status FilterAtomicNodesForMemoryAssign(map>> &normal_atomic_nodes_map, map> &connecting_output_atomic_nodes); ge::Status AssignContinuousInputMemory(const ge::NodePtr &node, int64_t &continuous_mem_start, - int64_t &continuous_mem_size, int64_t memory_type, uint32_t continuous_type); + int64_t &continuous_mem_size, int64_t memory_type, uint32_t continuous_type, + bool reverse_refresh = false); ge::Status AssignContinuousOutputMemory(const ge::NodePtr &node, int64_t memory_type, uint32_t continuous_type);