Using EventGuard in decorator form

Author: 0x3878f

python/paddle/jit/dy2static/pir_partial_program.py

@@ -20,6 +20,7 @@
 from typing import TYPE_CHECKING
 
 import numpy as np
+from sot.profiler import event_register
 
 import paddle
 import paddle.pir.core as ir_static
@@ -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/sot/symbolic/compile_cache.py

@@ -20,7 +20,7 @@
 import paddle
 
 from ..infer_meta import convert_meta_to_input_spec
-from ..profiler import EventGuard
+from ..profiler import EventGuard, event_register
 from ..utils import (
     ENV_SOT_EXPORT,
     Cache,
@@ -210,7 +210,8 @@ def update_compile_time_info(self, SIR, partial_program_layer):
         ] += partial_program_layer._compile_time_counter.get_total_time()
 
     def __call__(self, *args, **kwargs):
-        with EventGuard(f"FallbackWrapper: {self.SIR.name}"):
+        @event_register(f"FallbackWrapper: {self.SIR.name}")
+        def call_fn():
             if StepInfoManager().need_back_trace:
                 trace_back_frames()
 
@@ -233,8 +234,7 @@ def __call__(self, *args, **kwargs):
                         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)
+            outputs = self.partial_program.sot_call(*args, **kwargs)
 
             clear_eager_tensor_name(outputs)
             log_do(
@@ -252,6 +252,8 @@ def __call__(self, *args, **kwargs):
             self.is_first_call = False
             return outputs
 
+        return call_fn()
+
 
 class CompileSIRCache(Cache, metaclass=Singleton):
     """