PaddlePaddle
diff --git a/‎paddle/fluid/pybind/jit.cc
Copy file name to clipboard
+73Lines changed: 73 additions & 0 deletions b/‎paddle/fluid/pybind/jit.cc
Copy file name to clipboard
+73Lines changed: 73 additions & 0 deletions
diff --git a/‎paddle/fluid/pybind/sot/frame_proxy.h
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions b/‎paddle/fluid/pybind/sot/frame_proxy.h
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions
diff --git a/‎paddle/fluid/pybind/sot/guards.cc
Copy file name to clipboard
+56Lines changed: 56 additions & 0 deletions b/‎paddle/fluid/pybind/sot/guards.cc
Copy file name to clipboard
+56Lines changed: 56 additions & 0 deletions
diff --git a/‎paddle/fluid/pybind/sot/guards.h
Copy file name to clipboard
+111Lines changed: 111 additions & 0 deletions b/‎paddle/fluid/pybind/sot/guards.h
Copy file name to clipboard
+111Lines changed: 111 additions & 0 deletions
diff --git a/‎python/paddle/jit/sot/opcode_translator/executor/executor_cache.py
Copy file name to clipboard
+3-1Lines changed: 3 additions & 1 deletion b/‎python/paddle/jit/sot/opcode_translator/executor/executor_cache.py
Copy file name to clipboard
+3-1Lines changed: 3 additions & 1 deletion
diff --git a/‎python/paddle/jit/sot/opcode_translator/executor/function_graph.py
Copy file name to clipboard
+11-1Lines changed: 11 additions & 1 deletion b/‎python/paddle/jit/sot/opcode_translator/executor/function_graph.py
Copy file name to clipboard
+11-1Lines changed: 11 additions & 1 deletion
@@ -28,6 +28,8 @@ limitations under the License. */
 #include "paddle/phi/common/data_type.h"
 #include "paddle/phi/common/place.h"
 #include "paddle/utils/pybind.h"
+#include "pybind11/pybind11.h"
+#include "pybind11/pytypes.h"
 
 namespace py = pybind11;
 
@@ -123,6 +125,76 @@ void BindGuard(pybind11::module *m) {
 #endif
 }
 
+void BindGuardTree(pybind11::module *m) {
+#if SOT_IS_SUPPORTED
+  py::class_<GuardTree, std::shared_ptr<GuardTree>>(
+      *m, "GuardTree", R"DOC(GuardTree Class.)DOC")
+      .def(py::init<
+               const std::vector<std::vector<std::shared_ptr<GuardNode>>> &>(),
+           py::arg("guard_nodes_list"))
+      .def(
+          "lookup",
+          [](GuardTree &self, py::object frame) {
+            return self.lookup(reinterpret_cast<FrameProxy *>(frame.ptr()));
+          },
+          py::arg("frame"));
+
+  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<GuardNode>> &,
+                    const std::optional<int> &>(),
+           py::arg("guard"),
+           py::arg("expr"),
+           py::arg("next_guard_nodes") = py::list(),
+           py::arg("return_cache_index") = py::none())
+      .def_property(
+          "return_cache_index",
+          [](GuardNode &self) { return self.return_cache_index; },
+          [](GuardNode &self, int index) { self.return_cache_index = index; })
+      .def(
+          "lookup",
+          [](GuardNode &self, py::object frame) {
+            return self.lookup(reinterpret_cast<FrameProxy *>(frame.ptr()));
+          },
+          py::arg("frame"));
+
+  py::class_<ExprNode, std::shared_ptr<ExprNode>>(
+      *m, "ExprNode", R"DOC(ExprNode Class.)DOC")
+      .def(
+          "eval",
+          [](ExprNode &self, py::object frame) {
+            return self.eval(reinterpret_cast<FrameProxy *>(frame.ptr()));
+          },
+          py::arg("frame"));
+
+  py::class_<ConstantExprNode, ExprNode, std::shared_ptr<ConstantExprNode>>(
+      *m, "ConstantExprNode", R"DOC(ConstantExprNode Class.)DOC")
+      .def(py::init<const py::object &>(), 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"));
+
+  py::class_<GlobalVarExprNode, ExprNode, std::shared_ptr<GlobalVarExprNode>>(
+      *m, "GlobalVarExprNode", R"DOC(GlobalVarExprNode Class.)DOC")
+      .def(py::init<const std::string &>(), py::arg("var_name"));
+
+  py::class_<AttributeExprNode, ExprNode, std::shared_ptr<AttributeExprNode>>(
+      *m, "AttributeExprNode", R"DOC(AttributeExprNode Class.)DOC")
+      .def(py::init<std::shared_ptr<ExprNode>, const std::string &>(),
+           py::arg("var_expr"),
+           py::arg("attr_name"));
+
+  py::class_<ItemExprNode, ExprNode, std::shared_ptr<ItemExprNode>>(
+      *m, "ItemExprNode", R"DOC(ItemExprNode Class.)DOC")
+      .def(py::init<std::shared_ptr<ExprNode>, std::shared_ptr<ExprNode>>(),
+           py::arg("var_expr"),
+           py::arg("key_expr"));
+#endif
+}
+
 void BindSot(pybind11::module *m) {
 #if SOT_IS_SUPPORTED
   PyInit__eval_frame();
@@ -188,6 +260,7 @@ void BindSot(pybind11::module *m) {
       },
       py::arg("py_codes"));
   BindGuard(m);
