diff --git a/ge/hybrid/node_executor/aicpu/aicpu_node_executor.cc b/ge/hybrid/node_executor/aicpu/aicpu_node_executor.cc index c2ebf654..90a8d3ef 100755 --- a/ge/hybrid/node_executor/aicpu/aicpu_node_executor.cc +++ b/ge/hybrid/node_executor/aicpu/aicpu_node_executor.cc @@ -669,14 +669,104 @@ Status AicpuTfNodeTask::TaskCallback(TaskContext &context) { return callback_ret; } +Status AicpuNodeTask::InitForDependComputeTask() { + if ((unknown_type_ != DEPEND_COMPUTE) || (node_item_->num_outputs == 0)) { + GELOGD("Node[%s] type[%s] unknown_type is %d, output num is %d.", + node_name_.c_str(), node_item_->node_type.c_str(), unknown_type_, node_item_->num_outputs); + return SUCCESS; + } + + output_summary_.resize(node_item_->num_outputs); + constexpr auto result_summary_size = sizeof(aicpu::FWKAdapter::ResultSummary); + for (auto i = 0; i < node_item_->num_outputs; ++i) { + GE_CHK_STATUS_RET(AllocTensorBuffer(result_summary_size, output_summary_[i]), + "[Alloc][TensorBuffer] failed for Node[%s] to copy result summary info, size=%zu.", + node_name_.c_str(), result_summary_size); + } + output_summary_host_.resize(node_item_->num_outputs); + + // init for mem copy task + // copy task need copy output_data and output_shape, max len is 2 * output_num + const size_t copy_input_buf_len = node_item_->num_outputs * 2 * sizeof(uint64_t); + GE_CHK_STATUS_RET(AllocTensorBuffer(copy_input_buf_len, copy_input_release_flag_dev_), + "[Alloc][TensorBuffer] failed for Node[%s] to copy task input release_flag, size=%zu", + node_name_.c_str(), copy_input_buf_len); + GE_CHK_STATUS_RET(AllocTensorBuffer(copy_input_buf_len, copy_input_data_size_dev_), + "[Alloc][TensorBuffer] failed for Node[%s] to copy task input data_size, size=%zu", + node_name_.c_str(), copy_input_buf_len); + GE_CHK_STATUS_RET(AllocTensorBuffer(copy_input_buf_len, copy_input_src_dev_), + "[Alloc][TensorBuffer] failed for Node[%s] to copy task input src, size=%zu", + node_name_.c_str(), copy_input_buf_len); + GE_CHK_STATUS_RET(AllocTensorBuffer(copy_input_buf_len, copy_input_dst_dev_), + "[Alloc][TensorBuffer] failed for Node[%s] to copy task input dst, size=%zu", + node_name_.c_str(), copy_input_buf_len); + + std::vector copy_io_addr; + copy_io_addr.emplace_back(reinterpret_cast(copy_input_release_flag_dev_->GetData())); + copy_io_addr.emplace_back(reinterpret_cast(copy_input_data_size_dev_->GetData())); + copy_io_addr.emplace_back(reinterpret_cast(copy_input_src_dev_->GetData())); + copy_io_addr.emplace_back(reinterpret_cast(copy_input_dst_dev_->GetData())); + + // mem copy op has 4 inputs and 0 output. + const auto copy_io_addr_size = sizeof(uint64_t) * copy_io_addr.size(); + + // can alloc in init, it can reuse + GE_CHK_STATUS_RET(AllocTensorBuffer(copy_io_addr_size, copy_ioaddr_dev_), + "[Alloc][TensorBuffer] failed for Node[%s] to copy task ioaddr, size=%zu", + node_name_.c_str(), copy_io_addr_size); + + GE_CHK_RT_RET(rtMemcpy(copy_ioaddr_dev_->GetData(), copy_io_addr_size, + ©_io_addr[0], copy_io_addr_size, RT_MEMCPY_HOST_TO_DEVICE)); + return SUCCESS; +} + +Status AicpuNodeTask::SetMemCopyTask(const domi::TaskDef &task_def) { + if (node_item_->num_outputs == 0) { + GELOGD("Node[%s] type[%s] has no output, no need set mem_copy task.", + node_name_.c_str(), node_item_->node_type.c_str()); + return SUCCESS; + } + + GELOGD("Start to set memcpy task for node[%s].", node_name_.c_str()); + const domi::KernelExDef &kernel_def = task_def.kernel_ex(); + if (kernel_def.args_size() > sizeof(STR_FWK_OP_KERNEL)) { + GELOGE(PARAM_INVALID, "[Check][Size]sizeof STR_FWK_OP_KERNEL is:%lu, but args_size:%d is bigger", + sizeof(STR_FWK_OP_KERNEL), kernel_def.args_size()); + REPORT_INNER_ERROR("E19999", "sizeof STR_FWK_OP_KERNEL is:%lu, but args_size:%d is bigger.", + sizeof(STR_FWK_OP_KERNEL), kernel_def.args_size()); + return PARAM_INVALID; + } + STR_FWK_OP_KERNEL aicpu_task = {0}; + auto sec_ret = memcpy_s(&aicpu_task, sizeof(STR_FWK_OP_KERNEL), + kernel_def.args().data(), kernel_def.args_size()); + if (sec_ret != EOK) { + GELOGE(FAILED, "[Update][aicpu_task] failed, ret: %d", sec_ret); + REPORT_CALL_ERROR("E19999", "update aicpu_task failed, ret: %d.", sec_ret); + return FAILED; + } + + GE_CHK_STATUS_RET(AllocTensorBuffer(kernel_def.task_info_size(), copy_workspace_buf_), + "[Alloc][TensorBuffer] for Node[%s] to copy task workspace buf, size=%u.", + node_name_.c_str(), kernel_def.task_info_size()); + + GE_CHK_RT_RET(rtMemcpy(copy_workspace_buf_->GetData(), kernel_def.task_info_size(), + kernel_def.task_info().data(), kernel_def.task_info_size(), RT_MEMCPY_HOST_TO_DEVICE)); + + aicpu_task.fwkKernelBase.fwk_kernel.inputOutputAddr = reinterpret_cast(copy_ioaddr_dev_->GetData()); + aicpu_task.fwkKernelBase.fwk_kernel.workspaceBaseAddr = reinterpret_cast(copy_workspace_buf_->GetData()); + aicpu_task.fwkKernelBase.fwk_kernel.extInfoAddr = 0; + aicpu_task.fwkKernelBase.fwk_kernel.extInfoLen = 0; + + GE_CHK_RT_RET(rtMemcpy(copy_task_args_buf_->GetData(), sizeof(STR_FWK_OP_KERNEL), + &aicpu_task, sizeof(STR_FWK_OP_KERNEL), RT_MEMCPY_HOST_TO_DEVICE)); + GELOGD("Set memcpy task for node[%s] successfully.", node_name_.c_str()); + return SUCCESS; +} + Status AicpuNodeTask::Init(const HybridModel &model) { auto node_name = node_name_; GELOGD("Node[%s] init start.", node_name.c_str()); - GE_CHK_BOOL_RET_STATUS(unknown_type_ != DEPEND_COMPUTE, FAILED, - "[Check][Type]Node[%s] unknown type[%d] is depend compute, it's not supported now.", - node_name.c_str(), unknown_type_); - GE_CHK_BOOL_RET_STATUS(task_def_.has_kernel(), FAILED, "[Check][task_def_]Node[%s] task def does not has kernel.", node_name.c_str()); auto &kernel_def = task_def_.kernel(); @@ -766,6 +856,10 @@ Status AicpuNodeTask::Init(const HybridModel &model) { GE_CHK_STATUS_RET(InitExtInfo(kernel_ext_info, ext_session_id), "[Init][ExtInfo] failed for Node[%s].", node_name.c_str()); + GE_CHK_STATUS_RET(InitForDependComputeTask(), + "[Init][DependComputeTask] failed for Node[%s].", + node_name_.c_str()); + if (ext_info_addr_dev_ == nullptr) { aicpu_param_head->extInfoLength = 0; aicpu_param_head->extInfoAddr = 0; @@ -774,6 +868,11 @@ Status AicpuNodeTask::Init(const HybridModel &model) { aicpu_param_head->extInfoAddr = reinterpret_cast(ext_info_addr_dev_->GetData()); } + auto task_defs = model.GetTaskDefs(node_item_->node); + GE_CHECK_NOTNULL(task_defs); + if (unknown_type_ == DEPEND_COMPUTE) { + GE_CHK_STATUS_RET_NOLOG(SetMemCopyTask((*task_defs)[1])); + } GELOGD("Node[%s] init end.", node_name.c_str()); return SUCCESS; } diff --git a/ge/hybrid/node_executor/aicpu/aicpu_node_executor.h b/ge/hybrid/node_executor/aicpu/aicpu_node_executor.h index 7577d486..02180b1d 100644 --- a/ge/hybrid/node_executor/aicpu/aicpu_node_executor.h +++ b/ge/hybrid/node_executor/aicpu/aicpu_node_executor.h @@ -166,6 +166,11 @@ class AicpuNodeTask : public AicpuNodeTaskBase { Status UpdateIoAddr(TaskContext &context) override; + private: + Status SetMemCopyTask(const domi::TaskDef &task_def); + + Status InitForDependComputeTask(); + protected: // host mem std::unique_ptr args_;