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

SigureMo · web-flow · commit a1f5cdb59239 · 2024-04-03T00:40:26.000+08:00
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()

Original file line number	Diff line number	Diff line change
`@@ -503,7 +503,8 @@ def sot_call(self, inputs):`
`503`	`503`	`self._cuda_graph_vec,`
`504`	`504`	`*attrs,`
`505`	`505`	`)`
`506`		`- return out_vars`
	`506`	`+ restored_nest_out = self._restore_out(out_vars)`
	`507`	`+ return restored_nest_out`
`507`	`508`
`508`	`509`	`@cached_property`
`509`	`510`	`def origin_runnable_program(self):`
Original file line number	Diff line number	Diff line change
`@@ -266,11 +266,16 @@ def infer_meta_for_layer(layer, args, *kwargs):`
`266`	`266`	`partial_program_layer,`
`267`	`267`	`) = layer.forward.get_concrete_program(args_, *kwargs_)`
`268`	`268`
	`269`	`+ if use_pir_api():`
	`270`	`+ output_values = partial_program_layer._outputs.var_list`
	`271`	`+ else:`
	`272`	`+ output_values = concrete_program.outputs`
	`273`	`+`
`269`	`274`	`out = partial_program_layer._restore_out(`
`270`	`275`	`[`
`271`	`276`	`x`
`272`	`277`	`for x in paddle.utils.flatten(`
`273`		`- convert_variable_to_meta_info(concrete_program.outputs)`
	`278`	`+ convert_variable_to_meta_info(output_values)`
`274`	`279`	`)`
`275`	`280`	`if isinstance(x, MetaInfo)`
`276`	`281`	`]`