!12920 Add more isolated nodes in Parser, and clear dependency nodes that have no side effect.

From: @zh_qh
Reviewed-by: 
Signed-off-by:

mindspore/ccsrc/pipeline/jit/parse/function_block.cc

@@ -75,10 +75,11 @@ void FunctionBlock::WriteVariable(const std::string &var_name, const AnfNodePtr
     auto is_used = iter->second.second;
     auto hidden_node = iter->second.first;
     auto is_isolated = CanBeIsolatedNode(var_name, hidden_node);
-    MS_LOG(INFO) << "Isolated node found(Hidden), hidden_node: " << hidden_node->DebugString(2) << " is hidden by "
-                 << node->DebugString(2) << " with the same name, var_name: " << var_name
-                 << ", is_isolated: " << is_isolated << ", !is_used: " << !is_used;
     if (!is_used && is_isolated) {
+      MS_LOG(INFO) << "Isolated node found(Hidden), hidden_node: " << hidden_node->DebugString(2) << " is hidden by "
+                   << node->DebugString(2) << " with the same name, var_name: " << var_name << ", block: " << this
+                   << "/" << (func_graph() ? func_graph()->ToString() : "FG(Null)")
+                   << ", Line: " << trace::GetDebugInfo(hidden_node->debug_info(), "", kSourceLineTipDiscard);
       AddIsolatedNode(hidden_node);
     }
     iter->second = std::make_pair(node, false);
@@ -335,8 +336,6 @@ void FunctionBlock::Jump(const FunctionBlockPtr &target_block, AnfNodePtr node)
   jumps_[target_block.get()] = jump;
   target_block->AddPrevBlock(shared_from_this());
   func_graph()->set_output(jump);
-  // Attach all isolated nodes.
-  AttachIsolatedNodesBeforeReturn();
 }
 
 // Perform a conditional jump using switch operation.
@@ -352,8 +351,6 @@ void FunctionBlock::ConditionalJump(AnfNodePtr condNode, const FunctionBlockPtr
                                    NewValueNode(false_block->func_graph())});
   CNodePtr switch_app_new = func_graph()->NewCNodeInOrder({switch_app});
   func_graph()->set_output(switch_app_new);
-  // Attach all isolated nodes.
-  AttachIsolatedNodesBeforeReturn();
 }
 
 // Create cnode for the assign statement like 'self.target = source'.
@@ -363,7 +360,9 @@ void FunctionBlock::SetStateAssign(const AnfNodePtr &target, const AnfNodePtr &s
   const std::string module_name("mindspore.ops.functional");
   ValueNodePtr assign_op = NewValueNode(prim::GetPythonOps(primitive_name, module_name, true));
   auto assign_node = func_graph_->NewCNodeInOrder({assign_op, target, source});
-  MS_LOG(DEBUG) << "Isolated node found(Assign), assign_node: " << assign_node->DebugString(2);
+  MS_LOG(DEBUG) << "Isolated node found(Assign), assign_node: " << assign_node->DebugString(2) << ", block: " << this
+                << "/" << (func_graph() ? func_graph()->ToString() : "FG(Null)")
+                << ", Line: " << trace::GetDebugInfo(assign_node->debug_info(), "", kSourceLineTipDiscard);
   AddIsolatedNode(assign_node);
 }
 
@@ -399,7 +398,10 @@ void FunctionBlock::FindIsolatedNodes() {
     if (used.find(node) == used.end() && CanBeIsolatedNode(var_name, node)) {
       // We don't call AddIsolatedNode(node) anymore.
       // If need, to call FindIsolatedNodes() in appropriate place.
-      MS_LOG(INFO) << "Isolated node found(NoUse), node: " << node->DebugString(2) << ", var_name: " << var_name;
+      MS_LOG(INFO) << "Isolated node found(NoUse), node: " << node->DebugString(2) << ", var_name: " << var_name
+                   << ", block: " << this << "/" << (func_graph() ? func_graph()->ToString() : "FG(Null)")
+                   << ", Line: " << trace::GetDebugInfo(node->debug_info(), "", kSourceLineTipDiscard);
+      AddIsolatedNode(node);
     }
   }
 }

mindspore/ccsrc/pipeline/jit/parse/parse.cc

@@ -212,12 +212,6 @@ FuncGraphPtr Parser::ParseFuncGraph() {
     return nullptr;
   }
 
-  // Add unused variables as isolate nodes.
-  for (auto &block : func_block_list_) {
-    // Find unused variables.
-    block->FindIsolatedNodes();
-  }
-
   RemoveUnnecessaryPhis();
 
   MS_EXCEPTION_IF_NULL(pFnBlock);
