convert disttensor for run program (#59143)

wanghuancoder · web-flow · commit 9608fc3e3007 · 2023-11-20T14:41:01.000+08:00
diff --git a/paddle/fluid/pybind/eager_legacy_custom_python_api.h b/paddle/fluid/pybind/eager_legacy_custom_python_api.h
@@ -31,6 +31,13 @@ static PyObject *eager_api_run_program(PyObject *self,  // TOREMOVE
     auto Out = GetTensorPtrListFromArgs("run_program", "Out", args, 2, true);
     auto OutScope =
         GetScopePtrListFromArgs("run_program", "OutScope", args, 3, false);
+    const phi::distributed::ProcessMesh *mesh = nullptr;
+    if (InputsContainDistTensor(&mesh, X, Params, Out)) {
+      X = GetTensorListFromArgs("run_program", "X", args, 0, true, mesh);
+      Params =
+          GetTensorListFromArgs("run_program", "Params", args, 1, true, mesh);
+      Out = GetTensorPtrListFromArgs("run_program", "Out", args, 2, true, mesh);
+    }
     framework::AttributeMap attrs;
     // TODO(zengjinle): support CUDA Graph on eager mode
     ConstructAttrMapFromPyArgs(
@@ -70,6 +77,13 @@ static PyObject *pir_eager_api_run_program(PyObject *self,
     auto Out = GetTensorPtrListFromArgs("run_program", "Out", args, 2, true);
     auto OutScope =
         GetScopePtrListFromArgs("run_program", "OutScope", args, 3, false);
+    const phi::distributed::ProcessMesh *mesh = nullptr;
+    if (InputsContainDistTensor(&mesh, X, Params, Out)) {
+      X = GetTensorListFromArgs("run_program", "X", args, 0, true, mesh);
+      Params =
+          GetTensorListFromArgs("run_program", "Params", args, 1, true, mesh);
+      Out = GetTensorPtrListFromArgs("run_program", "Out", args, 2, true, mesh);
+    }
     framework::AttributeMap attrs;
     // TODO(zengjinle): support CUDA Graph on eager mode
     VLOG(1) << "Start Pir ConstructAttrMapFromPyArgs";
diff --git a/paddle/fluid/pybind/eager_utils.cc b/paddle/fluid/pybind/eager_utils.cc
@@ -1586,7 +1586,8 @@ std::vector<paddle::Tensor*> GetTensorPtrListFromArgs(
     const std::string& arg_name,
     PyObject* args,
     ssize_t arg_idx,
-    bool dispensable) {
+    bool dispensable,
+    const phi::distributed::ProcessMesh* mesh) {
   PyObject* list = PyTuple_GET_ITEM(args, arg_idx);
 
   if (list == nullptr) {
@@ -1602,7 +1603,7 @@ std::vector<paddle::Tensor*> GetTensorPtrListFromArgs(
   }
 
   std::vector<paddle::Tensor*> result;
-  const phi::distributed::ProcessMesh* local_mesh = nullptr;
+  const phi::distributed::ProcessMesh* local_mesh = mesh;
   int mesh_start_index = -1;
 
   if (PyList_Check(list)) {
diff --git a/paddle/fluid/pybind/eager_utils.h b/paddle/fluid/pybind/eager_utils.h
@@ -370,7 +370,8 @@ std::vector<paddle::Tensor*> GetTensorPtrListFromArgs(
     const std::string& arg_name,
     PyObject* args,
     ssize_t arg_idx,
-    bool dispensable = false);
+    bool dispensable = false,
+    const phi::distributed::ProcessMesh* mesh = nullptr);
 
 std::vector<paddle::Tensor*> GetTensorPtrListFromPyObject(PyObject* obj);
 
