Fix task SetIoAddrs for known node.

ge/graph/load/new_model_manager/task_info/kernel_ex_task_info.cc

@@ -127,8 +127,9 @@ Status KernelExTaskInfo::Init(const domi::TaskDef &task_def, DavinciModel *davin
     GE_IF_BOOL_EXEC(rt_ret != RT_ERROR_NONE, GELOGE(RT_FAILED, "rtMemcpy error, ret: Ox%X", rt_ret);
                     return RT_ERROR_TO_GE_STATUS(rt_ret);)
 
+    SetIoAddrs(op_desc);
     GELOGI("KernelExTaskInfo knonw node Init Success.");
-    return SetIoAddrs(op_desc);
+    return SUCCESS;
   }
 
   // 3. Set workspaceaddr, inputOutputDataAddr
@@ -190,9 +191,9 @@ Status KernelExTaskInfo::Init(const domi::TaskDef &task_def, DavinciModel *davin
                   return RT_ERROR_TO_GE_STATUS(rt_ret);)
 
   davinci_model_->SetZeroCopyAddr(op_desc, io_addrs, io_addrs.data(), input_output_addr_, addrs_size, 0);
-
+  SetIoAddrs(op_desc);
   GELOGI("KernelExTaskInfo Init Success. session id: %lu", session_id);
-  return SetIoAddrs(op_desc);
+  return SUCCESS;
 }
 
 Status KernelExTaskInfo::CalculateArgs(const domi::TaskDef &task_def, DavinciModel *davinci_model) {
@@ -227,39 +228,36 @@ Status KernelExTaskInfo::CalculateArgs(const domi::TaskDef &task_def, DavinciMod
     davinci_model->SetTotalFixedAddrsSize(peer_input_name, tensor_size);
     GELOGI("Calculate stream switch task args , tensor size is %ld, fixed addr offset %ld", tensor_size,
            fixed_addr_offset_);
+
+    const RuntimeParam &rts_param = davinci_model_->GetRuntimeParam();
+    vector<void *> input_data_addrs = ModelUtils::GetInputDataAddrs(rts_param, op_desc);
+    vector<void *> output_data_addrs = ModelUtils::GetOutputDataAddrs(rts_param, op_desc);
+    if (output_index > output_data_addrs.size()) {
+      GELOGE(FAILED, "The output data addr size[%zu] and output index[%u] are inconsistent.",
+             output_data_addrs.size(), output_index);
+      return FAILED;
+    }
+    io_addrs_.insert(io_addrs_.end(), input_data_addrs.begin(), input_data_addrs.end());
+    for (size_t i = 0; i < output_data_addrs.size(); ++i) {
+      if (i == output_index) {
+        void *fixed_addr = davinci_model_->GetCurrentFixedAddr(fixed_addr_offset_);
+        io_addrs_.emplace_back(fixed_addr);
+        continue;
+      }
+      io_addrs_.emplace_back(output_data_addrs[i]);
+    }
   }
   return SUCCESS;
 }
 
