[SOT][Faster Guard] support NumpyArrayVariable (#71979)

gouzil · web-flow · commit 00dac2cdbf7c · 2025-04-01T14:29:17.000+08:00
diff --git a/paddle/fluid/pybind/jit.cc b/paddle/fluid/pybind/jit.cc
@@ -115,6 +115,13 @@ void BindGuard(pybind11::module *m) {
              std::shared_ptr<NumpyDtypeMatchGuard>>(
       *m, "NumpyDtypeMatchGuard", R"DOC(NumpyDtypeMatchGuard Class.)DOC")
       .def(py::init<const py::object &>(), py::arg("dtype"));
+  py::class_<NumPyArrayValueMatchGuard,
+             GuardBase,
+             std::shared_ptr<NumPyArrayValueMatchGuard>>(
+      *m,
+      "NumPyArrayValueMatchGuard",
+      R"DOC(NumPyArrayValueMatchGuard Class.)DOC")
+      .def(py::init<const py::object &>(), py::arg("array"));
 
   m->def(
       "merge_guard",
diff --git a/paddle/fluid/pybind/sot/guards.cc b/paddle/fluid/pybind/sot/guards.cc
@@ -152,6 +152,18 @@ bool NumpyDtypeMatchGuard::check(PyObject* value) {
   return expected_dtype.equal(py::handle(value).get_type());
 }
 
+bool NumPyArrayValueMatchGuard::check(PyObject* value) {
+  if (value == nullptr) {
+    return false;
+  }
+
+  py::object py_value = py::cast<py::object>(value);
+  return py::cast<py::object>(expected_)
+      .attr("__eq__")(py_value)
+      .attr("all")()
+      .cast<bool>();
+}
+
 PyObject* ConstantExprNode::eval(FrameProxy* frame) { return value_ptr_; }
 
 PyObject* LocalVarExprNode::eval(FrameProxy* frame) {
diff --git a/paddle/fluid/pybind/sot/guards.h b/paddle/fluid/pybind/sot/guards.h
@@ -21,6 +21,7 @@ limitations under the License. */
 #include "paddle/fluid/pybind/sot/macros.h"
 #include "paddle/phi/core/utils/data_type.h"
 #include "paddle/utils/pybind.h"
+#include "pybind11/numpy.h"
 #include "pybind11/pybind11.h"
 
 namespace py = pybind11;
@@ -220,6 +221,21 @@ class NumpyDtypeMatchGuard : public GuardBase {
   PyObject* expected_;
 };
 
+class NumPyArrayValueMatchGuard : public GuardBase {
+ public:
+  explicit NumPyArrayValueMatchGuard(const py::object& array)
+      : expected_(array.ptr()) {
+    Py_INCREF(expected_);
+  }
+
+  ~NumPyArrayValueMatchGuard() override { Py_DECREF(expected_); }
+
+  bool check(PyObject* value) override;
+
+ private:
+  PyObject* expected_;
+};
+
 class GuardTreeNode {};
 
 class AttributeExprNode;
@@ -328,4 +344,5 @@ class GuardTree {
 };
 
 std::string guard_tree_to_str(const GuardTree& guard_tree);
+
 #endif
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py b/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py
@@ -1494,16 +1494,22 @@ def make_stringified_guard(self) -> list[StringifiedExpression]:
         frame_value_tracer = self.tracker.trace_value_from_frame()
         obj_free_var_name = f"__{self.id}"
 
-        dtype_guard = StringifiedExpression(
+        dtype_guard = FasterStringifiedExpression(
             f"{{}}.dtype == {NumpyVariable.format_dtype(self.get_py_value().dtype)}",
+            paddle.framework.core.NumpyDtypeMatchGuard(
+                self.get_py_value().dtype
+            ),
             [frame_value_tracer],
             union_free_vars(frame_value_tracer.free_vars, {"np": np}),
         )
 
         return [
             dtype_guard,
-            StringifiedExpression(
+            FasterStringifiedExpression(
                 f"({{}} == {obj_free_var_name}).all()",
+                paddle.framework.core.NumPyArrayValueMatchGuard(
+                    self.get_py_value()
+                ),
                 [frame_value_tracer],
                 union_free_vars(
                     frame_value_tracer.free_vars,
diff --git a/test/sot/test_faster_guard.py b/test/sot/test_faster_guard.py
@@ -148,6 +148,50 @@ def test_numpy_dtype_match_guard(self):
             guard_numpy_bool_dtype.check(np.array(1, dtype=np.bool_))
         )
 
+    def test_numpy_array_match_guard(self):
+        np_array = paddle.framework.core.NumPyArrayValueMatchGuard(
+            np.array([1, 2, 3])
+        )
+        self.assertTrue(np_array.check(np.array([1, 2, 3])))
+        self.assertFalse(np_array.check(np.array([4, 5, 6])))
+
+        np_array_all_one = paddle.framework.core.NumPyArrayValueMatchGuard(
+            np.array([1, 1, 1])
+        )
+        self.assertTrue(np_array_all_one.check(np.array([1, 1, 1])))
+        self.assertFalse(np_array_all_one.check(np.array([1, 2, 3])))
+        self.assertTrue(np_array_all_one.check(np.array([1, 1, 1], dtype=bool)))
+        self.assertTrue(np_array_all_one.check(np.array([True, True, True])))
+        self.assertTrue(np_array_all_one.check(np.array([1, 1, 1], dtype=int)))
+
+        np_bool_array = paddle.framework.core.NumPyArrayValueMatchGuard(
+            np.array([True, False, True])
+        )
+        self.assertTrue(np_bool_array.check(np.array([True, False, True])))
+        self.assertFalse(np_bool_array.check(np.array([True, True, True])))
+        self.assertFalse(np_bool_array.check(np.array([False, False, False])))
+        self.assertFalse(np_bool_array.check(np.array([1, 2, 3])))
+        self.assertTrue(np_bool_array.check(np.array([True, False, 1])))
+        self.assertFalse(np_bool_array.check(np.array([1, 2, 3], dtype=bool)))
+        self.assertFalse(np_bool_array.check(np.array([1, 2, 3], dtype=int)))
+
+        np_bool_array_all_true = (
+            paddle.framework.core.NumPyArrayValueMatchGuard(
+                np.array([True, True, True])
+            )
+        )
+        self.assertTrue(
+            np_bool_array_all_true.check(np.array([True, True, True]))
+        )
+        self.assertFalse(
+            np_bool_array_all_true.check(np.array([True, False, True]))
+        )
+        self.assertFalse(
+            np_bool_array_all_true.check(np.array([False, False, False]))
+        )
+        self.assertFalse(np_bool_array_all_true.check(np.array([1, 2, 3])))
+        self.assertTrue(np_bool_array_all_true.check(np.array([True, True, 1])))
+
 
 class TestFasterGuardGroup(unittest.TestCase):
     def test_guard_group(self):
