[SOT][Faster Guard] Move all trackers implements to tracker.py (#72501)

zrr1999 · web-flow · commit 266e2a9ca344 · 2025-04-27T01:55:14.000+08:00
diff --git a/paddle/fluid/pybind/sot/guards.cc b/paddle/fluid/pybind/sot/guards.cc
@@ -521,7 +521,8 @@ std::string DummyGuardNode::stringify(int indent) {
 void GuardTree::add_guard_chain(
     const std::vector<std::shared_ptr<GuardNodeBase>>& guard_chain) {
   if (guard_chain.empty()) {
-    throw std::runtime_error("Empty guard chain");
+    PADDLE_THROW(phi::errors::InvalidArgument(
+        "Empty guard chain, please check the 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/opcode_inline_executor.py b/python/paddle/jit/sot/opcode_translator/executor/opcode_inline_executor.py
@@ -18,8 +18,6 @@
 import sys
 from typing import TYPE_CHECKING
 
-import paddle
-
 from ...utils import (
     BreakGraphError,
     DataDependencyControlFlowBreak,
@@ -28,9 +26,8 @@
 )
 from ..instruction_utils import Instruction
 from .dispatch_functions import generator_send
-from .guard import StringifiedExpression, union_free_vars
 from .opcode_executor import OpcodeExecutorBase, Stop
-from .tracker import DanglingTracker, Tracker
+from .tracker import DanglingTracker
 from .variables import (
     BuiltinVariable,
     ConstantVariable,
@@ -43,68 +40,9 @@
 
 if TYPE_CHECKING:
     from .function_graph import FunctionGraph
-    from .pycode_generator import PyCodeGen
-    from .variables import FunctionVariable
     from .virtual_frame import VirtualFrame
 
 
-class FunctionGlobalTracker(Tracker):
-    """
-    A tracker class that represents a function global variable.
-
-    Args:
-        fn: FunctionVariable object.
-        name: The name of the global variable.
-
-    """
-
-    def __init__(self, fn: FunctionVariable, name: str):
-        super().__init__([fn])
-        self.fn = fn
-        self.name = name
-
-    def gen_instructions(self, codegen: PyCodeGen):
-        """
-        Generate bytecode instructions in order to put the variables at the top of the stack.
-
-        Args:
-            codegen: The PyCodeGen object used to generate bytecode.
-
-        """
-        self.fn.tracker.gen_instructions(codegen)
-        codegen.gen_load_attr("__globals__")
-        codegen.gen_load_const(self.name)
-        codegen.gen_subscribe()
-
-    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(
-                fn_tracer,
-                "__globals__",
-            ),
-            paddle.framework.core.ConstantExprNode(self.name),
-        )
-
-    def trace_value_from_frame(self) -> StringifiedExpression:
-        """
-        Trace the value of the function global variable from the frame.
-
-        Returns:
-            StringifiedExpression: The traced value of the function global variable.
-
-        """
-        fn_tracer = self.fn.tracker.trace_value_from_frame()
-        return StringifiedExpression(
-            f"{{}}.__globals__['{self.name}']",
-            [fn_tracer],
-            union_free_vars(fn_tracer.free_vars),
-        )
-
-    def __repr__(self) -> str:
-        return f"FunctionGlobalTracker(fn={self.fn}, name={self.name})"
-
-
 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
@@ -29,7 +29,7 @@
 
     from ...utils.magic_methods import BinaryOp, UnaryOp
     from .pycode_generator import PyCodeGen
-    from .variables import VariableBase
+    from .variables import FunctionVariable, VariableBase
 
 
 class Tracker:
@@ -496,3 +496,121 @@ def trace_value_from_frame(self):
 
     def __repr__(self) -> str:
         return f"CreateLayerTracker(Layer={self.layer_class}, args={self.args}, kwargs={self.kwargs})"
+
+
+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.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.
+
+        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})"
+
+
+class FunctionGlobalTracker(Tracker):
+    """
+    A tracker class that represents a function global variable.
+
+    Args:
+        fn: FunctionVariable object.
+        name: The name of the global variable.
+
+    """
+
+    def __init__(self, fn: FunctionVariable, name: str):
+        super().__init__([fn])
+        self.fn = fn
+        self.name = name
+
+    def gen_instructions(self, codegen: PyCodeGen):
+        """
+        Generate bytecode instructions in order to put the variables at the top of the stack.
+
+        Args:
+            codegen: The PyCodeGen object used to generate bytecode.
+
+        """
+        self.fn.tracker.gen_instructions(codegen)
+        codegen.gen_load_attr("__globals__")
+        codegen.gen_load_const(self.name)
+        codegen.gen_subscribe()
+
+    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(
+                fn_tracer,
+                "__globals__",
+            ),
+            paddle.framework.core.ConstantExprNode(self.name),
+        )
+
+    def trace_value_from_frame(self) -> StringifiedExpression:
+        """
+        Trace the value of the function global variable from the frame.
+
+        Returns:
+            StringifiedExpression: The traced value of the function global variable.
+
+        """
+        fn_tracer = self.fn.tracker.trace_value_from_frame()
+        return StringifiedExpression(
+            f"{{}}.__globals__['{self.name}']",
+            [fn_tracer],
+            union_free_vars(fn_tracer.free_vars),
+        )
+
+    def __repr__(self) -> str:
+        return f"FunctionGlobalTracker(fn={self.fn}, name={self.name})"
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py b/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py
@@ -2021,7 +2021,7 @@ def __init__(
         self.fn = fn
 
     def proxy_getter(self, proxy: MutableDictLikeData, key: Any):
-        from ..opcode_inline_executor import FunctionGlobalTracker
+        from ..tracker import FunctionGlobalTracker
 
         if key not in proxy.original_data:
             return MutableDictLikeData.Empty()
diff --git a/python/paddle/jit/sot/opcode_translator/executor/virtual_frame.py b/python/paddle/jit/sot/opcode_translator/executor/virtual_frame.py
@@ -18,17 +18,14 @@
 import re
 from typing import TYPE_CHECKING, Any, NamedTuple
 
-import paddle
-
 from ...utils import log
-from .guard import StringifiedExpression, union_free_vars
 from .tracker import (
     BuiltinTracker,
     CellTracker,
     ConstTracker,
     DanglingTracker,
+    FunctionClosureTracker,
     LocalTracker,
-    Tracker,
 )
 from .variable_stack import VariableStack
 from .variables.base import VariableBase, VariableFactory, fn_bind_inputs
@@ -45,74 +42,12 @@
     from typing_extensions import TypeAlias
 
     from .function_graph import FunctionGraph
-    from .pycode_generator import PyCodeGen
     from .variables.callable import FunctionVariable
 
     # The type to represent the (*args, **kwargs) pack in the call.
     CallArgsPack: TypeAlias = tuple[tuple[Any, ...], dict[str, Any]]
 
 
-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.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.
-
-        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 validate_value(value):
     assert isinstance(
         value, VariableBase
