PaddlePaddle
diff --git a/‎paddle/fluid/pybind/jit.cc
Copy file name to clipboard
+75-29Lines changed: 75 additions & 29 deletions b/‎paddle/fluid/pybind/jit.cc
Copy file name to clipboard
+75-29Lines changed: 75 additions & 29 deletions
diff --git a/‎paddle/fluid/pybind/sot/guards.cc
Copy file name to clipboard
+175-1Lines changed: 175 additions & 1 deletion b/‎paddle/fluid/pybind/sot/guards.cc
Copy file name to clipboard
+175-1Lines changed: 175 additions & 1 deletion
@@ -142,9 +142,11 @@ 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>>> &>(),
+      .def(py::init<const std::vector<
+               std::vector<std::shared_ptr<GuardNodeBase>>> &>(),
            py::arg("guard_chain_list"))
+      .def_property_readonly(
+          "guard_nodes", [](GuardTree &self) { return self.get_guard_nodes(); })
       .def(
           "lookup",
           [](GuardTree &self, py::object frame) {
@@ -154,73 +156,117 @@ void BindGuardTree(pybind11::module *m) {
       .def(
           "add_guard_chain",
           [](GuardTree &self,
-             const std::vector<std::shared_ptr<GuardNode>> &guard_chain) {
+             const std::vector<std::shared_ptr<GuardNodeBase>> &guard_chain) {
             self.add_guard_chain(guard_chain);
           },
           py::arg("guard_chain"))
       .def("stringify", &GuardTree::stringify);
 
-  py::class_<GuardNode, std::shared_ptr<GuardNode>>(
-      *m, "GuardNode", R"DOC(GuardNode Class.)DOC")
-      .def(py::init<const std::shared_ptr<GuardBase> &,
-                    const std::vector<std::shared_ptr<ExprNode>> &,
-                    const std::vector<std::shared_ptr<GuardNode>> &,
-                    const std::optional<int> &>(),
-           py::arg("guard"),
-           py::arg("exprs"),
-           py::arg("next_guard_nodes") = py::list(),
-           py::arg("return_cache_index") = py::none())
+  py::class_<GuardNodeBase, std::shared_ptr<GuardNodeBase>>(
+      *m, "GuardNodeBase", R"DOC(GuardNodeBase Class.)DOC")
       .def_property(
           "return_cache_index",
-          [](GuardNode &self) { return self.return_cache_index; },
-          [](GuardNode &self, int index) { self.return_cache_index = index; })
+          [](GuardNodeBase &self) { return self.return_cache_index; },
+          [](GuardNodeBase &self, int index) {
+            self.return_cache_index = index;
+          })
+      .def_property(
+          "next_guard_nodes",
+          [](GuardNodeBase &self) { return self.next_guard_nodes; },
+          [](GuardNodeBase &self,
+             const std::vector<std::shared_ptr<GuardNodeBase>>
+                 &next_guard_nodes) {
+            self.next_guard_nodes = next_guard_nodes;
+          })
       .def(
           "lookup",
-          [](GuardNode &self, py::object frame) {
+          [](GuardNodeBase &self, py::object frame) {
             return self.lookup(reinterpret_cast<FrameProxy *>(frame.ptr()));
           },
           py::arg("frame"))
-      .def("stringify", &GuardNode::stringify, py::arg("indent") = 0);
+      .def("stringify", &GuardNodeBase::stringify, py::arg("indent") = 0);
 
-  py::class_<ExprNode, std::shared_ptr<ExprNode>>(
-      *m, "ExprNode", R"DOC(ExprNode Class.)DOC")
+  py::class_<ExprNodeBase, std::shared_ptr<ExprNodeBase>>(
+      *m, "ExprNodeBase", R"DOC(ExprNodeBase Class.)DOC")
       .def(
           "eval",
-          [](ExprNode &self, py::object frame) {
+          [](ExprNodeBase &self, py::object frame) {
             return self.eval(reinterpret_cast<FrameProxy *>(frame.ptr()));
           },
           py::arg("frame"))
-      .def("stringify", &ExprNode::stringify, py::arg("indent") = 0);
+      .def("stringify", &ExprNodeBase::stringify, py::arg("indent") = 0);
+
+  py::class_<GuardNode, GuardNodeBase, std::shared_ptr<GuardNode>>(
+      *m, "GuardNode", R"DOC(GuardNode Class.)DOC")
+      .def(py::init<const std::shared_ptr<GuardBase> &,
+                    const std::vector<std::shared_ptr<ExprNodeBase>> &,
+                    const std::vector<std::shared_ptr<GuardNodeBase>> &,
+                    const std::optional<int> &>(),
+           py::arg("guard"),
+           py::arg("exprs"),
+           py::arg("next_guard_nodes") = py::list(),
+           py::arg("return_cache_index") = py::none());
 
-  py::class_<ConstantExprNode, ExprNode, std::shared_ptr<ConstantExprNode>>(
+  py::class_<ExprGuardNode, GuardNodeBase, std::shared_ptr<ExprGuardNode>>(
+      *m, "ExprGuardNode", R"DOC(ExprGuardNode Class.)DOC")
+      .def(py::init<const std::shared_ptr<ExprNodeBase> &,
+                    const std::vector<std::shared_ptr<GuardNodeBase>> &,
+                    const std::optional<int> &>(),
+           py::arg("expr"),
+           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"));
 
-  py::class_<ExternVarExprNode, ExprNode, std::shared_ptr<ExternVarExprNode>>(
+  py::class_<ExternVarExprNode,
+             ExprNodeBase,
+             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>>(
+  py::class_<LocalVarExprNode, ExprNodeBase, 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>>(
+  py::class_<GlobalVarExprNode,
+             ExprNodeBase,
+             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>>(
+  py::class_<AttributeExprNode,
+             ExprNodeBase,
+             std::shared_ptr<AttributeExprNode>>(
       *m, "AttributeExprNode", R"DOC(AttributeExprNode Class.)DOC")
-      .def(py::init<std::shared_ptr<ExprNode>, const std::string &>(),
+      .def(py::init<std::shared_ptr<ExprNodeBase>, const std::string &>(),
            py::arg("var_expr"),
            py::arg("attr_name"));
 
-  py::class_<ItemExprNode, ExprNode, std::shared_ptr<ItemExprNode>>(
+  py::class_<ItemExprNode, ExprNodeBase, std::shared_ptr<ItemExprNode>>(
       *m, "ItemExprNode", R"DOC(ItemExprNode Class.)DOC")
-      .def(py::init<std::shared_ptr<ExprNode>, std::shared_ptr<ExprNode>>(),
+      .def(py::init<std::shared_ptr<ExprNodeBase>,
+                    std::shared_ptr<ExprNodeBase>>(),
            py::arg("var_expr"),
            py::arg("key_expr"));
+
+  py::class_<BinaryExprNode, ExprNodeBase, std::shared_ptr<BinaryExprNode>>(
+      *m, "BinaryExprNode", R"DOC(BinaryExprNode Class.)DOC")
+      .def(py::init<std::shared_ptr<ExprNodeBase>,
+                    std::shared_ptr<ExprNodeBase>,
+                    const std::string &>(),
+           py::arg("lhs"),
+           py::arg("rhs"),
+           py::arg("op"));
+
+  py::class_<UnaryExprNode, ExprNodeBase, std::shared_ptr<UnaryExprNode>>(
+      *m, "UnaryExprNode", R"DOC(UnaryExprNode Class.)DOC")
+      .def(py::init<std::shared_ptr<ExprNodeBase>, const std::string &>(),
+           py::arg("expr"),
+           py::arg("op"));
 #endif
 }
 
 
@@ -254,6 +254,88 @@ std::string ItemExprNode::stringify(int indent) {
   return ss.str();
 }
 
+PyObject* BinaryExprNode::eval(FrameProxy* frame) {
+  PyObject* lhs = lhs_->eval(frame);
+  PyObject* rhs = rhs_->eval(frame);
+
+  if (!lhs || !rhs) {
+    PyErr_Clear();
+    return Py_False;
+  }
+
+  PyObject* result = nullptr;
+  if (op_type_ == OpType::COMPARE) {
+    int bool_result = PyObject_RichCompareBool(lhs, rhs, op_code_);
+    if (bool_result == -1) {
+      PyErr_Clear();
+      return Py_False;
+    }
+    result = bool_result ? Py_True : Py_False;
+  } else {
+    PyNumberMethods* nb = Py_TYPE(lhs)->tp_as_number;
+    if (nb == nullptr) {
+      PyErr_SetString(PyExc_TypeError,
+                      "Object does not support number operations");
+      return Py_False;
+    }
+
+    switch (op_code_) {
+      case 0:  // +
+        result = nb->nb_add(lhs, rhs);
+        break;
+      case 1:  // -
+        result = nb->nb_subtract(lhs, rhs);
+        break;
+      case 2:  // *
+        result = nb->nb_multiply(lhs, rhs);
+        break;
+      case 3:  // /
+        result = nb->nb_true_divide(lhs, rhs);
+        break;
+      case 4:  // //
+        result = nb->nb_floor_divide(lhs, rhs);
+        break;
+      case 5:  // %
+        result = nb->nb_remainder(lhs, rhs);
+        break;
+      case 6:  // **
+        result = nb->nb_power(lhs, rhs, nullptr);
+        break;
+      case 7:  // <<
+        result = nb->nb_lshift(lhs, rhs);
+        break;
+      case 8:  // >>
+        result = nb->nb_rshift(lhs, rhs);
+        break;
+      case 9:  // &
+        result = nb->nb_and(lhs, rhs);
+        break;
+      case 10:  // |
+        result = nb->nb_or(lhs, rhs);
+        break;
+      case 11:  // ^
+        result = nb->nb_xor(lhs, rhs);
+        break;
+      default:
+        PyErr_SetString(PyExc_TypeError, "Unsupported operation");
+        return Py_False;
+    }
+
+    if (result == nullptr) {
+      PyErr_Clear();
+      return Py_False;
+    }
+  }
+
+  return result;
+}
+
+std::string BinaryExprNode::stringify(int indent) {
+  std::stringstream ss;
+  ss << lhs_->stringify() << " " << op_str_ << " " << rhs_->stringify();
+  return ss.str();
+}
+
 std::optional<int> GuardNode::lookup(FrameProxy* frame) {
   // TODO(zrr1999): support multiple exprs
   auto expr = exprs.back();
@@ -287,8 +369,38 @@ std::string GuardNode::stringify(int indent) {
   return ss.str();
 }
 
+std::optional<int> ExprGuardNode::lookup(FrameProxy* frame) {
+  auto expr = expr_;
+  auto value = expr->eval(frame);
+  if (PyObject_IsTrue(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::string ExprGuardNode::stringify(int indent) {
+  std::stringstream ss;
+  ss << std::string(indent, ' ');
+  ss << "(" << expr_->stringify() << ")";
+  if (!next_guard_nodes.empty()) {
+    ss << " |" << std::endl;
+    for (auto& next_guard_node : next_guard_nodes) {
+      ss << std::string(indent + 2, ' ');
+      ss << next_guard_node->stringify(indent + 2) << std::endl;
+    }
+  }
+  return ss.str();
+}
+
 void GuardTree::add_guard_chain(
-    const std::vector<std::shared_ptr<GuardNode>>& 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.
@@ -321,4 +433,66 @@ std::string GuardTree::stringify() {
   return ss.str();
 }
 
+std::vector<std::shared_ptr<GuardNodeBase>> GuardTree::get_guard_nodes() const {
+  return guard_nodes_;
+}
+
+PyObject* UnaryExprNode::eval(FrameProxy* frame) {
+  PyObject* value = expr_->eval(frame);
+  if (!value) {
+    PyErr_Clear();
+    return Py_False;
+  }
+
+  PyObject* result = nullptr;
+  if (op_type_ == OpType::NUMBER) {
+    PyNumberMethods* nb = Py_TYPE(value)->tp_as_number;
+    if (nb == nullptr) {
+      PyErr_SetString(PyExc_TypeError,
+                      "Object does not support number operations");
+      return Py_False;
+    }
+
+    switch (op_code_) {
+      case 0:  // +
+        result = nb->nb_positive(value);
+        break;
+      case 1:  // -
+        result = nb->nb_negative(value);
+        break;
+      case 2:  // ~
+        result = nb->nb_invert(value);
+        break;
+      default:
+        PyErr_SetString(PyExc_TypeError, "Unsupported operation");
+        return Py_False;
+    }
+  } else {  // LOGICAL
+    switch (op_code_) {
+      case 0:  // not or !
+        result = PyObject_IsTrue(value) ? Py_False : Py_True;
+        break;
+      case 1:  // bool
+        result = PyObject_IsTrue(value) ? Py_True : Py_False;
+        break;
+      default:
+        PyErr_SetString(PyExc_TypeError, "Unsupported operation");
+        return Py_False;
+    }
+  }
+
+  if (result == nullptr) {
+    PyErr_Clear();
+    return Py_False;
+  }
+
+  return result;
+}
+
+std::string UnaryExprNode::stringify(int indent) {
+  std::stringstream ss;
+  ss << op_str_ << "(" << expr_->stringify() << ")";
+  return ss.str();
+}
+
 #endif