[SOT][FasterGuard] add IsDenseTensorHoldAllocationMatchGuard (#72823)

gouzil · Copilot · web-flow · commit 05a7baac2e03 · 2025-05-22T09:12:06.000+08:00
---------

Co-authored-by: Copilot &lt;175728472+Copilot@users.noreply.github.com&gt;
diff --git a/paddle/fluid/pybind/jit.cc b/paddle/fluid/pybind/jit.cc
@@ -132,6 +132,13 @@ 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_<IsNotDenseTensorHoldAllocationMatchGuard,
+             GuardBase,
+             std::shared_ptr<IsNotDenseTensorHoldAllocationMatchGuard>>(
+      *m,
+      "IsNotDenseTensorHoldAllocationMatchGuard",
+      R"DOC(IsNotDenseTensorHoldAllocationMatchGuard Class.)DOC")
+      .def(py::init());
 
   m->def(
       "merge_guard",
diff --git a/paddle/fluid/pybind/sot/guards.cc b/paddle/fluid/pybind/sot/guards.cc
@@ -237,6 +237,42 @@ bool WeakRefMatchGuard::check(PyObject* value) {
 #endif
 }
 
+bool IsNotDenseTensorHoldAllocationMatchGuard::check(PyObject* value) {
+  auto tensor = GetTensorFromPyObject(value);
+  HANDLE_NULL_TENSOR(tensor);
+
+  if (!tensor->defined() ||
+      (!tensor->is_dense_tensor() && !tensor->is_dist_tensor()))
+    return true;
+
+  PyObject* method =
+      PyObject_GetAttrString(value, "_is_dense_tensor_hold_allocation");
+  if (!method) {
+    PyErr_Print();
+    return false;
+  }
+
+  if (!PyCallable_Check(method)) {
+    Py_DECREF(method);
+    PyErr_SetString(PyExc_TypeError, "Attribute is not callable");
+    return false;
+  }
+
+  PyObject* result = PyObject_CallOneArg(method, value);
+  Py_DECREF(method);
+  if (result == nullptr) {
+    PyErr_Print();
+    return false;
+  }
+  int truthy = PyObject_IsTrue(result);
+  Py_DECREF(result);
+  if (truthy == -1) {
+    PyErr_Print();
+    return false;
+  }
+  return !static_cast<bool>(truthy);
+}
+
 PyObject* ConstantExprNode::eval(FrameProxy* frame) { return value_ptr_; }
 std::string ConstantExprNode::stringify(int indent) {
   return py::str(value_ptr_);
diff --git a/paddle/fluid/pybind/sot/guards.h b/paddle/fluid/pybind/sot/guards.h
@@ -285,6 +285,14 @@ class WeakRefMatchGuard : public GuardBase {
   PyObject* expected_;
 };
 
+class IsNotDenseTensorHoldAllocationMatchGuard : public GuardBase {
+ public:
+  bool check(PyObject* value) override;
+  std::string get_guard_name() const override {
+    return "IsNotDenseTensorHoldAllocationMatchGuard";
+  }
+};
+
 class GuardTreeNodeBase {
  public:
   virtual ~GuardTreeNodeBase() = default;
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py b/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py
@@ -526,7 +526,10 @@ def make_faster_guard(self) -> list[paddle.framework.core.GuardNodeBase]:
         meta = self.origin_meta
         if meta.is_null():
             return [
-                # TODO: Implement uninitialized tensor guard
+                paddle.framework.core.GuardNode(
+                    paddle.framework.core.IsNotDenseTensorHoldAllocationMatchGuard(),
+                    [expr_node],
+                ),
             ]
         meta = meta.unwrap_unsafe()
         return [
diff --git a/test/sot/test_faster_guard.py b/test/sot/test_faster_guard.py
@@ -239,6 +239,13 @@ def test_func():
         self.assertFalse(guard_object.check(1))
         self.assertFalse(guard_object.check("1"))
 
+    def test_is_dense_tensor_hold_allocation(self):
+        tensor = paddle.to_tensor([1, 2, 3])
+        guard = paddle.framework.core.IsNotDenseTensorHoldAllocationMatchGuard()
+        self.assertEqual(
+            guard.check(tensor), not tensor._is_dense_tensor_hold_allocation()
+        )
+
 
 class TestFasterGuardGroup(unittest.TestCase):
     def test_guard_group(self):