-Status KernelExTaskInfo::SetIoAddrs(const OpDescPtr &op_desc) {
-  const RuntimeParam &rts_param = davinci_model_->GetRuntimeParam();
-  vector<void *> input_data_addrs = ModelUtils::GetInputDataAddrs(rts_param, op_desc);
-  vector<void *> output_data_addrs = ModelUtils::GetOutputDataAddrs(rts_param, op_desc);
+void KernelExTaskInfo::SetIoAddrs(const OpDescPtr &op_desc) {
   if (!op_desc->HasAttr(ATTR_DYNAMIC_SHAPE_FIXED_ADDR)) {
+    const RuntimeParam &rts_param = davinci_model_->GetRuntimeParam();
+    vector<void *> input_data_addrs = ModelUtils::GetInputDataAddrs(rts_param, op_desc);
+    vector<void *> output_data_addrs = ModelUtils::GetOutputDataAddrs(rts_param, op_desc);
     io_addrs_.insert(io_addrs_.end(), input_data_addrs.begin(), input_data_addrs.end());
     io_addrs_.insert(io_addrs_.end(), output_data_addrs.begin(), output_data_addrs.end());
-  } else {
-    string peer_input_name;
-    if (AttrUtils::GetStr(op_desc, ATTR_DYNAMIC_SHAPE_FIXED_ADDR, peer_input_name)) {
-      uint32_t output_index = davinci_model_->GetFixedAddrOutputIndex(peer_input_name);
-      if (output_index > output_data_addrs.size()) {
-        GELOGE(FAILED, "The output data addr size[%zu] and output index[%u] are inconsistent.",
-               output_data_addrs.size(), output_index);
-        return FAILED;
-      }
-      io_addrs_.insert(io_addrs_.end(), input_data_addrs.begin(), input_data_addrs.end());
-      for (size_t i = 0; i < output_data_addrs.size(); ++i) {
-        if (i == output_index) {
-          void *fixed_addr = davinci_model_->GetCurrentFixedAddr(fixed_addr_offset_);
-          io_addrs_.emplace_back(fixed_addr);
-          continue;
-        }
-        io_addrs_.emplace_back(output_data_addrs[i]);
-      }
-    }
   }
-
-  return SUCCESS;
 }
 
 Status KernelExTaskInfo::UpdateArgs() {

ge/graph/load/new_model_manager/task_info/kernel_ex_task_info.h

@@ -59,7 +59,7 @@ class KernelExTaskInfo : public TaskInfo {
   };
  private:
   Status CopyTaskInfo(const domi::KernelExDef &kernel_def, const RuntimeParam &rts_param, const OpDescPtr &op_desc);
-  Status SetIoAddrs(const OpDescPtr &op_desc);
+  void SetIoAddrs(const OpDescPtr &op_desc);
 
   uint32_t task_id_;
   uint32_t stream_id_;

ge/graph/load/new_model_manager/task_info/memcpy_async_task_info.cc

@@ -45,9 +45,9 @@ Status MemcpyAsyncTaskInfo::Init(const domi::TaskDef &task_def, DavinciModel *da
     dst_ = reinterpret_cast<uint8_t *>(reinterpret_cast<uintptr_t>(src_) + sizeof(void *));
     // for zero copy
     kind_ = RT_MEMCPY_ADDR_DEVICE_TO_DEVICE;
+    SetIoAddrs(op_desc);
     GELOGI("MemcpyAsyncTaskInfo op name %s, src_ %p, dst_ %p, args_offset %u.",
            op_desc->GetName().c_str(), src_, dst_, args_offset_);
-    SetIoAddrs(op_desc);
     return SUCCESS;
   }
 
@@ -73,12 +73,12 @@ Status MemcpyAsyncTaskInfo::Init(const domi::TaskDef &task_def, DavinciModel *da
     }
   }
 
-  GELOGI("MemcpyAsyncTaskInfo Init Success, logic[0x%lx, 0x%lx], src:%p, dst:%p, max:%lu, count:%lu",
-         memcpy_async.src(), memcpy_async.dst(), src_, dst_, dst_max_, count_);
-
   davinci_model_->DisableZeroCopy(src_);
   davinci_model_->DisableZeroCopy(dst_);
   SetIoAddrs(op_desc);
+
+  GELOGI("MemcpyAsyncTaskInfo Init Success, logic[0x%lx, 0x%lx], src:%p, dst:%p, max:%lu, count:%lu",
+         memcpy_async.src(), memcpy_async.dst(), src_, dst_, dst_max_, count_);
   return SUCCESS;
 }
 
@@ -112,16 +112,16 @@ Status MemcpyAsyncTaskInfo::CalculateArgs(const domi::TaskDef &task_def, Davinci
     int64_t tensor_size = 0;
     GE_CHK_STATUS(TensorUtils::GetSize(tensor_desc, tensor_size));
     davinci_model->SetTotalFixedAddrsSize(peer_input_name, tensor_size);
+
+    void *fixed_addr = davinci_model_->GetCurrentFixedAddr(fixed_addr_offset_);
+    io_addrs_.emplace_back(fixed_addr);
   }
   return SUCCESS;
 }
 
 void MemcpyAsyncTaskInfo::SetIoAddrs(const OpDescPtr &op_desc) {
   io_addrs_.emplace_back(reinterpret_cast<void *>(src_));
-  if (op_desc->HasAttr(ATTR_DYNAMIC_SHAPE_FIXED_ADDR)) {
-    void *fixed_addr = davinci_model_->GetCurrentFixedAddr(fixed_addr_offset_);
-    io_addrs_.emplace_back(fixed_addr);
-  } else {
+  if (!op_desc->HasAttr(ATTR_DYNAMIC_SHAPE_FIXED_ADDR)) {
     io_addrs_.emplace_back(reinterpret_cast<void *>(dst_));
   }
 }
@@ -130,7 +130,6 @@ Status MemcpyAsyncTaskInfo::UpdateArgs() {
   GELOGI("MemcpyAsyncTaskInfo::UpdateArgs in.");
   GE_CHECK_NOTNULL(davinci_model_);
   davinci_model_->SetTotalIOAddrs(io_addrs_);
-
   GELOGI("MemcpyAsyncTaskInfo::UpdateArgs success.");
   return SUCCESS;
 }