!30549 fix the link bug of depend on depend node in the control flow

Merge pull request !30549 from limingqi107/new_actor_runtime
This commit is contained in:
i-robot 2022-02-25 07:48:08 +00:00 committed by Gitee
commit f52ca3053c
No known key found for this signature in database
GPG Key ID: 173E9B9CA92EEF8F
1 changed files with 7 additions and 0 deletions

View File

@ -69,9 +69,16 @@ void FetchRealDependNodeByAutoMonad(const AnfNodePtr &node, std::set<AnfNodePtr>
return;
}
const mindspore::HashSet<PrimitivePtr, PrimitiveHasher, PrimitiveEqual> recursion_prims = {
prim::kPrimDepend, prim::kPrimUpdateState, prim::kPrimLoad, prim::kPrimMakeTuple};
if (AnfAlgo::CheckPrimitiveType(real_node, prim::kPrimDepend) ||
AnfAlgo::CheckPrimitiveType(real_node, prim::kPrimLoad)) {
FetchRealDependNodeByAutoMonad(real_inputs[kDependAttachNodeIndex], depend_nodes);
// The real input may be this scene: depend/load --> load/depend, so need add the control arrow for real input
// node in this scene.
if (IsOneOfPrimitiveCNode(real_inputs[kRealInputIndexInDepend], recursion_prims)) {
FetchRealDependNodeByAutoMonad(real_inputs[kRealInputIndexInDepend], depend_nodes);
}
} else if (AnfAlgo::CheckPrimitiveType(real_node, prim::kPrimUpdateState)) {
for (size_t i = kUpdateStateRealInput; i < real_inputs.size(); ++i) {
FetchRealDependNodeByAutoMonad(real_inputs[i], depend_nodes);