@@ -338,6 +332,16 @@ FunctionBlockPtr Parser::ParseFunction(const py::object &node, const FunctionBlo
   py::object funcObj = python_adapter::GetPyObjAttr(node, "body");
   (void)ParseStatements(pFunBlock, funcObj);
 
+  // Add unused variables as isolate nodes.
+  for (auto &func_block : func_block_list_) {
+    if (func_block->func_graph()->get_return() != nullptr) {
+      // Find unused variables.
+      func_block->FindIsolatedNodes();
+      // Attach all isolated nodes.
+      func_block->AttachIsolatedNodesBeforeReturn();
+    }
+  }
+
   if (current_fg->get_return() == nullptr) {
     py::list ret = ast_->CallParserObjMethod(PYTHON_PARSE_GET_LOCATION, node);
     py::str desc = python_adapter::CallPyModFn(ast_->module(), PYTHON_MOD_GET_OBJECT_DESCRIPTION, node, ret[0], ret[1]);
@@ -356,8 +360,6 @@ FunctionBlockPtr Parser::ParseStatements(FunctionBlockPtr block, const py::objec
     block = ParseStatement(block, node);
     // Insert appropriate depended items for the function block if it has a return node
     if (block->func_graph()->get_return() != nullptr) {
-      // Attach all isolated nodes.
-      block->AttachIsolatedNodesBeforeReturn();
       // Skip statements after 'return' (or 'break', 'continue').
       break;
     }
@@ -436,7 +438,10 @@ FunctionBlockPtr Parser::ParseExpr(const FunctionBlockPtr &block, const py::obje
       // e.g.: print(x)
       // We save it as an isolated node.
       auto &no_return_node = call_node;
-      MS_LOG(INFO) << "Isolated node found(NoReturn), no_return_node: " << no_return_node->DebugString(2);
+      MS_LOG(INFO) << "Isolated node found(NoReturn), no_return_node: " << no_return_node->DebugString(2)
+                   << ", block: " << block << "/"
+                   << (block->func_graph() ? block->func_graph()->ToString() : "FG(Null)")
+                   << ", Line: " << trace::GetDebugInfo(no_return_node->debug_info(), "", kSourceLineTipDiscard);
       block->AddIsolatedNode(no_return_node);
     } else {
       // Expand the assign statement,

mindspore/ccsrc/pipeline/jit/static_analysis/auto_monad.cc

@@ -1013,6 +1013,10 @@ class AutoMonadConverter {
     if (HasSideEffects()) {
       HandleCNodes();
     }
+
+    // Safe to clear isolated nodes after handled side effect nodes.
+    ClearIsolatedNodes();
+
     // Clean up after conversion finished.
     func_graph_->ClearOrderList();
     return has_effect_cnodes_;
@@ -1077,6 +1081,22 @@ class AutoMonadConverter {
     }
   }
 
+  // Clean no side effect dependency nodes.
+  //   From:  output = Depend(output, StopGrad)
+  //          return output
+  //
+  //   To:    return output
+  void ClearIsolatedNodes() {
+    auto output = GetGraphOutput();
+    if (IsPrimitiveCNode(output, prim::kPrimDepend) &&
+        IsPrimitiveCNode(output->cast<CNodePtr>()->input(2), prim::kPrimStopGradient)) {
+      // Replace Depend(orig_output, StopGrad) node with orig_output.
+      // After that, nodes may be eliminated if have no side effects.
+      auto &orig_output = output->cast<CNodePtr>()->input(1);
+      func_graph_->set_output(orig_output);
+    }
+  }
+
   void HandleOuterNode(const CNodePtr &cnode, const EffectInfo &info) {
     if (info.memory || info.load) {
       (void)GetUniverse();
@@ -1248,16 +1268,19 @@ class AutoMonadConverter {
 
   void InsertStateDepend(const AnfNodePtr &state) {
     auto output = GetGraphOutput();
-    // It's safe to handle isolated nodes here:
-    //   Node: Depend(output, StopGrad)
+    auto depend = NewValueNode(prim::kPrimDepend);
+    // If isolated nodes dependencies exist.
     if (IsPrimitiveCNode(output, prim::kPrimDepend) &&
         IsPrimitiveCNode(output->cast<CNodePtr>()->input(2), prim::kPrimStopGradient)) {
-      // Replace Depend(orig_output, StopGrad) node with orig_output.
-      // After that, nodes may be eliminated if have no side effects.
-      output = output->cast<CNodePtr>()->input(1);
+      // Insert state Depend node into isolated Depend node.
+      auto isolated_depend = output->cast<CNodePtr>();
+      auto &orig_output = isolated_depend->input(1);
+      auto state_depend = func_graph_->NewCNode({depend, orig_output, state});
+      state_depend->set_abstract(orig_output->abstract());
+      manager_->SetEdge(isolated_depend, 1, state_depend);
+      return;
     }
-    // Insert Depend node and set it as output.
-    auto depend = NewValueNode(prim::kPrimDepend);
+    // Insert Depend node and set it as output, if no isolated nodes.
     auto depend_cnode = func_graph_->NewCNode({depend, output, state});
     depend_cnode->set_abstract(output->abstract());
     func_graph_->set_output(depend_cnode);