[SOT][Faster Guard] add ENV_SOT_ENABLE_FASTER_GUARD (#69263)

zrr1999 · web-flow · commit d466c1c1c86c · 2024-11-11T23:18:44.000+08:00
diff --git a/python/paddle/jit/sot/utils/__init__.py b/python/paddle/jit/sot/utils/__init__.py
@@ -18,17 +18,19 @@
     ENV_COST_MODEL,
     ENV_MIN_GRAPH_SIZE,
     ENV_SOT_ALLOW_DYNAMIC_SHAPE,
+    ENV_SOT_ENABLE_FASTER_GUARD,
     ENV_SOT_EXPORT,
     ENV_SOT_LOG_LEVEL,
     ENV_SOT_WITH_CONTROL_FLOW,
     ENV_STRICT_MODE,
+    allow_dynamic_shape_guard,
     cost_model_guard,
+    export_guard,
+    faster_guard_guard,
     min_graph_size_guard,
     sot_step_profiler_guard,
     strict_mode_guard,
-    with_allow_dynamic_shape_guard,
     with_control_flow_guard,
-    with_export_guard,
 )
 from .exceptions import (  # noqa: F401
     BreakGraphError,
diff --git a/python/paddle/jit/sot/utils/envs.py b/python/paddle/jit/sot/utils/envs.py
@@ -38,6 +38,10 @@
     # Enable SOT dynamic shape as default in PIR mode only
     paddle.framework.use_pir_api(),
 )
