diff --git a/ge/hybrid/node_executor/aicpu/aicpu_node_executor.cc b/ge/hybrid/node_executor/aicpu/aicpu_node_executor.cc index 51b587b8..030b6b48 100755 --- a/ge/hybrid/node_executor/aicpu/aicpu_node_executor.cc +++ b/ge/hybrid/node_executor/aicpu/aicpu_node_executor.cc @@ -234,7 +234,7 @@ Status AicpuNodeTaskBase::ExecuteAsync(TaskContext &context, std::functionnum_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); @@ -267,9 +267,15 @@ Status AicpuNodeTaskBase::InitForDependComputeTask() { node_name_.c_str(), copy_input_buf_len); // copy task args buf - 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)); + if (is_aicpu_kernel) { + GE_CHK_STATUS_RET(AllocTensorBuffer(sizeof(aicpu::AicpuParamHead), copy_task_args_buf_), + "[Alloc][TensorBuffer] failed for Node[%s] to copy task args, size=%zu", + node_name_.c_str(), sizeof(aicpu::AicpuParamHead)); + } else { + 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 copy_io_addr; copy_io_addr.emplace_back(reinterpret_cast(copy_input_release_flag_dev_->GetData())); @@ -331,7 +337,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(), "[Init][DependComputeTask] failed for Node[%s].", node_name_.c_str()); + GE_CHK_STATUS_RET(InitForDependComputeTask(false), "[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(), @@ -745,7 +751,7 @@ Status AicpuNodeTask::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(), + GE_CHK_STATUS_RET(InitForDependComputeTask(true), "[Init][DependComputeTask] failed for Node[%s].", node_name_.c_str()); if (ext_info_addr_dev_ == nullptr) { @@ -807,9 +813,9 @@ Status AicpuNodeTask::SetMemCopyTask(const domi::TaskDef &task_def) { "[Update][args] failed for Node[MemCopy], ret: %d", sec_ret); return sec_ret;); - - 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)); + + GE_CHK_RT_RET(rtMemcpy(copy_task_args_buf_->GetData(), sizeof(aicpu::AicpuParamHead), + &memcpy_task, sizeof(aicpu::AicpuParamHead), RT_MEMCPY_HOST_TO_DEVICE)); GELOGD("Set memcpy task for node[%s] successfully.", 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 45541d84..8843ed3b 100644 --- a/ge/hybrid/node_executor/aicpu/aicpu_node_executor.h +++ b/ge/hybrid/node_executor/aicpu/aicpu_node_executor.h @@ -57,7 +57,7 @@ class AicpuNodeTaskBase : public NodeTask { virtual Status LaunchTask(TaskContext &context) = 0; - Status InitForDependComputeTask(); + Status InitForDependComputeTask(bool is_aicpu_kernel); Status TaskCallback(TaskContext &context);