[SOT][Faster Guard] support NumpyDtypeMatchGuard (#71900)

gouzil · web-flow · commit e5b43407dfa6 · 2025-03-27T21:52:22.000+08:00
* [SOT][Faster Guard] support `NumpyDtypeMatchGuard`

* fix review

* fix review and test build

* test build

* fix build error

* use pyobj

* add todo
diff --git a/paddle/fluid/pybind/jit.cc b/paddle/fluid/pybind/jit.cc
@@ -108,6 +108,11 @@ void BindGuard(pybind11::module *m) {
              std::shared_ptr<InstanceCheckGuard>>(
       *m, "InstanceCheckGuard", R"DOC(InstanceCheckGuard Class.)DOC")
       .def(py::init<const py::object &>(), py::arg("isinstance_obj"));
+  py::class_<NumpyDtypeMatchGuard,
+             GuardBase,
+             std::shared_ptr<NumpyDtypeMatchGuard>>(
+      *m, "NumpyDtypeMatchGuard", R"DOC(NumpyDtypeMatchGuard Class.)DOC")
+      .def(py::init<const py::object &>(), py::arg("dtype"));
 
   m->def(
       "merge_guard",
diff --git a/paddle/fluid/pybind/sot/guards.cc b/paddle/fluid/pybind/sot/guards.cc
@@ -19,6 +19,7 @@ limitations under the License. */
 
 #include <Python.h>
 #include <frameobject.h>
+#include "pybind11/numpy.h"
 
 #if !defined(PyObject_CallOneArg) && !PY_3_9_PLUS
 static inline PyObject* PyObject_CallOneArg(PyObject* func, PyObject* arg) {
@@ -133,4 +134,20 @@ bool InstanceCheckGuard::check(PyObject* value) {
   return PyObject_IsInstance(value, expected_);
 }
 
+bool NumpyDtypeMatchGuard::check(PyObject* value) {
+  if (value == nullptr) {
+    return false;
+  }
+
+  // TODO(dev): encountered a compilation error: "declared with greater
+  // visibility than the type of its field", so had to put the conversion here
+  py::dtype expected_dtype = py::cast<py::dtype>(expected_);
+
+  if (py::isinstance<py::array>(value)) {
+    return py::cast<py::array>(value).dtype().is(expected_dtype);
+  }
+
+  return expected_dtype.equal(py::handle(value).get_type());
+}
+
 #endif
diff --git a/paddle/fluid/pybind/sot/guards.h b/paddle/fluid/pybind/sot/guards.h
@@ -202,4 +202,19 @@ class InstanceCheckGuard : public GuardBase {
   PyObject* expected_;
 };
 
+class NumpyDtypeMatchGuard : public GuardBase {
+ public:
+  explicit NumpyDtypeMatchGuard(const py::object& dtype)
+      : expected_(dtype.ptr()) {
+    Py_INCREF(expected_);
+  }
+
+  ~NumpyDtypeMatchGuard() override { Py_DECREF(expected_); }
+
+  bool check(PyObject* value) override;
+
+ private:
+  PyObject* expected_;
+};
+
 #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
@@ -1331,16 +1331,20 @@ def getattr(self, name: str, default=None):
     def make_stringified_guard(self) -> list[StringifiedExpression]:
         frame_value_tracer = self.tracker.trace_value_from_frame()
 
-        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"{{}} == {NumpyVariable.format_number(self.get_py_value())}",
+                paddle.framework.core.ValueMatchGuard(self.get_py_value()),
                 [frame_value_tracer],
                 union_free_vars(frame_value_tracer.free_vars, {"np": np}),
             ),
diff --git a/test/sot/test_faster_guard.py b/test/sot/test_faster_guard.py
@@ -17,6 +17,8 @@
 import unittest
 from collections import OrderedDict
 
+import numpy as np
+
 import paddle
 
 
@@ -125,6 +127,27 @@ def test_id_match_guard(self):
         self.assertTrue(guard_id.check(layer))
         self.assertFalse(guard_id.check(paddle.nn.Linear(10, 10)))
 
+    def test_numpy_dtype_match_guard(self):
+        np_array = np.array(1, dtype=np.int32)
+        guard_numpy_dtype = paddle.framework.core.NumpyDtypeMatchGuard(
+            np_array.dtype
+        )
+        self.assertTrue(guard_numpy_dtype.check(np_array))
+        self.assertTrue(guard_numpy_dtype.check(np.array(1, dtype=np.int32)))
+        self.assertTrue(guard_numpy_dtype.check(np.int32()))
+        self.assertFalse(guard_numpy_dtype.check(np.array(1, dtype=np.int64)))
+        self.assertFalse(guard_numpy_dtype.check(np.float32()))
+        self.assertFalse(guard_numpy_dtype.check(np.bool_()))
+
+        np_bool = np.bool_(1)
+        guard_numpy_bool_dtype = paddle.framework.core.NumpyDtypeMatchGuard(
+            np_bool.dtype
+        )
+        self.assertTrue(guard_numpy_bool_dtype.check(np.bool_()))
+        self.assertTrue(
+            guard_numpy_bool_dtype.check(np.array(1, dtype=np.bool_))
+        )
+
 
 class TestFasterGuardGroup(unittest.TestCase):
     def test_guard_group(self):