+ENV_SOT_ENABLE_FASTER_GUARD = BooleanEnvironmentVariable(
+    "SOT_ENABLE_FASTER_GUARD",
+    False,
+)
 ENV_SOT_EVENT_LEVEL = IntegerEnvironmentVariable("SOT_EVENT_LEVEL", 0)
 ENV_ENABLE_SOT_STEP_PROFILER = BooleanEnvironmentVariable(
     "ENABLE_SOT_STEP_PROFILER", False
@@ -69,17 +73,23 @@ def with_control_flow_guard(value: bool):
 
 
 @contextmanager
-def with_export_guard(value: str):
+def export_guard(value: str):
     with EnvironmentVariableGuard(ENV_SOT_EXPORT, value):
         yield
 
 
 @contextmanager
-def with_allow_dynamic_shape_guard(value: bool):
+def allow_dynamic_shape_guard(value: bool):
     with EnvironmentVariableGuard(ENV_SOT_ALLOW_DYNAMIC_SHAPE, value):
         yield
 
 
+@contextmanager
+def faster_guard_guard(value: bool):
+    with EnvironmentVariableGuard(ENV_SOT_ENABLE_FASTER_GUARD, value):
+        yield
+
+
 @contextmanager
 def sot_step_profiler_guard(value: bool):
     with EnvironmentVariableGuard(ENV_ENABLE_SOT_STEP_PROFILER, value):
diff --git a/test/sot/test_sot_dynamic_shape.py b/test/sot/test_sot_dynamic_shape.py
@@ -24,7 +24,7 @@
 import paddle
 from paddle.jit.sot.psdb import check_no_breakgraph
 from paddle.jit.sot.utils import (
-    with_allow_dynamic_shape_guard,
+    allow_dynamic_shape_guard,
 )
 
 
@@ -78,7 +78,7 @@ def forward(self, x):
 
 class TestOpcodeExecutorDynamicShapeCache(TestCaseBase):
     def test_dynamic_int_input_cache_hit_case1(self):
-        with with_allow_dynamic_shape_guard(
+        with allow_dynamic_shape_guard(
             True
         ), test_instruction_translator_cache_context() as ctx:
             self.assert_results(
@@ -92,7 +92,7 @@ def test_dynamic_int_input_cache_hit_case1(self):
                 self.assertEqual(ctx.translate_count, 2)
 
     def test_dynamic_int_input_cache_hit_case2(self):
-        with with_allow_dynamic_shape_guard(
+        with allow_dynamic_shape_guard(
             True
         ), test_instruction_translator_cache_context() as ctx:
             self.assert_results(
@@ -106,7 +106,7 @@ def test_dynamic_int_input_cache_hit_case2(self):
                 self.assertEqual(ctx.translate_count, 2)
 
     def test_dynamic_int_input_cache_hit_case3(self):
-        with with_allow_dynamic_shape_guard(
+        with allow_dynamic_shape_guard(
             True
         ), test_instruction_translator_cache_context() as ctx:
             for i in range(0, 6):
@@ -116,7 +116,7 @@ def test_dynamic_int_input_cache_hit_case3(self):
                 self.assertEqual(ctx.translate_count, i + 1)
 
     def test_dynamic_shape_input_cache_hit_case1(self):
-        with with_allow_dynamic_shape_guard(
+        with allow_dynamic_shape_guard(
             True
         ), test_instruction_translator_cache_context() as ctx:
             self.assert_results(
@@ -130,7 +130,7 @@ def test_dynamic_shape_input_cache_hit_case1(self):
                 self.assertEqual(ctx.translate_count, 2)
 
     def test_dynamic_shape_input_cache_hit_case2(self):
-        with with_allow_dynamic_shape_guard(
+        with allow_dynamic_shape_guard(
             True
         ), test_instruction_translator_cache_context() as ctx:
             self.assert_results(
@@ -145,7 +145,7 @@ def test_dynamic_shape_input_cache_hit_case2(self):
                 self.assertEqual(ctx.translate_count, 2)
 
     def test_dynamic_shape_cast(self):
-        with with_allow_dynamic_shape_guard(
+        with allow_dynamic_shape_guard(
             True
         ), test_instruction_translator_cache_context() as ctx:
             func1 = check_no_breakgraph(lambda n: bool(n))
@@ -156,7 +156,7 @@ def test_dynamic_shape_cast(self):
                 self.assert_results(func, 2)
 
     def test_dynamic_shape_in_list(self):
-        with with_allow_dynamic_shape_guard(
+        with allow_dynamic_shape_guard(
             True
         ), test_instruction_translator_cache_context() as ctx:
             self.assert_results(
@@ -174,7 +174,7 @@ def test_dynamic_shape_in_list(self):
                 self.assertEqual(ctx.translate_count, 2)
 
     def test_conv_dynamic_shape_fallback(self):
-        with with_allow_dynamic_shape_guard(
+        with allow_dynamic_shape_guard(
             True
         ), test_instruction_translator_cache_context() as ctx:
             for i in range(1, 5):
@@ -183,7 +183,7 @@ def test_conv_dynamic_shape_fallback(self):
                 self.assertEqual(ctx.translate_count, i)
 
     def test_pad_dynamic_shape_fallback(self):
-        with with_allow_dynamic_shape_guard(
+        with allow_dynamic_shape_guard(
             True
         ), test_instruction_translator_cache_context() as ctx:
             pad_func = check_no_breakgraph(
diff --git a/test/sot/test_sot_export.py b/test/sot/test_sot_export.py
@@ -17,7 +17,7 @@
 import unittest
 
 import paddle
-from paddle.jit.sot.utils import min_graph_size_guard, with_export_guard
+from paddle.jit.sot.utils import export_guard, min_graph_size_guard
 
 
 class Net(paddle.nn.Layer):
@@ -43,7 +43,7 @@ def test_basic(self):
         temp_dir_name = temp_dir.name
         net = Net()
         x = paddle.to_tensor([2, 3], dtype="float32", stop_gradient=True)
-        with with_export_guard(temp_dir_name):
+        with export_guard(temp_dir_name):
             y = paddle.jit.to_static(net)(x)
         assert os.path.exists(os.path.join(temp_dir_name, "SIR_0.py"))
         temp_dir.cleanup()
diff --git a/test/sot/test_trace_list_arg.py b/test/sot/test_trace_list_arg.py
@@ -22,7 +22,7 @@
 )
 
 import paddle
-from paddle.jit.sot.utils.envs import with_allow_dynamic_shape_guard
+from paddle.jit.sot.utils.envs import allow_dynamic_shape_guard
 
 
 def foo(x: list[paddle.Tensor], y: list[paddle.Tensor]):
@@ -47,7 +47,7 @@ def test_foo(self):
             self.assert_results(foo, [a], [c])  # Cache miss
             self.assertEqual(cache.translate_count, 2)
 
-    @with_allow_dynamic_shape_guard(False)
+    @allow_dynamic_shape_guard(False)
     def test_bar_static_shape(self):
         a = [paddle.to_tensor(1), paddle.to_tensor(2), paddle.to_tensor(3)]
         b = [paddle.to_tensor([2, 3]), paddle.to_tensor(4), paddle.to_tensor(5)]
@@ -60,7 +60,7 @@ def test_bar_static_shape(self):
             self.assert_results(bar, b, 1, 1)  # Cache hit
             self.assertEqual(cache.translate_count, 2)
 
-    @with_allow_dynamic_shape_guard(True)
+    @allow_dynamic_shape_guard(True)
     def test_bar_dynamic_shape(self):
         # TODO(zrr1999): mv to dynamic shape test
         a = [paddle.to_tensor(1), paddle.to_tensor(2), paddle.to_tensor(3)]
