PaddlePaddle
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/attribute_storage.h
Copy file name to clipboard
+6-6Lines changed: 6 additions & 6 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/attribute_storage.h
Copy file name to clipboard
+6-6Lines changed: 6 additions & 6 deletions
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_attribute.cc
Copy file name to clipboard
+10-31Lines changed: 10 additions & 31 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/dist_attribute.cc
Copy file name to clipboard
+10-31Lines changed: 10 additions & 31 deletions
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_attribute.h
Copy file name to clipboard
+13-15Lines changed: 13 additions & 15 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/dist_attribute.h
Copy file name to clipboard
+13-15Lines changed: 13 additions & 15 deletions
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_dialect.cc
Copy file name to clipboard
+6-74Lines changed: 6 additions & 74 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/dist_dialect.cc
Copy file name to clipboard
+6-74Lines changed: 6 additions & 74 deletions
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_tools.cc
Copy file name to clipboard
+61-8Lines changed: 61 additions & 8 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/dist_tools.cc
Copy file name to clipboard
+61-8Lines changed: 61 additions & 8 deletions
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_tools.h
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/dist_tools.h
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions
diff --git a/‎paddle/fluid/pir/dialect/op_generator/op_infermeta_func_gen.py
Copy file name to clipboard
+5-16Lines changed: 5 additions & 16 deletions b/‎paddle/fluid/pir/dialect/op_generator/op_infermeta_func_gen.py
Copy file name to clipboard
+5-16Lines changed: 5 additions & 16 deletions
@@ -125,8 +125,8 @@ class OperationDistAttrStorage : public pir::AttributeStorage {
                               std::vector<pir::Attribute>>;
   OperationDistAttrStorage(ParamKey&& param)  // NOLINT
       : mesh_attr(std::get<0>(param)),
-        operand_attrs(std::get<1>(param)),
-        result_attrs(std::get<2>(param)) {}
+        operands(std::get<1>(param)),
+        results(std::get<2>(param)) {}
 
   ///
   /// \brief Each derived TypeStorage must define a Construct method, which
@@ -156,13 +156,13 @@ class OperationDistAttrStorage : public pir::AttributeStorage {
   /// \brief Each derived TypeStorage needs to overload operator==.
   ///
   bool operator==(const ParamKey& key) const {
-    return mesh_attr == std::get<0>(key) && operand_attrs == std::get<1>(key) &&
-           result_attrs == std::get<2>(key);
+    return mesh_attr == std::get<0>(key) && operands == std::get<1>(key) &&
+           results == std::get<2>(key);
   }
 
   ProcessMeshAttribute mesh_attr;
-  std::vector<pir::Attribute> operand_attrs;
-  std::vector<pir::Attribute> result_attrs;
+  std::vector<pir::Attribute> operands;
+  std::vector<pir::Attribute> results;
 };
 
 }  // namespace dialect
 
