From f156eb242ce28370031a13cefd1e05728588dafa Mon Sep 17 00:00:00 2001
From: chuxing <chuxing@huawei.com>
Date: Fri, 11 Dec 2020 20:29:58 +0800
Subject: [PATCH] fix aicpu cc

---
 ge/single_op/task/aicpu_kernel_task_builder.cc | 13 ++++++++++---
 ge/single_op/task/aicpu_kernel_task_builder.h  |  4 ++--
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/ge/single_op/task/aicpu_kernel_task_builder.cc b/ge/single_op/task/aicpu_kernel_task_builder.cc
index fe02a88e..c676ccf8 100755
--- a/ge/single_op/task/aicpu_kernel_task_builder.cc
+++ b/ge/single_op/task/aicpu_kernel_task_builder.cc
@@ -17,12 +17,13 @@
 #include "single_op/task/aicpu_kernel_task_builder.h"
 #include "framework/common/taskdown_common.h"
 #include "graph/load/new_model_manager/model_manager.h"
+#include "build_task_utils.h"
 
 namespace ge {
 AiCpuCCTaskBuilder::AiCpuCCTaskBuilder(const OpDescPtr &op_desc, const domi::KernelDef &kernel_def)
     : op_desc_(op_desc), kernel_def_(kernel_def) {}
 
-Status AiCpuCCTaskBuilder::SetKernelArgs(AiCpuCCTask &task) {
+Status AiCpuCCTaskBuilder::SetKernelArgs(AiCpuCCTask &task, const SingleOpModelParam &param) {
   size_t aicpu_arg_size = kernel_def_.args_size();
   if (aicpu_arg_size <= sizeof(aicpu::AicpuParamHead)) {
     GELOGE(ACL_ERROR_GE_PARAM_INVALID, "aicpu_arg_size is invalid, value = %zu", aicpu_arg_size);
@@ -47,11 +48,17 @@ Status AiCpuCCTaskBuilder::SetKernelArgs(AiCpuCCTask &task) {
 
   task.SetIoAddr(reinterpret_cast<uintptr_t *>(aicpu_args.get() + sizeof(aicpu::AicpuParamHead)));
   task.SetKernelArgs(std::move(aicpu_args), aicpu_arg_size);
+
+  auto addresses = BuildTaskUtils::GetKernelArgs(op_desc_, param);
+  GE_CHECK_GE(addresses.size(), task.io_addr_num_);
+  for (size_t i = 0; i < task.io_addr_num_; ++i) {
+    task.io_addr_[i] = reinterpret_cast<uintptr_t>(addresses[i]);
+  }
   return SUCCESS;
 }
 
-Status AiCpuCCTaskBuilder::BuildTask(AiCpuCCTask &task, uint64_t kernel_id) {
-  auto ret = SetKernelArgs(task);
+Status AiCpuCCTaskBuilder::BuildTask(AiCpuCCTask &task, uint64_t kernel_id, const SingleOpModelParam &param) {
+  auto ret = SetKernelArgs(task, param);
   if (ret != SUCCESS) {
     return ret;
   }
diff --git a/ge/single_op/task/aicpu_kernel_task_builder.h b/ge/single_op/task/aicpu_kernel_task_builder.h
index e77e3c10..85d5034d 100755
--- a/ge/single_op/task/aicpu_kernel_task_builder.h
+++ b/ge/single_op/task/aicpu_kernel_task_builder.h
@@ -30,10 +30,10 @@ class AiCpuCCTaskBuilder {
   explicit AiCpuCCTaskBuilder(const OpDescPtr &op_desc, const domi::KernelDef &kernel_def);
   ~AiCpuCCTaskBuilder() = default;
 
-  Status BuildTask(AiCpuCCTask &task, uint64_t kernel_id);
+  Status BuildTask(AiCpuCCTask &task, uint64_t kernel_id, const SingleOpModelParam &param);
 
  private:
-  Status SetKernelArgs(AiCpuCCTask &task);
+  Status SetKernelArgs(AiCpuCCTask &task, const SingleOpModelParam &param);
   const OpDescPtr op_desc_;
   const domi::KernelDef &kernel_def_;
 };