fix

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<uint64_t> copy_io_addr;
+  copy_io_addr.emplace_back(reinterpret_cast<uintptr_t>(copy_input_release_flag_dev_->GetData()));
+  copy_io_addr.emplace_back(reinterpret_cast<uintptr_t>(copy_input_data_size_dev_->GetData()));
+  copy_io_addr.emplace_back(reinterpret_cast<uintptr_t>(copy_input_src_dev_->GetData()));
+  copy_io_addr.emplace_back(reinterpret_cast<uintptr_t>(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,
+                         &copy_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<uintptr_t>(copy_ioaddr_dev_->GetData());
+  aicpu_task.fwkKernelBase.fwk_kernel.workspaceBaseAddr = reinterpret_cast<uintptr_t>(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<uintptr_t>(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;
 }

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<uint8_t[]> args_;