vector of tensor, optional of tensor, optional of vector of tensor, support convert DistTensor (#59061)

wanghuancoder · web-flow · commit fa0c85857706 · 2023-11-17T14:53:37.000+08:00
diff --git a/paddle/fluid/eager/auto_code_generator/generator/python_c_gen.py b/paddle/fluid/eager/auto_code_generator/generator/python_c_gen.py
@@ -72,10 +72,22 @@ def FindParsingFunctionFromAttributeType(atype):
     "    auto {} = {}(\"{}\", \"{}\", args, {}, {});\n"
 )
 
-CONVERT_INPUT_TENSORS_TO_DIST_TENSOR_TEMPLATE = """
+CONVERT_TO_DISTTENSOR_AND_PARSE_PYTHON_C_TENSORS_TEMPLATE = (
+    "    {} = {}(\"{}\", \"{}\", args, {}, {}, mesh);\n"
+)
+
+CONVERT_INPUT_TENSORS_TO_DIST_TENSOR_WITH_SINGLE_TENSOR_TEMPLATE = """
     const phi::distributed::ProcessMesh* mesh = nullptr;
-    if (InputsContainDistTensor(&mesh{inputs})) {{
-      ConvertAllInputsToDistTensor(mesh{inputs});
+    if (InputsContainDistTensor(&mesh{input_names})) {{
+      ConvertAllInputsToDistTensor(mesh{input_single_tensor_names});
+      {optional_and_vector_convert_code}
+    }}
+"""
+
+CONVERT_INPUT_TENSORS_TO_DIST_TENSOR_WITHOUT_SINGLE_TENSOR_TEMPLATE = """
+    const phi::distributed::ProcessMesh* mesh = nullptr;
+    if (InputsContainDistTensor(&mesh{input_names})) {{
+      {optional_and_vector_convert_code}
     }}
 """
 
@@ -332,6 +344,7 @@ def GeneratePythonCFunction(self):
         # Generate Python-C Tensors Parsing Logic
         get_eager_tensor_str = ""
         input_names = ""
+        input_single_tensor_names = ""
         for name, (ttype, pos) in forward_inputs_position_map.items():
             input_names = input_names + ", " + name
             if forward_inplace_map and name in forward_inplace_map.keys():
@@ -373,6 +386,9 @@ def GeneratePythonCFunction(self):
                         )
                     )
                 else:
+                    input_single_tensor_names = (
+                        input_single_tensor_names + ", " + name
+                    )
                     get_eager_tensor_str += (
                         PARSE_PYTHON_C_TENSORS_TEMPLATE.format(
                             name,
@@ -385,12 +401,50 @@ def GeneratePythonCFunction(self):
                     )
         # No inputs, skip convert to DistTensor
         if len(input_names) > 0:
-            get_eager_tensor_str += (
-                CONVERT_INPUT_TENSORS_TO_DIST_TENSOR_TEMPLATE.format(
-                    inputs=input_names
-                )
-            )
+            optional_and_vector_convert_code = ""
+            for name, (ttype, pos) in forward_inputs_position_map.items():
+                is_optional = name in optional_inputs
+                if IsVectorTensorType(ttype):
+                    if is_optional:
+                        optional_and_vector_convert_code += CONVERT_TO_DISTTENSOR_AND_PARSE_PYTHON_C_TENSORS_TEMPLATE.format(
+                            name,
+                            "GetOptionalTensorListFromArgs",
+                            forward_api_name,
+                            name,
+                            pos,
+                            "true",
+                        )
+                    else:
+                        optional_and_vector_convert_code += CONVERT_TO_DISTTENSOR_AND_PARSE_PYTHON_C_TENSORS_TEMPLATE.format(
+                            name,
+                            "GetTensorListFromArgs",
+                            forward_api_name,
+                            name,
+                            pos,
+                            "false",
+                        )
+                else:
+                    if is_optional:
+                        optional_and_vector_convert_code += CONVERT_TO_DISTTENSOR_AND_PARSE_PYTHON_C_TENSORS_TEMPLATE.format(
+                            name,
+                            "GetOptionalTensorFromArgs",
+                            forward_api_name,
+                            name,
+                            pos,
+                            "true",
+                        )
 
+            if len(input_single_tensor_names) > 0:
+                get_eager_tensor_str += CONVERT_INPUT_TENSORS_TO_DIST_TENSOR_WITH_SINGLE_TENSOR_TEMPLATE.format(
+                    input_names=input_names,
+                    input_single_tensor_names=input_single_tensor_names,
+                    optional_and_vector_convert_code=optional_and_vector_convert_code,
+                )
+            else:
+                get_eager_tensor_str += CONVERT_INPUT_TENSORS_TO_DIST_TENSOR_WITHOUT_SINGLE_TENSOR_TEMPLATE.format(
+                    input_names=input_names,
+                    optional_and_vector_convert_code=optional_and_vector_convert_code,
+                )
         if forward_inplace_map:
             for name, (ttype, pos) in forward_outputs_position_map.items():
                 if name in forward_inplace_map.values():
diff --git a/paddle/fluid/pybind/eager_utils.cc b/paddle/fluid/pybind/eager_utils.cc
@@ -107,6 +107,35 @@ int TensorDtype2NumpyDtype(phi::DataType dtype) {
   }
 }
 
+void ConvertToDistTensor(Tensor* x, const phi::distributed::ProcessMesh* mesh) {
+  if (x->is_dist_tensor()) {
+    PADDLE_ENFORCE_EQ(
+        std::dynamic_pointer_cast<phi::distributed::DistTensor>(x->impl())
+            ->dist_attr()
+            .process_mesh(),
+        *mesh,
+        platform::errors::InvalidArgument(
+            "Input %s has different mesh. However all inputs should "
+            "have the same mesh.",
+            x->name()));
+    return;
+  } else {
+    PADDLE_ENFORCE_EQ(
+        phi::DenseTensor::classof(x->impl().get()),
+        true,
+        platform::errors::InvalidArgument(
+            "Failed to convert input %s impl to phi::distributed::DistTensor "
+            "as it's not phi::DenseTensor.",
+            x->name()));
+    phi::distributed::TensorDistAttr dist_attr(
+        phi::vectorize(x->impl()->dims()));
+    dist_attr.set_process_mesh(*mesh);
+    auto dense_t = std::static_pointer_cast<phi::DenseTensor>(x->impl());
+    x->set_impl(
+        std::make_shared<phi::distributed::DistTensor>(dense_t, dist_attr));
+  }
+}
+
 bool PyObject_CheckLongOrConvertToLong(PyObject** obj) {
   if (PyLong_Check(*obj) && !PyBool_Check(*obj)) {
     return true;
@@ -1172,7 +1201,8 @@ paddle::optional<paddle::Tensor> GetOptionalTensorFromArgs(
     const std::string& arg_name,
     PyObject* args,
     ssize_t arg_idx,
-    bool dispensable) {
+    bool dispensable,
+    const phi::distributed::ProcessMesh* mesh) {
   PyObject* obj = PyTuple_GET_ITEM(args, arg_idx);
 
   if (PyTuple_Check(obj)) {
@@ -1191,6 +1221,10 @@ paddle::optional<paddle::Tensor> GetOptionalTensorFromArgs(
   }
 
   if (PyObject_TypeCheck(obj, p_tensor_type)) {
+    if (mesh) {
+      ConvertToDistTensor(&(reinterpret_cast<TensorObject*>(obj)->tensor),
+                          mesh);
+    }
     return paddle::make_optional<paddle::Tensor>(
         reinterpret_cast<TensorObject*>(obj)->tensor);
   } else {
@@ -1255,11 +1289,13 @@ paddle::Tensor& GetTensorFromArgs(const std::string& op_type,
   return GetTensorFromPyObject(op_type, arg_name, obj, arg_idx, dispensable);
 }
 
-std::vector<paddle::Tensor> GetTensorListFromArgs(const std::string& op_type,
-                                                  const std::string& arg_name,
-                                                  PyObject* args,
-                                                  ssize_t arg_idx,
-                                                  bool dispensable) {
+std::vector<paddle::Tensor> GetTensorListFromArgs(
+    const std::string& op_type,
+    const std::string& arg_name,
+    PyObject* args,
+    ssize_t arg_idx,
+    bool dispensable,
+    const phi::distributed::ProcessMesh* mesh) {
   PyObject* list = PyTuple_GET_ITEM(args, arg_idx);
 
   if (list == nullptr) {
@@ -1288,6 +1324,11 @@ std::vector<paddle::Tensor> GetTensorListFromArgs(const std::string& op_type,
           arg_idx));
     }
     for (Py_ssize_t i = 0; i < len; i++) {
+      if (mesh) {
+        ConvertToDistTensor(
+            &(reinterpret_cast<TensorObject*>(PyList_GetItem(list, i))->tensor),
+            mesh);
+      }
       result.emplace_back(
           reinterpret_cast<TensorObject*>(PyList_GetItem(list, i))->tensor);
     }
@@ -1303,6 +1344,12 @@ std::vector<paddle::Tensor> GetTensorListFromArgs(const std::string& op_type,
           arg_idx));
     }
     for (Py_ssize_t i = 0; i < len; i++) {
+      if (mesh) {
+        ConvertToDistTensor(
+            &(reinterpret_cast<TensorObject*>(PyTuple_GetItem(list, i))
+                  ->tensor),
+            mesh);
+      }
       result.emplace_back(
           reinterpret_cast<TensorObject*>(PyTuple_GetItem(list, i))->tensor);
     }
@@ -1326,7 +1373,8 @@ paddle::optional<std::vector<paddle::Tensor>> GetOptionalTensorListFromArgs(
     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 || list == Py_None) {
@@ -1355,6 +1403,11 @@ paddle::optional<std::vector<paddle::Tensor>> GetOptionalTensorListFromArgs(
           arg_idx));
     }
     for (Py_ssize_t i = 0; i < len; i++) {
+      if (mesh) {
+        ConvertToDistTensor(
+            &(reinterpret_cast<TensorObject*>(PyList_GetItem(list, i))->tensor),
+            mesh);
+      }
       result.emplace_back(
           reinterpret_cast<TensorObject*>(PyList_GetItem(list, i))->tensor);
     }
@@ -1370,6 +1423,12 @@ paddle::optional<std::vector<paddle::Tensor>> GetOptionalTensorListFromArgs(
           arg_idx));
     }
     for (Py_ssize_t i = 0; i < len; i++) {
+      if (mesh) {
+        ConvertToDistTensor(
+            &(reinterpret_cast<TensorObject*>(PyTuple_GetItem(list, i))
+                  ->tensor),
+            mesh);
+      }
       result.emplace_back(
           reinterpret_cast<TensorObject*>(PyTuple_GetItem(list, i))->tensor);
     }
@@ -2164,34 +2223,7 @@ void DistTensorTypeParser::operator()(
   }
 }
 
-void DistTensorConverter::convert(Tensor* x) {
-  if (x->is_dist_tensor()) {
-    PADDLE_ENFORCE_EQ(
-        std::dynamic_pointer_cast<phi::distributed::DistTensor>(x->impl())
-            ->dist_attr()
-            .process_mesh(),
-        *mesh,
-        platform::errors::InvalidArgument(
-            "Input %s has different mesh. However all inputs should "
-            "have the same mesh.",
-            x->name()));
-    return;
-  } else {
-    PADDLE_ENFORCE_EQ(
-        phi::DenseTensor::classof(x->impl().get()),
-        true,
-        platform::errors::InvalidArgument(
-            "Failed to convert input %s impl to phi::distributed::DistTensor "
-            "as it's not phi::DenseTensor.",
-            x->name()));
-    phi::distributed::TensorDistAttr dist_attr(
-        phi::vectorize(x->impl()->dims()));
-    dist_attr.set_process_mesh(*mesh);
-    auto dense_t = std::static_pointer_cast<phi::DenseTensor>(x->impl());
-    x->set_impl(
-        std::make_shared<phi::distributed::DistTensor>(dense_t, dist_attr));
-  }
-}
+void DistTensorConverter::convert(Tensor* x) { ConvertToDistTensor(x, mesh); }
 
 void DistTensorConverter::operator()(Tensor* x) {
   DistTensorConverter::convert(x);
diff --git a/paddle/fluid/pybind/eager_utils.h b/paddle/fluid/pybind/eager_utils.h
@@ -334,7 +334,8 @@ paddle::optional<paddle::Tensor> GetOptionalTensorFromArgs(
     const std::string& arg_name,
     PyObject* args,
     ssize_t arg_idx,
-    bool dispensable = false);
+    bool dispensable = false,
+    const phi::distributed::ProcessMesh* mesh = nullptr);
 
 paddle::Tensor& GetTensorFromArgs(const std::string& op_type,
                                   const std::string& arg_name,
@@ -347,13 +348,16 @@ paddle::optional<std::vector<paddle::Tensor>> GetOptionalTensorListFromArgs(
     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> GetTensorListFromArgs(const std::string& op_type,
-                                                  const std::string& arg_name,
-                                                  PyObject* args,
-                                                  ssize_t arg_idx,
-                                                  bool dispensable = false);
+std::vector<paddle::Tensor> GetTensorListFromArgs(
+    const std::string& op_type,
+    const std::string& arg_name,
+    PyObject* args,
+    ssize_t arg_idx,
+    bool dispensable = false,
+    const phi::distributed::ProcessMesh* mesh = nullptr);
 
 paddle::Tensor* GetTensorPtrFromArgs(const std::string& op_type,
                                      const std::string& arg_name,
