add fetch_targets and open tests (#64532)

changeyoung98 · web-flow · commit 7b73a73a43fd · 2024-05-23T14:17:43.000+08:00
diff --git a/python/paddle/static/pir_io.py b/python/paddle/static/pir_io.py
@@ -74,12 +74,15 @@ def get_pir_parameters(program):
     return params, opts
 
 
-def get_pir_feed_names(program):
+def get_pir_feed_and_fetch(program):
     feed_name_list = []
+    fetch_targets = []
     for op in program.global_block().ops:
         if op.name() == "pd_op.data" or op.name() == "pd_op.feed":
             feed_name_list.append(op.attrs()["name"])
-    return feed_name_list
+        if op.name() == "pd_op.fetch":
+            fetch_targets.extend(op.operands_source())
+    return feed_name_list, fetch_targets
 
 
 def set_var(name, ndarray):
@@ -896,6 +899,5 @@ def load_pir_inference_model(path_prefix, executor, **kwargs):
                 filename=params_filename,
             )
 
-    feed_names = get_pir_feed_names(program)
-    # pir load program has fetch op, so if use exe.run to execute load program, don't need to set fetch_list
-    return [program, feed_names, []]
+    feed_names, fetch_targets = get_pir_feed_and_fetch(program)
+    return [program, feed_names, fetch_targets]
diff --git a/test/deprecated/legacy_test/test_save_model_without_var.py b/test/deprecated/legacy_test/test_save_model_without_var.py
@@ -17,9 +17,11 @@
 
 import paddle
 from paddle import base
+from paddle.pir_utils import test_with_pir_api
 
 
 class TestSaveModelWithoutVar(unittest.TestCase):
+    @test_with_pir_api
     def test_no_var_save(self):
         data = paddle.static.data(name='data', shape=[-1, 1], dtype='float32')
         data_plus = data + 1
diff --git a/test/deprecated/legacy_test/test_static_save_load.py b/test/deprecated/legacy_test/test_static_save_load.py
@@ -1993,11 +1993,13 @@ def setUp(self):
     def tearDown(self):
         self.temp_dir.cleanup()
 
+    @test_with_pir_api
     def test_no_params(self):
-        main_program = framework.Program()
-        with framework.program_guard(main_program):
+        main_program = paddle.static.Program()
+        with paddle.static.program_guard(main_program):
             x = paddle.static.data(name="x", shape=[10, 10], dtype='float32')
-            x.desc.set_need_check_feed(False)
+            if not in_pir_mode():
+                x.desc.set_need_check_feed(False)
             y = x + x
 
             place = paddle.CPUPlace()
@@ -2012,9 +2014,25 @@ def test_no_params(self):
             ] = paddle.static.load_inference_model(self.model_path, exe)
 
             self.assertEqual(feed_target_names, ['x'])
-            self.assertEqual(fetch_targets[0].shape, (10, 10))
-            ops = [op.type for op in inference_program.block(0).ops]
-            self.assertEqual(ops, ['feed', 'elementwise_add', 'scale', 'fetch'])
+            if in_pir_mode():
+                self.assertEqual(fetch_targets[0].shape, [10, 10])
+                ops = [op.name() for op in inference_program.global_block().ops]
+                self.assertEqual(
+                    ops,
+                    [
+                        'pd_op.data',
+                        'pd_op.add',
+                        'pd_op.full',
+                        'pd_op.scale',
+                        'pd_op.fetch',
+                    ],
+                )
+            else:
+                self.assertEqual(fetch_targets[0].shape, (10, 10))
+                ops = [op.type for op in inference_program.block(0).ops]
+                self.assertEqual(
+                    ops, ['feed', 'elementwise_add', 'scale', 'fetch']
+                )
 
 
 if __name__ == '__main__':
