CINN add fetch op for skip gc vars

paddle/fluid/framework/paddle2cinn/build_cinn_pass.cc

@@ -485,7 +485,8 @@ void AnalyseClusterVariables(
     GraphNodeSet* cluster_inputs,
     GraphNodeSet* cluster_outputs,
     GraphNodeSet* cluster_internals,
-    bool is_inference_stage) {
+    bool is_inference_stage,
+    const std::unordered_set<std::string>& skip_gc_var_names) {
   // collecting all input and output of op
   for (auto* op_node : cluster) {
     const auto& op_name = op_node->Name();
@@ -510,7 +511,8 @@ void AnalyseClusterVariables(
 
       // the internal node is must an output node of sub-graph,
       // but not any input node of out-graph.
-      bool is_only_used_internal = true;
+      // And should not in skip gc var
+      bool is_only_used_internal = !skip_gc_var_names.count(var_node->Name());
       for (auto* next_op_node : var_node->outputs) {
         is_only_used_internal &= (cluster.count(next_op_node) > 0);
       }
@@ -679,13 +681,20 @@ void SearchAllSubgraphs(Graph* graph, bool is_inference_stage) {
 
     auto deny_var_set = trans_info.GetDenyVarNames(cluster_set);
 
+    std::unordered_set<std::string> skip_gc_var_names;
+    if (graph->Has(kSkipGcVarNames)) {
+      skip_gc_var_names =
+          graph->Get<std::unordered_set<std::string>>(kSkipGcVarNames);
+    }
+
     GraphNodeSet cluster_inputs, cluster_outputs, cluster_internals;
     AnalyseClusterVariables(cluster_set,
                             deny_var_set,
                             &cluster_inputs,
                             &cluster_outputs,
                             &cluster_internals,
-                            is_inference_stage);
+                            is_inference_stage,
+                            skip_gc_var_names);
 
     VLOG(4) << "Cluster Ops: " << cluster_debug_info(cluster_set);
     VLOG(4) << "Cluster input vars: " << cluster_debug_info(cluster_inputs);

paddle/fluid/framework/paddle2cinn/build_cinn_pass_test.cc

@@ -670,6 +670,45 @@ TEST(BuildCinnPassTest, NoNeedBufferInput) {
       std::unordered_set<std::string>({"var5", "var6"}));
 }
 
+TEST(BuildCinnPassTest, TestSkipGcVars) {
+  auto g = BuildGraphWithOneCinnSubgraph();
+
+  std::unordered_set<std::string> all_skip_gc_vars = {"var1", "var3"};
+  g->SetNotOwned(kSkipGcVarNames, &all_skip_gc_vars);
+
+  auto pass =
+      paddle::framework::ir::PassRegistry::Instance().Get("build_cinn_pass");
+  pass->Apply(g.get());
+
+  // After search, the graph should as following
+  // fake1 --> v1 --
+  //                | --> kCinnLaunchOp --> v4 --> fake2
+  //           v2 --
+  const auto& nodes = g->Nodes();
+  ASSERT_EQ(nodes.size(), static_cast<size_t>(7));
+  ASSERT_TRUE(CheckGraphIndependence(nodes));
+
+  // A new op named kCinnLaunchOp should be added
+  ASSERT_TRUE(CheckNodeExisted(nodes, kCinnLaunchOp));
+
+  // After search, there should has just one cinn subgraph
+  // feed --> v1 --
+  //               | --> mul --> v3 --> relu --> v4 --> fetch
+  //          v2 --
+  auto compilation_keys = GetCompilationKeys(*g);
+  ASSERT_EQ(compilation_keys.size(), static_cast<size_t>(1));
+  auto* cinn_compiler = CinnCompiler::GetInstance();
+  const auto& subgraph = cinn_compiler->FindGraph(compilation_keys[0]);
+
+  const auto& subnodes = subgraph.Nodes();
+  ASSERT_EQ(subnodes.size(), static_cast<size_t>(10));
+  ASSERT_TRUE(CheckGraphIndependence(subnodes));
+
+  ASSERT_EQ(CountNode(subnodes, "feed"), 2);
+  // var3 and var4 should has fetch op
+  ASSERT_EQ(CountNode(subnodes, "fetch"), 2);
+}
+
 }  // namespace paddle2cinn
 }  // namespace framework
 }  // namespace paddle