mindspore2022/mindspore/ccsrc/device/ascend/dump/data_dumper.cc

/**
 * Copyright 2020 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#ifdef ENABLE_DATA_DUMP
#include "device/ascend/dump/data_dumper.h"

#include <map>
#include <memory>
#include <string>
#include "utility"
#include "session/anf_runtime_algorithm.h"
#include "runtime/mem.h"
#include "runtime/kernel.h"
#include "device/ascend/dump/ge_dump.h"
#include "proto/op_mapping_info.pb.h"
#include "utils/context/ms_context.h"
#include "debug/data_dump_parser.h"

constexpr uint32_t kAicpuLoadFlag = 1;
constexpr uint32_t kAicpuUnloadFlag = 0;
constexpr uint32_t kTupleTaskId = 0;
constexpr uint32_t kTupleStreamId = 1;
constexpr uint32_t kTupleArgs = 2;
constexpr uint32_t kCurrentStepTensorIndex = 0;
constexpr uint32_t kCurrentEpochTensorIndex = 1;
constexpr uint32_t kStepsPerEpochTensorIndex = 2;

namespace mindspore {
namespace device {
namespace ascend {
void DumpKernelOutput(const CNodePtr &kernel, void *args, NotNull<aicpu::dump::Task *> task);
void DumpKernelInput(const CNodePtr &kernel, void *args, NotNull<aicpu::dump::Task *> task);
void RtLoadDumpData(const aicpu::dump::OpMappingInfo &dump_info, void **ptr);

DataDumper::~DataDumper() {
  ReleaseDevMem(&dev_load_mem_);
  ReleaseDevMem(&dev_unload_mem_);
}

void DataDumper::LoadDumpInfo() {
  MS_LOG(INFO) << "[DataDump] LoadDumpInfo start";
  MS_EXCEPTION_IF_NULL(kernel_graph_);
  aicpu::dump::OpMappingInfo dump_info;
  SetOpMappingInfo(NOT_NULL(&dump_info));

  auto kernels = kernel_graph_->execution_order();
  for (const auto &kernel : kernels) {
    MS_EXCEPTION_IF_NULL(kernel);
    if (!KernelNeedDump(kernel)) {
      continue;
    }
    MS_LOG(INFO) << "[DataDump] LoadDumpInfo kernel:" << kernel->fullname_with_scope();
    dump_kernel_names_.emplace_back(kernel->fullname_with_scope());

    aicpu::dump::Task task;
    ConstructDumpTask(NOT_NULL(kernel), NOT_NULL(&task));
    MS_EXCEPTION_IF_NULL(dump_info.mutable_task());
    dump_info.mutable_task()->Add(std::move(task));
  }
  RtLoadDumpData(dump_info, &dev_load_mem_);
  load_flag_ = true;
  MS_LOG(INFO) << "[DataDump] LoadDumpInfo end";
}

void DataDumper::SetOpMappingInfo(NotNull<aicpu::dump::OpMappingInfo *> dump_info) const {
  auto context_ptr = MsContext::GetInstance();
  MS_EXCEPTION_IF_NULL(context_ptr);
  MS_EXCEPTION_IF_NULL(kernel_graph_);
  auto dump_path = DataDumpParser::GetInstance().GetDumpPath();
  if (!dump_path.has_value()) {
    MS_LOG(EXCEPTION) << "Dump path invalid";
  }
  auto device_id = context_ptr->device_id();
  dump_info->set_dump_path(dump_path.value() + "_" + std::to_string(device_id) + "/");
  MS_LOG(INFO) << "[DataDump] dump_path:" << dump_path.value();

  dump_info->set_model_name(DataDumpParser::GetInstance().net_name() + "_" + std::to_string(kernel_graph_->graph_id()));
  dump_info->set_dump_step(std::to_string(DataDumpParser::GetInstance().dump_step()));
  dump_info->set_model_id(kernel_graph_->graph_id());
  dump_info->set_flag(kAicpuLoadFlag);

  const auto &input_ctrl_tensors = kernel_graph_->input_ctrl_tensors();
  if (input_ctrl_tensors == nullptr || input_ctrl_tensors->size() < 3) {
    MS_LOG(INFO) << "[DataDump] Not data sink mode, input_ctrl_tensor";
    return;
  }
  const auto &current_step_tensor = input_ctrl_tensors->at(kCurrentStepTensorIndex);
  const auto &currnet_epoch_tensor = input_ctrl_tensors->at(kCurrentEpochTensorIndex);
  const auto &steps_per_epoch_tensor = input_ctrl_tensors->at(kStepsPerEpochTensorIndex);

  MS_EXCEPTION_IF_NULL(current_step_tensor);
  MS_EXCEPTION_IF_NULL(currnet_epoch_tensor);
  MS_EXCEPTION_IF_NULL(steps_per_epoch_tensor);
  MS_EXCEPTION_IF_NULL(current_step_tensor->device_address());
  MS_EXCEPTION_IF_NULL(currnet_epoch_tensor->device_address());
  MS_EXCEPTION_IF_NULL(steps_per_epoch_tensor->device_address());

  void *current_step = current_step_tensor->device_address()->ptr_;
  void *current_epoch = currnet_epoch_tensor->device_address()->ptr_;
  void *steps_per_epoch = steps_per_epoch_tensor->device_address()->ptr_;

  if (current_epoch != nullptr && current_step != nullptr && steps_per_epoch != nullptr) {
    dump_info->set_step_id_addr(reinterpret_cast<uint64_t>(current_epoch));
    dump_info->set_loop_cond_addr(reinterpret_cast<uint64_t>(current_step));
    dump_info->set_iterations_per_loop_addr(reinterpret_cast<uint64_t>(steps_per_epoch));
  } else {
    MS_LOG(INFO) << "Invalid ctrl tensor device address";
  }
}

bool DataDumper::KernelNeedDump(const CNodePtr &kernel) const {
  if (AnfAlgo::GetKernelType(kernel) != TBE_KERNEL && AnfAlgo::GetKernelType(kernel) != AICPU_KERNEL &&
      AnfAlgo::GetKernelType(kernel) != AKG_KERNEL) {
    return false;
  }
  MS_EXCEPTION_IF_NULL(kernel);
  const auto &kernel_set = DataDumpParser::GetInstance().kernel_set();
  return kernel_set.find(kernel->fullname_with_scope()) != kernel_set.end();
}

void DataDumper::UnloadDumpInfo() {
  if (!load_flag_) {
    MS_LOG(WARNING) << "Load not success, no need to unload";
    return;
  }
  MS_EXCEPTION_IF_NULL(kernel_graph_);
  MS_LOG(INFO) << "[DataDump] UnloadDumpInfo start. graphId:" << kernel_graph_->graph_id();

  aicpu::dump::OpMappingInfo op_mapping_info;
  op_mapping_info.set_model_id(kernel_graph_->graph_id());
  op_mapping_info.set_flag(kAicpuUnloadFlag);

  for (const auto &kernel_name : dump_kernel_names_) {
    aicpu::dump::Task task;
    auto iter = runtime_info_map_.find(kernel_name);
    if (iter == runtime_info_map_.end()) {
      MS_LOG(EXCEPTION) << "[DataDump] kernel name not found in runtime_info_map";
    }
    MS_EXCEPTION_IF_NULL(iter->second);
    auto task_id = std::get<kTupleTaskId>(*iter->second);
    task.set_task_id(task_id);
    MS_EXCEPTION_IF_NULL(op_mapping_info.mutable_task());
    op_mapping_info.mutable_task()->Add(std::move(task));
  }

  RtLoadDumpData(op_mapping_info, &dev_unload_mem_);
}

void DataDumper::ReleaseDevMem(void **ptr) const {
  if (ptr == nullptr) {
    return;
  }
  if (*ptr != nullptr) {
    rtError_t rt_error = rtFree(*ptr);
    if (rt_error != RT_ERROR_NONE) {
      MS_LOG(ERROR) << "[DataDump] Call rtFree failed, ret:" << rt_error;
    }
    *ptr = nullptr;
  }
}

void DataDumper::ConstructDumpTask(NotNull<const CNodePtr &> kernel, NotNull<aicpu::dump::Task *> dump_task) const {
  dump_task->set_end_graph(false);
  auto iter = runtime_info_map_.find(kernel->fullname_with_scope());
  if (iter == runtime_info_map_.end()) {
    MS_LOG(EXCEPTION) << "[DataDump] kernel name not found in runtime_info_map";
  }
  MS_EXCEPTION_IF_NULL(iter->second);
  auto task_id = std::get<kTupleTaskId>(*iter->second);
  auto stream_id = std::get<kTupleStreamId>(*iter->second);
  auto args = std::get<kTupleArgs>(*iter->second);
  MS_LOG(INFO) << "[DataDump] Get runtime info task_id:" << task_id << " stream_id:" << stream_id;

  dump_task->set_task_id(task_id);
  dump_task->set_stream_id(stream_id);
  MS_EXCEPTION_IF_NULL(dump_task->mutable_op());
  dump_task->mutable_op()->set_op_name(kernel->fullname_with_scope());
  dump_task->mutable_op()->set_op_type(AnfAlgo::GetCNodeName(kernel.get()));

  DumpKernelOutput(kernel, args, dump_task);
  DumpKernelInput(kernel, args, dump_task);
}

void RtLoadDumpData(const aicpu::dump::OpMappingInfo &dump_info, void **ptr) {
  std::string proto_str;
  size_t proto_size = dump_info.ByteSizeLong();
  bool ret = dump_info.SerializeToString(&proto_str);
  if (!ret || proto_size == 0) {
    MS_LOG(EXCEPTION) << "[DataDump] Protobuf SerializeToString failed, proto size %zu.";
  }

  rtError_t rt_ret = rtMalloc(ptr, proto_size, RT_MEMORY_HBM);
  if (rt_ret != RT_ERROR_NONE) {
    MS_LOG(EXCEPTION) << "[DataDump] Call rtMalloc failed";
  }

  if (ptr == nullptr) {
    MS_LOG(ERROR) << "[DataDump] rtMalloc failed, ptr is nullptr";
    return;
  }
  rt_ret = rtMemcpy(*ptr, proto_size, proto_str.c_str(), proto_size, RT_MEMCPY_HOST_TO_DEVICE);
  if (rt_ret != RT_ERROR_NONE) {
    MS_LOG(EXCEPTION) << "[DataDump] Call rtMemcpy failed";
  }

  MS_LOG(INFO) << "[DataDump] rtDatadumpInfoLoad start";
  rt_ret = rtDatadumpInfoLoad(*ptr, proto_size);
  if (rt_ret != RT_ERROR_NONE) {
    MS_LOG(EXCEPTION) << "[DataDump] Call rtDatadumpInfoLoad failed";
  }
}

void DumpKernelOutput(const CNodePtr &kernel, void *args, NotNull<aicpu::dump::Task *> task) {
  MS_LOG(INFO) << "[DataDump] DumpKernelOutput start. Kernel:" << kernel->fullname_with_scope();
  auto input_size = AnfAlgo::GetInputTensorNum(kernel);
  auto output_size = AnfAlgo::GetOutputTensorNum(kernel);
  uint64_t offset = sizeof(void *) * input_size;
  for (size_t i = 0; i < output_size; ++i) {
    auto data_type = AnfAlgo::GetOutputDeviceDataType(kernel, i);
    auto output_format = AnfAlgo::GetOutputFormat(kernel, i);
    auto output_shape = AnfAlgo::GetOutputDeviceShape(kernel, i);

    aicpu::dump::Output output;
    output.set_data_type(GetGeDataType(data_type));
    output.set_format(GetGeFormat(output_format, output_shape.size()));
    MS_EXCEPTION_IF_NULL(output.mutable_shape());
    for (auto dim : output_shape) {
      output.mutable_shape()->add_dim(dim);
    }
    output.set_original_output_format(GetGeFormat(output_format, output_shape.size()));
    output.set_address(static_cast<uint64_t>(reinterpret_cast<uintptr_t>(args)) + offset);
    MS_EXCEPTION_IF_NULL(task->mutable_output());
    task->mutable_output()->Add(std::move(output));
    offset += sizeof(void *);
  }
}

void DumpKernelInput(const CNodePtr &kernel, void *args, NotNull<aicpu::dump::Task *> task) {
  MS_LOG(INFO) << "[DataDump] DumpKernelInput start. Kernel:" << kernel->fullname_with_scope();
  auto input_size = AnfAlgo::GetInputTensorNum(kernel);
  uint64_t offset = 0;
  for (size_t i = 0; i < input_size; ++i) {
    aicpu::dump::Input input;
    auto input_node_with_index = AnfAlgo::GetPrevNodeOutput(kernel, i);
    auto input_node = input_node_with_index.first;
    auto input_index = input_node_with_index.second;
    std::string output_format = AnfAlgo::GetOutputFormat(input_node, input_index);
    auto output_type = AnfAlgo::GetOutputDeviceDataType(input_node, input_index);
    if (output_type == kTypeUnknown) {
      MS_LOG(WARNING) << "[DataDump] It is not suggested to use a lonely weight parameter as the output of graph";
      output_type = AnfAlgo::GetOutputInferDataType(input_node, input_index);
    }
    auto output_shape = AnfAlgo::GetOutputDeviceShape(input_node, input_index);

    input.set_data_type(GetGeDataType(output_type));
    input.set_format(GetGeFormat(output_format, output_shape.size()));
    MS_EXCEPTION_IF_NULL(input.mutable_shape());
    for (auto dim : output_shape) {
      input.mutable_shape()->add_dim(dim);
    }
    input.set_address(static_cast<uint64_t>(reinterpret_cast<uintptr_t>(args)) + offset);
    MS_EXCEPTION_IF_NULL(task->mutable_input());
    task->mutable_input()->Add(std::move(input));
    offset += sizeof(void *);
  }
}
}  // namespace ascend
}  // namespace device
}  // namespace mindspore
#endif