!19326 [MSLITE][DEVELOP] fix bug of npu+fp16

Merge pull request !19326 from yangruoqi713/npu
This commit is contained in:
i-robot 2021-07-03 08:27:43 +00:00 committed by Gitee
commit 77b96d7bce
2 changed files with 37 additions and 3 deletions

View File

@ -99,6 +99,11 @@ int Scheduler::Schedule(std::vector<kernel::LiteKernel *> *dst_kernels) {
}
}
FindAllInoutKernels(*dst_kernels);
ret = InitKernels(*dst_kernels);
if (ret != RET_OK) {
MS_LOG(ERROR) << "InitKernels failed.";
return ret;
}
auto src_kernel = *dst_kernels;
dst_kernels->clear();
@ -112,6 +117,36 @@ int Scheduler::Schedule(std::vector<kernel::LiteKernel *> *dst_kernels) {
return RET_OK;
}
int Scheduler::InitKernels(std::vector<kernel::LiteKernel *> dst_kernels) {
if (is_train_session_) {
return RET_OK;
}
for (auto kernel : dst_kernels) {
if (kernel->subgraph_type() != kernel::kNotSubGraph) {
auto subgraph_nodes = reinterpret_cast<kernel::SubGraphKernel *>(kernel)->nodes();
for (auto node : subgraph_nodes) {
auto ret = node->Init();
if (ret != RET_OK) {
MS_LOG(ERROR) << "Kernel " << node->name() << " Init failed.";
return ret;
}
}
continue;
}
// delegate graph kernel
if (kernel->desc().delegate != nullptr) {
continue;
}
// origin inner kernel
auto ret = kernel->Init();
if (ret != RET_OK) {
MS_LOG(ERROR) << "Kernel " << kernel->name() << " Init failed.";
return ret;
}
}
return RET_OK;
}
int Scheduler::ReplaceDelegateKernels(std::vector<kernel::LiteKernel *> *dst_kernels) {
std::vector<kernel::Kernel *> kernels;
for (size_t i = 0; i < dst_kernels->size(); i++) {
@ -665,9 +700,6 @@ int Scheduler::ScheduleSubGraphToKernels(size_t subgraph_index, std::vector<kern
kernel = SchedulePartialToKernel(node);
} else { // kernel
kernel = ScheduleNodeToKernel(node, prefer_data_type);
if ((kernel != nullptr) && (!is_train_session_)) {
ret = kernel->Init();
}
}
if (kernel == nullptr || ret != RET_OK) {
MS_LOG(ERROR) << "FindBackendKernel return nullptr, name: " << node->name_

View File

@ -71,6 +71,8 @@ class Scheduler {
int ReplaceDelegateKernels(std::vector<kernel::LiteKernel *> *dst_kernels);
int InitKernels(std::vector<kernel::LiteKernel *> dst_kernels);
// schedule a partial node to a subgraph_kernel
kernel::LiteKernel *SchedulePartialToKernel(const lite::Model::Node *src_node);
// schedule a node to a kernel