[SOT][Faster Guard] adapt to faster guard for more variables

paddle/fluid/pybind/jit.cc

@@ -87,12 +87,22 @@ void BindGuard(pybind11::module *m) {
       .def(py::init<const paddle::framework::proto::VarType &>(),
            py::arg("dtype"))
       .def(py::init<const phi::DataType &>(), py::arg("dtype"));
-  py::class_<LayerMatchGuard, GuardBase, std::shared_ptr<LayerMatchGuard>>(
-      *m, "LayerMatchGuard", R"DOC(LayerMatchGuard Class.)DOC")
-      .def(py::init<const py::object &>(), py::arg("layer_obj"));
+  py::class_<AttributeMatchGuard,
+             GuardBase,
+             std::shared_ptr<AttributeMatchGuard>>(
+      *m, "AttributeMatchGuard", R"DOC(AttributeMatchGuard Class.)DOC")
+      .def(py::init<const py::object &, const std::string &>(),
+           py::arg("obj"),
+           py::arg("attr_name"));
   py::class_<ShapeMatchGuard, GuardBase, std::shared_ptr<ShapeMatchGuard>>(
       *m, "ShapeMatchGuard", R"DOC(ShapeMatchGuard Class.)DOC")
       .def(py::init<const std::vector<py::object> &>(), py::arg("shape"));
+  py::class_<LayerMatchGuard, GuardBase, std::shared_ptr<LayerMatchGuard>>(
+      *m, "LayerMatchGuard", R"DOC(LayerMatchGuard Class.)DOC")
+      .def(py::init<const py::object &>(), py::arg("layer_obj"));
+  py::class_<RangeMatchGuard, GuardBase, std::shared_ptr<RangeMatchGuard>>(
+      *m, "RangeMatchGuard", R"DOC(RangeMatchGuard Class.)DOC")
+      .def(py::init<const py::object &>(), py::arg("range_obj"));
 
   m->def(
       "merge_guard",

paddle/fluid/pybind/sot/guards.cc

@@ -26,6 +26,22 @@ static inline PyObject* PyObject_CallOneArg(PyObject* func, PyObject* arg) {
 }
 #endif
 
+static inline bool PyObject_Equal(PyObject* a, PyObject* b) {
+  if (a == b) {
+    return true;
+  }
+  if (Py_TYPE(a) != Py_TYPE(b)) {
+    return false;
+  }
+  int result = PyObject_RichCompareBool(a, b, Py_EQ);
+  // Check for exception
+  if (result == -1) {
+    PyErr_Clear();
+    return false;
+  }
+  return result;
+}
+
 std::optional<paddle::Tensor> GetTensorFromPyObject(PyObject* obj) {
   if (!paddle::pybind::PyCheckTensor(obj)) {
     // TODO(zrr1999): PyCheckTensor only check if the object is a p_tensor_type.
@@ -59,19 +75,7 @@ bool TypeMatchGuard::check(PyObject* value) {
 }
 
 bool ValueMatchGuard::check(PyObject* value) {
-  if (value == expected_value_) {
-    return true;
-  }
-  if (Py_TYPE(value) != expected_type_) {
-    return false;
-  }
-  int result = PyObject_RichCompareBool(value, expected_value_, Py_EQ);
-  // Check for exception
-  if (result == -1) {
-    PyErr_Clear();
-    return false;
-  }
-  return result;
+  return PyObject_Equal(value, expected_value_);
 }
 
 bool LengthMatchGuard::check(PyObject* value) {
@@ -104,6 +108,11 @@ bool ShapeMatchGuard::check(PyObject* value) {
   return true;
 }
 
+bool AttributeMatchGuard::check(PyObject* value) {
+  PyObject* attr = PyObject_GetAttrString(value, attr_name_.c_str());
+  return PyObject_Equal(attr, attr_ptr_);
+}
+
 bool LayerMatchGuard::check(PyObject* value) {
   if (value != layer_ptr_) {
     return false;

paddle/fluid/pybind/sot/guards.h

@@ -72,9 +72,9 @@ class GuardGroup : public GuardBase {
 
 class TypeMatchGuard : public GuardBase {
  public:
+  explicit TypeMatchGuard(PyTypeObject* type_ptr) : expected_(type_ptr) {}
   explicit TypeMatchGuard(PyObject* type_ptr)
       : expected_(reinterpret_cast<PyTypeObject*>(type_ptr)) {}
-
   explicit TypeMatchGuard(const py::type& py_type)
       : expected_(reinterpret_cast<PyTypeObject*>(py_type.ptr())) {}
 
@@ -148,6 +148,19 @@ class ShapeMatchGuard : public GuardBase {
   std::vector<std::optional<int64_t>> expected_;
 };
 
+class AttributeMatchGuard : public GuardBase {
+ public:
+  AttributeMatchGuard(const py::object& obj, const std::string& attr_name)
+      : attr_ptr_(PyObject_GetAttrString(obj.ptr(), attr_name.c_str())),
+        attr_name_(attr_name) {}
+
+  bool check(PyObject* value);
+
+ private:
+  PyObject* attr_ptr_;
+  std::string attr_name_;
+};
+
 class LayerMatchGuard : public GuardBase {
  public:
   explicit LayerMatchGuard(PyObject* layer_ptr) : layer_ptr_(layer_ptr) {
@@ -164,4 +177,14 @@ class LayerMatchGuard : public GuardBase {
   bool training_;
 };
 
+class RangeMatchGuard : public GuardGroup {
+ public:
+  explicit RangeMatchGuard(const py::object& range_obj)
+      : GuardGroup({std::make_shared<TypeMatchGuard>(Py_TYPE(range_obj.ptr())),
+                    std::make_shared<AttributeMatchGuard>(range_obj, "start"),
+                    std::make_shared<AttributeMatchGuard>(range_obj, "stop"),
+                    std::make_shared<AttributeMatchGuard>(range_obj, "step")}) {
+  }
+};
+
 #endif

python/paddle/jit/sot/opcode_translator/executor/guard.py

@@ -98,7 +98,7 @@ def __init__(
         free_vars: dict[str, Any],
     ):
         self.faster_guard = faster_guard
-        if ENV_SOT_ENABLE_FASTER_GUARD:
+        if ENV_SOT_ENABLE_FASTER_GUARD.get():
             original_expr_template = expr_template
             guard_cls_name = faster_guard.__class__.__name__
             guard_name = f"{guard_cls_name}_{id(faster_guard)}"

python/paddle/jit/sot/opcode_translator/executor/variables/base.py

@@ -23,10 +23,19 @@
 import paddle
 
 from ....profiler import event_register
-from ....utils import NameGenerator, get_unbound_method, log
+from ....utils import (
+    NameGenerator,
+    get_unbound_method,
+    log,
+)
 from ....utils.exceptions import FallbackError, HasNoAttributeError
 from ..dispatcher import Dispatcher
-from ..guard import StringifiedExpression, check_guard, union_free_vars
+from ..guard import (
+    FasterStringifiedExpression,
+    StringifiedExpression,
+    check_guard,
+    union_free_vars,
+)
 from ..mutable_data import MutableDictLikeData
 from ..tracker import (
     DummyTracker,
@@ -364,18 +373,13 @@ def make_stringified_guard(self) -> list[StringifiedExpression]:
 
         # Get a ValueTracer object from the Tracker object associated with the variable
         frame_value_tracer = self.tracker.trace_value_from_frame()
-
         return [
-            StringifiedExpression(
-                f"id(type({{}})) == {id(self.get_py_type())}",
+            FasterStringifiedExpression(
+                f"id(type({{0}})) == {id(self.get_py_type())} and {{0}} == {self.get_py_value()!r}",
+                paddle.framework.core.ValueMatchGuard(self.get_py_value()),
                 [frame_value_tracer],
                 union_free_vars(frame_value_tracer.free_vars),
-            ),
-            StringifiedExpression(
-                f"{{}} == {self.get_py_value()!r}",
-                [frame_value_tracer],
-                union_free_vars(frame_value_tracer.free_vars),
-            ),
+            )
         ]
 
     def get_py_value(self, allow_tensor=False) -> Any:

python/paddle/jit/sot/opcode_translator/executor/variables/callable.py

@@ -48,6 +48,7 @@
 )
 from ..dispatcher import Dispatcher
 from ..guard import (
+    FasterStringifiedExpression,
     StringifiedExpression,
     check_guard,
     object_equal_stringified_guard,
@@ -445,13 +446,9 @@ def call_function(self, /, *args, **kwargs):
     def make_stringified_guard(self) -> list[StringifiedExpression]:
         frame_value_tracer = self.tracker.trace_value_from_frame()
         return [
-            StringifiedExpression(
-                f"id({{}}) == {id(self.get_py_value())}",
-                [frame_value_tracer],
-                union_free_vars(frame_value_tracer.free_vars),
-            ),
-            StringifiedExpression(
-                f"{{}}.training == {self.get_py_value().training}",
+            FasterStringifiedExpression(
+                f"id({{0}}) == {id(self.get_py_value())} and {{0}}.training == {self.get_py_value().training}",
+                paddle.framework.core.ValueMatchGuard(self.get_py_value()),
                 [frame_value_tracer],
                 union_free_vars(frame_value_tracer.free_vars),
             ),
@@ -503,13 +500,14 @@ def make_stringified_guard(self) -> list[StringifiedExpression]:
         if isinstance(self.value, PD_SEQ_CONTAINERS):
             frame_value_tracer = self.tracker.trace_value_from_frame()
 
-            len_guard = StringifiedExpression(
+            len_guard = FasterStringifiedExpression(
                 f"len({{}}) == {len(self.value)}",
+                paddle.framework.core.LengthMatchGuard(len(self.value)),
                 [frame_value_tracer],
                 frame_value_tracer.free_vars,
             )
 
-            guards = [len_guard]
+            guards: list[StringifiedExpression] = [len_guard]
             for idx, layer in enumerate(self.value):
                 layer_variable = VariableFactory.from_value(
                     layer, self.graph, GetItemTracker(self, idx)

test/sot/test_01_basic.py

@@ -14,7 +14,7 @@
 
 import unittest
 
-from test_case_base import TestCaseBase
+from test_case_base import TestCaseBase, test_with_faster_guard
 
 import paddle
 
@@ -24,6 +24,7 @@ def foo(x: int, y: paddle.Tensor):
 
 
 class TestBasic(TestCaseBase):
+    @test_with_faster_guard
     def test_simple(self):
         self.assert_results(foo, 1, paddle.to_tensor(2))
 

test/sot/test_12_for_loop.py

@@ -22,6 +22,7 @@
 from test_case_base import (
     TestCaseBase,
     test_instruction_translator_cache_context,
+    test_with_faster_guard,
 )
 
 import paddle
@@ -224,6 +225,7 @@ def test_for_without_zero_iter(self):
     def test_reconstruct_range_iter(self):
         self.assert_results(for_reconstruct_range_iter)
 
+    @test_with_faster_guard
     def test_layer_list(self):
         layers = paddle.nn.LayerList()
         for i in range(5):

test/sot/test_17_paddle_layer.py

@@ -14,7 +14,7 @@
 
 import unittest
 
-from test_case_base import TestCaseBase
+from test_case_base import TestCaseBase, test_with_faster_guard
 
 import paddle
 
@@ -66,6 +66,7 @@ def forward(self, x):
 
 
 class TestLayer(TestCaseBase):
+    @test_with_faster_guard
     def test_layer(self):
         x = paddle.rand((10,))
         y = paddle.rand((10, 10))
@@ -74,6 +75,7 @@ def test_layer(self):
         self.assert_results(net_call, y, net)
         self.assert_results(net_call_passed_by_user, x, net.forward)
 
+    @test_with_faster_guard
     def test_layer_with_sequential(self):
         x = paddle.rand((10,))
         y = paddle.rand((10, 10))
@@ -82,6 +84,7 @@ def test_layer_with_sequential(self):
         self.assert_results(net_call, y, net)
         self.assert_results(net_call_passed_by_user, x, net.forward)
 
+    @test_with_faster_guard
     def test_bound(self):
         x = paddle.rand((10,))
         y = paddle.rand((10, 10))

test/sot/test_faster_guard.py

@@ -19,7 +19,7 @@
 import paddle
 
 
-class TestFasterGuard(unittest.TestCase):
+class TestBasicFasterGuard(unittest.TestCase):
     def test_lambda_guard(self):
         guard_lambda = paddle.framework.core.LambdaGuard(lambda x: x == 1)
         self.assertTrue(guard_lambda.check(1))
@@ -69,6 +69,12 @@ def test_shape_match_guard(self):
         guard_shape = paddle.framework.core.ShapeMatchGuard([2, 3, 1])
         self.assertFalse(guard_shape.check(tensor))
 
+    def test_attribute_match_guard(self):
+        a = range(1, 10, 2)
+        guard_attribute = paddle.framework.core.AttributeMatchGuard(a, "start")
+        self.assertTrue(guard_attribute.check(a))
+        self.assertFalse(guard_attribute.check(range(10)))
+
     def test_layer_match_guard(self):
         layer = paddle.nn.Linear(10, 10)
         guard_layer = paddle.framework.core.LayerMatchGuard(layer)
@@ -90,7 +96,7 @@ def test_guard_group(self):
         self.assertTrue(guard_group.check(1))
         self.assertFalse(guard_group.check(2))
 
-    def test_negated_guard_group(self):
+    def test_nested_guard_group(self):
         guard_lambda = paddle.framework.core.LambdaGuard(lambda x: x == 1)
         guard_type_match = paddle.framework.core.TypeMatchGuard(int)
         guard_group = paddle.framework.core.GuardGroup(
@@ -103,6 +109,11 @@ def test_negated_guard_group(self):
         self.assertTrue(guard_group.check(1))
         self.assertFalse(guard_group.check(2))
 
+    def test_range_match_guard(self):
+        guard_range = paddle.framework.core.RangeMatchGuard(range(1, 10, 2))
+        self.assertTrue(guard_range.check(range(1, 10, 2)))
+        self.assertFalse(guard_range.check(range(11)))
+
 
 if __name__ == "__main__":
     unittest.main()