fix

ge/hybrid/node_executor/aicpu/aicpu_node_executor.cc

@@ -446,8 +446,8 @@ Status AicpuTfNodeTask::EnsureSessionCreated(uint64_t session_id) {
   return SUCCESS;
 }
 
-Status AicpuTfNodeTask::ReadResultSummaryAndPrepareMemory(TaskContext &context,
-                                                          std::vector<std::unique_ptr<TensorBuffer>> &out_shape_hbm) {
+Status AicpuNodeTaskBase::ReadResultSummaryAndPrepareMemory(TaskContext &context,
+                                                            std::vector<std::unique_ptr<TensorBuffer>> &out_shape_hbm) {
   for (auto i = 0; i < node_item_->num_outputs; ++i) {
     auto &result_summary = output_summary_host_[i];
     GE_CHK_RT_RET(rtMemcpy(&result_summary, sizeof(aicpu::FWKAdapter::ResultSummary),
@@ -484,6 +484,32 @@ Status AicpuTfNodeTask::CopyDataToHbm(TaskContext &context,
   RECORD_CALLBACK_EVENT(context.GetExecutionContext(), node_name_.c_str(), "[LaunchCopy] Start");
   GE_CHK_RT_RET(rtKernelLaunchEx(copy_task_args_buf_->GetData(), sizeof(STR_FWK_OP_KERNEL),
                                  RT_KERNEL_DEFAULT, context.GetStream()));
+
+  RECORD_CALLBACK_EVENT(context.GetExecutionContext(), node_name_.c_str(), "[LaunchCopy] End");
+
+  GE_CHK_RT_RET(rtStreamSynchronize(context.GetStream()));
+  RECORD_CALLBACK_EVENT(context.GetExecutionContext(), node_name_.c_str(), "[SynchronizeCopy] End");
+  return SUCCESS;
+}
+
+Status AicpuNodeTask::CopyDataToHbm(TaskContext &context,
+                                    const std::vector<std::unique_ptr<TensorBuffer>> &out_shape_hbm) {
+  GE_CHK_BOOL_RET_STATUS(out_shape_hbm.size() == static_cast<std::size_t>(node_item_->num_outputs),
+                         INTERNAL_ERROR,
+                         "[Check][Size]Node[%s] has %d outputs but out shape is %zu not equal.",
+                         node_name_.c_str(), node_item_->num_outputs, out_shape_hbm.size());
+
+  GE_CHK_STATUS_RET_NOLOG(PrepareCopyInputs(context, out_shape_hbm));
+
+  RECORD_CALLBACK_EVENT(context.GetExecutionContext(), node_name_.c_str(), "[LaunchCopy] Start");
+  uint32_t flag = RT_KERNEL_DEFAULT;
+  auto rt_ret = rtCpuKernelLaunchWithFlag(reinterpret_cast<const void *>(memcpy_so_name.c_str()),
+                                          reinterpret_cast<const void *>(memcpy_kernel_name.c_str()),
+                                          1, // default core dim is 1
+                                          memcpy_args_.get(), memcpy_args_size_,
+                                          nullptr, context.GetStream(), flag);
+  GE_CHK_RT_RET(rt_ret);
+
   RECORD_CALLBACK_EVENT(context.GetExecutionContext(), node_name_.c_str(), "[LaunchCopy] End");
 
   GE_CHK_RT_RET(rtStreamSynchronize(context.GetStream()));
@@ -491,8 +517,8 @@ Status AicpuTfNodeTask::CopyDataToHbm(TaskContext &context,
   return SUCCESS;
 }
 
-Status AicpuTfNodeTask::PrepareCopyInputs(const TaskContext &context,
-                                          const std::vector<std::unique_ptr<TensorBuffer>> &out_shape_hbm) {
+Status AicpuNodeTaskBase::PrepareCopyInputs(const TaskContext &context,
+                                            const std::vector<std::unique_ptr<TensorBuffer>> &out_shape_hbm) {
   std::vector<uint64_t> copy_input_release_flag;
   std::vector<uint64_t> copy_input_data_size;
   std::vector<uint64_t> copy_input_src;
@@ -533,8 +559,8 @@ Status AicpuTfNodeTask::PrepareCopyInputs(const TaskContext &context,
   return SUCCESS;
 }
 
-Status AicpuTfNodeTask::UpdateShapeByHbmBuffer(TaskContext &context,
-                                               const std::vector<std::unique_ptr<TensorBuffer>> &out_shape_hbm) {
+Status AicpuNodeTaskBase::UpdateShapeByHbmBuffer(TaskContext &context,
+                                                 const std::vector<std::unique_ptr<TensorBuffer>> &out_shape_hbm) {
   GE_CHK_BOOL_RET_STATUS(out_shape_hbm.size() == static_cast<std::size_t>(node_item_->num_outputs),
                          INTERNAL_ERROR,
                          "Node[%s] has %d outputs but out shape is %zu",
@@ -565,7 +591,7 @@ Status AicpuTfNodeTask::UpdateShapeByHbmBuffer(TaskContext &context,
   return SUCCESS;
 }
 
-Status AicpuTfNodeTask::UpdateShapeAndDataByResultSummary(TaskContext &context) {
+Status AicpuNodeTaskBase::UpdateShapeAndDataByResultSummary(TaskContext &context) {
   GELOGD("Node[%s] update shape and data by result summary begin.", node_name_.c_str());
 
   std::vector<std::unique_ptr<TensorBuffer>> out_shape_hbm;
@@ -680,51 +706,52 @@ Status AicpuNodeTask::SetMemCopyTask(const domi::TaskDef &task_def) {
   GELOGD("Start to set memcpy task for node[%s].", node_name_.c_str());
   const domi::KernelDef &kernel_def = task_def.kernel();
   auto &memcpy_args = kernel_def.args();
-  uint32_t memcpy_args_size = kernel_def.args_size();
-  GE_IF_BOOL_EXEC(memcpy_args.size() != memcpy_args_size,
+  memcpy_args_size_ = kernel_def.args_size();
+  memcpy_so_name_ = kernel_def.so_name();
+  memcpy_kernel_name_ = kernel_def.kernel_name();
+  GE_IF_BOOL_EXEC(memcpy_args.size() != memcpy_args_size_,
                   REPORT_INNER_ERROR("E19999", "MemCopy task def args.size=%zu, but args_size=%u not equal.",
-                                     memcpy_args.size(), memcpy_args_size);
+                                     memcpy_args.size(), memcpy_args_size_);
                   GELOGE(FAILED, "[Check][Size]MemCopy task def args.size=%zu, but args_size=%u not equal.",
-                         memcpy_args.size(), memcpy_args_size);
+                         memcpy_args.size(), memcpy_args_size_);
                   return FAILED;);
-  GE_IF_BOOL_EXEC(memcpy_args_size < sizeof(aicpu::AicpuParamHead),
+  GE_IF_BOOL_EXEC(memcpy_args_size_ < sizeof(aicpu::AicpuParamHead),
                   REPORT_INNER_ERROR("E19999",
                                      "Task def args_size=%u is less than aicpu param head len=%zu.",
-                                     memcpy_args_size, sizeof(aicpu::AicpuParamHead));
+                                     memcpy_args_size_, sizeof(aicpu::AicpuParamHead));
                   GELOGE(FAILED,
                          "[Check][Size] Task def args_size=%u is less than aicpu param head len=%zu.",
-                         memcpy_args_size, sizeof(aicpu::AicpuParamHead));
+                         memcpy_args_size_, sizeof(aicpu::AicpuParamHead));
                   return FAILED;);
 
-  std::unique_ptr<uint8_t[]> memcpy_task;
-  memcpy_task.reset(new(std::nothrow) uint8_t[memcpy_args_size]());
+  memcpy_args_.reset(new(std::nothrow) uint8_t[memcpy_args_size_]());
   GE_IF_BOOL_EXEC(memcpy_task == nullptr,
                   REPORT_INNER_ERROR("E19999", "new memory failed for Node[MemCopy], task_size[%u].",
-                                     memcpy_args_size);
+                                     memcpy_args_size_);
                   GELOGE(FAILED, "[Malloc][Memory] failed for Node[MemCopy], task_size[%u].",
-                         memcpy_args_size);
+                         memcpy_args_size_);
                   return FAILED;);
 
-  errno_t sec_ret = memcpy_s(memcpy_task.get(), memcpy_args_size, memcpy_args.c_str(), memcpy_args.size());
+  errno_t sec_ret = memcpy_s(memcpy_args_.get(), memcpy_args_size_, memcpy_args.c_str(), memcpy_args.size());
   GE_IF_BOOL_EXEC(sec_ret != EOK,
                   REPORT_INNER_ERROR("E19999",
                                      "memcpy_s argc_ failed for Node[MemCopy], ret: %d", sec_ret);
                   GELOGE(INTERNAL_ERROR,
                          "[Update][args] failed for Node[MemCopy], ret: %d", sec_ret);
                   return sec_ret;);
-  auto memcpy_param_head = reinterpret_cast<aicpu::AicpuParamHead *>(memcpy_task.get());
+  auto memcpy_param_head = reinterpret_cast<aicpu::AicpuParamHead *>(memcpy_args_.get());
   uint32_t memcpy_io_num = memcpy_param_head->ioAddrNum;
-  auto memcpy_io_addr = memcpy_task.get() + sizeof(aicpu::AicpuParamHead);
+  auto memcpy_io_addr = memcpy_args_.get() + sizeof(aicpu::AicpuParamHead);
   // if has input and output, need copy to ioaddr
-  int cpy_ret = memcpy_s(memcpy_io_addr, memcpy_args_size - sizeof(aicpu::AicpuParamHead),
+  int cpy_ret = memcpy_s(memcpy_io_addr, memcpy_args_size_ - sizeof(aicpu::AicpuParamHead),
                          &copy_ioaddr_dev_, sizeof(uint64_t) * memcpy_io_num);
   GE_IF_BOOL_EXEC(cpy_ret != 0,
                   REPORT_INNER_ERROR("E19999", "Node[Memcpoy] memcpy io addr to AicpuParamHead failed,"
                          "ret=%d, args_size=%u, io nums=%zu.",
-                         cpy_ret, memcpy_args_size, memcpy_io_num);
+                         cpy_ret, memcpy_args_size_, memcpy_io_num);
                   GELOGE(INTERNAL_ERROR, "[Update][io_addr]Node[MemCopy] memcpy io addr to AicpuParamHead failed,"
                          "ret=%d, args_size=%u, io nums=%zu.",
-                         cpy_ret, memcpy_args_size, memcpy_io_num);
+                         cpy_ret, memcpy_args_size_, memcpy_io_num);
                   return INTERNAL_ERROR;);
   GELOGD("Set memcpy task for node[MemCopy] successfully.");
   return SUCCESS;

ge/hybrid/node_executor/aicpu/aicpu_node_executor.h

@@ -61,11 +61,29 @@ class AicpuNodeTaskBase : public NodeTask {
 
   virtual Status UpdateIoAddr(TaskContext &context) = 0;
 
+  virtual Status CopyDataToHbm(TaskContext &context,
+                               const std::vector<std::unique_ptr<TensorBuffer>> &out_shape_hbm) = 0;
+
   static Status AllocTensorBuffer(size_t size, std::unique_ptr<TensorBuffer> &tensor_buffer);
 
  private:
   Status TaskCallback(TaskContext &context);
 
+  ///
+  /// read result summary and prepare copy task memory.
+  /// @param context task context
+  /// @param out_shape_hbm if scalar, TensorBuffer->data is null, size=0
+  /// @return SUCCESS:success other:failed
+  ///
+  Status ReadResultSummaryAndPrepareMemory(TaskContext &context,
+                                           std::vector<std::unique_ptr<TensorBuffer>> &out_shape_hbm);
+
+  Status UpdateShapeByHbmBuffer(TaskContext &context,
+                                const std::vector<std::unique_ptr<TensorBuffer>> &out_shape_hbm);
+
+  Status PrepareCopyInputs(const TaskContext &context,
+                           const std::vector<std::unique_ptr<TensorBuffer>> &out_shape_hbm);
+
  protected:
   const NodeItem *node_item_;
   // just reference.
@@ -111,26 +129,11 @@ class AicpuTfNodeTask : public AicpuNodeTaskBase {
 
   Status UpdateIoAddr(TaskContext &context) override;
 
+  Status CopyDataToHbm(TaskContext &context,
+                       const std::vector<std::unique_ptr<TensorBuffer>> &out_shape_hbm) override;
  private:
   Status SetMemCopyTask(const domi::TaskDef &task_def);
 
-  ///
-  /// read result summary and prepare copy task memory.
-  /// @param context task context
-  /// @param out_shape_hbm if scalar, TensorBuffer->data is null, size=0
-  /// @return SUCCESS:success other:failed
-  ///
-  Status ReadResultSummaryAndPrepareMemory(TaskContext &context,
-                                           std::vector<std::unique_ptr<TensorBuffer>> &out_shape_hbm);
-  Status CopyDataToHbm(TaskContext &context,
-                       const std::vector<std::unique_ptr<TensorBuffer>> &out_shape_hbm);
-
-  Status UpdateShapeByHbmBuffer(TaskContext &context,
-                                const std::vector<std::unique_ptr<TensorBuffer>> &out_shape_hbm);
-
-  Status PrepareCopyInputs(const TaskContext &context,
-                           const std::vector<std::unique_ptr<TensorBuffer>> &out_shape_hbm);
-
   static Status EnsureSessionCreated(uint64_t session_id);
   static uint64_t GetStepIdAddr(const HybridModel &model);
  private:
@@ -161,7 +164,8 @@ class AicpuNodeTask : public AicpuNodeTaskBase {
 
   Status LaunchTask(TaskContext &context) override;
 
-   override;
+  Status CopyDataToHbm(TaskContext &context,
+                       const std::vector<std::unique_ptr<TensorBuffer>> &out_shape_hbm) override;
 
   Status UpdateIoAddr(TaskContext &context) override;
 
@@ -171,9 +175,17 @@ class AicpuNodeTask : public AicpuNodeTaskBase {
  protected:
   // host mem
   std::unique_ptr<uint8_t[]> args_;
-
+  
+  // host memcpy mem
+  std::unique_ptr<uint8_t[]> memcpy_args_;
   // args size
   uint32_t args_size_ = 0;
+
+  std::string memcpy_so_name_;
+
+  std::string memcpy_kernel_name_;
+  // args size
+  uint32_t memcpy_args_size_ = 0;  
 };
 
 class AiCpuNodeExecutor : public NodeExecutor {