From e5762928620e8cd6e60eb8c074b4ec15dd9b8689 Mon Sep 17 00:00:00 2001
From: TinaMengtingZhang <tina.mengting.zhang@huawei.com>
Date: Mon, 14 Feb 2022 17:48:36 -0500
Subject: [PATCH] dump statistic file in kernel by kernel mode

---
 mindspore/ccsrc/debug/data_dump/e2e_dump.cc | 36 +++++++++++----------
 tests/st/dump/test_data_dump.py             | 18 +++++++++++
 2 files changed, 37 insertions(+), 17 deletions(-)

diff --git a/mindspore/ccsrc/debug/data_dump/e2e_dump.cc b/mindspore/ccsrc/debug/data_dump/e2e_dump.cc
index a5ea3fadd9e..3bf073f0972 100644
--- a/mindspore/ccsrc/debug/data_dump/e2e_dump.cc
+++ b/mindspore/ccsrc/debug/data_dump/e2e_dump.cc
@@ -182,17 +182,18 @@ void E2eDump::DumpOutputImpl(const CNodePtr &node, bool trans_flag, const std::s
     std::string file_path = dump_path + '/' + op_type + '.' + op_name + '.' + std::to_string(task_id) + '.' +
                             std::to_string(stream_id) + '.' + std::to_string(timestamp) + ".output." +
                             std::to_string(j);
-    if (IsDeviceTargetGPU()) {
-      if (DumpJsonParser::GetInstance().IsStatisticDump()) {
-        TensorStatDump stat_dump(op_type, op_name, task_id, stream_id, timestamp, false, j, j);
-        stat_dump.DumpTensorStatsToFile(GetKernelNodeName(node), dump_path, debugger);
-      }
-      if (DumpJsonParser::GetInstance().IsTensorDump()) {
+    if (DumpJsonParser::GetInstance().IsStatisticDump() &&
+        (IsDeviceTargetGPU() || Debugger::GetInstance()->GetAscendKernelByKernelFlag())) {
+      TensorStatDump stat_dump(op_type, op_name, task_id, stream_id, timestamp, false, j, j);
+      stat_dump.DumpTensorStatsToFile(GetKernelNodeName(node), dump_path, debugger);
+    }
+    if (DumpJsonParser::GetInstance().IsTensorDump()) {
+      if (IsDeviceTargetGPU()) {
         DumpGPUMemToFile(file_path, GetKernelNodeName(node), *addr, int_shapes, type, device_type, trans_flag, j,
                          debugger);
+      } else {
+        DumpMemToFile(file_path, *addr, int_shapes, type, trans_flag);
       }
-    } else {
-      DumpMemToFile(file_path, *addr, int_shapes, type, trans_flag);
     }
   }
 }
@@ -249,7 +250,7 @@ void E2eDump::DumpInputImpl(const CNodePtr &node, bool trans_flag, const std::st
 
     std::string tensor_name = GetKernelNodeName(node);
     size_t slot = j;
-    if (IsDeviceTargetGPU()) {
+    if (IsDeviceTargetGPU() || Debugger::GetInstance()->GetAscendKernelByKernelFlag()) {
       auto input_kernel = node->input(j + 1);
       std::string input_kernel_name = GetKernelNodeName(input_kernel);
       tensor_name = input_kernel_name;
@@ -267,16 +268,17 @@ void E2eDump::DumpInputImpl(const CNodePtr &node, bool trans_flag, const std::st
     std::string file_path = dump_path + '/' + op_type + '.' + op_name + '.' + std::to_string(task_id) + '.' +
                             std::to_string(stream_id) + '.' + std::to_string(timestamp) + ".input." + std::to_string(j);
     MS_EXCEPTION_IF_NULL(addr);
-    if (IsDeviceTargetGPU()) {
-      if (DumpJsonParser::GetInstance().IsStatisticDump()) {
-        TensorStatDump stat_dump(op_type, op_name, task_id, stream_id, timestamp, true, j, slot);
-        stat_dump.DumpTensorStatsToFile(tensor_name, dump_path, debugger);
-      }
-      if (DumpJsonParser::GetInstance().IsTensorDump()) {
+    if (DumpJsonParser::GetInstance().IsStatisticDump() &&
+        (IsDeviceTargetGPU() || Debugger::GetInstance()->GetAscendKernelByKernelFlag())) {
+      TensorStatDump stat_dump(op_type, op_name, task_id, stream_id, timestamp, true, j, slot);
+      stat_dump.DumpTensorStatsToFile(tensor_name, dump_path, debugger);
+    }
+    if (DumpJsonParser::GetInstance().IsTensorDump()) {
+      if (IsDeviceTargetGPU()) {
         DumpGPUMemToFile(file_path, tensor_name, *addr, int_shapes, type, device_type, trans_flag, slot, debugger);
+      } else {
+        DumpMemToFile(file_path, *addr, int_shapes, type, trans_flag);
       }
-    } else {
-      DumpMemToFile(file_path, *addr, int_shapes, type, trans_flag);
     }
   }
 }
diff --git a/tests/st/dump/test_data_dump.py b/tests/st/dump/test_data_dump.py
index 73a2b1546ea..eb83401a900 100644
--- a/tests/st/dump/test_data_dump.py
+++ b/tests/st/dump/test_data_dump.py
@@ -555,6 +555,24 @@ def test_ascend_statistic_dump():
     context.set_context(mode=context.GRAPH_MODE, device_target="Ascend")
     run_saved_data_dump_test('test_async_dump', 'statistic')
 
+
+@pytest.mark.level0
+@pytest.mark.platform_arm_ascend_training
+@pytest.mark.platform_x86_ascend_training
+@pytest.mark.env_onecard
+@security_off_wrap
+def test_ascend_statistic_dump_kernel_by_kernel():
+    """
+    Feature: Ascend Statistics Dump in kernel by kernel (mindRT) mode
+    Description: Test Ascend statistics dump
+    Expectation: Statistics are stored in statistic.csv files
+    """
+    # set env `GRAPH_OP_RUN`` to enable kernel-by-kernel mode.
+    os.environ['GRAPH_OP_RUN'] = "1"
+    context.set_context(mode=context.GRAPH_MODE, device_target="Ascend")
+    run_saved_data_dump_test('test_async_dump', 'statistic')
+    del os.environ['GRAPH_OP_RUN']
+
 @pytest.mark.level0
 @pytest.mark.platform_arm_ascend_training
 @pytest.mark.platform_x86_ascend_training