fuix

ge/hybrid/node_executor/aicpu/aicpu_node_executor.cc

@@ -234,7 +234,7 @@ Status AicpuNodeTaskBase::ExecuteAsync(TaskContext &context, std::function<void(
   return SUCCESS;
 }
 
-Status AicpuNodeTaskBase::InitForDependComputeTask() {
+Status AicpuNodeTaskBase::InitForDependComputeTask(bool is_aicpu_kernel) {
   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);
@@ -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<uint64_t> copy_io_addr;
   copy_io_addr.emplace_back(reinterpret_cast<uintptr_t>(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;
 }

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);