Browse Source

Save atomic kernel bin to model.

pull/1442/head
zhaozhixuan 4 years ago
parent
commit
4ee06ea0a6
1 changed files with 8 additions and 8 deletions
  1. +8
    -8
      ge/common/tbe_kernel_store.cc

+ 8
- 8
ge/common/tbe_kernel_store.cc View File

@@ -33,15 +33,15 @@ void TBEKernelStore::LoadTBEKernelBinToOpDesc(const std::shared_ptr<ge::OpDesc>
GE_IF_BOOL_EXEC(!op_desc->SetExtAttr(ge::OP_EXTATTR_NAME_TBE_KERNEL, kernel_bin), GE_IF_BOOL_EXEC(!op_desc->SetExtAttr(ge::OP_EXTATTR_NAME_TBE_KERNEL, kernel_bin),
GELOGW("LoadKernelTBEBinToOpDesc: SetExtAttr for kernel_bin failed");) GELOGW("LoadKernelTBEBinToOpDesc: SetExtAttr for kernel_bin failed");)
GELOGI("Load tbe kernel:%s, %zu", kernel_bin->GetName().c_str(), kernel_bin->GetBinDataSize()); GELOGI("Load tbe kernel:%s, %zu", kernel_bin->GetName().c_str(), kernel_bin->GetBinDataSize());
}


std::string atomic_kernel_name;
(void) AttrUtils::GetStr(op_desc, "ATOMIC_ATTR_TBE_KERNEL_NAME", atomic_kernel_name);
if (!atomic_kernel_name.empty()) {
GELOGI("Get atomic kernel name is %s", atomic_kernel_name.c_str());
auto atomic_kernel_bin = FindKernel(atomic_kernel_name);
GE_IF_BOOL_EXEC(!op_desc->SetExtAttr("EXT_ATTR_ATOMIC_TBE_KERNEL", atomic_kernel_bin),
GELOGW("LoadKernelTBEBinToOpDesc: SetExtAttr for atomic kernel_bin failed");)
std::string atomic_kernel_name;
(void) AttrUtils::GetStr(op_desc, "ATOMIC_ATTR_TBE_KERNEL_NAME", atomic_kernel_name);
if (!atomic_kernel_name.empty()) {
GELOGI("Get atomic kernel name is %s", atomic_kernel_name.c_str());
auto atomic_kernel_bin = FindKernel(atomic_kernel_name);
GE_IF_BOOL_EXEC(!op_desc->SetExtAttr("EXT_ATTR_ATOMIC_TBE_KERNEL", atomic_kernel_bin),
GELOGW("LoadKernelTBEBinToOpDesc: SetExtAttr for atomic kernel_bin failed");)
}
} }
} }
} }


Loading…
Cancel
Save