Browse Source

Fix bug of single_op kernel bin register.

tags/v1.5.1
zhaozhixuan 3 years ago
parent
commit
1942c5e74e
1 changed files with 2 additions and 2 deletions
  1. +2
    -2
      ge/hybrid/node_executor/aicore/aicore_op_task.cc

+ 2
- 2
ge/hybrid/node_executor/aicore/aicore_op_task.cc View File

@@ -81,7 +81,7 @@ Status AiCoreOpTask::Init(const OpDesc &op_desc, const domi::TaskDef &task_def)


Status AiCoreOpTask::RegisterTbeHandle(const OpDesc &op_desc) { Status AiCoreOpTask::RegisterTbeHandle(const OpDesc &op_desc) {
rtError_t rt_ret = rtQueryFunctionRegistered(stub_name_.c_str()); rtError_t rt_ret = rtQueryFunctionRegistered(stub_name_.c_str());
if (rt_ret != RT_ERROR_NONE || is_single_op_) {
if (rt_ret != RT_ERROR_NONE) {
auto op_desc_ptr = MakeShared<OpDesc>(op_desc); auto op_desc_ptr = MakeShared<OpDesc>(op_desc);
GE_CHECK_NOTNULL(op_desc_ptr); GE_CHECK_NOTNULL(op_desc_ptr);
auto tbe_kernel = op_desc_ptr->TryGetExtAttr(GetKeyForTbeKernel(), TBEKernelPtr()); auto tbe_kernel = op_desc_ptr->TryGetExtAttr(GetKeyForTbeKernel(), TBEKernelPtr());
@@ -194,7 +194,7 @@ Status AiCoreOpTask::RegisterKernelHandle(const OpDesc &op_desc) {
Status AiCoreOpTask::InitWithKernelDef(const OpDesc &op_desc, const domi::TaskDef &task_def) { Status AiCoreOpTask::InitWithKernelDef(const OpDesc &op_desc, const domi::TaskDef &task_def) {
const domi::KernelDef &kernel_def = task_def.kernel(); const domi::KernelDef &kernel_def = task_def.kernel();
const domi::KernelContext &context = kernel_def.context(); const domi::KernelContext &context = kernel_def.context();
stub_name_ = kernel_def.stub_func();
stub_name_ = is_single_op_ ? to_string(log_id_) + kernel_def.stub_func() : kernel_def.stub_func();
GE_CHK_STATUS_RET(RegisterTbeHandle(op_desc)); GE_CHK_STATUS_RET(RegisterTbeHandle(op_desc));
GE_CHK_RT_RET(rtGetFunctionByName(stub_name_.c_str(), &stub_func_)); GE_CHK_RT_RET(rtGetFunctionByName(stub_name_.c_str(), &stub_func_));
args_size_ = kernel_def.args_size(); args_size_ = kernel_def.args_size();


Loading…
Cancel
Save