diff --git a/ge/graph/build/memory/graph_mem_assigner.cc b/ge/graph/build/memory/graph_mem_assigner.cc index c7e3bd8d..947bf38e 100755 --- a/ge/graph/build/memory/graph_mem_assigner.cc +++ b/ge/graph/build/memory/graph_mem_assigner.cc @@ -2121,6 +2121,18 @@ void GraphMemoryAssigner::CheckNeedCalcDistAndUpdateVisitInfo( return; } +void GraphMemoryAssigner::UpdateMemBlockVisitInfo(const NodePtr &update_node, + const InDataAnchorPtr &in_data_anchor, + bool is_same_input, + size_t matched_mem_offset, + map>> &mem_block_visit_info) { + if (!is_same_input) { + mem_block_visit_info[matched_mem_offset].first = update_node; + mem_block_visit_info[matched_mem_offset].second.clear(); + } + mem_block_visit_info[matched_mem_offset].second.push_back(in_data_anchor->GetIdx()); +} + // calculate distance, update visit info, update prev_node input desc, update cur node input desc void GraphMemoryAssigner::CalcDistanceAndUpdateDesc(const map &node_index_in_stream, const InDataAnchorPtr &in_data_anchor, @@ -2132,6 +2144,7 @@ void GraphMemoryAssigner::CalcDistanceAndUpdateDesc(const map & auto prev_node = mem_block_visit_info[matched_mem_offset].first; auto prev_node_input_index_vec = mem_block_visit_info[matched_mem_offset].second; GE_IF_BOOL_EXEC(prev_node == nullptr, is_need_skip = true; return); + bool is_same_input = false; if (prev_node_input_index_vec.size() == 1 && prev_node_input_index_vec[0] == -1) { // prev_node is producer and the data is just be produced(not visited by other node) GE_IF_BOOL_EXEC(prev_node->GetOpDesc() == nullptr, is_need_skip = true; return); @@ -2145,24 +2158,15 @@ void GraphMemoryAssigner::CalcDistanceAndUpdateDesc(const map & distance = node_index_in_stream.at(node->GetName()) - iter->second - 1; } } - mem_block_visit_info[matched_mem_offset].first = node; - mem_block_visit_info[matched_mem_offset].second.clear(); - mem_block_visit_info[matched_mem_offset].second.push_back(in_data_anchor->GetIdx()); } else { // the data is visit by other customer just before. - if (prev_node_input_index_vec.empty()) { - GELOGW("Missing prev node[%s] input index.", prev_node->GetName().c_str()); - is_need_skip = true; - return; - } + GE_IF_BOOL_EXEC(prev_node_input_index_vec.empty(), + GELOGW("Miss prev node[%s] input idx", prev_node->GetName().c_str()); is_need_skip = true; return); if (prev_node == node) { // scene: multiple anchors of a node access the same data vector prev_next_distances; GE_IF_BOOL_EXEC(prev_node->GetOpDesc() == nullptr, is_need_skip = true; return); auto input_desc = prev_node->GetOpDesc()->GetInputDesc(prev_node_input_index_vec[0]); - if (!ge::AttrUtils::GetListInt(input_desc, ATTR_NAME_DATA_VISIT_DISTANCE, prev_next_distances)) { - GELOGW("Get ATTR_NAME_DATA_VISIT_DISTANCE failed."); - is_need_skip = true; - return; - } + GE_IF_BOOL_EXEC(!ge::AttrUtils::GetListInt(input_desc, ATTR_NAME_DATA_VISIT_DISTANCE, prev_next_distances), + GELOGW("Get ATTR_NAME_DATA_VISIT_DISTANCE failed."); is_need_skip = true; return); if (prev_next_distances.size() != kPrevNextDistanceNum) { GELOGW("Size of prev_next_distance is not %d.", kPrevNextDistanceNum); is_need_skip = true; @@ -2170,15 +2174,13 @@ void GraphMemoryAssigner::CalcDistanceAndUpdateDesc(const map & } else { distance = prev_next_distances[0]; // use the same prev_distance as previous anchor } - mem_block_visit_info[matched_mem_offset].second.push_back(in_data_anchor->GetIdx()); + is_same_input = true; } else { distance = node_index_in_stream.at(node->GetName()) - node_index_in_stream.at(prev_node->GetName()) - 1; UpdatePrevNodeInputDesc(prev_node, prev_node_input_index_vec, distance); - mem_block_visit_info[matched_mem_offset].first = node; - mem_block_visit_info[matched_mem_offset].second.clear(); - mem_block_visit_info[matched_mem_offset].second.push_back(in_data_anchor->GetIdx()); } } + UpdateMemBlockVisitInfo(node, in_data_anchor, is_same_input, matched_mem_offset, mem_block_visit_info); UpdateCurNodeInputDesc(node, in_data_anchor->GetIdx(), distance); } diff --git a/ge/graph/build/memory/graph_mem_assigner.h b/ge/graph/build/memory/graph_mem_assigner.h index a6a2a686..bec434e3 100755 --- a/ge/graph/build/memory/graph_mem_assigner.h +++ b/ge/graph/build/memory/graph_mem_assigner.h @@ -225,6 +225,12 @@ class GraphMemoryAssigner { map>> &mem_block_visit_info, bool &is_need_skip); + void UpdateMemBlockVisitInfo(const NodePtr &update_node, + const InDataAnchorPtr &in_data_anchor, + bool is_same_input, + size_t matched_mem_offset, + map>> &mem_block_visit_info); + void DeleteVisitInfoWhenLifecycleEnded(const NodePtr &node, const InDataAnchorPtr &in_data_anchor, size_t matched_mem_offset,