fix dump final ir

mindspore/ccsrc/backend/session/ascend_session.cc

@@ -483,7 +483,7 @@ GraphId AscendSession::CompileGraphImpl(NotNull<FuncGraphPtr> func_graph) {
   Load(root_graph);
   root_graph->SetInputNodes();
   root_graph->SetOptimizerFlag();
-  DumpAllGraphs(all_graphs);
+  DumpGraphs(all_graphs);
   // Save memory profiling data to proto file
 #ifndef ENABLE_SECURITY
   if (MemoryProfiling::GetInstance().IsMemoryProfilingInitialized()) {
@@ -569,7 +569,7 @@ void AscendSession::BuildGraphImpl(GraphId graph_id) {
   }
   // sync the initial const tensor to device
   SyncInitialTenosrToDevice();
-  DumpAllGraphs({graph});
+  DumpGraphs({graph});
   MS_LOG(INFO) << "End";
 }
 
@@ -1415,46 +1415,6 @@ void AscendSession::Dump(const std::shared_ptr<KernelGraph> &kernel_graph) const
 }
 #endif
 
-void AscendSession::DumpAllGraphs(const std::vector<KernelGraphPtr> &all_graphs) {
-#ifdef ENABLE_DUMP_IR
-  auto context_ptr = MsContext::GetInstance();
-  MS_EXCEPTION_IF_NULL(context_ptr);
-  bool save_graphs = context_ptr->get_param<bool>(MS_CTX_SAVE_GRAPHS_FLAG);
-  auto &json_parser = DumpJsonParser::GetInstance();
-  json_parser.Parse();
-  if (!save_graphs && !json_parser.e2e_dump_enabled() && !json_parser.async_dump_enabled() &&
-      !mindspore::RecorderManager::Instance().RdrEnable()) {
-    return;
-  }
-  auto kernel_runtime = device::KernelRuntimeManager::Instance().GetKernelRuntime(kAscendDevice, device_id_);
-  MS_EXCEPTION_IF_NULL(kernel_runtime);
-  for (auto &graph : all_graphs) {
-    MS_EXCEPTION_IF_NULL(graph);
-    std::string name = "graph_build." + std::to_string(graph->graph_id());
-    DumpGraphParams dump_params = {true, static_cast<int>(kWholeStack)};
-    (void)mindspore::RDR::RecordAnfGraph(SUBMODULE_ID, name, graph, dump_params, ".ir;.pb");
-    if (save_graphs) {
-      std::string file_name = "graph_build_" + std::to_string(graph->graph_id()) + ".ir";
-      DumpIR(file_name, graph, true, kWholeStack);
-      DumpIRProto(graph, "vm_build_" + std::to_string(graph->graph_id()));
-      DumpIR("trace_code_graph", graph, true, kWholeStack);
-    }
-    std::string final_graph = "trace_code_graph_" + std::to_string(graph->graph_id());
-    if (json_parser.e2e_dump_enabled() || json_parser.async_dump_enabled()) {
-      std::string root_dir = json_parser.path() + "/rank_" + std::to_string(rank_id_);
-      std::string target_dir = root_dir + "/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);
-      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());
-    }
-  }
-#endif
-}
-
 void AscendSession::LoadTensor(const std::shared_ptr<KernelGraph> &kernel_graph) const {
   MS_LOG(INFO) << "Start!";
   MS_EXCEPTION_IF_NULL(kernel_graph);

mindspore/ccsrc/backend/session/ascend_session.h

@@ -116,7 +116,6 @@ class AscendSession : public SessionBasic {
   void Dump(const std::shared_ptr<KernelGraph> &kernel_graph) const;
   void DumpSetup(const std::shared_ptr<KernelGraph> &kernel_graph) const;
 #endif
-  void DumpAllGraphs(const std::vector<KernelGraphPtr> &all_graphs);
   void LoadTensor(const std::shared_ptr<KernelGraph> &kernel_graph) const;
   // below functions are used for run op
   void RunOpHardwareOptimize(const std::shared_ptr<session::KernelGraph> &kernel_graph) const;

mindspore/ccsrc/backend/session/cpu_session.cc

@@ -155,7 +155,7 @@ GraphId CPUSession::CompileGraphImpl(const AnfNodePtrList &lst, const AnfNodePtr
   SetSummaryNodes(graph.get());
 #endif
   runtime_.IncreaseSummaryRefCount(graph->summary_nodes());
-  DumpGraph(graph);
+  DumpGraphs({graph});
   return graph_id;
 }
 

mindspore/ccsrc/backend/session/gpu_session.cc

@@ -507,7 +507,7 @@ GraphId GPUSession::CompileGraphImpl(const KernelGraphPtr &graph) {
     AllocateMemory(graph.get());
   }
 
-  DumpGraph(graph);
+  DumpGraphs({graph});
 
 #ifdef ENABLE_DEBUGGER
   if (debugger_ && debugger_->DebuggerBackendEnabled()) {

mindspore/ccsrc/backend/session/session_basic.cc

@@ -56,6 +56,21 @@
 #include "backend/session/pynative_task_manager.h"
 #include "pipeline/pynative/pynative_execute.h"
 #include "runtime/op_builder/op_lazy_builder.h"
+#ifdef ENABLE_DEBUGGER
+#include "debug/tensor_load.h"
+#include "debug/debugger/proto_exporter.h"
+#else
+#include "debug/debugger/proto_exporter_stub.h"
+#endif
+#ifdef ENABLE_DUMP_IR
+#include "debug/rdr/running_data_recorder.h"
+#include "debug/rdr/recorder_manager.h"
+#include "debug/rdr/graph_recorder.h"
+#endif
+#ifndef ENABLE_SECURITY
+#include "debug/data_dump/dump_json_parser.h"
+#include "debug/data_dump/e2e_dump.h"
+#endif
 
 namespace mindspore {
 namespace session {
@@ -2693,15 +2708,44 @@ void SessionBasic::FinalOptimize(const KernelGraphPtr &graph) const {
   MS_LOG(INFO) << "End FinalOptimize for graph: " << graph->graph_id();
 }
 
-void SessionBasic::DumpGraph(const std::shared_ptr<KernelGraph> &kernel_graph) {
+void SessionBasic::DumpGraphs(const std::vector<KernelGraphPtr> &graphs) {
 #ifdef ENABLE_DUMP_IR
   auto context_ptr = MsContext::GetInstance();
   MS_EXCEPTION_IF_NULL(context_ptr);
   bool save_graphs = context_ptr->get_param<bool>(MS_CTX_SAVE_GRAPHS_FLAG);
-  if (save_graphs) {
-    DumpIR("graph_build_" + std::to_string(kernel_graph->graph_id()) + ".ir", kernel_graph, true, kWholeStack);
-    DumpIRProto(kernel_graph, "vm_build_" + std::to_string(kernel_graph->graph_id()));
-    DumpIR("trace_code_graph", kernel_graph, true, kWholeStack);
+  auto &json_parser = DumpJsonParser::GetInstance();
+  json_parser.Parse();
+  if (!save_graphs && !json_parser.e2e_dump_enabled() && !json_parser.async_dump_enabled() &&
+      !mindspore::RecorderManager::Instance().RdrEnable()) {
+    return;
+  }
+  for (auto &graph : graphs) {
+    MS_EXCEPTION_IF_NULL(graph);
+    std::string name = "graph_build." + std::to_string(graph->graph_id());
+    DumpGraphParams dump_params = {true, static_cast<int>(kWholeStack)};
+    (void)mindspore::RDR::RecordAnfGraph(SUBMODULE_ID, name, graph, dump_params, ".ir;.pb");
+    if (save_graphs) {
+      std::string file_name = "graph_build_" + std::to_string(graph->graph_id()) + ".ir";
+      DumpIR(file_name, graph, true, kWholeStack);
+      DumpIRProto(graph, "vm_build_" + std::to_string(graph->graph_id()));
+      DumpIR("trace_code_graph", graph, true, kWholeStack);
+    }
+    if (context_ptr->get_param<std::string>(MS_CTX_DEVICE_TARGET) != kAscendDevice) {
+      // Here dump data only with Ascend.
+      continue;
+    }
+    std::string final_graph = "trace_code_graph_" + std::to_string(graph->graph_id());
+    if (json_parser.e2e_dump_enabled() || json_parser.async_dump_enabled()) {
+      std::string root_dir = json_parser.path() + "/rank_" + std::to_string(rank_id_);
+      std::string target_dir = root_dir + "/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);
+      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());
+    }
   }
 #endif
 }

mindspore/ccsrc/backend/session/session_basic.h

@@ -320,7 +320,7 @@ class SessionBasic : public std::enable_shared_from_this<SessionBasic> {
   void ClearAllBucket(const GraphId &graph_id);
   std::vector<uint32_t> GetAllReduceSplitIndex();
   virtual std::string GetCommWorldGroup() { return std::string(); }
-  void DumpGraph(const std::shared_ptr<KernelGraph> &kernel_graph);
+  void DumpGraphs(const std::vector<KernelGraphPtr> &graphs);
 #if ((defined ENABLE_CPU) && (!defined _WIN32) && !defined(__APPLE__))
   void CheckPSModeConsistence(const KernelGraphPtr &kernel_graph) const;
   void GetBatchElements(const AnfNodePtr &kernel_node) const;

mindspore/ccsrc/runtime/framework/graph_compiler.cc

@@ -372,6 +372,8 @@ GraphId GraphCompiler::CompileGraph(const GraphSegmentPtr &segment, const AnfNod
 
   auto graph_id = CompileGraphImpl(graph, device_context);
 
+  session_->DumpGraphs({graph});
+
   // Cache the backend graph output nodes to front nodes with output index.
   auto backend_node = graph->output();
   MS_EXCEPTION_IF_NULL(backend_node);
@@ -408,7 +410,7 @@ GraphId GraphCompiler::CompileGraph(const FuncGraphPtr &func_graph, const Device
   auto graph_id = CompileGraphImpl(root_graph, device_context);
 
   // dump all graphs.
-  device_context->DumpAllGraphs(all_graphs);
+  session_->DumpGraphs(all_graphs);
 
   // Cache the backend graph output nodes to front nodes with output index.
   auto output = func_graph->output();
@@ -496,7 +498,6 @@ GraphId GraphCompiler::CompileGraphImpl(const KernelGraphPtr &graph, const Devic
 #endif
 
   device_context->EnableRuntimeCache(graph);
-  session_->DumpGraph(graph);
   return graph->graph_id();
 }
 

mindspore/ccsrc/runtime/hardware/ascend/ascend_device_context.cc

@@ -234,44 +234,6 @@ void Dump(const KernelGraphPtr &graph, uint32_t rank_id) {
 }
 #endif
 
-void AscendDeviceContext::DumpAllGraphs(const std::vector<KernelGraphPtr> &all_graphs) const {
-#ifdef ENABLE_DUMP_IR
-  auto context_ptr = MsContext::GetInstance();
-  MS_EXCEPTION_IF_NULL(context_ptr);
-  bool save_graphs = context_ptr->get_param<bool>(MS_CTX_SAVE_GRAPHS_FLAG);
-  auto &json_parser = DumpJsonParser::GetInstance();
-  json_parser.Parse();
-  if (!save_graphs && !json_parser.e2e_dump_enabled() && !json_parser.async_dump_enabled() &&
-      !mindspore::RecorderManager::Instance().RdrEnable()) {
-    return;
-  }
-  for (auto &graph : all_graphs) {
-    MS_EXCEPTION_IF_NULL(graph);
-    std::string name = "graph_build." + std::to_string(graph->graph_id());
-    DumpGraphParams dump_params = {true, static_cast<int>(kWholeStack)};
-    (void)mindspore::RDR::RecordAnfGraph(SUBMODULE_ID, name, graph, dump_params, ".ir;.pb");
-    if (save_graphs) {
-      std::string file_name = "graph_build_" + std::to_string(graph->graph_id()) + ".ir";
-      DumpIR(file_name, graph, true, kWholeStack);
-      DumpIRProto(graph, "vm_build_" + std::to_string(graph->graph_id()));
-      DumpIR("trace_code_graph", graph, true, kWholeStack);
-    }
-    std::string final_graph = "trace_code_graph_" + std::to_string(graph->graph_id());
-    if (json_parser.e2e_dump_enabled() || json_parser.async_dump_enabled()) {
-      std::string root_dir = json_parser.path() + "/rank_" + std::to_string(rank_id_);
-      std::string target_dir = root_dir + "/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);
-      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());
-    }
-  }
-#endif
-}
-
 void AscendDeviceContext::Initialize() {
   MS_LOG(INFO) << "Status record: Enter Initialize...";
   if (initialized_) {

mindspore/ccsrc/runtime/hardware/ascend/ascend_device_context.h

@@ -129,9 +129,6 @@ class AscendDeviceContext : public DeviceContext {
   // set rt_context_ to this thread to control device
   void BindDeviceToCurrentThread() const;
 
-  // dump all graphs.
-  void DumpAllGraphs(const std::vector<KernelGraphPtr> &all_graphs) const override;
-
  private:
   // Graph loader interface
   void AllocateGraphMemory(const NotNull<KernelGraphPtr> &root_graph) const;

mindspore/ccsrc/runtime/hardware/device_context.h

@@ -27,6 +27,10 @@
 #include "backend/session/kernel_graph.h"
 #include "backend/session/anf_runtime_algorithm.h"
 #include "backend/optimizer/common/common_backend_optimization.h"
+#ifdef ENABLE_DUMP_IR
+#include "debug/anf_ir_dump.h"
+#include "debug/dump_proto.h"
+#endif
 
 namespace mindspore {
 namespace device {
@@ -155,10 +159,6 @@ class DeviceContext {
   // Return collective communication object for caller to access
   CollectiveCommunicationLib *collective_comm_lib() const { return collective_comm_lib_; }
 
-  // TODO(jiaorui): will be delete
-  // Dump all graphs.
-  virtual void DumpAllGraphs(const std::vector<KernelGraphPtr> &all_graphs) const {}
-
   void EnableRuntimeCache(const KernelGraphPtr &graph) const {
     auto node_list = graph->TopoSort(graph->get_return());
     for (auto &node : node_list) {