[Dy2St][CINN] Dont access cpp flag FLAGS_use_cinn via os.getenv (#71817)

SigureMo · web-flow · commit 054ee47bfb30 · 2025-03-21T19:55:00.000+08:00
diff --git a/python/paddle/base/framework.py b/python/paddle/base/framework.py
@@ -27,6 +27,7 @@
 import traceback
 import warnings
 from collections.abc import Iterable
+from contextlib import contextmanager
 from types import FunctionType, MethodType
 from typing import TYPE_CHECKING, Callable, TypeVar, overload
 
@@ -200,6 +201,16 @@ def get_flags(flags: str | Sequence[str]) -> dict[str, bool | str | float]:
     return flags_value
 
 
+@contextmanager
+def flag_guard(flag_name, flag_value):
+    old_value = paddle.get_flags(flag_name)[flag_name]
+    paddle.set_flags({flag_name: flag_value})
+    try:
+        yield
+    finally:
+        paddle.set_flags({flag_name: old_value})
+
+
 # use thread local to create thread save global variables.
 class GlobalThreadLocal(threading.local):
     def __init__(self):
@@ -432,8 +443,11 @@ def in_cinn_mode() -> bool:
         bool: Whether paddle runs in cinn mode.
 
     """
-    flag = str(os.environ.get("FLAGS_use_cinn")).lower()
-    return flag in ("true", "1")
+    CINN_FLAG_NAME = "FLAGS_use_cinn"
+    # NOTE: This flag only available when compiled with CINN
+    if not is_compiled_with_cinn():
+        return False
+    return paddle.get_flags(CINN_FLAG_NAME)[CINN_FLAG_NAME]
 
 
 global_ipu_index = -1
diff --git a/python/paddle/jit/dy2static/utils.py b/python/paddle/jit/dy2static/utils.py
@@ -662,9 +662,7 @@ def cinn_is_enabled(build_strategy, backend):
         return True
     if build_strategy is not None and build_strategy.build_cinn_pass:
         return True
-
-    value = os.getenv('FLAGS_use_cinn')
-    if value is not None and value.lower() in ['true', '1']:
+    if paddle.base.framework.in_cinn_mode():
         return True
     return False
 
diff --git a/test/ir/pir/test_cse_pass.py b/test/ir/pir/test_cse_pass.py
@@ -16,6 +16,7 @@
 from contextlib import contextmanager
 
 import paddle
+from paddle.base.framework import flag_guard
 
 paddle.enable_static()
 
@@ -31,16 +32,6 @@ def program_scope_guard():
             yield main_program
 
 
-@contextmanager
-def flag_guard(flag_name, flag_value):
-    old_value = paddle.get_flags(flag_name)[flag_name]
-    paddle.set_flags({flag_name: flag_value})
-    try:
-        yield
-    finally:
-        paddle.set_flags({flag_name: old_value})
-
-
 def walk_block(block, fn):
     for op in block.ops:
         fn(op)
@@ -412,11 +403,11 @@ def get_cond(loop_var):
     "This case only works when compiled with CINN",
 )
 class TestCSEDenyFullInCinn(unittest.TestCase, AssertOpCountEqualMixin):
-    CINN_FLAGS_NAME = "FLAGS_use_cinn"
+    CINN_FLAG_NAME = "FLAGS_use_cinn"
 
     def test_replace_full_without_cinn(self):
         with flag_guard(
-            self.CINN_FLAGS_NAME, False
+            self.CINN_FLAG_NAME, False
         ), program_scope_guard() as main_program:
             # Inputs
             x1 = paddle.full([2], 1.0, dtype="float32")
@@ -428,7 +419,7 @@ def test_replace_full_without_cinn(self):
 
     def test_replace_full_with_cinn(self):
         with flag_guard(
-            self.CINN_FLAGS_NAME, True
+            self.CINN_FLAG_NAME, True
         ), program_scope_guard() as main_program:
             # Inputs
             x1 = paddle.full([2], 1.0, dtype="float32")
diff --git a/test/ir/pir/test_pir_executor_flag.py b/test/ir/pir/test_pir_executor_flag.py
@@ -15,7 +15,8 @@
 import os
 import unittest
 
-from paddle.base.framework import in_cinn_mode, in_pir_executor_mode
+import paddle
+from paddle.base.framework import flag_guard, in_cinn_mode, in_pir_executor_mode
 
 
 class TestPIRModeFlags(unittest.TestCase):
@@ -27,9 +28,13 @@ def test_pir_mode_flags(self):
 
 class TestCinnModeFlags(unittest.TestCase):
     def test_cinn_mode_flags(self):
+        CINN_FLAG_NAME = "FLAGS_use_cinn"
+        if not paddle.is_compiled_with_cinn():
+            self.assertFalse(in_cinn_mode())
+            return
         self.assertFalse(in_cinn_mode())
-        os.environ["FLAGS_use_cinn"] = "true"
-        self.assertTrue(in_cinn_mode())
+        with flag_guard(CINN_FLAG_NAME, True):
+            self.assertTrue(in_cinn_mode())
 
 
 if __name__ == '__main__':
