[Core] Fix memory_optmize_pass for reshape/reshape2 op with inplace=True (#3045)

lite/core/mir/memory_optimize_pass.cc

@@ -39,52 +39,79 @@ void MemoryOptimizePass::CollectLifeCycleByDevice(
   auto is_host = [](TargetType x) -> bool {
     return x == TARGET(kHost) || x == TARGET(kX86) || x == TARGET(kARM);
   };
-  // The vars which inputs or outputs are invalid op will not be reused.
-  auto valid_var = [&](Node* node) -> bool {
-    std::set<std::string> invalid_op = {"while",
-                                        "conditional_block",
-                                        "conditional_block_infer",
-                                        "merge_lod_tensor_infer",
-                                        "merge_lod_tensor",
-                                        "equal",
-                                        "lod_reset",
-                                        "concat",
-                                        "yolo_box",
-                                        "subgraph",
-                                        "feed",
-                                        "fetch"};
-    for (auto* tmp : node->inlinks) {
-      CHECK(tmp->IsStmt());
-      std::string op_type = tmp->AsStmt().op_info()->Type();
-      if (std::find(invalid_op.begin(), invalid_op.end(), op_type) !=
-          invalid_op.end()) {
-        return false;
+
+  // Collect the invalid input and output variables that will not be reused.
+  std::unordered_set<std::string> invalid_var_names;
+  for (auto& op_node : graph->StmtTopologicalOrder()) {
+    if (!op_node->IsStmt()) continue;
+    auto op_info = op_node->AsStmt().op_info();
+    auto op_type = op_info->Type();
+    // The all of input and output variables of the Ops will not be reused.
+    std::unordered_set<std::string> invalid_op_nodes = {
+        "while",
+        "conditional_block",
+        "conditional_block_infer",
+        "merge_lod_tensor_infer",
+        "merge_lod_tensor",
+        "equal",
+        "lod_reset",
+        "concat",
+        "yolo_box",
+        "subgraph",
+        "feed",
+        "fetch"};
+    auto invalid_op_node = invalid_op_nodes.find(op_type);
+    if (invalid_op_node != invalid_op_nodes.end()) {
+      for (auto in_var_node : op_node->inlinks) {
+        CHECK(in_var_node->IsArg());
+        invalid_var_names.insert(in_var_node->AsArg().name);
       }
+      for (auto out_var_node : op_node->outlinks) {
+        CHECK(out_var_node->IsArg());
+        invalid_var_names.insert(out_var_node->AsArg().name);
+      }
+      continue;
     }
-    for (auto* tmp : node->outlinks) {
-      CHECK(tmp->IsStmt());
-      std::string op_type = tmp->AsStmt().op_info()->Type();
-      if (std::find(invalid_op.begin(), invalid_op.end(), op_type) !=
-          invalid_op.end()) {
-        return false;
+    // The specified input and output variables of the Ops whose 'inplace' attr
+    // is true will not be reused, such as reshape/reshape2's X and Out
+    // variables
+    std::unordered_map<std::string,
+                       std::pair<std::unordered_set<std::string>,
+                                 std::unordered_set<std::string>>>
+        inplace_op_nodes = {{"reshape", {{"X"}, {"Out"}}},
+                            {"reshape2", {{"X"}, {"Out"}}}};
+    auto inplace_op_node = inplace_op_nodes.find(op_type);
+    if (inplace_op_node != inplace_op_nodes.end()) {
+      bool inplace = false;
+      if (op_info->HasAttr("inplace")) {
+        inplace = op_info->GetAttr<bool>("inplace");
+      }
+      if (inplace) {
+        for (auto& in_param_name : inplace_op_node->second.first) {
+          const auto& in_arg_names = op_info->Input(in_param_name);
+          invalid_var_names.insert(in_arg_names.begin(), in_arg_names.end());
+        }
+        for (auto& out_param_name : inplace_op_node->second.second) {
+          const auto& out_arg_names = op_info->Output(out_param_name);
+          invalid_var_names.insert(out_arg_names.begin(), out_arg_names.end());
+        }
       }
     }
-    return true;
-  };
+  }
 
   for (auto& op_node : graph->StmtTopologicalOrder()) {
     if (op_node->IsStmt()) {
-      auto inputs = op_node->inlinks;
-      auto outputs = op_node->outlinks;
-      std::vector<Node*> requires(inputs.begin(), inputs.end());
-      requires.insert(requires.end(), outputs.begin(), outputs.end());
-      for (Node* node : requires) {
-        CHECK(node->IsArg());
-        auto& arg = node->AsArg();
+      std::vector<Node*> var_nodes(op_node->inlinks.begin(),
+                                   op_node->inlinks.end());
+      var_nodes.insert(
+          var_nodes.end(), op_node->outlinks.begin(), op_node->outlinks.end());
+      for (auto* var_node : var_nodes) {
+        CHECK(var_node->IsArg());
+        auto& arg = var_node->AsArg();
         if (arg.is_weight || arg.is_persist) continue;
-        if (!valid_var(node)) continue;
         std::string var_name = arg.name;
-        TargetType target_type = node->AsArg().type->target();
+        if (invalid_var_names.count(var_name)) continue;
+        TargetType target_type = arg.type->target();
         if (is_host(target_type)) target_type = TARGET(kHost);
 
         if (!(*lifecycles)[TargetToStr(target_type)].count(var_name)) {