【pir_save_load】paddle.load and paddle.static.load_inference_model support load pdmodel translate to pir (Compatibility) (#65671)

xiaoguoguo626807 · changeyoung98 · web-flow · commit 38720019a843 · 2024-07-05T14:05:27.000+08:00
* save old load new

* bug fix

* add pdmodel translate pir

* modify flag

* adjust test order aviod segmentfault

* modify pir py3

---------

Co-authored-by: chenzhiyang &lt;1792266893@qq.com&gt;
diff --git a/paddle/fluid/pir/serialize_deserialize/src/ir_deserialize.cc b/paddle/fluid/pir/serialize_deserialize/src/ir_deserialize.cc
@@ -112,7 +112,8 @@ pir::Operation* ProgramReader::ReadParameterOp(Json* op_json) {
       attrs_json.size(),
       4,
       common::errors::InvalidArgument(
-          "builtin ParameterOp has %d 's  attributes, which should be 4"));
+          "builtin ParameterOp has %d 's  attributes, which should be 4",
+          attrs_json.size()));
   pir::AttributeMap attributes;
   pir::IrContext* ctx = pir::IrContext::Instance();
   attributes.insert(
diff --git a/python/paddle/framework/__init__.py b/python/paddle/framework/__init__.py
@@ -54,6 +54,7 @@
     get_flags,
     in_dygraph_mode as in_dynamic_mode,
     in_dynamic_or_pir_mode,
+    in_pir_executor_mode,
     in_pir_mode,
     set_flags,
     switch_main_program,
diff --git a/python/paddle/framework/io.py b/python/paddle/framework/io.py
@@ -1226,6 +1226,32 @@ def load(path, **configs):
                             program = Program.parse_from_string(
                                 program_desc_str
                             )
+                            if paddle.framework.in_pir_executor_mode():
+                                with paddle.pir_utils.IrGuard():
+                                    program = paddle.pir.translate_to_pir(
+                                        program.desc
+                                    )
+                                    block = program.global_block()
+                                    remove_op_list = []
+                                    for op in block.ops:
+                                        if op.name() == "pd_op.feed":
+                                            var_name = op.attrs()["name"]
+                                            org_value = op.result(0)
+                                            with block:
+                                                value = paddle.static.data(
+                                                    name=var_name,
+                                                    shape=org_value.shape,
+                                                    dtype=org_value.dtype,
+                                                )
+                                                org_value.replace_all_uses_with(
+                                                    value
+                                                )
+                                                value.get_defining_op().move_before(
+                                                    op
+                                                )
+                                            remove_op_list.append(op)
+                                    for op in remove_op_list:
+                                        block.remove_op(op)
                             return program
                     except:
                         raise ValueError(
diff --git a/python/paddle/static/io.py b/python/paddle/static/io.py
@@ -945,13 +945,37 @@ def load_inference_model(path_prefix, executor, **kwargs):
                 predicate=is_persistable,
                 filename=params_filename,
             )
-
     feed_target_names = program.desc.get_feed_target_names()
-    fetch_target_names = program.desc.get_fetch_target_names()
-    fetch_targets = [
-        program.global_block().var(name) for name in fetch_target_names
-    ]
+    if paddle.framework.in_pir_executor_mode():
+        with paddle.pir_utils.IrGuard():
+            program = paddle.pir.translate_to_pir(program.desc)
+            block = program.global_block()
+            remove_op_list = []
+            fetch_targets = []
+            for op in block.ops:
+                if op.name() == "pd_op.feed":
+                    var_name = op.attrs()["name"]
+                    org_value = op.result(0)
+                    with block:
+                        value = paddle.static.data(
+                            name=var_name,
+                            shape=org_value.shape,
+                            dtype=org_value.dtype,
+                        )
+                        org_value.replace_all_uses_with(value)
+                        value.get_defining_op().move_before(op)
+                    remove_op_list.append(op)
+            for op in remove_op_list:
+                block.remove_op(op)
+            for op in block.ops:
+                if op.name() == "pd_op.fetch":
+                    fetch_targets.append(op.operand_source(0))
 
+    else:
+        fetch_target_names = program.desc.get_fetch_target_names()
+        fetch_targets = [
+            program.global_block().var(name) for name in fetch_target_names
+        ]
     return [program, feed_target_names, fetch_targets]
 
 
diff --git a/test/ir/pir/test_ir_load_oldir.py b/test/ir/pir/test_ir_load_oldir.py
@@ -0,0 +1,132 @@
+# 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 os
+import tempfile
+import unittest
+
+import numpy as np
+
+import paddle
+
+
+class TestALoadPdmodelTranslatePir(unittest.TestCase):
+    def setUp(self):
+        paddle.seed(2022)
+        self.temp_dir = tempfile.TemporaryDirectory()
+        self.save_path = os.path.join(self.temp_dir.name, 'saveload')
+        self.place = (
+            paddle.CUDAPlace(0)
+            if paddle.is_compiled_with_cuda()
+            else paddle.CPUPlace()
+        )
+
+    def tearDown(self):
+        self.temp_dir.cleanup()
+
+    def test_load_inference_model(self):
+        paddle.enable_static()
+        if paddle.framework.in_pir_mode():
+            return
+        np_x = np.random.randn(9, 10, 11).astype('float32')
+        main_prog = paddle.static.Program()
+        startup_prog = paddle.static.Program()
+        with paddle.static.program_guard(main_prog, startup_prog):
+            x = paddle.static.data(shape=np_x.shape, name='x', dtype=np_x.dtype)
+            linear = paddle.nn.Linear(np_x.shape[-1], np_x.shape[-1])
+            linear_out = linear(x)
+            relu_out = paddle.nn.functional.relu(linear_out)
+            axis = paddle.full([1], 2, dtype='int64')
+            out = paddle.cumsum(relu_out, axis=axis)
+            loss = paddle.mean(out)
+            sgd = paddle.optimizer.SGD(learning_rate=0.0)
+            sgd.minimize(paddle.mean(out))
+
+            exe = paddle.static.Executor(self.place)
+            exe.run(startup_prog)
+            out_old = exe.run(feed={'x': np_x}, fetch_list=[out])
+
+            # run infer
+            paddle.static.save_inference_model(
+                self.save_path, [x], [out], exe, program=main_prog
+            )
+
+            exe = paddle.static.Executor(self.place)
+
+            load_program, _, _ = paddle.static.load_inference_model(
+                self.save_path, exe
+            )
+
+        with paddle.pir_utils.IrGuard():
+            startup_prog = paddle.static.Program()
+            with paddle.static.program_guard(load_program, startup_prog):
+                exe.run(startup_prog)
+                out_new = exe.run(load_program, feed={'x': np_x}, fetch_list=[])
+                np.testing.assert_allclose(out_old, out_new)
+
+        load_program = paddle.load(
+            self.save_path + '.pdmodel',
+        )
+
+        with paddle.pir_utils.IrGuard():
+            startup_prog = paddle.static.Program()
+            with paddle.static.program_guard(load_program, startup_prog):
+                exe.run(startup_prog)
+                out_new = exe.run(load_program, feed={'x': np_x}, fetch_list=[])
+                np.testing.assert_allclose(out_old, out_new)
+
+
+class TestJitSaveOp(unittest.TestCase):
+    def setUp(self):
+        self.temp_dir = tempfile.TemporaryDirectory()
+        self.model_path = os.path.join(self.temp_dir.name, "pir_save_load")
+        paddle.disable_static()
+        linear = paddle.nn.Linear(10, 10)
+        path = os.path.join(self.model_path, "linear")
+
+        paddle.jit.save(
+            linear,
+            path,
+            input_spec=[paddle.static.InputSpec([10, 10], 'float32', 'x')],
+        )
+
+    def tearDown(self):
+        self.temp_dir.cleanup()
+
+    def test_with_pir(self):
+        paddle.enable_static()
+        if paddle.framework.in_pir_mode():
+            return
+        place = (
+            paddle.CUDAPlace(0)
+            if paddle.is_compiled_with_cuda()
+            else paddle.CPUPlace()
+        )
+
+        exe = paddle.static.Executor(place)
+
+        [
+            inference_program,
+            feed_target_names,
+            fetch_targets,
+        ] = paddle.static.io.load_inference_model(
+            self.model_path,
+            executor=exe,
+            model_filename="linear.pdmodel",
+            params_filename="linear.pdiparams",
+        )
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/test/ir/pir/test_standalone_pir.py b/test/ir/pir/test_standalone_pir.py
@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
-import tempfile
 import unittest
 
 import numpy as np
@@ -287,47 +285,6 @@ def test_with_pir(self):
         np.testing.assert_array_equal(out[0], gold_res)
 
 
-class TestJitSaveOp(unittest.TestCase):
-    def setUp(self):
-        self.temp_dir = tempfile.TemporaryDirectory()
-        self.model_path = os.path.join(self.temp_dir.name, "pir_save_load")
-
-    def tearDown(self):
-        self.temp_dir.cleanup()
-
-    def test_with_pir(self):
-        with paddle.pir_utils.OldIrGuard():
-            paddle.disable_static()
-            linear = paddle.nn.Linear(10, 10)
-            path = os.path.join(self.model_path, "linear")
-
-            paddle.jit.save(
-                linear,
-                path,
-                input_spec=[paddle.static.InputSpec([10, 10], 'float32', 'x')],
-            )
-
-            paddle.enable_static()
-            place = (
-                paddle.CUDAPlace(0)
-                if paddle.is_compiled_with_cuda()
-                else paddle.CPUPlace()
-            )
-
-            exe = paddle.static.Executor(place)
-
-            [
-                inference_program,
-                feed_target_names,
-                fetch_targets,
-            ] = paddle.static.io.load_inference_model(
-                self.model_path,
-                executor=exe,
-                model_filename="linear.pdmodel",
-                params_filename="linear.pdiparams",
-            )
-
-
 class TestPirConcatDygraph(unittest.TestCase):
     def test_with_pir(self):
         paddle.disable_static()
