[SOT][Faster Guard] Implement DummyGuardNode (#72503)

zrr1999 · web-flow · commit 38af65b85340 · 2025-04-27T00:51:15.000+08:00
diff --git a/paddle/fluid/pybind/jit.cc b/paddle/fluid/pybind/jit.cc
@@ -223,6 +223,15 @@ void BindGuardTree(pybind11::module *m) {
            py::arg("next_guard_nodes") = py::list(),
            py::arg("return_cache_index") = py::none());
 
+  py::class_<DummyGuardNode, GuardNodeBase, std::shared_ptr<DummyGuardNode>>(
+      *m, "DummyGuardNode", R"DOC(DummyGuardNode Class.)DOC")
+      .def(py::init<bool,
+                    const std::vector<std::shared_ptr<GuardNodeBase>> &,
+                    const std::optional<int> &>(),
+           py::arg("return_true") = true,
+           py::arg("next_guard_nodes") = py::list(),
+           py::arg("return_cache_index") = py::none());
+
   py::class_<ConstantExprNode, ExprNodeBase, std::shared_ptr<ConstantExprNode>>(
       *m, "ConstantExprNode", R"DOC(ConstantExprNode Class.)DOC")
       .def(py::init<const py::object &>(), py::arg("value_ptr"));
diff --git a/paddle/fluid/pybind/sot/guards.cc b/paddle/fluid/pybind/sot/guards.cc
@@ -433,6 +433,8 @@ std::optional<int> GuardNode::lookup(FrameProxy* frame) {
   }();
 
   if (guard->check(value)) {
+    // TODO(zrr1999): To extract the reusable code, we need to add a new method
+    // to GuardNodeBase<N>
     if (return_cache_index.has_value()) {
       Py_DECREF(value);
       return return_cache_index.value();
@@ -467,6 +469,8 @@ std::optional<int> ExprGuardNode::lookup(FrameProxy* frame) {
   auto expr = expr_;
   auto value = expr->eval(frame);
   if (PyObject_IsTrue(value)) {
+    // TODO(zrr1999): To extract the reusable code, we need to add a new method
+    // to GuardNodeBase<N>
     if (return_cache_index.has_value()) {
       return return_cache_index.value();
     }
@@ -493,6 +497,27 @@ std::string ExprGuardNode::stringify(int indent) {
   return ss.str();
 }
 
+std::optional<int> DummyGuardNode::lookup(FrameProxy* frame) {
+  if (return_true_) {
+    // TODO(zrr1999): To extract the reusable code, we need to add a new method
+    // to GuardNodeBase
+    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::string DummyGuardNode::stringify(int indent) {
+  return std::string(indent, ' ') + "DummyGuard(" +
+         (return_true_ ? "True" : "False") + ")";
+}
+
 void GuardTree::add_guard_chain(
     const std::vector<std::shared_ptr<GuardNodeBase>>& guard_chain) {
   if (guard_chain.empty()) {
diff --git a/paddle/fluid/pybind/sot/guards.h b/paddle/fluid/pybind/sot/guards.h
@@ -323,19 +323,6 @@ class ExprNodeBase : public GuardTreeNodeBase,
   virtual ~ExprNodeBase() = default;
 };
 
-class GuardNodeBase : public GuardTreeNodeBase {
- public:
-  std::vector<std::shared_ptr<GuardNodeBase>> next_guard_nodes;
-  // return_cache_index is used to record the index of the guard list
-  std::optional<int> return_cache_index;
-  GuardNodeBase(std::vector<std::shared_ptr<GuardNodeBase>> next_guard_nodes,
-                std::optional<int> return_cache_index)
-      : next_guard_nodes(next_guard_nodes),
-        return_cache_index(return_cache_index) {}
-  virtual ~GuardNodeBase() = default;
-  virtual std::optional<int> lookup(FrameProxy* frame) = 0;
-};
-
 class ConstantExprNode : public ExprNodeBase {
  public:
   explicit ConstantExprNode(PyObject* value_ptr) : value_ptr_(value_ptr) {}
@@ -511,6 +498,19 @@ class UnaryExprNode : public ExprNodeBase {
   int op_code_;
 };
 
+class GuardNodeBase : public GuardTreeNodeBase {
+ public:
+  std::vector<std::shared_ptr<GuardNodeBase>> next_guard_nodes;
+  // return_cache_index is used to record the index of the guard list
+  std::optional<int> return_cache_index;
+  GuardNodeBase(std::vector<std::shared_ptr<GuardNodeBase>> next_guard_nodes,
+                std::optional<int> return_cache_index)
+      : next_guard_nodes(next_guard_nodes),
+        return_cache_index(return_cache_index) {}
+  virtual ~GuardNodeBase() = default;
+  virtual std::optional<int> lookup(FrameProxy* frame) = 0;
+};
+
 class ExprGuardNode : public GuardNodeBase {
  public:
   explicit ExprGuardNode(
@@ -543,6 +543,22 @@ class GuardNode : public GuardNodeBase {
   std::optional<int> lookup(FrameProxy* frame) override;
 };
 
+class DummyGuardNode : public GuardNodeBase {
+ public:
+  explicit DummyGuardNode(
+      bool return_true,
+      std::vector<std::shared_ptr<GuardNodeBase>> next_guard_nodes,
+      std::optional<int> return_cache_index)
+      : GuardNodeBase(next_guard_nodes, return_cache_index),
+        return_true_(return_true) {}
+  virtual ~DummyGuardNode() = default;
+  std::string stringify(int indent = 0) override;
+  std::optional<int> lookup(FrameProxy* frame) override;
+
+ private:
+  bool return_true_;
+};
+
 class GuardTree {
  public:
   GuardTree(const std::vector<std::vector<std::shared_ptr<GuardNodeBase>>>&
diff --git a/python/paddle/jit/sot/opcode_translator/executor/executor_cache.py b/python/paddle/jit/sot/opcode_translator/executor/executor_cache.py
@@ -375,13 +375,7 @@ def start_translate(
             )
         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)],
-                )
-            ]
+            guard_chain: GuardChain = [paddle.framework.core.DummyGuardNode()]
         return new_custom_code, guard_fn, guard_chain
     # TODO(0x45f): handle BreakGraphError to trigger fallback
     except BreakGraphError as e:
@@ -400,13 +394,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: GuardChain = [
-            # TODO(zrr1999): GuardNode should support zero-expr constructor
-            paddle.framework.core.GuardNode(
-                paddle.framework.core.DummyGuard(),
-                [paddle.framework.core.ConstantExprNode(True)],
-            )
-        ]
+        dummy_guard_chain: GuardChain = [paddle.framework.core.DummyGuardNode()]
         guard, guard_chain = dummy_guard, dummy_guard_chain
 
         if isinstance(e, ConditionalFallbackError):
diff --git a/python/paddle/jit/sot/opcode_translator/executor/tracker.py b/python/paddle/jit/sot/opcode_translator/executor/tracker.py
@@ -460,7 +460,7 @@ def gen_instructions(self, codegen: PyCodeGen):
             codegen.gen_call_function_ex(has_kwargs=True)
 
     def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
-        # TODO(zrr1999): implement LayerExprNode
+        # TODO(zrr1999): implement LayerExprNode.guard_tree_expr_node
         raise NotImplementedError("LayerExprNode is not implemented")
 
     def trace_value_from_frame(self):
diff --git a/test/sot/test_instruction_translator_cache.py b/test/sot/test_instruction_translator_cache.py
@@ -102,35 +102,17 @@ def mock_start_translate(frame: FrameType, **kwargs):
         FRAME_1: (
             CustomCode(FRAME_2.f_code, False),
             GuardCode(True),
-            # TODO(zrr1999): GuardNode should support zero-expr constructor
-            [
-                paddle.framework.core.GuardNode(
-                    paddle.framework.core.DummyGuard(),
-                    [paddle.framework.core.ConstantExprNode(True)],
-                )
-            ],
+            [paddle.framework.core.DummyGuardNode()],
         ),
         FRAME_3: (
             CustomCode(FRAME_4.f_code, False),
             GuardCode(False),
-            # TODO(zrr1999): GuardNode should support zero-expr constructor
-            [
-                paddle.framework.core.GuardNode(
-                    paddle.framework.core.DummyGuard(),
-                    [paddle.framework.core.ConstantExprNode(False)],
-                )
-            ],
+            [paddle.framework.core.DummyGuardNode(False)],
         ),  # Always re-compile
         FRAME_5: (
             CustomCode(None, False),
             lambda frame: True,
-            # TODO(zrr1999): GuardNode should support zero-expr constructor
-            [
-                paddle.framework.core.GuardNode(
-                    paddle.framework.core.DummyGuard(),
-                    [paddle.framework.core.ConstantExprNode(True)],
-                )
-            ],
+            [paddle.framework.core.DummyGuardNode()],
         ),
     }
     return translate_map[frame]
