diff --git a/ge/single_op/task/op_task.cc b/ge/single_op/task/op_task.cc index 7117368e..c0c33d5e 100755 --- a/ge/single_op/task/op_task.cc +++ b/ge/single_op/task/op_task.cc @@ -420,6 +420,18 @@ AiCpuBaseTask::~AiCpuBaseTask() { if (ext_info_addr_dev_ != nullptr) { (void)rtFree(ext_info_addr_dev_); } + + FreeHbm(copy_ioaddr_dev_); + FreeHbm(copy_input_release_flag_dev_); + FreeHbm(copy_input_data_size_dev_); + FreeHbm(copy_input_src_dev_); + FreeHbm(copy_input_dst_dev_); + for (auto summary : output_summary_) { + FreeHbm(summary); + } + for (auto out_shape : out_shape_hbm_) { + FreeHbm(out_shape); + } } Status AiCpuBaseTask::SetExtInfoAndType(const std::string &kernel_ext_info, uint64_t kernel_id) { @@ -623,18 +635,7 @@ AiCpuTask::~AiCpuTask() { FreeHbm(io_addr_); FreeHbm(workspace_addr_); FreeHbm(copy_workspace_buf_); - FreeHbm(copy_ioaddr_dev_); - FreeHbm(copy_input_release_flag_dev_); - FreeHbm(copy_input_data_size_dev_); - FreeHbm(copy_input_src_dev_); - FreeHbm(copy_input_dst_dev_); FreeHbm(copy_task_args_buf_); - for (auto summary : output_summary_) { - FreeHbm(summary); - } - for (auto out_shape : out_shape_hbm_) { - FreeHbm(out_shape); - } } Status AiCpuTask::LaunchKernel(rtStream_t stream) { @@ -664,7 +665,7 @@ Status AiCpuTask::LaunchKernel(rtStream_t stream) { return SUCCESS; } -Status AiCpuTask::PrepareCopyInputs(vector &outputs) { +Status AiCpuBaseTask::PrepareCopyInputs(vector &outputs) { std::vector copy_input_release_flag; std::vector copy_input_data_size; std::vector copy_input_src; @@ -705,7 +706,7 @@ Status AiCpuTask::PrepareCopyInputs(vector &outputs) { return SUCCESS; } -Status AiCpuTask::ReadResultSummaryAndPrepareMemory() { +Status AiCpuBaseTask::ReadResultSummaryAndPrepareMemory() { for (size_t i = 0; i < num_outputs_; ++i) { auto &result_summary = output_summary_host_[i]; @@ -732,7 +733,20 @@ Status AiCpuTask::CopyDataToHbm(vector &outputs, return SUCCESS; } -Status AiCpuTask::UpdateShapeByHbmBuffer(vector &output_desc) { +Status AiCpuCCTask::CopyDataToHbm(vector &outputs, + rtStream_t stream) { + GE_CHK_STATUS_RET_NOLOG(PrepareCopyInputs(outputs)); + + auto ret = rtCpuKernelLaunchWithFlag(static_cast(memcpy_so_name_.data()), + static_cast(memcpy_kernel_name_.data()), + block_dim_, memcpy_args_.get(), static_cast(memcpy_arg_size_), + nullptr, stream, dump_flag_); + GE_CHK_RT_RET(ret); + GE_CHK_RT_RET(rtStreamSynchronize(stream)); + return SUCCESS; +} + +Status AiCpuBaseTask::UpdateShapeByHbmBuffer(vector &output_desc) { for (size_t i = 0; i < num_outputs_; ++i) { const auto &result_summary = output_summary_host_[i]; std::vector shape_dims; diff --git a/ge/single_op/task/op_task.h b/ge/single_op/task/op_task.h index 546b1681..6f1da40a 100644 --- a/ge/single_op/task/op_task.h +++ b/ge/single_op/task/op_task.h @@ -146,7 +146,6 @@ class AiCpuBaseTask : public OpTask { rtStream_t stream); Status ReadResultSummaryAndPrepareMemory(); - Status CopyDataToHbm(vector &outputs, rtStream_t stream); Status PrepareCopyInputs(vector &outputs); Status UpdateShapeByHbmBuffer(vector &output_desc); @@ -190,6 +189,8 @@ class AiCpuTask : public AiCpuBaseTask { // for copy task. Status InitForSummaryAndCopy(); + Status CopyDataToHbm(vector &outputs, rtStream_t stream); + private: friend class AiCpuTaskBuilder; void *workspace_addr_ = nullptr; std::string task_info_; @@ -232,7 +233,9 @@ class AiCpuCCTask : public AiCpuBaseTask { std::vector &output_buffers, rtStream_t stream) override; private: - Status InitForSummaryAndCopy(); + Status InitForSummaryAndCopy(); + + Status CopyDataToHbm(vector &outputs, rtStream_t stream); private: friend class AiCpuCCTaskBuilder; std::string so_name_;