fix heterogeneous kernel by kernel

This commit is contained in:
baihuawei 2022-01-29 17:04:56 +08:00
parent 380fb6982a
commit dc383dd573
1 changed files with 13 additions and 0 deletions

View File

@ -700,6 +700,19 @@ AnfNodePtr SessionBasic::CreateNewParameterFromCNode(const AnfNodePtr &anf, Kern
MS_EXCEPTION_IF_NULL(anf);
MS_EXCEPTION_IF_NULL(graph);
MS_LOG(INFO) << "Create a new parameter from cnode[" << anf->DebugString() << "]";
if (IsPrimitiveCNode(anf, prim::kPrimLoad)) {
auto input = AnfAlgo::GetInputNode(anf->cast<CNodePtr>(), 0);
MS_EXCEPTION_IF_NULL(input);
if (input->isa<Parameter>()) {
auto new_param = CreateNewParameterFromParameter(input, graph);
auto context_ptr = MsContext::GetInstance();
MS_EXCEPTION_IF_NULL(context_ptr);
if (context_ptr->get_param<bool>(MS_CTX_ENABLE_MINDRT) == true) {
graph->CacheInternalParameterToFrontNode(new_param, {anf, 0});
}
return new_param;
}
}
return CreateParameterFromTuple(anf, graph);
}