From c2b91d9ef41fda64f8e5a3aaa74c50701a2233e0 Mon Sep 17 00:00:00 2001
From: unknown <zhang.xiaokun@huawei.com>
Date: Thu, 12 Nov 2020 15:49:16 +0800
Subject: [PATCH] Fix Label for dynamic graph.

---
 ge/graph/build/label_allocator.cc | 31 +++++++++++++++++++++----------
 ge/graph/build/model_builder.cc   |  4 ++--
 2 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/ge/graph/build/label_allocator.cc b/ge/graph/build/label_allocator.cc
index fad7d0c2..51f31003 100644
--- a/ge/graph/build/label_allocator.cc
+++ b/ge/graph/build/label_allocator.cc
@@ -32,11 +32,6 @@ Status LabelAllocator::AssignFunctionalLabels() {
     return INTERNAL_ERROR;
   }
 
-  if (compute_graph_->GetGraphUnknownFlag()) {
-    GELOGD("Graph[%s] is unknown graph, skip label allocator.", compute_graph_->GetName().c_str());
-    return SUCCESS;
-  }
-
   // Add label task for sub graph.
   GELOGI("AssignFunctionalLabels start: %s.", compute_graph_->GetName().c_str());
   std::set<NodePtr> functional_nodes;
@@ -62,7 +57,7 @@ Status LabelAllocator::AssignFunctionalLabels() {
   }
 
   (void)AttrUtils::SetInt(*compute_graph_, ATTR_MODEL_LABEL_NUM, label_index);
-  GELOGI("AssignFunctionalLabels success.");
+  GELOGI("AssignFunctionalLabels success, Num: %u.", label_index);
   return SUCCESS;
 }
 
@@ -72,13 +67,29 @@ bool LabelAllocator::CollectFunctionalNode(ComputeGraphPtr &graph, std::set<Node
     return false;
   }
 
-  NodePtr parent = graph->GetParentNode();
-  if (parent == nullptr) {
-    GELOGE(INTERNAL_ERROR, "ComputeGraph owner not set: %s.", graph->GetName().c_str());
+  if (graph->GetGraphUnknownFlag()) {
+    GELOGD("Graph[%s] is unknown graph, skip label allocator.", graph->GetName().c_str());
+    return true;
+  }
+
+  NodePtr func_node = graph->GetParentNode();
+  if (func_node == nullptr) {
+    GELOGE(INTERNAL_ERROR, "Parent functional node not set: %s.", graph->GetName().c_str());
     return false;
   }
 
-  (void)functional_nodes.insert(parent);  // unique functional node.
+  ComputeGraphPtr owner_graph = func_node->GetOwnerComputeGraph();
+  if (owner_graph == nullptr) {
+    GELOGE(INTERNAL_ERROR, "ComputeGraph owner not set: %s.", func_node->GetName().c_str());
+    return false;
+  }
+
+  if (owner_graph->GetGraphUnknownFlag()) {
+    GELOGD("Graph[%s] is unknown graph, skip label allocator.", owner_graph->GetName().c_str());
+    return true;
+  }
+
+  (void)functional_nodes.insert(func_node);  // unique functional node.
   return true;
 }
 }  // namespace ge
diff --git a/ge/graph/build/model_builder.cc b/ge/graph/build/model_builder.cc
index 56a5b4dc..f382c24a 100755
--- a/ge/graph/build/model_builder.cc
+++ b/ge/graph/build/model_builder.cc
@@ -690,8 +690,8 @@ Status ModelBuilder::BuildModelForGetTask(ge::Model &model) {
   GE_TIMESTAMP_END(AssignLogicalStreams, "GraphBuilder::AssignLogicalStreams");
 
   // Assign functional op labels.
-  label_num_ = 0;
-  (void)AttrUtils::GetInt(*compute_graph_, ATTR_MODEL_LABEL_NUM, label_num_);
+  auto root_graph = GraphUtils::FindRootGraph(compute_graph_);
+  (void)AttrUtils::GetInt(*root_graph, ATTR_MODEL_LABEL_NUM, label_num_);
 
   GE_TIMESTAMP_START(AssignMemory);
   MemoryAssigner mem_assigner(compute_graph_);