[SOT] Reduce eval frame callback launch overhead (#68153)

SigureMo · web-flow · commit 97d2de7132cf · 2024-09-12T19:52:16.000+08:00
diff --git a/python/paddle/jit/sot/opcode_translator/eval_frame_callback.py b/python/paddle/jit/sot/opcode_translator/eval_frame_callback.py
@@ -15,13 +15,10 @@
 from __future__ import annotations
 
 import dis
-from functools import partial
 from typing import TYPE_CHECKING
 
-from paddle.base.dygraph.base import sot_simulation_mode_guard
-
 from ..profiler import EventGuard
-from ..utils import log_do, log_format
+from ..utils import log_do, log_enabled, log_format
 from .executor.executor_cache import OpcodeExecutorCache
 
 if TYPE_CHECKING:
@@ -58,31 +55,38 @@ def eval_frame_callback(frame, **kwargs) -> CustomCode:
     with EventGuard(
         f"eval_frame_callback: {frame.f_code.co_name}", event_level=2
     ):
-        log_format(
-            2, "[eval_frame_callback] start to translate: {}\n", frame.f_code
-        )
-        log_do(4, partial(print_locals, frame))
-
-        log_format(
-            3, "[eval_frame_callback] OriginCode: {}\n", frame.f_code.co_name
-        )
-        log_do(3, lambda: dis.dis(frame.f_code))
-
-        with sot_simulation_mode_guard(True):
-            custom_code = OpcodeExecutorCache()(frame, **kwargs)
-
-        if custom_code.code is None:
+        # A quick way to check if the log is enabled
+        need_log = log_enabled(1)
+        if need_log:
             log_format(
-                3,
-                "[eval_frame_callback] NewCode (same as origin code): {}\n",
-                frame.f_code.co_name,
+                2,
+                "[eval_frame_callback] start to translate: {}\n",
+                frame.f_code,
             )
-        else:
+            log_do(4, lambda: print_locals(frame))
+
             log_format(
                 3,
-                "[eval_frame_callback] NewCode: {}\n",
-                custom_code.code.co_name,
+                "[eval_frame_callback] OriginCode: {}\n",
+                frame.f_code.co_name,
             )
-            log_do(3, lambda: dis.dis(custom_code.code))
+            log_do(3, lambda: dis.dis(frame.f_code))
+
+        custom_code = OpcodeExecutorCache()(frame, **kwargs)
+
+        if need_log:
+            if custom_code.code is None:
+                log_format(
+                    3,
+                    "[eval_frame_callback] NewCode (same as origin code): {}\n",
+                    frame.f_code.co_name,
+                )
+            else:
+                log_format(
+                    3,
+                    "[eval_frame_callback] NewCode: {}\n",
+                    custom_code.code.co_name,
+                )
+                log_do(3, lambda: dis.dis(custom_code.code))
 
         return custom_code
diff --git a/python/paddle/jit/sot/opcode_translator/executor/executor_cache.py b/python/paddle/jit/sot/opcode_translator/executor/executor_cache.py
@@ -18,6 +18,8 @@
 import traceback
 from typing import TYPE_CHECKING, List, Tuple
 
+from paddle.base.dygraph.base import sot_simulation_mode_guard
+
 from ...profiler import EventGuard, event_register
 from ...psdb import NO_FALLBACK_CODES
 from ...utils import (
@@ -217,7 +219,8 @@ def start_translate(
     simulator = OpcodeExecutor(frame, **kwargs)
     try:
         simulator.check_code_simulatable()
-        new_custom_code, guard_fn = simulator.transform()
+        with sot_simulation_mode_guard(True):
+            new_custom_code, guard_fn = simulator.transform()
         if not simulator._graph.need_cache:
             return (
                 CustomCode(None, True),
diff --git a/python/paddle/jit/sot/profiler/profiler.py b/python/paddle/jit/sot/profiler/profiler.py
@@ -37,9 +37,8 @@ def disable(self):
 
 @contextmanager
 def EventGuard(event_name, event_level=1):
+    need_pop = False
     try:
-        global _event_level
-        need_pop = False
         if ENV_SOT_EVENT_LEVEL.get_with_cache() >= event_level:
             core.nvprof_nvtx_push(event_name)
             need_pop = True
@@ -61,7 +60,7 @@ def call_with_event(*args, **kwargs):
     def do_nothing(func):
         return func
 
-    if ENV_SOT_EVENT_LEVEL.get() >= event_level:
+    if ENV_SOT_EVENT_LEVEL.get_with_cache() >= event_level:
         return event_wrapper
     else:
         return do_nothing
diff --git a/python/paddle/jit/sot/utils/__init__.py b/python/paddle/jit/sot/utils/__init__.py
@@ -73,6 +73,7 @@
     list_find_index_by_id,
     log,
     log_do,
+    log_enabled,
     log_format,
     map_if,
     map_if_extend,
diff --git a/python/paddle/jit/sot/utils/utils.py b/python/paddle/jit/sot/utils/utils.py
@@ -129,21 +129,25 @@ def next(self):
 
 
 def log(level, *args):
-    cur_level = ENV_SOT_LOG_LEVEL.get()
+    cur_level = ENV_SOT_LOG_LEVEL.get_with_cache()
     if level <= cur_level:
-        print(*args, end="")
+        print(*args, end="", flush=True)
 
 
 def log_do(level, fn):
-    cur_level = ENV_SOT_LOG_LEVEL.get()
+    cur_level = ENV_SOT_LOG_LEVEL.get_with_cache()
     if level <= cur_level:
         fn()
 
 
 def log_format(level, str, *args):
-    cur_level = ENV_SOT_LOG_LEVEL.get()
+    cur_level = ENV_SOT_LOG_LEVEL.get_with_cache()
     if level <= cur_level:
-        print(str.format(*args), end="")
+        print(str.format(*args), end="", flush=True)
+
+
+def log_enabled(level):
+    return level <= ENV_SOT_LOG_LEVEL.get_with_cache()
 
 
 def no_eval_frame(func):
