[PIR] refine pir guard (#63336)

wanghuancoder · web-flow · commit 81b52b61958c · 2024-04-10T12:54:12.000+08:00
* pir guard
diff --git a/python/paddle/pir_utils.py b/python/paddle/pir_utils.py
@@ -19,99 +19,114 @@
 from paddle.framework.dtype import bind_datatype, bind_vartype
 
 
-class IrGuard:
-    def __init__(self):
-        self.in_dygraph_outside = False
-        old_flag = paddle.base.framework.get_flags("FLAGS_enable_pir_api")
-        paddle.base.framework.set_flags({"FLAGS_enable_pir_api": False})
-        paddle.base.framework.global_var._use_pir_api_ = False
-        if not paddle.base.framework.get_flags("FLAGS_enable_pir_api")[
-            "FLAGS_enable_pir_api"
-        ]:
-            self.old_Program = paddle.static.Program
-            self.old_program_guard = paddle.base.program_guard
-            self.old_default_main_program = paddle.static.default_main_program
-            self.old_default_startup_program = (
-                paddle.static.default_startup_program
-            )
-        else:
-            raise RuntimeError(
-                "IrGuard only init when paddle.framework.in_pir_mode(): is false, \
-                please set FLAGS_enable_pir_api = false"
-            )
-        paddle.base.framework.set_flags(old_flag)
-        paddle.base.framework.global_var._use_pir_api_ = old_flag[
-            "FLAGS_enable_pir_api"
-        ]
+def _switch_to_pir_():
+    paddle.base.framework.global_var._use_pir_api_ = True
+    paddle.framework.set_flags({"FLAGS_enable_pir_in_executor": True})
+    paddle.pir.register_paddle_dialect()
+    # TODO find a better place to init the registion of dist dialect.
+    paddle.pir.register_dist_dialect()
 
+    paddle.base.Program = paddle.pir.Program
+    paddle.base.program_guard = paddle.pir.core.program_guard
+    paddle.base.default_main_program = paddle.pir.core.default_main_program
+    paddle.base.default_startup_program = (
+        paddle.pir.core.default_startup_program
+    )
+    paddle.static.Program = paddle.pir.Program
+    paddle.static.program_guard = paddle.pir.core.program_guard
+    paddle.static.default_main_program = paddle.pir.core.default_main_program
+    paddle.static.default_startup_program = (
+        paddle.pir.core.default_startup_program
+    )
+
+
+def _switch_to_old_ir_():
+    paddle.base.framework.global_var._use_pir_api_ = False
+    paddle.framework.set_flags({"FLAGS_enable_pir_in_executor": False})
+
+    paddle.base.Program = paddle.base.framework.Program
+    paddle.base.program_guard = paddle.base.framework.program_guard
+    paddle.base.default_main_program = (
+        paddle.base.framework.default_main_program
+    )
+    paddle.base.default_startup_program = (
+        paddle.base.framework.default_startup_program
+    )
+    paddle.static.Program = paddle.base.framework.Program
+    paddle.static.program_guard = paddle.base.framework.program_guard
+    paddle.static.default_main_program = (
+        paddle.base.framework.default_main_program
+    )
+    paddle.static.default_startup_program = (
+        paddle.base.framework.default_startup_program
+    )
+
+
+class IrGuard:
     def __enter__(self):
         self.in_dygraph_outside = paddle.base.framework.in_dygraph_mode()
+        self.old_flag = paddle.base.framework.get_flags("FLAGS_enable_pir_api")[
+            "FLAGS_enable_pir_api"
+        ]
         if self.in_dygraph_outside:
             paddle.enable_static()
-        paddle.framework.set_flags({"FLAGS_enable_pir_api": True})
-        paddle.base.framework.global_var._use_pir_api_ = True
-        bind_datatype()
-        self._switch_to_pir()
+        if not self.old_flag:
+            paddle.framework.set_flags({"FLAGS_enable_pir_api": True})
+            paddle.base.framework.global_var._use_pir_api_ = True
+            bind_datatype()
+            self._switch_to_pir()
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        paddle.framework.set_flags({"FLAGS_enable_pir_api": False})
-        paddle.base.framework.global_var._use_pir_api_ = False
-        bind_vartype()
-        self._switch_to_old_ir()
         if self.in_dygraph_outside:
             paddle.disable_static()
+        if not self.old_flag:
+            paddle.framework.set_flags({"FLAGS_enable_pir_api": False})
+            paddle.base.framework.global_var._use_pir_api_ = False
+            bind_vartype()
+            self._switch_to_old_ir()
 
     def _switch_to_pir(self):
         if paddle.base.framework.get_flags("FLAGS_enable_pir_api")[
             "FLAGS_enable_pir_api"
         ]:
-            paddle.framework.set_flags({"FLAGS_enable_pir_in_executor": True})
-            paddle.pir.register_paddle_dialect()
-            # TODO find a better place to init the registion of dist dialect.
-            paddle.pir.register_dist_dialect()
-
-            paddle.base.Program = paddle.pir.Program
-            paddle.base.program_guard = paddle.pir.core.program_guard
-            # paddle.base.default_main_program = (
-            #     paddle.pir.core.default_main_program
-            # )
-            # paddle.base.default_startup_program = (
-            #     paddle.pir.core.default_startup_program
-            # )
-            paddle.static.Program = paddle.pir.Program
-            paddle.static.program_guard = paddle.pir.core.program_guard
-            paddle.static.default_main_program = (
-                paddle.pir.core.default_main_program
-            )
-            paddle.static.default_startup_program = (
-                paddle.pir.core.default_startup_program
-            )
+            _switch_to_pir_()
 
     def _switch_to_old_ir(self):
         if not paddle.base.framework.get_flags("FLAGS_enable_pir_api")[
             "FLAGS_enable_pir_api"
         ]:
-            paddle.framework.set_flags({"FLAGS_enable_pir_in_executor": False})
-
-            paddle.base.Program = self.old_Program
-            paddle.base.program_guard = self.old_program_guard
-            # paddle.base.default_main_program = self.old_default_main_program
-            # paddle.base.default_startup_program = (
-            #     self.old_default_startup_program
-            # )
-            paddle.static.Program = self.old_Program
-            paddle.static.program_guard = self.old_program_guard
-            paddle.static.default_main_program = self.old_default_main_program
-            paddle.static.default_startup_program = (
-                self.old_default_startup_program
-            )
+            _switch_to_old_ir_()
         else:
             raise RuntimeError(
                 "IrGuard._switch_to_old_ir only work when paddle.framework.in_pir_mode() is false, \
                 please set FLAGS_enable_pir_api = false"
             )
 
 
+class OldIrGuard:
+    def __enter__(self):
+        self.in_dygraph_outside = paddle.base.framework.in_dygraph_mode()
+        self.old_flag = paddle.base.framework.get_flags("FLAGS_enable_pir_api")[
+            "FLAGS_enable_pir_api"
+        ]
+        if self.in_dygraph_outside:
+            paddle.enable_static()
+        if self.old_flag:
+            paddle.framework.set_flags({"FLAGS_enable_pir_api": False})
+            paddle.base.framework.global_var._use_pir_api_ = False
+            bind_vartype()
+            _switch_to_old_ir_()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.in_dygraph_outside:
+            paddle.disable_static()
+        if self.old_flag:
+            paddle.framework.set_flags({"FLAGS_enable_pir_api": True})
+            paddle.base.framework.global_var._use_pir_api_ = True
+            bind_datatype()
+            _switch_to_pir_()
+
+
 def test_with_pir_api(func):
     @wraps(func)
     def impl(*args, **kwargs):
@@ -120,3 +135,12 @@ def impl(*args, **kwargs):
             func(*args, **kwargs)
 
     return impl
+
+
+def test_with_old_ir_only(func):
+    @wraps(func)
+    def impl(*args, **kwargs):
+        with OldIrGuard():
+            func(*args, **kwargs)
+
+    return impl
diff --git a/test/legacy_test/op_test.py b/test/legacy_test/op_test.py
@@ -96,7 +96,7 @@ def check_out_dtype(api_fn, in_specs, expect_dtypes, target_index=0, **configs):
         check_out_dtype(base.layers.pad_constant_like, [([2,3,2,3], 'float64'), ([1, 3, 1,3], )], ['float32', 'float64', 'int64'], target_index=1, pad_value=0.)
 
     """
-    with paddle_static_guard():
+    with paddle.pir_utils.OldIrGuard():
         for i, expect_dtype in enumerate(expect_dtypes):
             with paddle.static.program_guard(paddle.static.Program()):
                 input_t = []
@@ -1141,7 +1141,7 @@ def cal_python_api(python_api, args, kernel_sig):
             return result
 
         with base.dygraph.base.guard(place=place):
-            block = base.default_main_program().global_block()
+            block = base.framework.default_main_program().global_block()
             op_proto = OpProtoHolder.instance().get_op_proto(self.op_type)
             # prepare input variable
             dygraph_tensor_inputs = (
@@ -1208,7 +1208,7 @@ def _calc_dygraph_output(
             self.op_type
         )  # for ci check, please not delete it for now
         with base.dygraph.base.guard(place=place):
-            block = base.default_main_program().global_block()
+            block = base.framework.default_main_program().global_block()
 
             op_proto = OpProtoHolder.instance().get_op_proto(self.op_type)
 
@@ -1246,7 +1246,7 @@ def _calc_dygraph_output(
 
     def get_kernel_signature(self, place, egr_inps=None, egr_oups=None):
         with base.dygraph.base.guard(place=place):
-            block = base.default_main_program().global_block()
+            block = base.framework.default_main_program().global_block()
             op_proto = OpProtoHolder.instance().get_op_proto(self.op_type)
             # prepare input variable
             dygraph_tensor_inputs = (
@@ -1525,7 +1525,7 @@ def _calc_output(
         for_inplace_test=None,
         check_cinn=False,
     ):
-        with static_guard():
+        with paddle.pir_utils.OldIrGuard():
             program = Program()
             block = program.global_block()
             op = self._append_ops(block)
@@ -1590,6 +1590,7 @@ def _calc_output(
                 program = compiled_prog
 
             executor = Executor(place)
+
             outs = executor.run(
                 program,
                 feed=feed_map,
@@ -1685,37 +1686,38 @@ def _construct_grad_program_from_forward(
         Returns:
             grad_program (program): The program which contains the grad_op.
         """
-        grad_program = Program()
-        grad_block = grad_program.global_block()
-        new_op_desc = grad_block.desc.append_op()
-        new_op_desc.copy_from(grad_op_desc)
-        grad_program._sync_with_cpp()
-
-        # Create grad vars based on fwd vars (shape and dtype)
-        for arg in (
-            grad_op_desc.input_arg_names() + grad_op_desc.output_arg_names()
-        ):
-            fwd_var_name = op_grad_to_var.get(arg, None)
-            if fwd_var_name is None:
-                fwd_var_name = arg
-            fwd_var = fwd_program.global_block().vars.get(fwd_var_name)
-            assert fwd_var is not None, f"{fwd_var_name} cannot be found"
-            grad_var = grad_block.create_var(
-                name=arg,
-                dtype=fwd_var.dtype,
-                shape=fwd_var.shape,
-                type=fwd_var.type,
-                persistable=False,
-            )
+        with paddle.pir_utils.OldIrGuard():
+            grad_program = Program()
+            grad_block = grad_program.global_block()
+            new_op_desc = grad_block.desc.append_op()
+            new_op_desc.copy_from(grad_op_desc)
+            grad_program._sync_with_cpp()
+
+            # Create grad vars based on fwd vars (shape and dtype)
+            for arg in (
+                grad_op_desc.input_arg_names() + grad_op_desc.output_arg_names()
+            ):
+                fwd_var_name = op_grad_to_var.get(arg, None)
+                if fwd_var_name is None:
+                    fwd_var_name = arg
+                fwd_var = fwd_program.global_block().vars.get(fwd_var_name)
+                assert fwd_var is not None, f"{fwd_var_name} cannot be found"
+                grad_var = grad_block.create_var(
+                    name=arg,
+                    dtype=fwd_var.dtype,
+                    shape=fwd_var.shape,
+                    type=fwd_var.type,
+                    persistable=False,
+                )
 
-            # Some variables' tensors hold no buffer (tensor's _holder is NULL), like XShape in reshape2 op,
-            # and the shapes of those variables contain 0 (eg. Xshape.shape = [0, 2, 5]).
-            # Set persistable for those variables in order to get them from global_scope for inplace grad test directly other than feed them,
-            # since feed op calls check_memory_size() which fails when tensor's holder_ is NULL.
-            if 0 in grad_var.shape:
-                grad_var.persistable = True
-        grad_program._sync_with_cpp()
-        return grad_program
+                # Some variables' tensors hold no buffer (tensor's _holder is NULL), like XShape in reshape2 op,
+                # and the shapes of those variables contain 0 (eg. Xshape.shape = [0, 2, 5]).
+                # Set persistable for those variables in order to get them from global_scope for inplace grad test directly other than feed them,
+                # since feed op calls check_memory_size() which fails when tensor's holder_ is NULL.
+                if 0 in grad_var.shape:
+                    grad_var.persistable = True
+            grad_program._sync_with_cpp()
+            return grad_program
 
     def _construct_grad_feed_map_from_forward(
         self, place, fwd_res, grad_op_desc, op_grad_to_var
@@ -3412,7 +3414,7 @@ def _get_dygraph_grad(
             check_dygraph = False
 
         with base.dygraph.base.guard(place=place):
-            block = base.default_main_program().global_block()
+            block = base.framework.default_main_program().global_block()
 
             op_proto = OpProtoHolder.instance().get_op_proto(self.op_type)
 
@@ -3637,7 +3639,7 @@ def _get_gradient(
         parallel=False,
         check_cinn=False,
     ):
-        with static_guard():
+        with paddle.pir_utils.OldIrGuard():
             prog = Program()
             scope = core.Scope()
             ir_scope = core.Scope()
diff --git a/test/legacy_test/test_where_op.py b/test/legacy_test/test_where_op.py
@@ -784,7 +784,6 @@ def test_where_condition(self):
 
 
 class TestWhereOpError(unittest.TestCase):
-    @test_with_pir_api
     def test_errors(self):
         with paddle.static.program_guard(
             paddle.static.Program(), paddle.static.Program()
@@ -805,15 +804,20 @@ def test_Value():
             self.assertRaises(TypeError, test_Value)
 
             def test_type():
-                x = paddle.static.data(name='x', shape=[-1, 4], dtype='bool')
-                x.desc.set_need_check_feed(False)
-                y = paddle.static.data(name='y', shape=[-1, 4], dtype='float16')
-                y.desc.set_need_check_feed(False)
-                cond = paddle.static.data(
-                    name='cond', shape=[-1, 4], dtype='int32'
-                )
-                cond.desc.set_need_check_feed(False)
-                paddle.where(cond, x, y)
+                with paddle.pir_utils.OldIrGuard():
+                    x = paddle.static.data(
+                        name='x', shape=[-1, 4], dtype='bool'
+                    )
+                    x.desc.set_need_check_feed(False)
+                    y = paddle.static.data(
+                        name='y', shape=[-1, 4], dtype='float16'
+                    )
+                    y.desc.set_need_check_feed(False)
+                    cond = paddle.static.data(
+                        name='cond', shape=[-1, 4], dtype='int32'
+                    )
+                    cond.desc.set_need_check_feed(False)
+                    paddle.where(cond, x, y)
 
             self.assertRaises(TypeError, test_type)
 
