[SOT] Restore short circuit mechanism in guard (#67890)

SigureMo · web-flow · commit 8efbe545703e · 2024-09-02T16:33:14.000+08:00
diff --git a/python/paddle/jit/sot/opcode_translator/executor/executor_cache.py b/python/paddle/jit/sot/opcode_translator/executor/executor_cache.py
@@ -42,7 +42,7 @@
 
 dummy_guard: Guard = lambda frame: True
 dummy_guard.expr = "lambda frame: True"
-dummy_guard.lambda_expr = "lambda frame: True"
+dummy_guard.inlined_expr = "lambda frame: True"
 
 
 class OpcodeExecutorCache(metaclass=Singleton):
@@ -133,6 +133,11 @@ def lookup(
                     )
             except Exception as e:
                 log(2, f"[Cache]: Guard function error: {e}\n")
+                log(
+                    2,
+                    f"[Cache]: Guard string is: {getattr(guard_fn, 'expr', 'None')}\n",
+                )
+
                 continue
 
         log(2, "[Cache]: all guards missed\n")
@@ -173,7 +178,7 @@ def inner():
 
     def analyse_guard_error(self, guard_fn, frame):
         def inner():
-            guard_expr = guard_fn.lambda_expr
+            guard_expr = guard_fn.inlined_expr
             lambda_head = "lambda frame: "
             guard_expr = guard_expr.replace(lambda_head, "")
             guards = guard_expr.split(" and ")
diff --git a/python/paddle/jit/sot/opcode_translator/executor/function_graph.py b/python/paddle/jit/sot/opcode_translator/executor/function_graph.py
@@ -48,7 +48,7 @@
     log,
     log_do,
     map_if,
-    tmp_name_guard,
+    switch_symbol_registry,
 )
 from ..instruction_utils import get_instructions
 from .guard import Guard, StringifiedExpression, make_guard
@@ -299,7 +299,7 @@ def collect(inp):
     @property
     @event_register("guard_fn")
     def guard_fn(self) -> Guard:
