[SOT][Faster Guard] support TensorVariable Dist check (#72327)

gouzil · web-flow · commit 026a8df4f865 · 2025-04-26T15:48:32.000+08:00
diff --git a/paddle/fluid/pybind/jit.cc b/paddle/fluid/pybind/jit.cc
@@ -128,6 +128,13 @@ void BindGuard(pybind11::module *m) {
   py::class_<DummyGuard, GuardBase, std::shared_ptr<DummyGuard>>(
       *m, "DummyGuard", R"DOC(DummyGuard Class.)DOC")
       .def(py::init<>());
+  py::class_<TensorDistMetaMatchGuard,
+             GuardBase,
+             std::shared_ptr<TensorDistMetaMatchGuard>>(
+      *m,
+      "TensorDistMetaMatchGuard",
+      R"DOC(TensorDistMetaMatchGuard Class.)DOC")
+      .def(py::init<const py::object &>(), py::arg("dist_info"));
 
   m->def(
       "merge_guard",
diff --git a/paddle/fluid/pybind/sot/guards.cc b/paddle/fluid/pybind/sot/guards.cc
@@ -33,10 +33,31 @@ static inline PyObject* PyObject_CallOneArg(PyObject* func, PyObject* arg) {
 #define Py_IsNone(x) ((x) == Py_None)
 #endif
 
+// check if the tensor is null, tensor is std::optional<paddle::Tensor>
+#define HANDLE_NULL_TENSOR(tensor) \
+  {                                \
+    if (!tensor) {                 \
+      return false;                \
+    }                              \
+  }
+
+// check if the value is null and decref it
+#define HANDLE_NULL_VALUE_DECREF(value) \
+  {                                     \
+    if ((value) == NULL) {              \
+      Py_DECREF(value);                 \
+      PyErr_Clear();                    \
+      return false;                     \
+    }                                   \
+  }
+
+// check if the value is null
 #define HANDLE_NULL_VALUE(value) \
-  if ((value) == NULL) {         \
-    PyErr_Clear();               \
-    return false;                \
+  {                              \
+    if ((value) == NULL) {       \
+      PyErr_Clear();             \
+      return false;              \
+    }                            \
   }
 
 static inline bool PyObject_Equal(PyObject* a, PyObject* b) {
@@ -106,19 +127,15 @@ bool LengthMatchGuard::check(PyObject* value) {
 
 bool DtypeMatchGuard::check(PyObject* value) {
   auto tensor = GetTensorFromPyObject(value);
-  if (!tensor) {
-    return false;
-  }
+  HANDLE_NULL_TENSOR(tensor);
   auto dtype = tensor->type();
   return phi::TransToProtoVarType(dtype) == expected_;
 }
 
 bool ShapeMatchGuard::check(PyObject* value) {
   HANDLE_NULL_VALUE(value);
   auto tensor = GetTensorFromPyObject(value);
-  if (!tensor) {
-    return false;
-  }
+  HANDLE_NULL_TENSOR(tensor);
   auto shape = tensor->shape();
   if (shape.size() != expected_.size()) {
     return false;
@@ -201,6 +218,60 @@ bool WeakRefMatchGuard::check(PyObject* value) {
 #endif
 }
 
+bool TensorDistMetaMatchGuard::check(PyObject* value) {
+  HANDLE_NULL_VALUE(value);
+
+  PyObject* expr = PyTuple_GetItem(value, 0);
+  HANDLE_NULL_VALUE(expr);
+
+  auto tensor = GetTensorFromPyObject(expr);
+  HANDLE_NULL_TENSOR(tensor);
+
+  if (tensor->is_dist_tensor() == false && is_dist_ == false) return true;
+  if (tensor->is_dist_tensor() != is_dist_) {
+    return false;
+  }
+
+  PyObject* dist_info_from_tensor_func = PyTuple_GetItem(value, 1);
+  HANDLE_NULL_VALUE(dist_info_from_tensor_func);
+
+  PyObject* dist_info = PyObject_CallOneArg(dist_info_from_tensor_func, expr);
+  HANDLE_NULL_VALUE_DECREF(dist_info);
+
+  PyObject* mesh = PyObject_GetAttrString(dist_info, "mesh");
+  HANDLE_NULL_VALUE_DECREF(mesh);
+
+  PyObject* mesh_shape = PyObject_GetAttrString(mesh, "shape");
+  HANDLE_NULL_VALUE_DECREF(mesh_shape);
+  PyObject* process_ids = PyObject_GetAttrString(mesh, "process_ids");
+  HANDLE_NULL_VALUE_DECREF(process_ids);
+  PyObject* dims_mapping = PyObject_GetAttrString(dist_info, "dims_mapping");
+  HANDLE_NULL_VALUE_DECREF(dims_mapping);
+  PyObject* local_shape = PyObject_GetAttrString(dist_info, "local_shape");
+  HANDLE_NULL_VALUE_DECREF(local_shape);
+
+  if (py::handle(mesh_shape).cast<std::vector<int>>() != mesh_shape_expected_ ||
+      py::handle(process_ids).cast<std::vector<int>>() !=
+          mesh_process_ids_expected_.value() ||
+      !PyObject_Equal(dims_mapping, dims_mapping_expected_.value()) ||
+      !PyObject_Equal(local_shape, local_shape_expected_.value())) {
+    Py_DECREF(mesh);
+    Py_DECREF(mesh_shape);
+    Py_DECREF(process_ids);
+    Py_DECREF(dims_mapping);
+    Py_DECREF(local_shape);
+    PyErr_Clear();
+    return false;
+  }
+
+  Py_DECREF(mesh);
+  Py_DECREF(mesh_shape);
+  Py_DECREF(process_ids);
+  Py_DECREF(dims_mapping);
+  Py_DECREF(local_shape);
+  return true;
+}
+
 PyObject* ConstantExprNode::eval(FrameProxy* frame) { return value_ptr_; }
 std::string ConstantExprNode::stringify(int indent) {
   return py::str(value_ptr_);
@@ -338,15 +409,38 @@ std::string BinaryExprNode::stringify(int indent) {
 
 std::optional<int> GuardNode::lookup(FrameProxy* frame) {
   // TODO(zrr1999): support multiple exprs
-  auto expr = exprs.back();
-  auto value = expr->eval(frame);
+  PyObject* value = [this, frame]() {
+    if (exprs.size() == 1) {
+      PyObject* v = exprs.back()->eval(frame);
+      if (v) {
+        // TODO(dev): DECREF v.
+        Py_INCREF(v);
+      }
+      return v;
+    }
+    auto values = std::vector<PyObject*>(exprs.size());
+    for (size_t i = 0; i < exprs.size(); ++i) {
+      values[i] = exprs[i]->eval(frame);
+      if (values[i]) {
+        Py_INCREF(values[i]);
+      }
+    }
+    auto packed_value = PyTuple_New(exprs.size());
+    for (size_t i = 0; i < exprs.size(); ++i) {
+      PyTuple_SetItem(packed_value, i, values[i]);
+    }
+    return packed_value;
+  }();
+
   if (guard->check(value)) {
     if (return_cache_index.has_value()) {
+      Py_DECREF(value);
       return return_cache_index.value();
     }
     for (auto& next_guard_node : next_guard_nodes) {
       auto ret = next_guard_node->lookup(frame);
       if (ret.has_value()) {
+        Py_DECREF(value);
         return ret.value();
       }
     }
diff --git a/paddle/fluid/pybind/sot/guards.h b/paddle/fluid/pybind/sot/guards.h
@@ -15,6 +15,7 @@ limitations under the License. */
 
 #include <Python.h>
 #include <memory>
+#include <vector>
 #include "paddle/fluid/framework/data_type.h"
 #include "paddle/fluid/pybind/sot/eval_frame_tools.h"
 #include "paddle/fluid/pybind/sot/frame_proxy.h"
@@ -23,6 +24,7 @@ limitations under the License. */
 #include "paddle/utils/pybind.h"
 #include "pybind11/numpy.h"
 #include "pybind11/pybind11.h"
+#include "pybind11/pytypes.h"
 
 namespace py = pybind11;
 #define PYBIND11_DETAILED_ERROR_MESSAGES
@@ -266,6 +268,42 @@ class WeakRefMatchGuard : public GuardBase {
   PyObject* expected_;
 };
 
+class TensorDistMetaMatchGuard : public GuardBase {
+ public:
+  explicit TensorDistMetaMatchGuard(const py::object& obj) {
+    if (obj != py::none()) {
+      mesh_shape_expected_ =
+          obj.attr("mesh").attr("shape").cast<std::vector<int>>();
+      mesh_process_ids_expected_ =
+          obj.attr("mesh").attr("process_ids").cast<std::vector<int>>();
+      dims_mapping_expected_ = obj.attr("dims_mapping").ptr();
+      local_shape_expected_ = obj.attr("local_shape").ptr();
+
+      is_dist_ = true;
+      Py_INCREF(dims_mapping_expected_.value());
+      Py_INCREF(local_shape_expected_.value());
+    }
+  }
+
+  ~TensorDistMetaMatchGuard() override {
+    if (is_dist_) {
+      Py_DECREF(dims_mapping_expected_.value());
+      Py_DECREF(local_shape_expected_.value());
+    }
+  }
+  bool check(PyObject* value) override;
+  std::string get_guard_name() const override {
+    return "TensorDistMetaMatchGuard";
+  }
+
+ private:
+  bool is_dist_ = false;
+  std::optional<std::vector<int>> mesh_shape_expected_;
+  std::optional<std::vector<int>> mesh_process_ids_expected_;
+  std::optional<PyObject*> dims_mapping_expected_;
+  std::optional<PyObject*> local_shape_expected_;
+};
+
 class DummyGuard : public GuardBase {
  public:
   bool check(PyObject* value) override { return true; }
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py b/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py
@@ -540,7 +540,18 @@ def make_faster_guard(self) -> list[paddle.framework.core.GuardNodeBase]:
                     )
                 ],
             ),
-            # TODO(zrr1999): use TensorMetaMatchGuard to support dist_info check
+            # Check dist info
+            paddle.framework.core.GuardNode(
+                paddle.framework.core.TensorDistMetaMatchGuard(
+                    self.meta.dist_info
+                ),
+                [
+                    expr_node,
+                    paddle.framework.core.ExternVarExprNode(
+                        "___dist_info_from_tensor", DistInfo.from_tensor
+                    ),
+                ],
+            ),
         ]
 
     @check_guard
diff --git a/test/sot/test_faster_guard.py b/test/sot/test_faster_guard.py
@@ -20,6 +20,8 @@
 import numpy as np
 
 import paddle
+import paddle.distributed as dist
+from paddle.jit.sot.infer_meta import DistInfo
 
 
 class TestBasicFasterGuard(unittest.TestCase):
@@ -202,6 +204,35 @@ def test_func():
         self.assertFalse(guard_object.check(1))
         self.assertFalse(guard_object.check("1"))
 
+    @unittest.skipIf(
+        not paddle.is_compiled_with_distribute(),
+        reason='Not compiled with distribute.',
+    )
+    def test_tensor_dist_meta_guard(self):
+        x = paddle.ones([2, 2])
+        y = paddle.ones([1, 2])
+        x.stop_gradient = False
+        y.stop_gradient = False
+        mesh1 = dist.ProcessMesh([0, 1], dim_names=['x'])
+        mesh2 = dist.ProcessMesh([0, 1], dim_names=['y'])
+        dist_x1 = dist.shard_tensor(
+            x, mesh1, [dist.Replicate()], stop_gradient=False
+        )
+        dist_y1 = dist.shard_tensor(
+            y, mesh2, [dist.Replicate()], stop_gradient=False
+        )
+        guard_tensor_is_dist = paddle.framework.core.TensorDistMetaMatchGuard(
+            DistInfo.from_tensor(dist_x1)
+        )
+        self.assertTrue(
+            guard_tensor_is_dist.check((dist_x1, DistInfo.from_tensor))
+        )
+        self.assertFalse(
+            guard_tensor_is_dist.check((dist_y1, DistInfo.from_tensor))
+        )
+        self.assertFalse(guard_tensor_is_dist.check((x, DistInfo.from_tensor)))
+        self.assertFalse(guard_tensor_is_dist.check((y, DistInfo.from_tensor)))
+
 
 class TestFasterGuardGroup(unittest.TestCase):
     def test_guard_group(self):
