PaddlePaddle
diff --git a/‎paddle/fluid/pybind/jit.cc
Copy file name to clipboard
+47-25Lines changed: 47 additions & 25 deletions b/‎paddle/fluid/pybind/jit.cc
Copy file name to clipboard
+47-25Lines changed: 47 additions & 25 deletions
diff --git a/‎paddle/fluid/pybind/sot/guards.cc
Copy file name to clipboard
+68-129Lines changed: 68 additions & 129 deletions b/‎paddle/fluid/pybind/sot/guards.cc
Copy file name to clipboard
+68-129Lines changed: 68 additions & 129 deletions
@@ -125,16 +125,6 @@ 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<>());
-  py::class_<TensorDistMetaMatchGuard,
-             GuardBase,
-             std::shared_ptr<TensorDistMetaMatchGuard>>(
-      *m,
-      "TensorDistMetaMatchGuard",
-      R"DOC(TensorDistMetaMatchGuard Class.)DOC")
-      .def(py::init<const py::object &>(), py::arg("dist_info"));
 
   m->def(
       "merge_guard",
@@ -190,28 +180,50 @@ void BindGuardTree(pybind11::module *m) {
           [](GuardNodeBase &self, py::object frame) {
             return self.lookup(reinterpret_cast<FrameProxy *>(frame.ptr()));
           },
-          py::arg("frame"))
-      .def("stringify", &GuardNodeBase::stringify, py::arg("indent") = 0);
-
-  py::class_<ExprNodeBase, std::shared_ptr<ExprNodeBase>>(
-      *m, "ExprNodeBase", R"DOC(ExprNodeBase Class.)DOC")
-      .def(
-          "eval",
-          [](ExprNodeBase &self, py::object frame) {
-            return self.eval(reinterpret_cast<FrameProxy *>(frame.ptr()));
-          },
-          py::arg("frame"))
-      .def("stringify", &ExprNodeBase::stringify, py::arg("indent") = 0);
-
-  py::class_<GuardNode, GuardNodeBase, std::shared_ptr<GuardNode>>(
+          py::arg("frame"));
+
+  py::class_<CheckGuardNode<1>,
+             GuardNodeBase,
+             std::shared_ptr<CheckGuardNode<1>>>(*m, "CheckGuardNode1")
+      .def_property_readonly("exprs",
+                             [](CheckGuardNode<1> &self) { return self.exprs; })
+      .def("get_guard_name", &CheckGuardNode<1>::get_guard_name);
+
+  py::class_<CheckGuardNode<2>,
+             GuardNodeBase,
+             std::shared_ptr<CheckGuardNode<2>>>(*m, "CheckGuardNode2")
+      .def_property_readonly("exprs",
+                             [](CheckGuardNode<2> &self) { return self.exprs; })
+      .def("get_guard_name", &CheckGuardNode<2>::get_guard_name);
+
+  py::class_<LegacyGuardNode,
+             CheckGuardNode<1>,
+             std::shared_ptr<LegacyGuardNode>>(
       *m, "GuardNode", R"DOC(GuardNode Class.)DOC")
       .def(py::init<const std::shared_ptr<GuardBase> &,
-                    const std::vector<std::shared_ptr<ExprNodeBase>> &,
+                    const std::array<std::shared_ptr<ExprNodeBase>, 1> &,
                     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())
+      .def_property_readonly("guard",
+                             [](LegacyGuardNode &self) { return self.guard; });
+
+  py::class_<TensorDistMetaMatchGuardNode,
+             CheckGuardNode<2>,
+             std::shared_ptr<TensorDistMetaMatchGuardNode>>(
+      *m,
+      "TensorDistMetaMatchGuardNode",
+      R"DOC(TensorDistMetaMatchGuardNode Class.)DOC")
+      .def(py::init<const py::object &,
+                    const std::array<std::shared_ptr<ExprNodeBase>, 2> &,
+                    const std::vector<std::shared_ptr<GuardNodeBase>> &,
+                    const std::optional<int> &>(),
+           py::arg("dist_info"),
+           py::arg("exprs"),
+           py::arg("next_guard_nodes") = py::list(),
            py::arg("return_cache_index") = py::none());
 
   py::class_<ExprGuardNode, GuardNodeBase, std::shared_ptr<ExprGuardNode>>(
@@ -232,6 +244,16 @@ void BindGuardTree(pybind11::module *m) {
            py::arg("next_guard_nodes") = py::list(),
            py::arg("return_cache_index") = py::none());
 
+  py::class_<ExprNodeBase, std::shared_ptr<ExprNodeBase>>(
+      *m, "ExprNodeBase", R"DOC(ExprNodeBase Class.)DOC")
+      .def(
+          "eval",
+          [](ExprNodeBase &self, py::object frame) {
+            return self.eval(reinterpret_cast<FrameProxy *>(frame.ptr()));
+          },
+          py::arg("frame"))
+      .def("stringify", &ExprNodeBase::stringify, py::arg("indent") = 0);
+
   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"));
 
@@ -218,60 +218,6 @@ bool WeakRefMatchGuard::check(PyObject* value) {
 #endif
 }
 
-bool TensorDistMetaMatchGuard::check(PyObject* value) {
-  HANDLE_NULL_VALUE(value);
-
-  PyObject* expr = PyTuple_GetItem(value, 0);
-  HANDLE_NULL_VALUE(expr);
-
-  auto tensor = GetTensorFromPyObject(expr);
-  HANDLE_NULL_TENSOR(tensor);
-
-  if (tensor->is_dist_tensor() == false && is_dist_ == false) return true;
-  if (tensor->is_dist_tensor() != is_dist_) {
-    return false;
-  }
-
-  PyObject* dist_info_from_tensor_func = PyTuple_GetItem(value, 1);
-  HANDLE_NULL_VALUE(dist_info_from_tensor_func);
-
-  PyObject* dist_info = PyObject_CallOneArg(dist_info_from_tensor_func, expr);
-  HANDLE_NULL_VALUE_DECREF(dist_info);
-
-  PyObject* mesh = PyObject_GetAttrString(dist_info, "mesh");
-  HANDLE_NULL_VALUE_DECREF(mesh);
-
-  PyObject* mesh_shape = PyObject_GetAttrString(mesh, "shape");
-  HANDLE_NULL_VALUE_DECREF(mesh_shape);
-  PyObject* process_ids = PyObject_GetAttrString(mesh, "process_ids");
-  HANDLE_NULL_VALUE_DECREF(process_ids);
-  PyObject* dims_mapping = PyObject_GetAttrString(dist_info, "dims_mapping");
-  HANDLE_NULL_VALUE_DECREF(dims_mapping);
-  PyObject* local_shape = PyObject_GetAttrString(dist_info, "local_shape");
-  HANDLE_NULL_VALUE_DECREF(local_shape);
-
-  if (py::handle(mesh_shape).cast<std::vector<int>>() != mesh_shape_expected_ ||
-      py::handle(process_ids).cast<std::vector<int>>() !=
-          mesh_process_ids_expected_.value() ||
-      !PyObject_Equal(dims_mapping, dims_mapping_expected_.value()) ||
-      !PyObject_Equal(local_shape, local_shape_expected_.value())) {
-    Py_DECREF(mesh);
-    Py_DECREF(mesh_shape);
-    Py_DECREF(process_ids);
-    Py_DECREF(dims_mapping);
-    Py_DECREF(local_shape);
-    PyErr_Clear();
-    return false;
-  }
-
-  Py_DECREF(mesh);
-  Py_DECREF(mesh_shape);
-  Py_DECREF(process_ids);
-  Py_DECREF(dims_mapping);
-  Py_DECREF(local_shape);
-  return true;
-}
-
 PyObject* ConstantExprNode::eval(FrameProxy* frame) { return value_ptr_; }
 std::string ConstantExprNode::stringify(int indent) {
   return py::str(value_ptr_);
@@ -407,86 +353,89 @@ std::string BinaryExprNode::stringify(int indent) {
   return ss.str();
 }
 
-std::optional<int> GuardNode::lookup(FrameProxy* frame) {
-  // TODO(zrr1999): support multiple exprs
-  PyObject* value = [this, frame]() {
-    if (exprs.size() == 1) {
-      PyObject* v = exprs.back()->eval(frame);
-      if (v) {
-        // TODO(dev): DECREF v.
-        Py_INCREF(v);
-      }
-      return v;
-    }
-    auto values = std::vector<PyObject*>(exprs.size());
-    for (size_t i = 0; i < exprs.size(); ++i) {
-      values[i] = exprs[i]->eval(frame);
-      if (values[i]) {
-        Py_INCREF(values[i]);
-      }
-    }
-    auto packed_value = PyTuple_New(exprs.size());
-    for (size_t i = 0; i < exprs.size(); ++i) {
-      PyTuple_SetItem(packed_value, i, values[i]);
-    }
-    return packed_value;
-  }();
-
-  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();
-    }
-    for (auto& next_guard_node : next_guard_nodes) {
-      auto ret = next_guard_node->lookup(frame);
-      if (ret.has_value()) {
-        Py_DECREF(value);
-        return ret.value();
-      }
+std::optional<int> GuardNodeBase::lookup_next(FrameProxy* frame) {
+  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 GuardNode::stringify(int indent) {
-  std::stringstream ss;
-  // TODO(zrr1999): support multiple exprs
-  auto expr = exprs.back();
-  ss << std::string(indent, ' ') << guard->get_guard_name();
-  ss << "(" << exprs.back()->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;
-    }
+
+bool TensorDistMetaMatchGuardNode::check(std::array<PyObject*, 2> values) {
+  PyObject* expr = values[0];
+  HANDLE_NULL_VALUE(expr);
+
+  auto tensor = GetTensorFromPyObject(expr);
+  HANDLE_NULL_TENSOR(tensor);
+
+  if (tensor->is_dist_tensor() == false && is_dist_ == false) return true;
+  if (tensor->is_dist_tensor() != is_dist_) {
+    return false;
   }
-  return ss.str();
+
+  PyObject* dist_info_from_tensor_func = values[1];
+  HANDLE_NULL_VALUE(dist_info_from_tensor_func);
+
+  PyObject* dist_info = PyObject_CallOneArg(dist_info_from_tensor_func, expr);
+  HANDLE_NULL_VALUE_DECREF(dist_info);
+
+  PyObject* mesh = PyObject_GetAttrString(dist_info, "mesh");
+  HANDLE_NULL_VALUE_DECREF(mesh);
+
+  PyObject* mesh_shape = PyObject_GetAttrString(mesh, "shape");
+  HANDLE_NULL_VALUE_DECREF(mesh_shape);
+  PyObject* process_ids = PyObject_GetAttrString(mesh, "process_ids");
+  HANDLE_NULL_VALUE_DECREF(process_ids);
+  PyObject* dims_mapping = PyObject_GetAttrString(dist_info, "dims_mapping");
+  HANDLE_NULL_VALUE_DECREF(dims_mapping);
+  PyObject* local_shape = PyObject_GetAttrString(dist_info, "local_shape");
+  HANDLE_NULL_VALUE_DECREF(local_shape);
+
+  if (py::handle(mesh_shape).cast<std::vector<int>>() != mesh_shape_expected_ ||
+      py::handle(process_ids).cast<std::vector<int>>() !=
+          mesh_process_ids_expected_.value() ||
+      !PyObject_Equal(dims_mapping, dims_mapping_expected_.value()) ||
+      !PyObject_Equal(local_shape, local_shape_expected_.value())) {
+    Py_DECREF(mesh);
+    Py_DECREF(mesh_shape);
+    Py_DECREF(process_ids);
+    Py_DECREF(dims_mapping);
+    Py_DECREF(local_shape);
+    PyErr_Clear();
+    return false;
+  }
+
+  Py_DECREF(mesh);
+  Py_DECREF(mesh_shape);
+  Py_DECREF(process_ids);
+  Py_DECREF(dims_mapping);
+  Py_DECREF(local_shape);
+  return true;
+}
+
+bool LegacyGuardNode::check(std::array<PyObject*, 1> values) {
+  // TODO(zrr1999): support multiple exprs
+  PyObject* value = values[0];
+  HANDLE_NULL_VALUE(value);
+  return guard->check(value);
 }
 
 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();
-    }
-    for (auto& next_guard_node : next_guard_nodes) {
-      auto ret = next_guard_node->lookup(frame);
-      if (ret.has_value()) {
-        return ret.value();
-      }
-    }
+    return lookup_next(frame);
   }
   return std::nullopt;
 }
 std::string ExprGuardNode::stringify(int indent) {
   std::stringstream ss;
   ss << std::string(indent, ' ');
-  ss << "(" << expr_->stringify() << ")";
+  ss << "(" << expr->stringify() << ")";
   if (!next_guard_nodes.empty()) {
     ss << " |" << std::endl;
     for (auto& next_guard_node : next_guard_nodes) {
@@ -499,17 +448,7 @@ std::string ExprGuardNode::stringify(int indent) {
 
 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 lookup_next(frame);
   }
   return std::nullopt;
 }