-        with tmp_name_guard():
+        with switch_symbol_registry():
             guards: list[StringifiedExpression] = []
             with EventGuard("guard_fn: find vars and make stringified guard"):
                 for variable in find_traceable_vars(
diff --git a/python/paddle/jit/sot/opcode_translator/executor/guard.py b/python/paddle/jit/sot/opcode_translator/executor/guard.py
@@ -16,12 +16,13 @@
 
 import types
 import weakref
+from functools import cached_property
 from typing import TYPE_CHECKING, Any, Callable, TypeVar
 
 import paddle
 
 from ...profiler import EventGuard
-from ...utils import current_tmp_name_records, log, log_do
+from ...utils import current_symbol_registry, log, log_do
 
 Guard = Callable[[types.FrameType], bool]
 
@@ -44,14 +45,37 @@ class StringifiedExpression:
     Used to store string based expressions for generating Guard.
     """
 
-    def __init__(self, str_expr, sub_exprs, free_vars):
-        expr = str_expr.format(*[arg.expr for arg in sub_exprs])
-        self.expr = current_tmp_name_records().add_tmp_var(expr)
-        self.inlined_expr = str_expr.format(
-            *[arg.inlined_expr for arg in sub_exprs]
+    def __init__(
+        self,
+        expr_template: str,
+        sub_exprs: list[StringifiedExpression],
+        free_vars: dict[str, Any],
+    ):
+        self.expr_template = expr_template
+        expr = self.expr_template.format(
+            *[sub_expr.symbol for sub_expr in sub_exprs]
         )
+        self.registered_expr = expr
+        self.symbol = current_symbol_registry().request_symbol(expr)
+        self.sub_exprs = sub_exprs
         self.free_vars = free_vars
 
+    @cached_property
+    def inlined_expr(self):
+        return self.expr_template.format(
+            *[sub_expr.inlined_expr for sub_expr in self.sub_exprs]
+        )
+
+    def gen_expr(self):
+        def gen_expr_fn():
+            return self.expr_template.format(
+                *[sub_expr.gen_expr() for sub_expr in self.sub_exprs]
+            )
+
+        return current_symbol_registry().gen_expr(
+            self.registered_expr, gen_expr_fn
+        )
+
     def __hash__(self):
         if self.free_vars:
             return hash((self.inlined_expr, id(self)))
@@ -77,43 +101,26 @@ def make_guard(stringified_guards: list[StringifiedExpression]) -> Guard:
         if not num_guards:
             guard = lambda frame: True
             guard.expr = "lambda frame: True"
+            guard.original_guard = guard
             return guard
 
-        def analyse_expressions(stringified_exprs, tmp_names):
-            func_string = "def built_guard_fn(frame):\n"
-            lambda_string = "lambda frame: "
-            free_vars = {}
-
-            for k, v in tmp_names.items():
-                func_string += f"    {v} = {k}\n"
+        free_vars = {}
+        for expr in stringified_guards:
+            free_vars = union_free_vars(free_vars, expr.free_vars)
 
-            func_result = ""
-            for str_expr in stringified_exprs:
-                func_result += str_expr.expr + " and "
-                lambda_string += str_expr.inlined_expr + " and "
-                free_vars = union_free_vars(free_vars, str_expr.free_vars)
-
-            func_string += f"    return {func_result[:-5]}"
-
-            return func_string, free_vars, lambda_string[:-5]
-
-        (
-            func_string,
-            free_vars,
-            lambda_string,
-        ) = analyse_expressions(
-            stringified_guards, current_tmp_name_records().tmp_names_record
+        inlined_guard_expr = "lambda frame: " + " and ".join(
+            [expr.inlined_expr for expr in stringified_guards]
         )
-
-        exec(
-            func_string,
-            free_vars,
+        guard_expr: str = "lambda frame: " + " and ".join(
+            [expr.gen_expr() for expr in stringified_guards]
         )
 
-        guard = free_vars['built_guard_fn']
-        log(3, f"[Guard]: {lambda_string}\n")
-        guard.lambda_expr = lambda_string
-        guard.expr = func_string
+        guard = eval(guard_expr, free_vars)
+
+        log(3, f"[Guard]: {inlined_guard_expr}\n")
+        guard.inlined_expr = inlined_guard_expr
+        guard.expr = guard_expr
+
         assert callable(guard), "guard must be callable."
 
         return guard
@@ -136,7 +143,7 @@ def wrapper(self: CheckGuardInputT) -> list[StringifiedExpression]:
         def guard_log():
             frame_value_tracer = self.tracker.trace_value_from_frame()
             print(
-                f"[Guard]: guard_fn for {self}, tracker={self.tracker.__class__.__name__}, value={frame_value_tracer.expr}"
+                f"[Guard]: guard_fn for {self}, tracker={self.tracker.__class__.__name__}, value={frame_value_tracer.registered_expr}"
             )
 
         log_do(4, guard_log)
diff --git a/python/paddle/jit/sot/profiler/profiler.py b/python/paddle/jit/sot/profiler/profiler.py
@@ -40,7 +40,7 @@ def EventGuard(event_name, event_level=1):
     try:
         global _event_level
         need_pop = False
-        if ENV_SOT_EVENT_LEVEL.get() >= event_level:
+        if ENV_SOT_EVENT_LEVEL.get_with_cache() >= event_level:
             core.nvprof_nvtx_push(event_name)
             need_pop = True
         yield
diff --git a/python/paddle/jit/sot/utils/__init__.py b/python/paddle/jit/sot/utils/__init__.py
@@ -56,7 +56,7 @@
     StepInfoManager,
     StepState,
     count_if,
-    current_tmp_name_records,
+    current_symbol_registry,
     execute_time,
     flatten,
     flatten_extend,
@@ -80,5 +80,5 @@
     meta_str,
     no_eval_frame,
     printable,
-    tmp_name_guard,
+    switch_symbol_registry,
 )
diff --git a/python/paddle/jit/sot/utils/utils.py b/python/paddle/jit/sot/utils/utils.py
@@ -77,38 +77,46 @@ def match_name(self, name: str) -> bool:
         return name.startswith(self.prefix)
 
 
-class TmpNameRecords:
+class SymbolRegistry:
     def __init__(self):
-        self.name_generator = NameGenerator(prefix="_sot_tmp_")
+        self.symbol_generator = NameGenerator(prefix="___t_")
         self.tmp_names_record = OrderedDict()
+        self.declared_symbols: set[str] = set()
+        self.symbol_table = {}
 
-    def next_name(self):
-        return self.name_generator.next()
+    def next_symbol(self) -> str:
+        return self.symbol_generator.next()
 
-    def add_tmp_var(self, expr):
-        if expr in self.tmp_names_record:
-            return self.tmp_names_record[expr]
-        else:
-            tmp_name = self.next_name()
-            self.tmp_names_record[expr] = tmp_name
-            return tmp_name
+    def request_symbol(self, expr: str) -> str:
+        if expr in self.symbol_table:
+            return self.symbol_table[expr]
+        symbol = self.next_symbol()
+        self.symbol_table[expr] = symbol
+        return symbol
+
+    def gen_expr(self, expr: str, gen_expr_fn):
+        symbol = self.symbol_table[expr]
+        if symbol in self.declared_symbols:
+            return symbol
+        self.declared_symbols.add(symbol)
+        return f"({symbol} := ({gen_expr_fn()}))"
 
 
-_tmp_name_records = TmpNameRecords()
+_symbol_registry = SymbolRegistry()
 
 
 @contextmanager
-def tmp_name_guard():
-    global _tmp_name_records
-    old = _tmp_name_records
-    _tmp_name_records = TmpNameRecords()
+def switch_symbol_registry():
+    global _symbol_registry
+    original_registry = _symbol_registry
+    _symbol_registry = SymbolRegistry()
     yield
-    _tmp_name_records = old
+    _symbol_registry = original_registry
 
 
-def current_tmp_name_records():
-    global _tmp_name_records
-    return _tmp_name_records
+def current_symbol_registry():
+    global _symbol_registry
+    return _symbol_registry
 
 
 class ResumeFnNameFactory(metaclass=Singleton):
diff --git a/python/paddle/utils/environments.py b/python/paddle/utils/environments.py
@@ -39,6 +39,12 @@ class EnvironmentVariable(Generic[T]):
     def __init__(self, name: str, default: T):
         self.name = name
         self.default = default
+        self._cached_value: T | None = None
+
+    def get_with_cache(self) -> T:
+        if self._cached_value is None:
+            self._cached_value = self.get()
+        return self._cached_value
 
     def get(self) -> T:
         raise NotImplementedError
