fix env_getitem has not been eliminated

mindspore/ccsrc/backend/session/ascend_auto_monad.cc

@@ -802,6 +802,24 @@ class ExecuteOrderGenerator {
     graph_->set_execution_order(std::move(execution_order));
   }
 
+  std::set<CNodePtr> GetAllNodes() {
+    auto &all_graphs = context_.visited_graphs();
+    std::set<CNodePtr> all_nodes;
+    for (auto &graph : all_graphs) {
+      auto out = graph->get_return();
+      MS_EXCEPTION_IF_NULL(out);
+      auto nodes = TopoSort(out);
+      for (auto &node : nodes) {
+        MS_EXCEPTION_IF_NULL(node);
+        auto cnode = node->cast<CNodePtr>();
+        if (cnode != nullptr) {
+          all_nodes.insert(cnode);
+        }
+      }
+    }
+    return all_nodes;
+  }
+
   static const AnfNodePtr &GetRealNode(const AnfNodePtr &input) {
     if (IsPrimitiveCNode(input, prim::kPrimLoad) || IsPrimitiveCNode(input, prim::kPrimDepend)) {
       return input->cast<CNodePtr>()->inputs().at(1);
@@ -813,6 +831,7 @@ class ExecuteOrderGenerator {
   void EraseParameter() {
     // Copy out execution order list.
     auto exec_order = graph_->execution_order();
+    std::set<CNodePtr> all_nodes = GetAllNodes();
 
     // Remove assigns that target and source are same.
     for (auto iter = exec_order.begin(); iter != exec_order.end();) {
@@ -844,6 +863,18 @@ class ExecuteOrderGenerator {
           auto kg = target->func_graph()->cast<KernelGraphPtr>();
           MS_EXCEPTION_IF_NULL(kg);
           kg->ReplaceNode(NOT_NULL(target), NOT_NULL(source));
+
+          // replace parameter in node
+          for (auto &iter_node : all_nodes) {
+            for (size_t i = 0; i < iter_node->size(); ++i) {
+              if (iter_node->input(i) == target) {
+                MS_LOG(INFO) << "Replace " << iter_node->DebugString() << " input " << i << " by "
+                             << source->DebugString();
+                iter_node->set_input(i, source);
+              }
+            }
+          }
+
           // replace parameter in graph input
           auto &all_graphs = context_.visited_graphs();
           for (auto &g : all_graphs) {

mindspore/ccsrc/frontend/optimizer/irpass/updatestate_eliminate.cc

@@ -84,23 +84,27 @@ AnfNodePtr EliminateUpdateStateOnlyUsedNode(const CNodePtr &update_state, const
 //   user(u)
 AnfNodePtr EliminateUpdateStateForPureNode(const CNodePtr &update_state, const AnfNodePtr &attach) {
   if (IsPrimitiveCNode(attach, prim::kPrimTupleGetItem)) {
-    // Skip tuple_getitem.
-    return nullptr;
-  }
-  auto cnode = dyn_cast<CNode>(attach);
-  if (cnode == nullptr) {
-    // Skip value node or parameter.
-    return nullptr;
-  }
-  if (cnode->size() > 1) {
-    // If the last input is a monad, means the attach node has side-effect and
-    // we should keep UpdateState; otherwise, we will remove the UpdateState.
-    if (HasAbstractMonad(cnode->inputs().back())) {
+    auto tuple_getitem_cnode = attach->cast<CNodePtr>();
+    auto mgr = GetManager(attach);
+    if (mgr == nullptr) {
+      return nullptr;
+    }
+    if (!OnlyUpdateStateUse(update_state, attach)) {
+      // Skip if UpdateState is not the only user of cnode.
       return nullptr;
     }
+    auto &node_users = mgr->node_users();
+    auto iter = node_users.find(tuple_getitem_cnode->input(1));
+    if (iter == node_users.end()) {
+      return nullptr;
+    }
+    auto &partial_users = iter->second;
+    if (partial_users.size() > 1) {
+      // Remove UpdateState by replace it with its input monad.
+      return update_state->input(kInputIndex);
+    }
   }
-  // Remove UpdateState by replace it with its input monad.
-  return update_state->input(kInputIndex);
+  return nullptr;
 }
 
 // Eliminate redundant UpdateState/Depend pair nodes caused by inline.