diff --git a/ge/graph/load/new_model_manager/davinci_model.cc b/ge/graph/load/new_model_manager/davinci_model.cc index 3e24a952..eb4960c1 100755 --- a/ge/graph/load/new_model_manager/davinci_model.cc +++ b/ge/graph/load/new_model_manager/davinci_model.cc @@ -3021,7 +3021,14 @@ Status DavinciModel::MallocKnownArgs() { GELOGE(RT_FAILED, "Call rtMalloc failed, ret: 0x%X", rt_ret); return RT_ERROR_TO_GE_STATUS(rt_ret); } - + // malloc dynamic and static hybrid memory + if (total_hybrid_args_size_ != 0) { + rt_ret = rtMalloc(&hybrid_addrs_, total_hybrid_args_size_, RT_MEMORY_HBM); + if (rt_ret != RT_ERROR_NONE) { + GELOGE(RT_FAILED, "Call rtMalloc failed, ret: 0x%X", rt_ret); + return RT_ERROR_TO_GE_STATUS(rt_ret); + } + } // malloc fixed addr memory, eg: rts op if (total_fixed_addr_size_ != 0) { GELOGI("Begin to allocate fixed addr."); diff --git a/ge/graph/load/new_model_manager/davinci_model.h b/ge/graph/load/new_model_manager/davinci_model.h index 0969559e..300245b0 100755 --- a/ge/graph/load/new_model_manager/davinci_model.h +++ b/ge/graph/load/new_model_manager/davinci_model.h @@ -498,6 +498,14 @@ class DavinciModel { return cur_args; } void SetTotalIOAddrs(const vector &io_addrs); + void SetHybridArgsSize(uint32_t args_size) { total_hybrid_args_size_ += args_size; } + uint32_t GetHybridArgsSize() { + return total_hybrid_args_size_; + } + void *GetCurrentHybridArgsAddr(uint32_t offset) { + void *cur_args = static_cast(hybrid_addrs_) + offset; + return cur_args; + } void SetTotalFixedAddrsSize(string tensor_name, int64_t fix_addr_size); int64_t GetFixedAddrsSize(string tensor_name); void *GetCurrentFixedAddr(int64_t offset) const { @@ -975,6 +983,8 @@ class DavinciModel { void *args_ = nullptr; void *args_host_ = nullptr; void *fixed_addrs_ = nullptr; + void *hybrid_addrs_ = nullptr; + uint32_t total_hybrid_args_size_ = 0; int64_t total_fixed_addr_size_ = 0; map known_input_data_info_; map known_output_data_info_; diff --git a/ge/graph/load/new_model_manager/task_info/kernel_task_info.cc b/ge/graph/load/new_model_manager/task_info/kernel_task_info.cc index 0c8507bc..fb351f89 100755 --- a/ge/graph/load/new_model_manager/task_info/kernel_task_info.cc +++ b/ge/graph/load/new_model_manager/task_info/kernel_task_info.cc @@ -444,7 +444,6 @@ void KernelTaskInfo::SetIoAddrs(const OpDescPtr &op_desc) { vector workspace_data_addrs = ModelUtils::GetWorkspaceDataAddrs(rts_param, op_desc); io_addrs_.insert(io_addrs_.end(), workspace_data_addrs.begin(), workspace_data_addrs.end()); } - return SUCCESS; } Status KernelTaskInfo::UpdateArgs() { @@ -897,6 +896,12 @@ Status KernelTaskInfo::InitAicpuTask(uint32_t op_index, const domi::KernelDef &k aicpu_param_head->extInfoAddr = reinterpret_cast(aicpu_ext_info_addr_); aicpu_param_head->extInfoLength = static_cast(ext_info.size()); + + if (davinci_model_->IsKnownNode()) { + args_ = davinci_model_->GetCurrentHybridArgsAddr(hybrid_args_offset_); + InitDumpTask(sizeof(aicpu::AicpuParamHead)); + return SUCCESS; + } const RuntimeParam &rts_param = davinci_model_->GetRuntimeParam(); vector input_addrs = ModelUtils::GetInputDataAddrs(rts_param, op_desc); vector output_addrs = ModelUtils::GetOutputDataAddrs(rts_param, op_desc);