+  BindGuardTree(m);
 #endif
 }
 
 
@@ -51,13 +51,13 @@ typedef struct PyInterpreterFrameProxy {
 PyInterpreterFrameProxy *PyInterpreterFrameProxy_New(
     _PyInterpreterFrame *frame);
 PyMODINIT_FUNC PyInit__frame_proxy();
-
+typedef PyInterpreterFrameProxy FrameProxy;
 #else
 typedef PyFrameObject FrameObject;
+typedef PyFrameObject FrameProxy;
 #endif
 
 #endif
-
 #ifdef __cplusplus
 }
 #endif
@@ -13,12 +13,14 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/fluid/pybind/sot/guards.h"
+#include <optional>
 #include "paddle/phi/api/include/tensor.h"
 
 #if SOT_IS_SUPPORTED
 
 #include <Python.h>
 #include <frameobject.h>
+#include <object.h>
 #include "pybind11/numpy.h"
 
 #if !defined(PyObject_CallOneArg) && !PY_3_9_PLUS
@@ -150,4 +152,58 @@ bool NumpyDtypeMatchGuard::check(PyObject* value) {
   return expected_dtype.equal(py::handle(value).get_type());
 }
 
+PyObject* ConstantExprNode::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());
+#elif PY_3_11_PLUS
+  return PyDict_GetItemString(frame->frame->f_locals, var_name_.c_str());
+#else
+  return PyDict_GetItemString(frame->f_locals, var_name_.c_str());
+#endif
+}
+PyObject* GlobalVarExprNode::eval(FrameProxy* frame) {
+#if PY_3_11_PLUS
+  return PyDict_GetItemString(frame->frame->f_globals, var_name_.c_str());
+#else
+  return PyDict_GetItemString(frame->f_globals, var_name_.c_str());
+#endif
+}
+PyObject* AttributeExprNode::eval(FrameProxy* frame) {
+  PyObject* var = var_expr_->eval(frame);
+  return PyObject_GetAttrString(var, attr_name_.c_str());
+}
+PyObject* ItemExprNode::eval(FrameProxy* frame) {
+  PyObject* var = var_expr_->eval(frame);
+  PyObject* key = key_expr_->eval(frame);
+  return PyObject_GetItem(var, key);
+}
+
+std::optional<int> GuardNode::lookup(FrameProxy* frame) {
+  auto value = expr->eval(frame);
+  if (guard->check(value)) {
+    if (return_cache_index.has_value()) {
+      return return_cache_index.value();
+    }
+    for (auto& next_guard_node : next_guard_nodes) {
+      auto ret = next_guard_node->lookup(frame);
+      if (ret.has_value()) {
+        return ret.value();
+      }
+    }
+  }
+  return std::nullopt;
+}
+
+std::optional<int> GuardTree::lookup(FrameProxy* frame) {
+  for (auto& guard_node : guard_nodes_) {
+    auto ret = guard_node->lookup(frame);
+    if (ret.has_value()) {
+      return ret.value();
+    }
+  }
+  return std::nullopt;
+}
+
 #endif
@@ -14,7 +14,10 @@ limitations under the License. */
 #pragma once
 
 #include <Python.h>
+#include <memory>
 #include "paddle/fluid/framework/data_type.h"
+#include "paddle/fluid/pybind/sot/eval_frame_tools.h"
+#include "paddle/fluid/pybind/sot/frame_proxy.h"
 #include "paddle/fluid/pybind/sot/macros.h"
 #include "paddle/phi/core/utils/data_type.h"
 #include "paddle/utils/pybind.h"
@@ -217,4 +220,112 @@ class NumpyDtypeMatchGuard : public GuardBase {
   PyObject* expected_;
 };
 
