|
|
@@ -49,14 +49,14 @@ bool TbeHandleRegistry::AddHandle(std::unique_ptr<TbeHandleHolder> &&holder) { |
|
|
|
return ret.second; |
|
|
|
} |
|
|
|
|
|
|
|
Status AiCoreOpTask::Init(const OpDesc &op_desc, const domi::TaskDef &task_def) { |
|
|
|
Status AiCoreOpTask::Init(HybridModel &hybrid_model, const OpDesc &op_desc, const domi::TaskDef &task_def) { |
|
|
|
log_name_ = op_desc.GetName() + "_tvmbin"; |
|
|
|
log_id_ = log_id++; |
|
|
|
auto op_desc_ptr = MakeShared<OpDesc>(op_desc); |
|
|
|
GE_CHECK_NOTNULL(op_desc_ptr); |
|
|
|
auto task_info = BuildTaskUtils::GetTaskInfo(op_desc_ptr); |
|
|
|
GELOGI("[TASK_INFO] %lu/%s %s.", log_id_, log_name_.c_str(), task_info.c_str()); |
|
|
|
GE_CHK_STATUS_RET_NOLOG(InitWithTaskDef(op_desc, task_def)); |
|
|
|
GE_CHK_STATUS_RET_NOLOG(InitWithTaskDef(hybrid_model, op_desc, task_def)); |
|
|
|
GE_CHK_STATUS_RET_NOLOG(InitTilingInfo(op_desc)); |
|
|
|
|
|
|
|
GE_CHECK_LE(op_desc.GetOutputsSize(), static_cast<size_t>(INT_MAX)); |
|
|
@@ -78,7 +78,7 @@ Status AiCoreOpTask::Init(const OpDesc &op_desc, const domi::TaskDef &task_def) |
|
|
|
return SUCCESS; |
|
|
|
} |
|
|
|
|
|
|
|
Status AiCoreOpTask::RegisterTbeHandle(const OpDesc &op_desc) { |
|
|
|
Status AiCoreOpTask::RegisterTbeHandle(HybridModel &hybrid_model, const OpDesc &op_desc) { |
|
|
|
rtError_t rt_ret = rtQueryFunctionRegistered(stub_name_.c_str()); |
|
|
|
if (rt_ret != RT_ERROR_NONE || is_single_op_) { |
|
|
|
auto op_desc_ptr = MakeShared<OpDesc>(op_desc); |
|
|
@@ -133,7 +133,11 @@ Status AiCoreOpTask::RegisterTbeHandle(const OpDesc &op_desc) { |
|
|
|
GELOGI("TBE: binfile_key=%s, kernel_name=%s", stub_name_.c_str(), kernel_name.c_str()); |
|
|
|
GE_CHK_RT_RET(rtFunctionRegister(bin_handle, stub_name_.c_str(), |
|
|
|
stub_name_.c_str(), kernel_name.c_str(), 0)); |
|
|
|
hybrid_model.SetUsedTbeHandleMap(stub_name_.c_str()); |
|
|
|
return SUCCESS; |
|
|
|
} |
|
|
|
// Kernel registed, Increase used num in store. |
|
|
|
hybrid_model.StoreTbeHandle(stub_name_.c_str()); |
|
|
|
return SUCCESS; |
|
|
|
} |
|
|
|
|
|
|
@@ -190,11 +194,11 @@ Status AiCoreOpTask::RegisterKernelHandle(const OpDesc &op_desc) { |
|
|
|
return SUCCESS; |
|
|
|
} |
|
|
|
|
|
|
|
Status AiCoreOpTask::InitWithKernelDef(const OpDesc &op_desc, const domi::TaskDef &task_def) { |
|
|
|
Status AiCoreOpTask::InitWithKernelDef(HybridModel &hybrid_model, const OpDesc &op_desc, const domi::TaskDef &task_def) { |
|
|
|
const domi::KernelDef &kernel_def = task_def.kernel(); |
|
|
|
const domi::KernelContext &context = kernel_def.context(); |
|
|
|
stub_name_ = kernel_def.stub_func(); |
|
|
|
GE_CHK_STATUS_RET(RegisterTbeHandle(op_desc)); |
|
|
|
GE_CHK_STATUS_RET(RegisterTbeHandle(hybrid_model, op_desc)); |
|
|
|
GE_CHK_RT_RET(rtGetFunctionByName(stub_name_.c_str(), &stub_func_)); |
|
|
|
args_size_ = kernel_def.args_size(); |
|
|
|
block_dim_ = kernel_def.block_dim(); |
|
|
@@ -304,7 +308,7 @@ Status AiCoreOpTask::InitWithKernelDefWithHandle(const OpDesc &op_desc, const do |
|
|
|
return SUCCESS; |
|
|
|
} |
|
|
|
|
|
|
|
Status AiCoreOpTask::InitWithTaskDef(const OpDesc &op_desc, const domi::TaskDef &task_def) { |
|
|
|
Status AiCoreOpTask::InitWithTaskDef(HybridModel &hybrid_model, const OpDesc &op_desc, const domi::TaskDef &task_def) { |
|
|
|
|
|
|
|
auto rt_ret = ValidateTaskDef(task_def); |
|
|
|
if (rt_ret != SUCCESS) { |
|
|
@@ -316,7 +320,7 @@ Status AiCoreOpTask::InitWithTaskDef(const OpDesc &op_desc, const domi::TaskDef |
|
|
|
} |
|
|
|
|
|
|
|
if (task_def.type() != RT_MODEL_TASK_ALL_KERNEL) { |
|
|
|
GE_CHK_STATUS_RET(InitWithKernelDef(op_desc, task_def)); |
|
|
|
GE_CHK_STATUS_RET(InitWithKernelDef(hybrid_model, op_desc, task_def)); |
|
|
|
} else { |
|
|
|
GE_CHK_STATUS_RET(InitWithKernelDefWithHandle(op_desc, task_def)); |
|
|
|
} |
|
|
@@ -558,8 +562,8 @@ std::string AiCoreOpTask::GetKeyForKernelName(const OpDesc &op_desc) const { |
|
|
|
return op_desc.GetName() + "_kernelname"; |
|
|
|
} |
|
|
|
|
|
|
|
Status AtomicAddrCleanOpTask::Init(const OpDesc &op_desc, const domi::TaskDef &task_def) { |
|
|
|
GE_CHK_STATUS_RET_NOLOG(AiCoreOpTask::Init(op_desc, task_def)); |
|
|
|
Status AtomicAddrCleanOpTask::Init(HybridModel &hybrid_model, const OpDesc &op_desc, const domi::TaskDef &task_def) { |
|
|
|
GE_CHK_STATUS_RET_NOLOG(AiCoreOpTask::Init(hybrid_model, op_desc, task_def)); |
|
|
|
return InitAtomicAddrCleanIndices(op_desc); |
|
|
|
} |
|
|
|
|
|
|
|