Pre Merge pull request !1699 from lichun/master

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<size_t, pair<NodePtr, vector<int64_t>>> &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<string, int64_t> &node_index_in_stream,
                                                     const InDataAnchorPtr &in_data_anchor,
@@ -2132,6 +2144,7 @@ void GraphMemoryAssigner::CalcDistanceAndUpdateDesc(const map<string, int64_t> &
   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<string, int64_t> &
         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<int64_t> 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<string, int64_t> &
       } 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);
 }
 

ge/graph/build/memory/graph_mem_assigner.h

@@ -225,6 +225,12 @@ class GraphMemoryAssigner {
                                  map<size_t, pair<NodePtr, vector<int64_t>>> &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<size_t, pair<NodePtr, vector<int64_t>>> &mem_block_visit_info);
+
   void DeleteVisitInfoWhenLifecycleEnded(const NodePtr &node,
                                          const InDataAnchorPtr &in_data_anchor,
                                          size_t matched_mem_offset,