[PIR]fix compare_legacy_with_pt_bug (#65932)

risemeup1 · web-flow · commit d94e15455ba9 · 2024-07-12T11:18:27.000+08:00
* fix

* fix
diff --git a/test/legacy_test/test_sparse_conv_op.py b/test/legacy_test/test_sparse_conv_op.py
@@ -16,6 +16,7 @@
 import unittest
 
 import numpy as np
+from utils import compare_legacy_with_pt
 
 import paddle
 from paddle import sparse
@@ -314,6 +315,7 @@ def test_Conv3D_bias(self):
 
 
 class TestStatic(unittest.TestCase):
+    @compare_legacy_with_pt
     def test(self):
         paddle.enable_static()
         main = paddle.static.Program()
@@ -382,6 +384,7 @@ def test(self):
             self.assertTrue(out_indices.dtype == paddle.int32)
         paddle.disable_static()
 
+    @compare_legacy_with_pt
     def test_cpu(self):
         paddle.enable_static()
         main = paddle.static.Program()
@@ -451,6 +454,7 @@ def test_cpu(self):
             self.assertTrue(out_indices.dtype == paddle.int32)
         paddle.disable_static()
 
+    @compare_legacy_with_pt
     def test2D(self):
         paddle.enable_static()
         main = paddle.static.Program()
@@ -514,6 +518,7 @@ def test2D(self):
             self.assertTrue(out_indices.dtype == paddle.int32)
         paddle.disable_static()
 
+    @compare_legacy_with_pt
     def test2D_cpu(self):
         paddle.enable_static()
         main = paddle.static.Program()
diff --git a/test/legacy_test/test_sparse_isnan_op.py b/test/legacy_test/test_sparse_isnan_op.py
@@ -15,6 +15,7 @@
 import unittest
 
 import numpy as np
+from utils import compare_legacy_with_pt
 
 import paddle
 
@@ -63,6 +64,7 @@ def test_isnan_dtype(self):
 
 
 class TestStatic(unittest.TestCase):
+    @compare_legacy_with_pt
     def test(self):
         paddle.enable_static()
         main_program = paddle.static.Program()
diff --git a/test/legacy_test/test_sparse_slice_op.py b/test/legacy_test/test_sparse_slice_op.py
@@ -15,6 +15,7 @@
 import unittest
 
 import numpy as np
+from utils import compare_legacy_with_pt
 
 import paddle
 
@@ -206,26 +207,32 @@ def check_result_with_list(self, x, axes, starts, ends, format='coo'):
         if format == 'coo':
             self._check_result_coo(np_x, axes, starts, ends)
 
+    @compare_legacy_with_pt
     def test_coo_5d(self):
         for item in data_5d:
             self.check_result_with_shape(*item, format='coo')
 
+    @compare_legacy_with_pt
     def test_coo_4d(self):
         for item in data_4d:
             self.check_result_with_shape(*item, format='coo')
 
+    @compare_legacy_with_pt
     def test_coo_3d(self):
         for item in data_3d:
             self.check_result_with_shape(*item, format='coo')
 
+    @compare_legacy_with_pt
     def test_coo_2d(self):
         for item in data_2d:
             self.check_result_with_shape(*item, format='coo')
 
+    @compare_legacy_with_pt
     def test_coo_1d(self):
         x = [-49, 55, -5, 0, 3, 0, 0, -60, -21, 0, 0, 0]
         self.check_result_with_list(x, [0], [3], [5], format='coo')
 
+    @compare_legacy_with_pt
     def test_coo_1d_zero(self):
         x = [-49, 55, -5, 0, 3, 0, 0, -60, -21, 0, 0, 0]
         self.check_result_with_list(x, [0], [-3], [-1], format='coo')
diff --git a/test/legacy_test/test_sparse_sum_op.py b/test/legacy_test/test_sparse_sum_op.py
@@ -15,6 +15,7 @@
 import unittest
 
 import numpy as np
+from utils import compare_legacy_with_pt
 
 import paddle
 
@@ -176,6 +177,7 @@ def check_result_coo(self, x_shape, dims, keepdim, dtype=None):
                 )
             paddle.disable_static()
 
+    @compare_legacy_with_pt
     def test_sum(self):
         # 1d
         self.check_result_coo([5], None, False)
diff --git a/test/legacy_test/utils.py b/test/legacy_test/utils.py
@@ -22,6 +22,11 @@
 from paddle.base import core
 from paddle.base.framework import _dygraph_guard
 from paddle.base.wrapped_decorator import signature_safe_contextmanager
+from paddle.pir_utils import DygraphOldIrGuard
+from paddle.utils.environments import (
+    BooleanEnvironmentVariable,
+    EnvironmentVariableGuard,
+)
 
 __all__ = ['DyGraphProgramDescTracerTestHelper', 'is_equal_program']
 
@@ -169,22 +174,37 @@ def pir_executor_guard():
         set_flags({"FLAGS_enable_pir_in_executor": tmp_cpp})
 
 
+ENV_ENABLE_PIR_WITH_PT = BooleanEnvironmentVariable(
+    "FLAGS_enable_pir_in_executor", False
+)
+
+
 def to_pir_pt_test(fn):
     @wraps(fn)
     def impl(*args, **kwargs):
-        ir_outs = None
-        if os.environ.get('FLAGS_use_stride_kernel', False):
-            return
-        with static.scope_guard(static.Scope()):
-            with static.program_guard(static.Program()):
-                pir_flag = 'FLAGS_enable_pir_in_executor'
-                try:
-                    os.environ[pir_flag] = 'True'
-                    set_flags({pir_flag: True})
-                    ir_outs = fn(*args, **kwargs)
-                finally:
-                    del os.environ[pir_flag]
-                    set_flags({pir_flag: False})
+        with DygraphOldIrGuard():
+            pt_flag = ENV_ENABLE_PIR_WITH_PT.name
+            original_flag_value = get_flags(pt_flag)[pt_flag]
+            if os.environ.get('FLAGS_use_stride_kernel', False):
+                return
+            with static.scope_guard(static.Scope()):
+                with static.program_guard(static.Program()):
+                    print(
+                        "before:",
+                        os.environ.get(ENV_ENABLE_PIR_WITH_PT.name),
+                        flush=True,
+                    )
+                    with EnvironmentVariableGuard(ENV_ENABLE_PIR_WITH_PT, True):
+                        try:
+                            set_flags({pt_flag: True})
+                            ir_outs = fn(*args, **kwargs)
+                        finally:
+                            set_flags({pt_flag: original_flag_value})
+                    print(
+                        "after:",
+                        os.environ.get(ENV_ENABLE_PIR_WITH_PT.name),
+                        flush=True,
+                    )
         return ir_outs
 
     return impl
