diff --git a/mindspore/ccsrc/pipeline/jit/parse/function_block.cc b/mindspore/ccsrc/pipeline/jit/parse/function_block.cc index 657d1e4017..88c55e0685 100644 --- a/mindspore/ccsrc/pipeline/jit/parse/function_block.cc +++ b/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); } } } diff --git a/mindspore/ccsrc/pipeline/jit/parse/parse.cc b/mindspore/ccsrc/pipeline/jit/parse/parse.cc index b29cce39fe..64347e9d51 100644 --- a/mindspore/ccsrc/pipeline/jit/parse/parse.cc +++ b/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, diff --git a/mindspore/ccsrc/pipeline/jit/static_analysis/auto_monad.cc b/mindspore/ccsrc/pipeline/jit/static_analysis/auto_monad.cc index a30df55512..cbda67bbad 100644 --- a/mindspore/ccsrc/pipeline/jit/static_analysis/auto_monad.cc +++ b/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()->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()->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()->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()->input(1); + // Insert state Depend node into isolated Depend node. + auto isolated_depend = output->cast(); + 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);