diff --git a/ge/graph/build/memory/graph_mem_assigner.cc b/ge/graph/build/memory/graph_mem_assigner.cc index b098a5f5..dddba2cd 100755 --- a/ge/graph/build/memory/graph_mem_assigner.cc +++ b/ge/graph/build/memory/graph_mem_assigner.cc @@ -2120,6 +2120,18 @@ void GraphMemoryAssigner::CheckNeedCalcDistAndUpdateVisitInfo( return; } +void GraphMemoryAssigner::UpdateMemBlockVisitInfo(const NodePtr &update_node, + const InDataAnchorPtr &in_data_anchor, + bool need_update_node, + size_t matched_mem_offset, + map>> &mem_block_visit_info) { + if (need_update_node) { + 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()); +} + // 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, @@ -2131,6 +2143,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 need_update_node = true; 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); @@ -2144,24 +2157,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; @@ -2169,15 +2173,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()); + need_update_node = false; } 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, matched_mem_offset, mem_block_visit_info, need_update_node); 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..f0e8a7e2 100755 --- a/ge/graph/build/memory/graph_mem_assigner.h +++ b/ge/graph/build/memory/graph_mem_assigner.h @@ -225,6 +225,11 @@ class GraphMemoryAssigner { map>> &mem_block_visit_info, bool &is_need_skip); + void UpdateMemBlockVisitInfo(const NodePtr &update_node, + const InDataAnchorPtr &in_data_anchor, + 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,