[Dy2St] Mark VariableRefArray inner vars as no need buffer to avoid memery leak (#63395)

SigureMo · web-flow · commit d2033d0e91a2 · 2024-04-11T19:53:42.000+08:00
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -1283,6 +1283,10 @@ static auto GetNoNeedBufferValue(const ::pir::Block *whole_block,
   std::unordered_set<::pir::Value> no_need_buffer_values;
   range_block_do(
       whole_block, range, [&need_buffer_values](::pir::Operation *op) {
+        // NOTE(SigureMo): We should process the CombineOp in it's users.
+        if (op->isa<pir::CombineOp>()) {
+          return;
+        }
         if (op->HasInterface<paddle::dialect::OpYamlInfoInterface>() == false) {
           // not a OpYamlInfoInterface, can't have no_need_buffer.
           for (const auto &operand : op->operands_source()) {
@@ -1293,8 +1297,16 @@ static auto GetNoNeedBufferValue(const ::pir::Block *whole_block,
               op->dyn_cast<paddle::dialect::OpYamlInfoInterface>().GetOpInfo();
           int counter = 0;
           for (const auto &op_input_info : std::get<0>(opinfo)) {
+            auto value = op->operand_source(counter);
             if (!op_input_info.no_need_buffer) {
-              need_buffer_values.insert(op->operand_source(counter));
+              need_buffer_values.insert(value);
+              if (!IsFakeValue(value) && value.defining_op() &&
+                  value.defining_op()->isa<pir::CombineOp>()) {
+                for (const auto &combine_value :
+                     value.defining_op()->operands_source()) {
+                  need_buffer_values.insert(combine_value);
+                }
+              }
             }
             counter += 1;
           }
diff --git a/test/dygraph_to_static/test_no_need_buffer.py b/test/dygraph_to_static/test_no_need_buffer.py
@@ -0,0 +1,61 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+from dygraph_to_static_utils import (
+    Dy2StTestBase,
+    test_ast_only,
+    test_pir_only,
+)
+
+import paddle
+
+
+# The input of concat_grad is no_need_buffer
+def concat_net(x):
+    y = x + 1
+    z = paddle.concat([y, y], axis=0)
+    return z
+
+
+class TestNoNeedBuffer(Dy2StTestBase):
+    @test_ast_only
+    @test_pir_only
+    def test_no_need_buffer(self):
+        input = paddle.to_tensor([1, 2])
+        input.stop_gradient = False
+        static_fn = paddle.jit.to_static(concat_net)
+        static_res = static_fn(input)
+        dygraph_res = concat_net(input)
+        np.testing.assert_allclose(static_res.numpy(), dygraph_res.numpy())
+
+        _, partial_program_layer = static_fn.get_concrete_program(input)
+        no_need_buffers = partial_program_layer.program.program_attr[
+            "no_need_buffers"
+        ]
+        for no_need_buffer_value in no_need_buffers:
+            defining_op = no_need_buffer_value.get_defining_op()
+            # y = x + 1, it's defining op is `pd_op.scale`
+            if defining_op is not None and defining_op.name() == 'pd_op.scale':
+                break
+        else:
+            raise AssertionError(
+                "middle var `y` should be no_need_buffer value"
+            )
+
+
+if __name__ == '__main__':
+    unittest.main()
