From 8d96ee8b464b68a9ea9c02712624a3603005f3b2 Mon Sep 17 00:00:00 2001
From: Parastoo Ashtari <parastoo.ashtari@huawei.com>
Date: Fri, 14 Jan 2022 17:56:55 -0500
Subject: [PATCH] Move DumpConstantData to DebugOnStepEnd for ascend MindRT and
 add check for kGraphOpRun

---
 mindspore/ccsrc/backend/session/session_basic.cc  |  5 ++++-
 .../ccsrc/debug/data_dump/dump_json_parser.cc     |  7 +++++++
 mindspore/ccsrc/debug/debugger/debugger.cc        | 15 +++++++++++++++
 mindspore/ccsrc/debug/debugger/debugger.h         |  2 ++
 4 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/mindspore/ccsrc/backend/session/session_basic.cc b/mindspore/ccsrc/backend/session/session_basic.cc
index 08d968d508c..8da95a0176c 100644
--- a/mindspore/ccsrc/backend/session/session_basic.cc
+++ b/mindspore/ccsrc/backend/session/session_basic.cc
@@ -2775,7 +2775,10 @@ void SessionBasic::DumpGraphs(const std::vector<KernelGraphPtr> &graphs) {
       std::string cst_file_dir = GenerateDumpPath(graph->root_graph_id(), rank_id, true);
       std::string ir_file_path = target_dir + "/" + "ms_output_" + final_graph + ".ir";
       DumpIRProtoWithSrcInfo(graph, final_graph, target_dir, kDebugWholeStack);
-      DumpConstantInfo(graph, cst_file_dir);
+      if (!MsContext::GetInstance()->get_param<bool>(MS_CTX_ENABLE_MINDRT)) {
+        // Dump constant data for old runtime ascend.
+        DumpConstantInfo(graph, cst_file_dir);
+      }
       DumpIR("trace_code_graph", graph, true, kWholeStack, ir_file_path);
       DumpGraphExeOrder("ms_execution_order_graph_" + std::to_string(graph->graph_id()) + ".csv", root_dir,
                         graph->execution_order());
diff --git a/mindspore/ccsrc/debug/data_dump/dump_json_parser.cc b/mindspore/ccsrc/debug/data_dump/dump_json_parser.cc
index cf0c82b69ae..2fadf926e69 100644
--- a/mindspore/ccsrc/debug/data_dump/dump_json_parser.cc
+++ b/mindspore/ccsrc/debug/data_dump/dump_json_parser.cc
@@ -65,10 +65,17 @@ std::string GetIfstreamString(const std::ifstream &ifstream) {
 }
 
 bool DumpJsonParser::IsDumpEnabled() {
+  auto single_op = common::GetEnv(kGraphOpRun);
   auto config_path = common::GetEnv(kMindsporeDumpConfig);
   if (config_path.empty()) {
     return false;
   }
+  // Dump is supported with Ascend kernel-by-kernel mode (mindRT) when kGraphOpRun is set.
+  if (!single_op.empty() && single_op == "1" && !MsContext::GetInstance()->get_param<bool>(MS_CTX_ENABLE_MINDRT)) {
+    MS_LOG(WARNING) << "Dump is not supported when task is not sink. Please set env GRAPH_OP_RUN to 0 to enable task "
+                       "sink, so that the data can be dumped.";
+    return false;
+  }
   MS_LOG(INFO) << "Dump config path is " << config_path;
 
   auto context = MsContext::GetInstance();
diff --git a/mindspore/ccsrc/debug/debugger/debugger.cc b/mindspore/ccsrc/debug/debugger/debugger.cc
index 841f1f6ae32..cecb377ecbd 100644
--- a/mindspore/ccsrc/debug/debugger/debugger.cc
+++ b/mindspore/ccsrc/debug/debugger/debugger.cc
@@ -441,12 +441,26 @@ void Debugger::Dump(const KernelGraphPtr &kernel_graph) const {
   if (debugger_ && debugger_->DebuggerBackendEnabled()) {
     MS_EXCEPTION_IF_NULL(kernel_graph);
     (void)E2eDump::DumpParametersData(kernel_graph.get(), rank_id, debugger_.get());
+    // Dump constant data for GPU mindRT.
     E2eDump::DumpConstantData(kernel_graph.get(), rank_id, debugger_.get());
   } else {
     DumpJsonParser::GetInstance().UpdateDumpIter();
   }
 }
 
+void Debugger::DumpConstantDataAscend(const KernelGraphPtr &graph) {
+  if (device_target_ != kAscendDevice) {
+    return;
+  }
+  auto &json_parser = DumpJsonParser::GetInstance();
+  if (json_parser.e2e_dump_enabled() || json_parser.async_dump_enabled()) {
+    // Dump constant data for ascend mindRT, for old runtime constant data is dumped in session_basic.
+    uint32_t rank_id = GetRankID();
+    std::string cst_file_dir = GenerateDumpPath(graph->root_graph_id(), rank_id, true);
+    DumpConstantInfo(graph, cst_file_dir);
+  }
+}
+
 void Debugger::DumpSingleNode(const CNodePtr &node, uint32_t graph_id) {
   if (debugger_ && debugger_->DebuggerBackendEnabled()) {
     uint32_t rank_id = GetRankID();
@@ -498,6 +512,7 @@ void Debugger::PostExecuteGraphDebugger() {
     // Dump Parameters and consts
     for (auto graph : graph_ptr_step_vec_) {
       debugger_->Dump(graph);
+      DumpConstantDataAscend(graph);
       if (!debugger_->debugger_enabled()) {
         debugger_->ClearCurrentData();
       }
diff --git a/mindspore/ccsrc/debug/debugger/debugger.h b/mindspore/ccsrc/debug/debugger/debugger.h
index 49ef0581b79..20352a76db1 100644
--- a/mindspore/ccsrc/debug/debugger/debugger.h
+++ b/mindspore/ccsrc/debug/debugger/debugger.h
@@ -99,6 +99,8 @@ class Debugger : public std::enable_shared_from_this<Debugger> {
 
   void Dump(const KernelGraphPtr &kernel_graph) const;
 
+  void DumpConstantDataAscend(const KernelGraphPtr &graph);
+
   void DumpSingleNode(const CNodePtr &node, uint32_t graph_id);
 
   void DumpSetup(const KernelGraphPtr &kernel_graph) const;