From d6ce5c65fbfd9975a22a0ce3cf2ef8f4099f5dbd Mon Sep 17 00:00:00 2001
From: zhou_chao1993 <zhouchao46@huawei.com>
Date: Sat, 23 Jan 2021 17:59:26 +0800
Subject: [PATCH] dump known shape

---
 ge/graph/load/model_manager/data_dumper.cc               | 2 +-
 ge/graph/load/model_manager/davinci_model.cc             | 9 ++++++---
 ge/graph/load/model_manager/davinci_model.h              | 7 +++++++
 .../load/model_manager/task_info/kernel_ex_task_info.cc  | 2 +-
 .../compiledsubgraph/known_node_executor.cc              | 6 ++++++
 5 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/ge/graph/load/model_manager/data_dumper.cc b/ge/graph/load/model_manager/data_dumper.cc
index 947aac1d..235cffa9 100644
--- a/ge/graph/load/model_manager/data_dumper.cc
+++ b/ge/graph/load/model_manager/data_dumper.cc
@@ -820,6 +820,7 @@ Status DataDumper::UnloadDumpInfo() {
   for (const auto &op_iter : op_list_) {
     aicpu::dump::Task task;
     task.set_task_id(op_iter.task_id);
+    task.set_stream_id(op_iter.stream_id);
     op_mapping_info.mutable_task()->Add(std::move(task));
   }
   auto ret = ExecuteUnLoadDumpInfo(op_mapping_info);
@@ -834,7 +835,6 @@ void DataDumper::DumpShrink() {
   compute_graph_.reset();
   input_map_.clear();
   ref_info_.clear();
-  op_list_.clear();
 }
 
 void DataDumper::PrintCheckLog(string &dump_list_key) {
diff --git a/ge/graph/load/model_manager/davinci_model.cc b/ge/graph/load/model_manager/davinci_model.cc
index 136a041c..cdf934c0 100755
--- a/ge/graph/load/model_manager/davinci_model.cc
+++ b/ge/graph/load/model_manager/davinci_model.cc
@@ -3959,8 +3959,11 @@ void DavinciModel::SetDataDumperArgs(const ComputeGraphPtr &graph, const map<str
   }
   data_dumper_.SetDeviceId(device_id);
 
-  // set loop count addr
-  auto get_var_addr = [&](const string &name) -> void *{
+  if (known_node_) {
+    data_dumper_.SetLoopAddr(known_shape_global_step_, nullptr, nullptr);
+  } else {
+    // set loop count addr
+    auto get_var_addr = [&](const string &name) -> void *{
     const auto it = variable_by_name.find(name);
     if (it != variable_by_name.end()) {
       const auto output_sizes = ModelUtils::GetOutputSize(it->second);
@@ -3973,10 +3976,10 @@ void DavinciModel::SetDataDumperArgs(const ComputeGraphPtr &graph, const map<str
     GELOGD("op: %s is null.", name.c_str());
     return nullptr;
   };
-
   data_dumper_.SetLoopAddr(get_var_addr(NODE_NAME_GLOBAL_STEP),
                            get_var_addr(NODE_NAME_FLOWCTRL_LOOP_PER_ITER),
                            get_var_addr(NODE_NAME_FLOWCTRL_LOOP_COND));
+  }
 }
 
 uint32_t DavinciModel::GetFlowctrlIndex(uint32_t op_index) {
diff --git a/ge/graph/load/model_manager/davinci_model.h b/ge/graph/load/model_manager/davinci_model.h
index 22f3bcab..53e9cd4d 100755
--- a/ge/graph/load/model_manager/davinci_model.h
+++ b/ge/graph/load/model_manager/davinci_model.h
@@ -470,6 +470,10 @@ class DavinciModel {
     data_dumper_.SaveDumpTask(task_id, stream_id, op_desc, args);
   }
 
+  void SetKnownShapeGlobalStep(void *global_step) {
+    known_shape_global_step_ = global_step;
+  }
+
   void DumperShrink() {
     data_dumper_.DumpShrink();
   }
@@ -1057,6 +1061,9 @@ class DavinciModel {
   vector<uint32_t> input_formats_;
   vector<InputOutputDescInfo> output_descs_;
   vector<uint32_t> output_formats_;
+
+  // known shape node for dump
+  void *known_shape_global_step_;
 };
 }  // namespace ge
 #endif  // GE_GRAPH_LOAD_NEW_MODEL_MANAGER_DAVINCI_MODEL_H_
diff --git a/ge/graph/load/model_manager/task_info/kernel_ex_task_info.cc b/ge/graph/load/model_manager/task_info/kernel_ex_task_info.cc
index c34a4e9a..6da1bf63 100644
--- a/ge/graph/load/model_manager/task_info/kernel_ex_task_info.cc
+++ b/ge/graph/load/model_manager/task_info/kernel_ex_task_info.cc
@@ -192,7 +192,7 @@ void KernelExTaskInfo::InitDumpTask(void *addr, const OpDescPtr &op_desc) {
   if (davinci_model_->GetDumpProperties().IsLayerNeedDump(davinci_model_->Name(), davinci_model_->OmName(),
                                                           op_desc->GetName())) {
     dump_flag_ = RT_KERNEL_DUMPFLAG;
-    dump_args_ = input_output_addr_;
+    dump_args_ = addr;
   }
 }
 
diff --git a/ge/hybrid/node_executor/compiledsubgraph/known_node_executor.cc b/ge/hybrid/node_executor/compiledsubgraph/known_node_executor.cc
index 2bca3e06..0837ffff 100755
--- a/ge/hybrid/node_executor/compiledsubgraph/known_node_executor.cc
+++ b/ge/hybrid/node_executor/compiledsubgraph/known_node_executor.cc
@@ -126,6 +126,12 @@ Status KnownNodeTask::Init(TaskContext &context) {
     auto dump_properties = context.GetDumpProperties();
     if (dump_properties.IsDumpOpen()) {
       davinci_model_->SetDumpProperties(dump_properties);
+      void *global_step = nullptr;
+      TensorValue *varible_global_step = context.GetVariable(NODE_NAME_GLOBAL_STEP);
+      if (varible_global_step != nullptr) {
+        global_step = varible_global_step->MutableData();
+      }
+      davinci_model_->SetKnownShapeGlobalStep(global_step);
     }
     int32_t device_id = 0;
     rtError_t rt_ret = rtGetDevice(&device_id);