@@ -78,47 +78,27 @@ TensorDistAttribute TensorDistAttribute::get(
 ProcessMeshAttribute OperationDistAttribute::process_mesh_attr() const {
   return storage()->mesh_attr;
 }
-const std::vector<pir::Attribute>& OperationDistAttribute::operand_attrs()
-    const {
-  return storage()->operand_attrs;
-}
-TensorDistAttribute OperationDistAttribute::operand_dist_attr(
-    uint32_t index) const {
-  return operand_attrs().at(index).dyn_cast<TensorDistAttribute>();
-}
-
-pir::ArrayAttribute OperationDistAttribute::operand_array_attr(
-    uint32_t index) const {
-  return operand_attrs().at(index).dyn_cast<pir::ArrayAttribute>();
+const std::vector<pir::Attribute>& OperationDistAttribute::operands() const {
+  return storage()->operands;
 }
 
 uint32_t OperationDistAttribute::num_operands() const {
-  return operand_attrs().size();
+  return operands().size();
 }
 
-const std::vector<pir::Attribute>& OperationDistAttribute::result_attrs()
-    const {
-  return storage()->result_attrs;
-}
-TensorDistAttribute OperationDistAttribute::result_dist_attr(
-    uint32_t index) const {
-  return result_attrs().at(index).dyn_cast<TensorDistAttribute>();
-}
-
-pir::ArrayAttribute OperationDistAttribute::result_array_attr(
-    uint32_t index) const {
-  return result_attrs().at(index).dyn_cast<pir::ArrayAttribute>();
+const std::vector<pir::Attribute>& OperationDistAttribute::results() const {
+  return storage()->results;
 }
 
 uint32_t OperationDistAttribute::num_results() const {
-  return result_attrs().size();
+  return results().size();
 }
 
 OperationDistAttribute OperationDistAttribute::get(
     pir::IrContext* ctx,
     ProcessMeshAttribute mesh,
-    const std::vector<pir::Attribute>& operand_attrs,
-    const std::vector<pir::Attribute>& result_attrs) {
+    const std::vector<pir::Attribute>& operands,
+    const std::vector<pir::Attribute>& results) {
   auto check_dist_attr = [=](pir::Attribute attr) {
     auto dist_attr = attr.dyn_cast<TensorDistAttribute>();
     auto ids = mesh.process_ids();
@@ -132,8 +112,7 @@ OperationDistAttribute OperationDistAttribute::get(
                             mesh));
     }
   };
-
-  for (auto attr : operand_attrs) {
+  for (auto attr : operands) {
     // NOTE: The operand dist attr maybe empty while the corresponding input is
     // optional.
     if (!attr) continue;
@@ -145,7 +124,7 @@ OperationDistAttribute OperationDistAttribute::get(
       check_dist_attr(attr);
     }
   }
-  return Base::get(ctx, mesh, operand_attrs, result_attrs);
+  return Base::get(ctx, mesh, operands, results);
 }
 
 }  // namespace dialect
 
@@ -105,28 +105,26 @@ class OperationDistAttribute : public pir::AttrBase<OperationDistAttribute,
   using Base::Base;
   ProcessMeshAttribute process_mesh_attr() const;
 
-  const std::vector<Attribute>& operand_attrs() const;
-  TensorDistAttribute operand_dist_attr(uint32_t index) const;
-  pir::ArrayAttribute operand_array_attr(uint32_t index) const;
+  const std::vector<Attribute>& operands() const;
+  pir::Attribute operand(uint32_t index) const { return operands().at(index); }
   uint32_t num_operands() const;
 
-  const std::vector<Attribute>& result_attrs() const;
-  TensorDistAttribute result_dist_attr(uint32_t index) const;
-  pir::ArrayAttribute result_array_attr(uint32_t index) const;
+  const std::vector<Attribute>& results() const;
+
+  pir::Attribute result(uint32_t index) const { return results().at(index); }
+
   uint32_t num_results() const;
 
   static OperationDistAttribute get(pir::IrContext* ctx,
                                     ProcessMeshAttribute mesh,
-                                    const std::vector<Attribute>& operand_attrs,
-                                    const std::vector<Attribute>& result_attrs);
+                                    const std::vector<Attribute>& operands,
+                                    const std::vector<Attribute>& results);
 
-  static OperationDistAttribute get(
-      pir::IrContext* ctx,
-      const phi::distributed::ProcessMesh& mesh,
-      const std::vector<Attribute>& operand_attrs,
-      const std::vector<Attribute>& result_attrs) {
-    return get(
-        ctx, ProcessMeshAttribute::get(ctx, mesh), operand_attrs, result_attrs);
+  static OperationDistAttribute get(pir::IrContext* ctx,
+                                    const phi::distributed::ProcessMesh& mesh,
+                                    const std::vector<Attribute>& operands,
+                                    const std::vector<Attribute>& results) {
+    return get(ctx, ProcessMeshAttribute::get(ctx, mesh), operands, results);
   }
 };
 
 
