[SOT] Using EventGuard in decorator form (#73891)

* Using EventGuard in decorator form

* fix: import EventGuard, SotProfiler, event_register in paddle.jit.sot.profiler package

Author: 0x3878f

python/paddle/jit/dy2static/pir_partial_program.py

@@ -32,6 +32,7 @@
 from paddle.base.dygraph.base import switch_to_static_graph
 from paddle.pir import Value, fake_value, get_fake_value_name, is_fake_value
 
+from ..profiler import event_register
 from .logging_utils import TranslatorLogger
 from .utils import (
     RETURN_NO_VALUE_MAGIC_NUM,
@@ -779,6 +780,7 @@ def __call__(self, inputs):
         restored_nest_out = self._restore_out(out)
         return self._remove_no_value(restored_nest_out)
 
+    @event_register("sot call partial_program")
     def sot_call(self, inputs):
         """
         In sot, inputs and outputs of partial program only contain tensors, so we can skip some step to speed up

python/paddle/jit/dy2static/utils.py

@@ -48,6 +48,7 @@
 from paddle.utils import flatten, gast
 from paddle.utils.environments import (
     BooleanEnvironmentVariable,
+    IntegerEnvironmentVariable,
 )
 
 from .ast_utils import ast_to_source_code
@@ -83,6 +84,7 @@
     core.VarDesc.VarType.FETCH_LIST,
 ]
 
+ENV_SOT_EVENT_LEVEL = IntegerEnvironmentVariable("SOT_EVENT_LEVEL", 0)
 ENV_ENABLE_SOT = BooleanEnvironmentVariable("ENABLE_FALL_BACK", True)
 ENV_ENABLE_CINN_IN_DY2ST = BooleanEnvironmentVariable(
     "ENABLE_CINN_IN_DY2ST", True

python/paddle/jit/sot/profiler/profiler.py renamed to python/paddle/jit/profiler.py

@@ -12,12 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 from contextlib import contextmanager
 from functools import wraps
+from typing import Callable, TypeVar
+
+from typing_extensions import ParamSpec
 
 from paddle.framework import core
 
-from ..utils.envs import ENV_SOT_EVENT_LEVEL
+from .dy2static.utils import ENV_SOT_EVENT_LEVEL
+
+P = ParamSpec("P")
+T = TypeVar("T")
 
 
 class SotProfiler:
@@ -48,16 +56,23 @@ def EventGuard(event_name, event_level=1):
             core.nvprof_nvtx_pop()
 
 
-def event_register(event_name, event_level=1):
-    def event_wrapper(func):
+def event_register(
+    event_name_formatter: Callable[P, str] | str, event_level=1
+) -> Callable[[Callable[P, T]], Callable[P, T]]:
+    def event_wrapper(func: Callable[P, T]) -> Callable[P, T]:
         @wraps(func)
-        def call_with_event(*args, **kwargs):
+        def call_with_event(*args: P.args, **kwargs: P.kwargs):
+            event_name = (
+                event_name_formatter(*args, **kwargs)
+                if callable(event_name_formatter)
+                else event_name_formatter
+            )
             with EventGuard(event_name, event_level=event_level):
                 return func(*args, **kwargs)
 
         return call_with_event
 
-    def do_nothing(func):
+    def do_nothing(func: Callable[P, T]) -> Callable[P, T]:
         return func
 
     if ENV_SOT_EVENT_LEVEL.get() >= event_level:

python/paddle/jit/sot/profiler/__init__.py

@@ -12,9 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .kernel_stats import SotStepProfilerGuard as SotStepProfilerGuard
-from .profiler import (
+from paddle.jit.profiler import (
     EventGuard as EventGuard,
     SotProfiler as SotProfiler,
     event_register as event_register,
 )
+
+from .kernel_stats import SotStepProfilerGuard as SotStepProfilerGuard

python/paddle/jit/sot/symbolic/compile_cache.py

@@ -18,9 +18,9 @@
 from typing import TYPE_CHECKING
 
 import paddle
+from paddle.jit.profiler import EventGuard, event_register
 
 from ..infer_meta import convert_meta_to_input_spec
-from ..profiler import EventGuard
 from ..utils import (
     ENV_SOT_EXPORT,
     Cache,
@@ -209,48 +209,49 @@ def update_compile_time_info(self, SIR, partial_program_layer):
             code
         ] += partial_program_layer._compile_time_counter.get_total_time()
 
+    @event_register(
+        lambda self, *args, **kwargs: f"FallbackWrapper: {self.SIR.name}"
+    )
     def __call__(self, *args, **kwargs):
-        with EventGuard(f"FallbackWrapper: {self.SIR.name}"):
-            if StepInfoManager().need_back_trace:
-                trace_back_frames()
+        if StepInfoManager().need_back_trace:
+            trace_back_frames()
 
-            log_do(
-                2,
-                lambda: print("[FallbackWrapper] start run SIR: \n", self.SIR),
-            )
-            log_do(
-                4,
-                lambda: print(
-                    self.compiled_fn.get_concrete_program(*args, **kwargs)[
-                        1
-                    ].train_program
-                ),
-            )
-            if self.partial_program is None:
-                with EventGuard("FallbackWrapper: get_concrete_program"):
-                    (
-                        self.concrete_program,
-                        self.partial_program,
-                    ) = self.compiled_fn.get_concrete_program(*args, **kwargs)
-                    self.partial_program.training = self.is_training
-            with EventGuard("FallbackWrapper: sot call partial_program"):
-                outputs = self.partial_program.sot_call(*args, **kwargs)
-
-            clear_eager_tensor_name(outputs)
-            log_do(
-                4,
-                lambda: print("[CompileCache] run sir forward success."),
-            )
-            self.collect_new_symbol_hit_rate(args, outputs)
-            self.collect_subgraph_relation(args, outputs, self.partial_program)
-            self.collect_subgraph_info(self.concrete_program.main_program)
-            self.update_compile_time_info(self.SIR, self.partial_program)
-            if ENV_SOT_EXPORT.get() != "" and not self.exported:
-                export(self.SIR, ENV_SOT_EXPORT.get())
-                self.exported = True
-
-            self.is_first_call = False
-            return outputs
+        log_do(
+            2,
+            lambda: print("[FallbackWrapper] start run SIR: \n", self.SIR),
+        )
+        log_do(
+            4,
+            lambda: print(
+                self.compiled_fn.get_concrete_program(*args, **kwargs)[
+                    1
+                ].train_program
+            ),
+        )
+        if self.partial_program is None:
+            with EventGuard("FallbackWrapper: get_concrete_program"):
+                (
+                    self.concrete_program,
+                    self.partial_program,
+                ) = self.compiled_fn.get_concrete_program(*args, **kwargs)
+                self.partial_program.training = self.is_training
+        outputs = self.partial_program.sot_call(*args, **kwargs)
+
+        clear_eager_tensor_name(outputs)
+        log_do(
+            4,
+            lambda: print("[CompileCache] run sir forward success."),
+        )
+        self.collect_new_symbol_hit_rate(args, outputs)
+        self.collect_subgraph_relation(args, outputs, self.partial_program)
+        self.collect_subgraph_info(self.concrete_program.main_program)
+        self.update_compile_time_info(self.SIR, self.partial_program)
+        if ENV_SOT_EXPORT.get() != "" and not self.exported:
+            export(self.SIR, ENV_SOT_EXPORT.get())
+            self.exported = True
+
+        self.is_first_call = False
+        return outputs
 
 
 class CompileSIRCache(Cache, metaclass=Singleton):

python/paddle/jit/sot/utils/envs.py

@@ -135,7 +135,6 @@ def parse_parameterized_key(input_str: str) -> dict[str, list[str]]:
     "SOT_ENABLE_GUARD_TREE",
     False,
 )
-ENV_SOT_EVENT_LEVEL = IntegerEnvironmentVariable("SOT_EVENT_LEVEL", 0)
 ENV_ENABLE_SOT_STEP_PROFILER = BooleanEnvironmentVariable(
     "ENABLE_SOT_STEP_PROFILER", False
 )