+class GuardTreeNode {};
+
+class AttributeExprNode;
+class ItemExprNode;
+class ExprNode : public GuardTreeNode,
+                 public std::enable_shared_from_this<ExprNode> {
+ public:
+  virtual PyObject* eval(FrameProxy* frame) = 0;
+};
+class ConstantExprNode : public ExprNode {
+ public:
+  explicit ConstantExprNode(PyObject* value_ptr) : value_ptr_(value_ptr) {}
+  explicit ConstantExprNode(const py::object& value_obj)
+      : value_ptr_(value_obj.ptr()) {
+    Py_INCREF(value_ptr_);
+  }
+  ~ConstantExprNode() { Py_DECREF(value_ptr_); }
+  PyObject* eval(FrameProxy* frame);
+
+ private:
+  PyObject* value_ptr_;
+};
+
+class LocalVarExprNode : public ExprNode {
+ public:
+  explicit LocalVarExprNode(const std::string& var_name)
+      : var_name_(var_name) {}
+
+  PyObject* eval(FrameProxy* frame);
+
+ private:
+  std::string var_name_;
+};
+class GlobalVarExprNode : public ExprNode {
+ public:
+  explicit GlobalVarExprNode(const std::string& var_name)
+      : var_name_(var_name) {}
+
+  PyObject* eval(FrameProxy* frame);
+
+ private:
+  std::string var_name_;
+};
+class AttributeExprNode : public ExprNode {
+ public:
+  explicit AttributeExprNode(std::shared_ptr<ExprNode> var_expr,
+                             const std::string& attr_name)
+      : var_expr_(var_expr), attr_name_(attr_name) {}
+
+  PyObject* eval(FrameProxy* frame);
+
+ private:
+  std::shared_ptr<ExprNode> var_expr_;
+  std::string attr_name_;
+};
+class ItemExprNode : public ExprNode {
+ public:
+  explicit ItemExprNode(std::shared_ptr<ExprNode> var_expr,
+                        std::shared_ptr<ExprNode> key_expr)
+      : var_expr_(var_expr), key_expr_(key_expr) {}
+
+  PyObject* eval(FrameProxy* frame);
+
+ private:
+  std::shared_ptr<ExprNode> var_expr_;
+  std::shared_ptr<ExprNode> key_expr_;
+};
+
+class GuardNode : public GuardTreeNode {
+ public:
+  std::shared_ptr<GuardBase> guard;
+  std::shared_ptr<ExprNode> expr;
+  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<GuardNode>> next_guard_nodes,
+            std::optional<int> return_cache_index)
+      : guard(guard),
+        expr(expr),
+        next_guard_nodes(next_guard_nodes),
+        return_cache_index(return_cache_index) {}
+
+  std::optional<int> lookup(FrameProxy* frame);
+};
+
+class GuardTree {
+ public:
+  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());
+    }
+  }
+
+  std::optional<int> lookup(FrameProxy* frame);
+
+ private:
+  std::vector<std::shared_ptr<GuardNode>> guard_nodes_;
+};
+
+std::string guard_tree_to_str(const GuardTree& guard_tree);
 #endif
@@ -24,6 +24,7 @@
 from ...psdb import NO_FALLBACK_CODES
 from ...utils import (
     ENV_SOT_ALLOW_DYNAMIC_SHAPE,
+    ENV_SOT_ENABLE_GUARD_TREE,
     ENV_SOT_ENABLE_STRICT_GUARD_CHECK,
     BreakGraphError,
     CompileCountInfo,
@@ -157,7 +158,8 @@ def lookup(
                         f"[Cache] Cache hit, Guard is \n{getattr(guard_fn, 'expr', 'None')}\n",
                     )
                     return custom_code
-                else:
+                elif not ENV_SOT_ENABLE_GUARD_TREE.get():
+                    # TODO(zrr1999): remove condition after faster guard tree support error analysis
                     log_do(
                         4,
                         self.analyse_guard_global_object(guard_fn),
 
@@ -41,6 +41,7 @@
 from ...symbolic.symbolic_context import SymbolicTraceContext
 from ...utils import (
     ENV_SOT_ALLOW_DYNAMIC_SHAPE,
+    ENV_SOT_ENABLE_GUARD_TREE,
     NameGenerator,
     SotUndefinedVar,
     inner_error_default_handler,
@@ -58,7 +59,7 @@
     SotExtraInfo,
 )
 from ..instruction_utils import get_instructions
-from .guard import Guard, StringifiedExpression, make_guard
+from .guard import Guard, StringifiedExpression, make_faster_guard, make_guard
 from .mutable_data import MutationDel, MutationNew, MutationSet
 from .pycode_generator import PyCodeGen
 from .side_effects import (
@@ -316,6 +317,15 @@ 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"):
+                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)
+
         with switch_symbol_registry():
             guards: list[StringifiedExpression] = []
             with EventGuard("guard_fn: find vars and make stringified guard"):