[SOT][Faster Guard] adapt to faster guard for more variables (#69672)

zrr1999 · web-flow · commit c76d2dbaf196 · 2024-11-29T01:08:34.000+08:00
diff --git a/paddle/fluid/pybind/jit.cc b/paddle/fluid/pybind/jit.cc
@@ -76,6 +76,9 @@ void BindGuard(pybind11::module *m) {
   py::class_<TypeMatchGuard, GuardBase, std::shared_ptr<TypeMatchGuard>>(
       *m, "TypeMatchGuard", R"DOC(TypeMatchGuard Class.)DOC")
       .def(py::init<const py::type &>(), py::arg("py_type"));
+  py::class_<IdMatchGuard, GuardBase, std::shared_ptr<IdMatchGuard>>(
+      *m, "IdMatchGuard", R"DOC(IdMatchGuard Class.)DOC")
+      .def(py::init<const py::object &>(), py::arg("py_obj"));
   py::class_<LengthMatchGuard, GuardBase, std::shared_ptr<LengthMatchGuard>>(
       *m, "LengthMatchGuard", R"DOC(LengthMatchGuard Class.)DOC")
       .def(py::init<const Py_ssize_t &>(), py::arg("length"));
diff --git a/paddle/fluid/pybind/sot/guards.cc b/paddle/fluid/pybind/sot/guards.cc
@@ -74,6 +74,8 @@ bool TypeMatchGuard::check(PyObject* value) {
   return Py_TYPE(value) == expected_;
 }
 
+bool IdMatchGuard::check(PyObject* value) { return value == expected_; }
+
 bool ValueMatchGuard::check(PyObject* value) {
   return PyObject_Equal(value, expected_value_);
 }
diff --git a/paddle/fluid/pybind/sot/guards.h b/paddle/fluid/pybind/sot/guards.h
@@ -84,6 +84,19 @@ class TypeMatchGuard : public GuardBase {
   PyTypeObject* expected_;
 };
 
+class IdMatchGuard : public GuardBase {
+ public:
+  explicit IdMatchGuard(PyObject* obj_ptr)
+      : expected_(reinterpret_cast<PyObject*>(obj_ptr)) {}
+  explicit IdMatchGuard(const py::object& py_obj)
+      : expected_(reinterpret_cast<PyObject*>(py_obj.ptr())) {}
+
+  bool check(PyObject* value);
+
+ private:
+  PyObject* expected_;
+};
+
 class ValueMatchGuard : public GuardBase {
  public:
   explicit ValueMatchGuard(PyObject* value_ptr)
diff --git a/test/sot/test_faster_guard.py b/test/sot/test_faster_guard.py
@@ -85,6 +85,14 @@ def test_layer_match_guard(self):
         layer.train()
         self.assertTrue(guard_layer.check(layer))
 
+    def test_id_match_guard(self):
+        layer = paddle.nn.Linear(10, 10)
+        guard_id = paddle.framework.core.IdMatchGuard(layer)
+        self.assertTrue(guard_id.check(layer))
+        layer.eval()
+        self.assertTrue(guard_id.check(layer))
+        self.assertFalse(guard_id.check(paddle.nn.Linear(10, 10)))
+
 
 class TestFasterGuardGroup(unittest.TestCase):
     def test_guard_group(self):

Original file line number	Diff line number	Diff line change
`@@ -74,6 +74,8 @@ bool TypeMatchGuard::check(PyObject* value) {`
`74`	`74`	`return Py_TYPE(value) == expected_;`
`75`	`75`	`}`
`76`	`76`
	`77`	`+bool IdMatchGuard::check(PyObject* value) { return value == expected_; }`
	`78`	`+`
`77`	`79`	`bool ValueMatchGuard::check(PyObject* value) {`
`78`	`80`	`return PyObject_Equal(value, expected_value_);`
`79`	`81`	`}`