[SOT][Faster Guard][3.13] implement faster lookup (#71994)

Author: zrr1999

paddle/fluid/pybind/jit.cc

@@ -125,6 +125,9 @@ void BindGuard(pybind11::module *m) {
   py::class_<WeakRefMatchGuard, GuardBase, std::shared_ptr<WeakRefMatchGuard>>(
       *m, "WeakRefMatchGuard", R"DOC(WeakRefMatchGuard Class.)DOC")
       .def(py::init<const py::object &>(), py::arg("func"));
+  py::class_<DummyGuard, GuardBase, std::shared_ptr<DummyGuard>>(
+      *m, "DummyGuard", R"DOC(DummyGuard Class.)DOC")
+      .def(py::init<>());
 
   m->def(
       "merge_guard",
@@ -147,16 +150,23 @@ void BindGuardTree(pybind11::module *m) {
           [](GuardTree &self, py::object frame) {
             return self.lookup(reinterpret_cast<FrameProxy *>(frame.ptr()));
           },
-          py::arg("frame"));
+          py::arg("frame"))
+      .def(
+          "add_guard_chain",
+          [](GuardTree &self,
+             const std::vector<std::shared_ptr<GuardNode>> &guard_chain) {
+            self.add_guard_chain(guard_chain);
+          },
+          py::arg("guard_chain"));
 
   py::class_<GuardNode, std::shared_ptr<GuardNode>>(
       *m, "GuardNode", R"DOC(GuardNode Class.)DOC")
       .def(py::init<const std::shared_ptr<GuardBase> &,
-                    const std::shared_ptr<ExprNode> &,
+                    const std::vector<std::shared_ptr<ExprNode>> &,
                     const std::vector<std::shared_ptr<GuardNode>> &,
                     const std::optional<int> &>(),
            py::arg("guard"),
-           py::arg("expr"),
+           py::arg("exprs"),
            py::arg("next_guard_nodes") = py::list(),
            py::arg("return_cache_index") = py::none())
       .def_property(
@@ -183,6 +193,12 @@ void BindGuardTree(pybind11::module *m) {
       *m, "ConstantExprNode", R"DOC(ConstantExprNode Class.)DOC")
       .def(py::init<const py::object &>(), py::arg("value_ptr"));
 
+  py::class_<ExternVarExprNode, ExprNode, std::shared_ptr<ExternVarExprNode>>(
+      *m, "ExternVarExprNode", R"DOC(ExternVarExprNode Class.)DOC")
+      .def(py::init<const std::string &, const py::object &>(),
+           py::arg("var_name"),
+           py::arg("value_ptr"));
+
   py::class_<LocalVarExprNode, ExprNode, std::shared_ptr<LocalVarExprNode>>(
       *m, "LocalVarExprNode", R"DOC(LocalVarExprNode Class.)DOC")
       .def(py::init<const std::string &>(), py::arg("var_name"));

paddle/fluid/pybind/sot/guards.cc

@@ -195,6 +195,8 @@ bool WeakRefMatchGuard::check(PyObject* value) {
 
 PyObject* ConstantExprNode::eval(FrameProxy* frame) { return value_ptr_; }
 
+PyObject* ExternVarExprNode::eval(FrameProxy* frame) { return value_ptr_; }
+
 PyObject* LocalVarExprNode::eval(FrameProxy* frame) {
 #if PY_3_13_PLUS
   return PyDict_GetItemString(frame->locals, var_name_.c_str());
@@ -222,6 +224,8 @@ PyObject* ItemExprNode::eval(FrameProxy* frame) {
 }
 
 std::optional<int> GuardNode::lookup(FrameProxy* frame) {
+  // TODO(zrr1999): support multiple exprs
+  auto expr = exprs.back();
   auto value = expr->eval(frame);
   if (guard->check(value)) {
     if (return_cache_index.has_value()) {

paddle/fluid/pybind/sot/guards.h

@@ -250,6 +250,11 @@ class WeakRefMatchGuard : public GuardBase {
   PyObject* expected_;
 };
 
+class DummyGuard : public GuardBase {
+ public:
+  bool check(PyObject* value) override { return true; }
+};
+
 class GuardTreeNode {};
 
 class AttributeExprNode;
@@ -272,6 +277,21 @@ class ConstantExprNode : public ExprNode {
  private:
   PyObject* value_ptr_;
 };
+class ExternVarExprNode : public ExprNode {
+ public:
+  explicit ExternVarExprNode(const std::string& var_name,
+                             const py::object& value_obj)
+      : value_ptr_(value_obj.ptr()), var_name_(var_name) {
+    Py_INCREF(value_ptr_);
+  }
+
+  ~ExternVarExprNode() { Py_DECREF(value_ptr_); }
+  PyObject* eval(FrameProxy* frame);
+
+ private:
+  PyObject* value_ptr_;
+  std::string var_name_;
+};
 
 class LocalVarExprNode : public ExprNode {
  public:
@@ -321,16 +341,16 @@ class ItemExprNode : public ExprNode {
 class GuardNode : public GuardTreeNode {
  public:
   std::shared_ptr<GuardBase> guard;
-  std::shared_ptr<ExprNode> expr;
+  std::vector<std::shared_ptr<ExprNode>> exprs;
   std::vector<std::shared_ptr<GuardNode>> next_guard_nodes;
   // return_cache_index is used to record the index of the guard list
   std::optional<int> return_cache_index;
   GuardNode(std::shared_ptr<GuardBase> guard,
-            std::shared_ptr<ExprNode> expr,
+            std::vector<std::shared_ptr<ExprNode>> exprs,
             std::vector<std::shared_ptr<GuardNode>> next_guard_nodes,
             std::optional<int> return_cache_index)
       : guard(guard),
-        expr(expr),
+        exprs(exprs),
         next_guard_nodes(next_guard_nodes),
         return_cache_index(return_cache_index) {}
 
@@ -342,13 +362,21 @@ class GuardTree {
   GuardTree(const std::vector<std::vector<std::shared_ptr<GuardNode>>>&
                 guard_nodes_list) {
     for (size_t index = 0; index < guard_nodes_list.size(); ++index) {
-      const auto& guard_nodes = guard_nodes_list[index];
-      for (size_t i = 1; i < guard_nodes.size(); ++i) {
-        guard_nodes[i - 1]->next_guard_nodes.push_back(guard_nodes[i]);
-      }
-      guard_nodes.back()->return_cache_index = index;
-      guard_nodes_.push_back(guard_nodes.front());
+      add_guard_chain(guard_nodes_list[index]);
+    }
+  }
+  void add_guard_chain(
+      const std::vector<std::shared_ptr<GuardNode>>& guard_chain) {
+    if (guard_chain.empty()) {
+      // TODO(zrr1999): empty guard nodes means that some
+      // tracker.make_faster_guard is not implemented.
+      return;
+    }
+    for (size_t i = 1; i < guard_chain.size(); ++i) {
+      guard_chain[i - 1]->next_guard_nodes.push_back(guard_chain[i]);
     }
+    guard_chain.back()->return_cache_index = guard_nodes_.size();
+    guard_nodes_.push_back(guard_chain.front());
   }
 
   std::optional<int> lookup(FrameProxy* frame);

python/paddle/jit/sot/opcode_translator/executor/executor_cache.py

@@ -18,6 +18,7 @@
 import traceback
 from typing import TYPE_CHECKING, List, Tuple
 
+import paddle
 from paddle.base.dygraph.base import sot_simulation_mode_guard
 
 from ...profiler import EventGuard, event_register
@@ -47,6 +48,8 @@
 
 GuardedFunction = Tuple[CustomCode, Guard]
 GuardedFunctions = List[GuardedFunction]
+GuardChain = List[paddle.framework.core.GuardNode]
+GuardChainList = List[GuardChain]
 
 dummy_guard: Guard = lambda frame: True
 dummy_guard.expr = "lambda frame: True"
@@ -66,7 +69,9 @@ class OpcodeExecutorCache(metaclass=Singleton):
     """
 
     MAX_CACHE_SIZE = 20
-    cache: dict[types.CodeType, GuardedFunctions]
+    cache: dict[
+        types.CodeType, tuple[GuardedFunctions, paddle.framework.core.GuardTree]
+    ]
     translate_count: int
     code_symbolic_inputs: dict[types.CodeType, dict[str, None | dict[int, int]]]
 
@@ -105,16 +110,25 @@ def __call__(self, frame: types.FrameType, **kwargs) -> CustomCode:
         code: types.CodeType = frame.f_code
         if code not in self.cache:
             log(2, f"[Cache]: Firstly call {code}\n")
-            new_custom_code, guard_fn = self.translate(frame, **kwargs)
+            new_custom_code, guard_fn, guard_chain = self.translate(
+                frame, **kwargs
+            )
             assert guard_fn is not None
-            self.cache[code] = [(new_custom_code, guard_fn)]
+            assert guard_chain is not None
+            self.cache[code] = [
+                (new_custom_code, guard_fn)
+            ], paddle.framework.core.GuardTree([guard_chain])
             return new_custom_code
-        guarded_fns = self.cache[code]
-        return self.lookup(frame, guarded_fns, **kwargs)
+        guarded_fns, guard_tree = self.cache[code]
+        return self.lookup(frame, guarded_fns, guard_tree, **kwargs)
 
     @event_register("lookup")
     def lookup(
-        self, frame: types.FrameType, guarded_fns: GuardedFunctions, **kwargs
+        self,
+        frame: types.FrameType,
+        guarded_fns: GuardedFunctions,
+        guard_tree: paddle.framework.core.GuardTree,
+        **kwargs,
     ) -> CustomCode:
         """
         Looks up the cache for a matching code object and returns a custom code object if a matching guard function is found, otherwise None.
@@ -132,8 +146,17 @@ def lookup(
             return CustomCode(None, False)
 
         enable_strict_guard = ENV_SOT_ENABLE_STRICT_GUARD_CHECK.get()
+        enable_guard_tree = ENV_SOT_ENABLE_GUARD_TREE.get()
+
+        cache_index = None
+        if enable_strict_guard or enable_guard_tree:
+            cache_index = guard_tree.lookup(frame)
 
-        for custom_code, guard_fn in guarded_fns:
+        if not enable_strict_guard and cache_index is not None:
+            # TODO(zrr1999): add a mapping between custom_code and cache_index
+            return guarded_fns[cache_index][0]
+
+        for index, (custom_code, guard_fn) in enumerate(guarded_fns):
             if enable_strict_guard:
                 mirror_guard_error = None
                 try:
@@ -157,9 +180,12 @@ def lookup(
                         2,
                         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})"
                     return custom_code
-                elif not ENV_SOT_ENABLE_GUARD_TREE.get():
-                    # TODO(zrr1999): remove condition after faster guard tree support error analysis
+                else:
                     log_do(
                         4,
                         self.analyse_guard_global_object(guard_fn),
@@ -192,9 +218,11 @@ def lookup(
                     )
 
         log(2, "[Cache]: all guards missed\n")
-        new_custom_code, guard_fn = self.translate(frame, **kwargs)
+        new_custom_code, guard_fn, guard_chain = self.translate(frame, **kwargs)
         if guard_fn is not None:
+            assert guard_chain is not None
             guarded_fns.append((new_custom_code, guard_fn))
+            guard_tree.add_guard_chain(guard_chain)
         return new_custom_code
 
     def before_translate_hook(self, frame: types.FrameType):
@@ -203,7 +231,7 @@ def before_translate_hook(self, frame: types.FrameType):
 
     def translate(
         self, frame: types.FrameType, **kwargs
-    ) -> tuple[CustomCode, Guard | None]:
+    ) -> tuple[CustomCode, Guard | None, GuardChain | None]:
         """
         Translates the given frame's code object and returns the cache getter function and a guarded function for the translated code object.
 
@@ -215,8 +243,10 @@ def translate(
         """
         self.before_translate_hook(frame)
         self.translate_count += 1
-        custom_new_code, guard_fn = start_translate(frame, **kwargs)
-        return custom_new_code, guard_fn
+        custom_new_code, guard_fn, guard_chain = start_translate(
+            frame, **kwargs
+        )
+        return custom_new_code, guard_fn, guard_chain
 
     def analyse_guard_global_object(self, guard_fn):
         def inner():
@@ -255,15 +285,15 @@ def inner():
 def start_translate(
     frame: types.FrameType,
     **kwargs,
-) -> tuple[CustomCode, Guard | None]:
+) -> tuple[CustomCode, Guard | None, GuardChain | None]:
     """
-    Starts the translation process for the given frame and returns the translated code object and its guard function, or None if translation fails.
+    Starts the translation process for the given frame and returns the translated code object, its guard function and its guard tree node, or None if translation fails.
 
     Args:
         frame: The frame to be translated.
 
     Returns:
-        tuple[CustomCode, Guard | None]: The translated code object and its guard function, or None if translation fails.
+        tuple[CustomCode, Guard | None, GuardChain | None]: The translated code object, its guard function and its guard tree node, or None if translation fails.
     """
     graph = FunctionGraph(frame.f_code, frame.f_globals, **kwargs)
     vframe = VirtualFrame.from_real_frame(frame, graph)
@@ -280,8 +310,10 @@ def start_translate(
             return (
                 CustomCode(None, True),
                 None,
+                None,
             )
-        return new_custom_code, guard_fn
+        guard_chain = simulator.guard_chain
+        return new_custom_code, guard_fn, guard_chain
     # TODO(0x45f): handle BreakGraphError to trigger fallback
     except BreakGraphError as e:
         raise RuntimeError(
@@ -299,9 +331,18 @@ 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 = [
+            # TODO(zrr1999): GuardNode should support zero-expr constructor
+            paddle.framework.core.GuardNode(
+                paddle.framework.core.DummyGuard(),
+                [paddle.framework.core.ConstantExprNode(True)],
+            )
+        ]
         return (
             CustomCode(None, e.disable_eval_frame),
             dummy_guard,
+            dummy_guard_chain,
         )
     except Exception as e:
         raise InnerError(OpcodeExecutorBase.error_message_summary(e)) from e

python/paddle/jit/sot/opcode_translator/executor/function_graph.py

@@ -43,6 +43,7 @@
 from ...utils import (
     ENV_SOT_ALLOW_DYNAMIC_SHAPE,
     ENV_SOT_ENABLE_GUARD_TREE,
+    ENV_SOT_ENABLE_STRICT_GUARD_CHECK,
     NameGenerator,
     SotUndefinedVar,
     inner_error_default_handler,
@@ -60,7 +61,7 @@
     SotExtraInfo,
 )
 from ..instruction_utils import get_instructions
-from .guard import Guard, StringifiedExpression, make_faster_guard, make_guard
+from .guard import Guard, StringifiedExpression, make_guard
 from .mutable_data import MutationDel, MutationNew, MutationSet
 from .pycode_generator import PyCodeGen
 from .side_effects import (
@@ -316,17 +317,30 @@ def collect(inp):
         )
 
     @property
-    @event_register("guard_fn")
-    def guard_fn(self) -> Guard:
-        if ENV_SOT_ENABLE_GUARD_TREE.get():
-            guard_nodes: list[paddle.framework.core.GuardNode] = []
-            with EventGuard("guard_fn: find vars and make faster guard"):
+    @event_register("guard_chain")
+    def guard_chain(self) -> list[paddle.framework.core.GuardNode]:
+        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.GuardNode] = []
+
+        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_nodes.extend(variable.make_faster_guard())
-            return make_faster_guard(guard_nodes)
+                    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 []
+        return guard_chain
 
+    @property
+    @event_register("guard_fn")
+    def guard_fn(self) -> Guard:
         with switch_symbol_registry():
             guards: list[StringifiedExpression] = []
             with EventGuard("guard_fn: find vars and make stringified guard"):