[Dy2St][PIR] Add restore_out in PIR sot_call (#63190)

PaddlePaddle · Apr 2, 2024 · a1f5cdb · a1f5cdb
1 parent 5dfe454
commit a1f5cdb
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 7 deletions.
diff --git a/python/paddle/jit/dy2static/pir_partial_program.py b/python/paddle/jit/dy2static/pir_partial_program.py
@@ -503,7 +503,8 @@ def sot_call(self, inputs):
             self._cuda_graph_vec,
             *attrs,
         )
-        return out_vars
+        restored_nest_out = self._restore_out(out_vars)
+        return restored_nest_out
 
     @cached_property
     def origin_runnable_program(self):

diff --git a/python/paddle/jit/sot/infer_meta.py b/python/paddle/jit/sot/infer_meta.py
@@ -266,11 +266,16 @@ def infer_meta_for_layer(layer, *args, **kwargs):
         partial_program_layer,
     ) = layer.forward.get_concrete_program(*args_, **kwargs_)
 
+    if use_pir_api():
+        output_values = partial_program_layer._outputs.var_list
+    else:
+        output_values = concrete_program.outputs
+
     out = partial_program_layer._restore_out(
         [
             x
             for x in paddle.utils.flatten(
-                convert_variable_to_meta_info(concrete_program.outputs)
+                convert_variable_to_meta_info(output_values)
             )
             if isinstance(x, MetaInfo)
         ]

diff --git a/test/dygraph_to_static/test_duplicate_output.py b/test/dygraph_to_static/test_duplicate_output.py
@@ -24,11 +24,6 @@
 
 np.random.seed(1)
 
-if paddle.base.is_compiled_with_cuda():
-    place = paddle.base.CUDAPlace(0)
-else:
-    place = paddle.base.CPUPlace()
-
 
 class SimpleNet(paddle.nn.Layer):
     def __init__(self):
@@ -41,6 +36,17 @@ def forward(self, x):
         return x, x
 
 
+class DuplicateOutputInPaddleLayer(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+        # In GRUCell, the output is a tuple (h, h)
+        self.layer = paddle.nn.GRUCell(10, 20)
+
+    def forward(self, x):
+        x = self.layer(x)
+        return x
+
+
 class TestDuplicateOutput(Dy2StTestBase):
     def _run_static(self):
         net = paddle.jit.to_static(SimpleNet())
@@ -58,5 +64,19 @@ def test_ast_to_func(self):
         self._run_static()
 
 
+class TestDuplicateOutputInPaddleLayer(Dy2StTestBase):
+    def check_dygraph_and_static_result(self, net, x):
+        static_net = paddle.jit.to_static(net)
+        dy_out = net(x)
+        st_out = static_net(x)
+        np.testing.assert_allclose(dy_out, st_out)
+
+    @test_legacy_and_pt_and_pir
+    def test_ast_to_func(self):
+        net = DuplicateOutputInPaddleLayer()
+        x = paddle.randn([10, 10])
+        self.check_dygraph_and_static_result(net, x)
+
+
 if __name__ == '__main__':
     unittest.main()