[SOT][Faster Guard] Implement moreguard_tree_expr_node and TensorDtypeVariable.make_faster_guard (#72463)

zrr1999 · web-flow · commit 99a2dc9a4809 · 2025-04-26T03:58:51.000+08:00
diff --git a/paddle/fluid/pybind/sot/guards.cc b/paddle/fluid/pybind/sot/guards.cc
@@ -402,9 +402,7 @@ std::string ExprGuardNode::stringify(int indent) {
 void GuardTree::add_guard_chain(
     const std::vector<std::shared_ptr<GuardNodeBase>>& guard_chain) {
   if (guard_chain.empty()) {
-    // TODO(zrr1999): empty guard nodes means that some
-    // tracker.make_faster_guard is not implemented.
-    return;
+    throw std::runtime_error("Empty guard chain");
   }
   for (size_t i = 1; i < guard_chain.size(); ++i) {
     guard_chain[i - 1]->next_guard_nodes.push_back(guard_chain[i]);
diff --git a/python/paddle/jit/sot/opcode_translator/executor/executor_cache.py b/python/paddle/jit/sot/opcode_translator/executor/executor_cache.py
@@ -124,9 +124,10 @@ def __call__(self, frame: types.FrameType, **kwargs) -> CustomCode:
             )
             assert guard_fn is not None
             assert guard_chain is not None
-            self.cache[code] = [
-                (new_custom_code, guard_fn)
-            ], paddle.framework.core.GuardTree([guard_chain])
+            self.cache[code] = (
+                [(new_custom_code, guard_fn)],
+                paddle.framework.core.GuardTree([guard_chain]),
+            )
             return new_custom_code
         guarded_fns, guard_tree = self.cache[code]
         compile_time_for_code = self.compile_time_stats.get(code, 0)
@@ -228,9 +229,9 @@ def lookup(
                         f"[Cache] Cache hit, Guard is \n{getattr(guard_fn, 'expr', 'None')}\n",
                     )
                     # TODO(zrr1999): cache_index should be equal to index when enable_strict_guard.
-                    # assert (
-                    #     cache_index is None or index == cache_index
-                    # ), f"cache_index({cache_index}) is not equal to index({index})"
+                    assert (
+                        cache_index is None or index == cache_index
+                    ), f"cache_index({cache_index}) is not equal to index({index})"
                     return custom_code
                 else:
                     log_do(
@@ -373,6 +374,14 @@ def start_translate(
                 None,
             )
         guard_chain = simulator.guard_chain
+        if len(guard_chain) == 0:
+            # TODO(zrr1999): GuardNode should support zero-expr constructor, to implement DummyGuardNode
+            guard_chain: GuardChain = [
+                paddle.framework.core.GuardNode(
+                    paddle.framework.core.DummyGuard(),
+                    [paddle.framework.core.ConstantExprNode(True)],
+                )
+            ]
         return new_custom_code, guard_fn, guard_chain
     # TODO(0x45f): handle BreakGraphError to trigger fallback
     except BreakGraphError as e:
@@ -391,7 +400,7 @@ def start_translate(
             f"Unsupported Frame is {frame.f_code}, error message is: \n"
             + "".join(traceback.format_exception(type(e), e, e.__traceback__)),
         )
-        dummy_guard_chain = [
+        dummy_guard_chain: GuardChain = [
             # TODO(zrr1999): GuardNode should support zero-expr constructor
             paddle.framework.core.GuardNode(
                 paddle.framework.core.DummyGuard(),
diff --git a/python/paddle/jit/sot/opcode_translator/executor/function_graph.py b/python/paddle/jit/sot/opcode_translator/executor/function_graph.py
@@ -42,8 +42,6 @@
 from ...symbolic_shape.operators import SYMBOLIC_BINARY_OPS, SYMBOLIC_UNARY_OPS
 from ...utils import (
     ENV_SOT_ALLOW_DYNAMIC_SHAPE,
-    ENV_SOT_ENABLE_GUARD_TREE,
-    ENV_SOT_ENABLE_STRICT_GUARD_CHECK,
     NameGenerator,
     SIRToCodeMap,
     SotUndefinedVar,
@@ -320,23 +318,13 @@ def collect(inp):
     @property
     @event_register("guard_chain")
     def guard_chain(self) -> list[paddle.framework.core.GuardNodeBase]:
-        enable_strict_guard = ENV_SOT_ENABLE_STRICT_GUARD_CHECK.get()
-        enable_guard_tree = ENV_SOT_ENABLE_GUARD_TREE.get()
-
-        if not enable_strict_guard and not enable_guard_tree:
-            return []
         guard_chain: list[paddle.framework.core.GuardNodeBase] = []
 
         with EventGuard("guard_fn: find vars and make faster guard"):
-            try:
-                for variable in find_traceable_vars(
-                    self.input_variables + list(self._global_guarded_variables)
-                ):
-                    guard_chain.extend(variable.make_faster_guard())
-            except NotImplementedError as e:
-                log(2, f"[Guard] make faster guard nodes error: {e}\n")
-                # TODO(zrr1999): empty list means that some tracker.make_faster_guard is not implemented.
-                return []
+            for variable in find_traceable_vars(
+                self.input_variables + list(self._global_guarded_variables)
+            ):
+                guard_chain.extend(variable.make_faster_guard())
         return guard_chain
 
     @property
diff --git a/python/paddle/jit/sot/opcode_translator/executor/opcode_inline_executor.py b/python/paddle/jit/sot/opcode_translator/executor/opcode_inline_executor.py
@@ -76,7 +76,7 @@ def gen_instructions(self, codegen: PyCodeGen):
         codegen.gen_load_const(self.name)
         codegen.gen_subscribe()
 
-    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNode:
+    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
         fn_tracer = self.fn.tracker.guard_tree_expr_node()
         return paddle.framework.core.ItemExprNode(
             paddle.framework.core.AttributeExprNode(
@@ -105,58 +105,6 @@ def __repr__(self) -> str:
         return f"FunctionGlobalTracker(fn={self.fn}, name={self.name})"
 
 
-class FunctionClosureTracker(Tracker):
-    """
-    A tracker class that represents a function closure variable.
-
-    Args:
-        fn: The FunctionVariable object.
-        idx: The index of the closure variable.
-
-    """
-
-    def __init__(self, fn: FunctionVariable, idx: int):
-        super().__init__([fn])
-        self.fn = fn
-        self.idx = idx
-
-    def gen_instructions(self, codegen: PyCodeGen):
-        """
-        Generate bytecode instructions to trace the value of the function closure variable.
-
-        Args:
-            codegen: The PyCodeGen object used to generate bytecode.
-
-        """
-        self.fn.tracker.gen_instructions(codegen)
-        codegen.gen_load_attr("__closure__")
-        codegen.gen_load_const(self.idx)
-        codegen.gen_subscribe()
-        codegen.gen_load_attr("cell_contents")
-
-    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNode:
-        # TODO(zrr1999): implement FunctionClosureExprNode
-        raise NotImplementedError("FunctionClosureExprNode is not implemented")
-
-    def trace_value_from_frame(self):
-        """
-        Trace the value of the function closure variable from the frame.
-
-        Returns:
-            The traced value of the function closure variable.
-
-        """
-        fn_tracer = self.fn.tracker.trace_value_from_frame()
-        return StringifiedExpression(
-            f"{{}}.__closure__[{self.idx}].cell_contents",
-            [fn_tracer],
-            union_free_vars(fn_tracer.free_vars),
-        )
-
-    def __repr__(self) -> str:
-        return f"FunctionClosureTracker(fn={self.fn}, idx={self.idx})"
-
-
 def inline_for_iter_impl(exe: OpcodeExecutorBase, instr: Instruction):
     iterator = exe.stack.top
     assert isinstance(iterator, IterVariable)
diff --git a/python/paddle/jit/sot/opcode_translator/executor/tracker.py b/python/paddle/jit/sot/opcode_translator/executor/tracker.py
@@ -61,7 +61,7 @@ def gen_instructions(self, codegen: PyCodeGen) -> None:
         """
         raise NotImplementedError
 
-    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNode:
+    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
         raise NotImplementedError(
             f"{self.__class__.__name__} has no guard_tree_expr_node"
         )
@@ -197,7 +197,7 @@ def __init__(self, name: str):
     def gen_instructions(self, codegen: PyCodeGen) -> None:
         codegen.gen_load_fast(self.name)
 
-    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNode:
+    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
         return paddle.framework.core.LocalVarExprNode(self.name)
 
     def trace_value_from_frame(self) -> StringifiedExpression:
@@ -211,7 +211,7 @@ class CellTracker(LocalTracker):
     def gen_instructions(self, codegen: PyCodeGen):
         codegen.gen_load_deref(self.name)
 
-    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNode:
+    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
         return paddle.framework.core.LocalVarExprNode(self.name)
 
     def trace_value_from_frame(self):
@@ -236,7 +236,7 @@ def __init__(self, name: str):
     def gen_instructions(self, codegen: PyCodeGen) -> None:
         codegen.gen_load_global(self.name, push_null=False)
 
-    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNode:
+    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
         return paddle.framework.core.GlobalVarExprNode(self.name)
 
     def trace_value_from_frame(self) -> StringifiedExpression:
@@ -261,7 +261,7 @@ def __init__(self, name: str):
     def gen_instructions(self, codegen: PyCodeGen) -> None:
         codegen.gen_load_global(self.name, push_null=False)
 
-    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNode:
+    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
         return paddle.framework.core.ConstantExprNode(
             getattr(builtins, self.name)
         )
@@ -290,7 +290,7 @@ def __init__(self, value):
     def gen_instructions(self, codegen: PyCodeGen):
         codegen.gen_load_const(self.value)
 
-    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNode:
+    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
         return paddle.framework.core.ConstantExprNode(self.value)
 
     def trace_value_from_frame(self):
@@ -324,7 +324,7 @@ def gen_instructions(self, codegen: PyCodeGen):
         self.obj.tracker.gen_instructions(codegen)
         codegen.gen_load_attr(self.attr)
 
-    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNode:
+    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
         obj_tracer = self.obj.tracker.guard_tree_expr_node()
         return paddle.framework.core.AttributeExprNode(
             obj_tracer,
@@ -377,7 +377,7 @@ def gen_instructions(self, codegen: PyCodeGen):
             codegen.gen_load_const(self.key)
         codegen.gen_subscribe()
 
-    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNode:
+    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
         container_tracer = self.container.tracker.guard_tree_expr_node()
         return paddle.framework.core.ItemExprNode(
             container_tracer,
@@ -418,7 +418,7 @@ def gen_instructions(self, codegen: PyCodeGen):
         self.iter_source.tracker.gen_instructions(codegen)
         codegen.add_instr("GET_ITER")
 
-    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNode:
+    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
         # TODO(zrr1999): implement IterExprNode
         raise NotImplementedError("IterExprNode is not implemented")
 
@@ -459,7 +459,7 @@ def gen_instructions(self, codegen: PyCodeGen):
             codegen.gen_build_map(len(self.kwargs))
             codegen.gen_call_function_ex(has_kwargs=True)
 
-    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNode:
+    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
         # TODO(zrr1999): implement LayerExprNode
         raise NotImplementedError("LayerExprNode is not implemented")
 
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py b/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py
@@ -335,7 +335,19 @@ def __init__(self, value, graph, tracker):
 
     @check_faster_guard
     def make_faster_guard(self) -> list[paddle.framework.core.GuardNodeBase]:
-        raise NotImplementedError
+        if isinstance(self.tracker, GetAttrTracker) and isinstance(
+            self.tracker.obj, TensorVariable
+        ):
+            expr_node = self.tracker.obj.tracker.guard_tree_expr_node()
+            assert paddle.framework.use_pir_api(), "Only support PIR"
+            return [
+                paddle.framework.core.GuardNode(
+                    paddle.framework.core.DtypeMatchGuard(self.value),
+                    [expr_node],
+                )
+            ]
+        else:
+            return object_equal_faster_guard(self)
 
     @check_guard
     def make_stringified_guard(self) -> list[StringifiedExpression]:
@@ -1240,6 +1252,8 @@ def make_stringified_guard(self) -> list[StringifiedExpression]:
         assert ENV_SOT_ALLOW_DYNAMIC_SHAPE.get()
         from ..executor_cache import OpcodeExecutorCache
 
+        # NOTE(zrr1999): SymbolicVariable is not supported in faster guard mode
+
         frame_value_tracer = self.tracker.trace_value_from_frame()
         symbolic_inputs = OpcodeExecutorCache().get_symbolic_inputs(
             self.graph.pycode_gen._origin_code
@@ -1267,7 +1281,6 @@ def make_stringified_guard(self) -> list[StringifiedExpression]:
                 [frame_value_tracer],
                 union_free_vars(frame_value_tracer.free_vars),
             ),
-            # TODO(zrr1999): replace it with FasterStringifiedExpression
             *constraint_guards,
         ]
         return guards
diff --git a/python/paddle/jit/sot/opcode_translator/executor/virtual_frame.py b/python/paddle/jit/sot/opcode_translator/executor/virtual_frame.py
@@ -20,6 +20,8 @@
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, Any, NamedTuple
 
+import paddle
+
 from ...utils import log
 from .guard import StringifiedExpression, union_free_vars
 from .tracker import (
@@ -82,6 +84,19 @@ def gen_instructions(self, codegen: PyCodeGen):
         codegen.gen_subscribe()
         codegen.gen_load_attr("cell_contents")
 
+    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
+        fn_tracer = self.fn.tracker.guard_tree_expr_node()
+        return paddle.framework.core.AttributeExprNode(
+            paddle.framework.core.ItemExprNode(
+                paddle.framework.core.AttributeExprNode(
+                    fn_tracer,
+                    "__closure__",
+                ),
+                paddle.framework.core.ConstantExprNode(self.idx),
+            ),
+            "cell_contents",
+        )
+
     def trace_value_from_frame(self):
         """
         Trace the value of the function closure variable from the frame.