@@ -102,81 +102,13 @@ void DistDialect::PrintAttribute(pir::Attribute attr, std::ostream &os) const {
               phi::distributed::auto_parallel::str_join(
                   op_dist_attr.process_mesh_attr().process_ids()) +
               "]}";
-    auto num_operand_dist_attrs = op_dist_attr.num_operands();
-    for (uint32_t i = 0; i < num_operand_dist_attrs; ++i) {
-      auto dist_attr = op_dist_attr.operand_dist_attr(i);
-      os << ",operand(" + std::to_string(i) + "):{";
-      if (!dist_attr) {
-        os << "null}";
-        continue;
-      }
-      if (dist_attr.process_mesh_attr() != op_dist_attr.process_mesh_attr()) {
-        os << "mesh_shape:[" +
-                  phi::distributed::auto_parallel::str_join(
-                      dist_attr.process_mesh_attr().shape()) +
-                  "],";
-        os << "process_ids:[" +
-                  phi::distributed::auto_parallel::str_join(
-                      dist_attr.process_mesh_attr().process_ids()) +
-                  "],";
-      }
-      os << "dims_maping:[" +
-                phi::distributed::auto_parallel::str_join(
-                    dist_attr.dims_mapping()) +
-                "]";
-      if (dist_attr.partial_status().size() > 0) {
-        std::vector<std::string> partial_status_strs;
-        for (auto &itr : dist_attr.partial_status()) {
-          std::string s = "partial(" + std::to_string(itr.first) + "," +
-                          phi::ReduceTypeStrings[static_cast<int>(itr.second)] +
-                          ")";
-          partial_status_strs.emplace_back(s);
-        }
-        os << "," +
-                  phi::distributed::auto_parallel::str_join(
-                      partial_status_strs) +
-                  "}";
-      } else {
-        os << "}";
-      }
+    for (uint32_t i = 0; i < op_dist_attr.num_operands(); ++i) {
+      os << ",operand(" + std::to_string(i) + "):{" << op_dist_attr.operand(i)
+         << "}";
     }
-    auto num_result_dist_attrs = op_dist_attr.num_results();
-    for (uint32_t i = 0; i < num_result_dist_attrs; ++i) {
-      auto dist_attr = op_dist_attr.result_dist_attr(i);
-      os << ",result(" + std::to_string(i) + "):{";
-      if (!dist_attr) {
-        os << "null}";
-        continue;
-      }
-      if (dist_attr.process_mesh_attr() != op_dist_attr.process_mesh_attr()) {
-        os << "mesh_shape:[" +
-                  phi::distributed::auto_parallel::str_join(
-                      dist_attr.process_mesh_attr().shape()) +
-                  "],";
-        os << "process_ids:[" +
-                  phi::distributed::auto_parallel::str_join(
-                      dist_attr.process_mesh_attr().process_ids()) +
-                  "],";
-      }
-      os << "dims_maping:[" +
-                phi::distributed::auto_parallel::str_join(
-                    dist_attr.dims_mapping()) +
-                "]";
-      if (dist_attr.partial_status().size() > 0) {
-        std::vector<std::string> partial_status_strs;
-        for (auto &itr : dist_attr.partial_status()) {
-          std::string s = "partial(" + std::to_string(itr.first) + "," +
-                          phi::ReduceTypeStrings[static_cast<int>(itr.second)] +
-                          ")";
-          partial_status_strs.emplace_back(s);
-        }
-        os << "," +
-                  phi::distributed::auto_parallel::str_join(
-                      partial_status_strs) +
-                  "}";
-      } else {
-        os << "}";
-      }
+    for (uint32_t i = 0; i < op_dist_attr.num_results(); ++i) {
+      os << ",result(" + std::to_string(i) + "):{" << op_dist_attr.result(i)
+         << "}";
     }
     os << "}";
   } else {
 
@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include "paddle/fluid/pir/dialect/distributed/ir/dist_tools.h"
+#include "glog/logging.h"
 #include "paddle/common/enforce.h"
 #include "paddle/pir/include/core/operation.h"
 
@@ -76,14 +77,66 @@ phi::distributed::DistMetaTensor CvtToDistMetaTensor(DistDenseTensorType type) {
   return phi::distributed::DistMetaTensor(type.global_ddim(), phi_attr);
 }
 
-TensorDistAttribute CvtToPirDistAttr(
-    const phi::distributed::ArgDistAttr& dist_attr) {
-  auto& attr = PADDLE_GET_CONST(phi::distributed::TensorDistAttr, dist_attr);
-  if (attr.process_mesh().empty()) return nullptr;
-  return TensorDistAttribute::get(pir::IrContext::Instance(),
-                                  attr.process_mesh(),
-                                  attr.dims_mapping(),
-                                  attr.partial_status());
+pir::Attribute CvtToPirAttr(const phi::distributed::ArgDistAttr& dist_attr) {
+  auto ctx = pir::IrContext::Instance();
+  if (holds_alternative<phi::distributed::TensorDistAttr>(dist_attr)) {
+    auto& attr = PADDLE_GET_CONST(phi::distributed::TensorDistAttr, dist_attr);
+    return TensorDistAttribute::get(
+        ctx, attr.process_mesh(), attr.dims_mapping(), attr.partial_status());
+  } else {
+    auto& vec = PADDLE_GET_CONST(std::vector<phi::distributed::TensorDistAttr>,
+                                 dist_attr);
+    std::vector<pir::Attribute> array;
+    for (auto& attr : vec) {
+      array.push_back(TensorDistAttribute::get(ctx,
+                                               attr.process_mesh(),
+                                               attr.dims_mapping(),
+                                               attr.partial_status()));
+    }
+    return pir::ArrayAttribute::get(ctx, array);
+  }
+}
+
+pir::Type CvtToPirDistType(pir::Type prim_type, pir::Attribute dist_attr) {
+  if (!prim_type) return nullptr;
+  auto ctx = pir::IrContext::Instance();
+  if (auto dense_tensor_type = prim_type.dyn_cast<pir::DenseTensorType>()) {
+    auto tensor_dist_attr = dist_attr.dyn_cast<TensorDistAttribute>();
+    if (!tensor_dist_attr) {
+      VLOG(0) << "Convert dense tensor type to dist type with attribute {"
+              << dist_attr << "}.";
+      PADDLE_THROW(common::errors::InvalidArgument(
+          "Only allowed convert a densor tensor type to dist dense tensor type "
+          "with non-empty TensorDistAttr"));
+    }
+    return DistDenseTensorType::get(ctx, dense_tensor_type, tensor_dist_attr);
+  } else if (auto vec_type = prim_type.dyn_cast<pir::VectorType>()) {
+    auto array_attr = dist_attr.dyn_cast<pir::ArrayAttribute>();
+    if (!array_attr) {
+      VLOG(0) << "Convert vector type to dist type with attribute {"
+              << dist_attr << "}.";
+      PADDLE_THROW(common::errors::InvalidArgument(
+          "Only allowed convert a vector type to dist vector type with "
+          "non-empty ArrayAttr"));
+    }
+    PADDLE_ENFORCE_EQ(
+        vec_type.size(),
+        array_attr.size(),
+        common::errors::InvalidArgument(
+            "The vector type size must equal to array attribute size."));
+    std::vector<pir::Type> dist_vec_type;
+    for (size_t idx = 0; idx < vec_type.size(); ++idx) {
+      dist_vec_type.push_back(CvtToPirDistType(vec_type[idx], array_attr[idx]));
+    }
+    return pir::VectorType::get(ctx, dist_vec_type);
+  } else {
+    VLOG(0) << "Convert type{" << prim_type << "} to dist type with attribute {"
+            << dist_attr << "}.";
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "Currently only support convert dense_tensor_type r vector type to "
+        "dist."));
+  }
+  return nullptr;
 }
 
 void CopyLeafOpToMesh(pir::Value value, ProcessMeshAttribute mesh_attr) {
 
@@ -28,8 +28,8 @@ void CvtAllInputsToDist(const std::vector<pir::Value>& inputs,
                         ProcessMeshAttribute mesh_attr);
 
 phi::distributed::DistMetaTensor CvtToDistMetaTensor(DistDenseTensorType type);
-TensorDistAttribute CvtToPirDistAttr(
-    const phi::distributed::ArgDistAttr& dist_attr);
+pir::Attribute CvtToPirAttr(const phi::distributed::ArgDistAttr& dist_attr);
+pir::Type CvtToPirDistType(pir::Type prim_type, pir::Attribute dist_attr);
 
 ///
 /// When the following conditions are met:
 
@@ -664,7 +664,7 @@ def GenDistBranch(args, op_info):
     PADDLE_ENFORCE_EQ(spmd_info.first.size(), {input_size}u, common::errors::Unavailable(
         "Size of spmd_info.first for op[{op_name}]is unexpected."));
     for(auto& arg_dist : spmd_info.first) {{
-        dist_operand_attrs.push_back(CvtToPirDistAttr(arg_dist));
+        dist_operand_attrs.push_back(CvtToPirAttr(arg_dist));
     }}
 """
     dist_branch_str += TEMPLATE.format(
@@ -691,23 +691,12 @@ def GenDistBranch(args, op_info):
         )
 
     for idx, output_name in enumerate(op_info.output_name_list):
-        # is a vector<Tensor>
-        if 'pir::VectorType' in op_info.output_type_list[idx]:
-            TEMPLATE = """
-    auto dist_attr_{name} = CvtToPirDistAttr(spmd_info.second[{idx}]);
-    dist_result_attrs.push_back(dist_attr_{name});
-    argument_outputs.push_back(DistDenseTensorType::get(ctx, {name}_type.dyn_cast<pir::DenseTensorType>(), dist_attr_{name}));
-"""
-            # Todo: support vector<Tensor> case
-            dist_branch_str += ""
-        # is a Tensor
-        else:
-            TEMPLATE = """
-    auto dist_attr_{name} = CvtToPirDistAttr(spmd_info.second[{idx}]);
+        TEMPLATE = """
+    auto dist_attr_{name} = CvtToPirAttr(spmd_info.second[{idx}]);
     dist_result_attrs.push_back(dist_attr_{name});
-    argument_outputs.push_back(DistDenseTensorType::get(ctx, {name}_type.dyn_cast<pir::DenseTensorType>(), dist_attr_{name}));
+    argument_outputs.push_back(CvtToPirDistType({name}_type, dist_attr_{name}));
 """
-            dist_branch_str += TEMPLATE.format(idx=idx, name=output_name)
+        dist_branch_str += TEMPLATE.format(idx=idx, name=output_name)
     TEMPLATE = """
     attributes[kAttrOpDistAttr] = OperationDistAttribute::get(
         ctx,