fix

ge/hybrid/node_executor/aicpu/aicpu_node_executor.cc

@@ -230,7 +230,7 @@ Status AicpuNodeTaskBase::ExecuteAsync(TaskContext &context, std::function<void(
   return SUCCESS;
 }
 
-Status AicpuNodeTaskBase::InitForDependComputeTask(bool is_tfkernel) {
+Status AicpuTfNodeTask::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);
@@ -262,11 +262,9 @@ Status AicpuNodeTaskBase::InitForDependComputeTask(bool is_tfkernel) {
                     "[Alloc][TensorBuffer] failed for Node[%s] to copy task input dst, size=%zu",
                     node_name_.c_str(), copy_input_buf_len);
   // copy task args buf
-  if (is_tfkernel) {
-    GE_CHK_STATUS_RET(AllocTensorBuffer(sizeof(STR_FWK_OP_KERNEL), copy_task_args_buf_),
-                     "[Alloc][TensorBuffer] failed for Node[%s] to copy task args, size=%zu",
-                      node_name_.c_str(), sizeof(STR_FWK_OP_KERNEL));
-  }
+  GE_CHK_STATUS_RET(AllocTensorBuffer(sizeof(STR_FWK_OP_KERNEL), copy_task_args_buf_),
+                    "[Alloc][TensorBuffer] failed for Node[%s] to copy task args, size=%zu",
+                    node_name_.c_str(), sizeof(STR_FWK_OP_KERNEL));
 
   std::vector<uint64_t> copy_io_addr;
   copy_io_addr.emplace_back(reinterpret_cast<uintptr_t>(copy_input_release_flag_dev_->GetData()));
@@ -328,7 +326,7 @@ Status AicpuTfNodeTask::Init(const HybridModel &model) {
   uint64_t ext_session_id = model.GetSessionId();
   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(true), "[Init][DependComputeTask] failed for Node[%s].", node_name_.c_str());
+  GE_CHK_STATUS_RET(InitForDependComputeTask(), "[Init][DependComputeTask] failed for Node[%s].", node_name_.c_str());
 
   // build fwk_op_kernel.
   GE_IF_BOOL_EXEC(sizeof(STR_FWK_OP_KERNEL) < kernel_ex_def.args_size(),
@@ -753,6 +751,57 @@ Status AicpuNodeTask::SetMemCopyTask(const domi::TaskDef &task_def) {
   return SUCCESS;
 }
 
+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::Init(const HybridModel &model) {
   auto node_name = node_name_;
   GELOGD("Node[%s] init start.", node_name.c_str());
@@ -846,7 +895,7 @@ 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(false),
+  GE_CHK_STATUS_RET(InitForDependComputeTask(),
                     "[Init][DependComputeTask] failed for Node[%s].",
                     node_name_.c_str());
 

ge/hybrid/node_executor/aicpu/aicpu_node_executor.h

@@ -57,7 +57,7 @@ class AicpuNodeTaskBase : public NodeTask {
   
   virtual Status UpdateShapeAndDataByResultSummary(TaskContext &context);
 
-  Status InitForDependComputeTask(bool is_tfkernel);
+  virtual Status InitForDependComputeTask() = 0;
 
   Status TaskCallback(TaskContext &context);
 
@@ -130,6 +130,8 @@ class AicpuTfNodeTask : public AicpuNodeTaskBase {
   
   Status CopyDataToHbm(TaskContext &context,
                        const std::vector<std::unique_ptr<TensorBuffer>> &out_shape_hbm) override;
+
+  Status InitForDependComputeTask() override;
  private:
   Status SetMemCopyTask(const domi::TaskDef &task_def);
 
@@ -168,6 +170,7 @@ class AicpuNodeTask : public AicpuNodeTaskBase {
 
   Status UpdateIoAddr(TaskContext &context) override;
 
+  Status InitForDependComputeTask() override;
 private:
   Status SetMemCopyTask(const domi::TaskDef &task_def);