PaddlePaddle
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/attribute_storage.h
Copy file name to clipboard
+8-9Lines changed: 8 additions & 9 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/attribute_storage.h
Copy file name to clipboard
+8-9Lines changed: 8 additions & 9 deletions
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_attribute.cc
Copy file name to clipboard
+44-24Lines changed: 44 additions & 24 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/dist_attribute.cc
Copy file name to clipboard
+44-24Lines changed: 44 additions & 24 deletions
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_attribute.h
Copy file name to clipboard
+15-15Lines changed: 15 additions & 15 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/dist_attribute.h
Copy file name to clipboard
+15-15Lines changed: 15 additions & 15 deletions
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_dialect.cc
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/dist_dialect.cc
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_interface.h
Copy file name to clipboard
+22-7Lines changed: 22 additions & 7 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/dist_interface.h
Copy file name to clipboard
+22-7Lines changed: 22 additions & 7 deletions
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_op.cc
Copy file name to clipboard
+27-31Lines changed: 27 additions & 31 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/dist_op.cc
Copy file name to clipboard
+27-31Lines changed: 27 additions & 31 deletions
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_type.h
Copy file name to clipboard
+5-2Lines changed: 5 additions & 2 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/dist_type.h
Copy file name to clipboard
+5-2Lines changed: 5 additions & 2 deletions
@@ -121,12 +121,12 @@ class OperationDistAttrStorage : public pir::AttributeStorage {
   /// \brief Declare ParamKey according to parameter type.
   ///
   using ParamKey = std::tuple<ProcessMeshAttribute,
-                              std::vector<TensorDistAttribute>,
-                              std::vector<TensorDistAttribute>>;
+                              std::vector<pir::Attribute>,
+                              std::vector<pir::Attribute>>;
   OperationDistAttrStorage(ParamKey&& param)  // NOLINT
       : mesh_attr(std::get<0>(param)),
-        operand_dist_attrs(std::get<1>(param)),
-        result_dist_attrs(std::get<2>(param)) {}
+        operand_attrs(std::get<1>(param)),
+        result_attrs(std::get<2>(param)) {}
 
   ///
   /// \brief Each derived TypeStorage must define a Construct method, which
@@ -156,14 +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_dist_attrs == std::get<1>(key) &&
-           result_dist_attrs == std::get<2>(key);
+    return mesh_attr == std::get<0>(key) && operand_attrs == std::get<1>(key) &&
+           result_attrs == std::get<2>(key);
   }
 
   ProcessMeshAttribute mesh_attr;
-  std::vector<TensorDistAttribute> operand_dist_attrs;
-  std::vector<TensorDistAttribute> result_dist_attrs;
+  std::vector<pir::Attribute> operand_attrs;
+  std::vector<pir::Attribute> result_attrs;
 };
 
 }  // namespace dialect
 
@@ -78,48 +78,68 @@ TensorDistAttribute TensorDistAttribute::get(
 ProcessMeshAttribute OperationDistAttribute::process_mesh_attr() const {
   return storage()->mesh_attr;
 }
-const std::vector<TensorDistAttribute>&
-OperationDistAttribute::operand_dist_attrs() const {
-  return storage()->operand_dist_attrs;
+const std::vector<pir::Attribute>& OperationDistAttribute::operand_attrs()
+    const {
+  return storage()->operand_attrs;
 }
 TensorDistAttribute OperationDistAttribute::operand_dist_attr(
     uint32_t index) const {
-  return operand_dist_attrs().at(index);
+  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>();
 }
-uint32_t OperationDistAttribute::num_operand_dist_attrs() const {
-  return operand_dist_attrs().size();
+
+uint32_t OperationDistAttribute::num_operands() const {
+  return operand_attrs().size();
 }
 
-const std::vector<TensorDistAttribute>&
-OperationDistAttribute::result_dist_attrs() const {
-  return storage()->result_dist_attrs;
+const std::vector<pir::Attribute>& OperationDistAttribute::result_attrs()
+    const {
+  return storage()->result_attrs;
 }
 TensorDistAttribute OperationDistAttribute::result_dist_attr(
     uint32_t index) const {
-  return result_dist_attrs().at(index);
+  return result_attrs().at(index).dyn_cast<TensorDistAttribute>();
 }
-uint32_t OperationDistAttribute::num_result_dist_attrs() const {
-  return result_dist_attrs().size();
+
+pir::ArrayAttribute OperationDistAttribute::result_array_attr(
+    uint32_t index) const {
+  return result_attrs().at(index).dyn_cast<pir::ArrayAttribute>();
 }
+
+uint32_t OperationDistAttribute::num_results() const {
+  return result_attrs().size();
+}
+
 OperationDistAttribute OperationDistAttribute::get(
     pir::IrContext* ctx,
     ProcessMeshAttribute mesh,
-    const std::vector<TensorDistAttribute>& operand_dist_attrs,
-    const std::vector<TensorDistAttribute>& result_dist_attrs) {
-  for (const auto& iter : operand_dist_attrs) {
+    const std::vector<pir::Attribute>& operand_attrs,
+    const std::vector<pir::Attribute>& result_attrs) {
+  auto check_dist_attr = [=](pir::Attribute attr) {
+    auto dist_attr = attr.dyn_cast<TensorDistAttribute>();
+    PADDLE_ENFORCE_EQ(mesh,
+                      dist_attr.process_mesh_attr(),
+                      common::errors::PreconditionNotMet(
+                          "operand_dist_attrs element's mesh(%s) not equal "
+                          "to input mesh(%s)"));
+  };
+  for (auto attr : operand_attrs) {
     // NOTE: The operand dist attr maybe empty while the corresponding input is
     // optional.
-    if (iter) {
-      PADDLE_ENFORCE_EQ(mesh,
-                        iter.process_mesh_attr(),
-                        common::errors::PreconditionNotMet(
-                            "operand_dist_attrs element's mesh(%s) not equal "
-                            "to input mesh(%s)",
-                            iter.process_mesh_attr(),
-                            mesh));
+    if (!attr) continue;
+    if (auto array_attr = attr.dyn_cast<pir::ArrayAttribute>()) {
+      for (size_t i = 0; i < array_attr.size(); ++i) {
+        check_dist_attr(array_attr[i]);
+      }
+    } else {
+      check_dist_attr(attr);
     }
   }
-  return Base::get(ctx, mesh, operand_dist_attrs, result_dist_attrs);
+  return Base::get(ctx, mesh, operand_attrs, result_attrs);
 }
 
 }  // namespace dialect
 
@@ -17,6 +17,7 @@
 #include "paddle/phi/common/reduce_type.h"
 #include "paddle/phi/core/distributed/auto_parallel/process_mesh.h"
 #include "paddle/pir/include/core/attribute.h"
+#include "paddle/pir/include/core/builtin_attribute.h"
 #include "paddle/pir/include/core/builtin_attribute_storage.h"
 #include "paddle/pir/include/core/utils.h"
 #include "paddle/utils/flat_hash_map.h"
@@ -104,29 +105,28 @@ class OperationDistAttribute : public pir::AttrBase<OperationDistAttribute,
   using Base::Base;
   ProcessMeshAttribute process_mesh_attr() const;
 
-  const std::vector<TensorDistAttribute>& operand_dist_attrs() const;
+  const std::vector<Attribute>& operand_attrs() const;
   TensorDistAttribute operand_dist_attr(uint32_t index) const;
-  uint32_t num_operand_dist_attrs() const;
+  pir::ArrayAttribute operand_array_attr(uint32_t index) const;
+  uint32_t num_operands() const;
 
-  const std::vector<TensorDistAttribute>& result_dist_attrs() const;
+  const std::vector<Attribute>& result_attrs() const;
   TensorDistAttribute result_dist_attr(uint32_t index) const;
-  uint32_t num_result_dist_attrs() const;
+  pir::ArrayAttribute result_array_attr(uint32_t index) const;
+  uint32_t num_results() const;
 
-  static OperationDistAttribute get(
-      pir::IrContext* ctx,
-      ProcessMeshAttribute mesh,
-      const std::vector<TensorDistAttribute>& operand_dist_attrs,
-      const std::vector<TensorDistAttribute>& result_dist_attrs);
+  static OperationDistAttribute get(pir::IrContext* ctx,
+                                    ProcessMeshAttribute mesh,
+                                    const std::vector<Attribute>& operand_attrs,
+                                    const std::vector<Attribute>& result_attrs);
 
   static OperationDistAttribute get(
       pir::IrContext* ctx,
       const phi::distributed::ProcessMesh& mesh,
-      const std::vector<TensorDistAttribute>& operand_dist_attrs,
-      const std::vector<TensorDistAttribute>& result_dist_attrs) {
-    return get(ctx,
-               ProcessMeshAttribute::get(ctx, mesh),
-               operand_dist_attrs,
-               result_dist_attrs);
+      const std::vector<Attribute>& operand_attrs,
+      const std::vector<Attribute>& result_attrs) {
+    return get(
+        ctx, ProcessMeshAttribute::get(ctx, mesh), operand_attrs, result_attrs);
   }
 };
 
 
@@ -98,7 +98,7 @@ 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_operand_dist_attrs();
+    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) + "):{";
@@ -132,7 +132,7 @@ void DistDialect::PrintAttribute(pir::Attribute attr, std::ostream &os) const {
         os << "}";
       }
     }
-    auto num_result_dist_attrs = op_dist_attr.num_result_dist_attrs();
+    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) + "):{";
 
@@ -26,19 +26,24 @@ class IR_API DistTypeInterface
  public:
   struct Concept {
     /// Defined these methods with the interface.
-    explicit Concept(pir::Type (*local_type)(pir::Type),
-                     ProcessMeshAttribute (*process_mesh_attr)(pir::Type),
-                     TensorDistAttribute (*tensor_dist_attr)(pir::Type),
-                     pir::Type (*copy_with_new_mesh)(pir::Type,
-                                                     ProcessMeshAttribute mesh))
+    explicit Concept(
+        pir::Type (*local_type)(pir::Type),
+        ProcessMeshAttribute (*process_mesh_attr)(pir::Type),
+        TensorDistAttribute (*tensor_dist_attr)(pir::Type),
+        pir::Type (*copy_with_new_mesh)(pir::Type, ProcessMeshAttribute mesh),
+        pir::Type (*copy_with_new_dist_attr)(pir::Type,
+                                             TensorDistAttribute dist_attr))
         : local_type(local_type),
           process_mesh_attr(process_mesh_attr),
           tensor_dist_attr(tensor_dist_attr),
-          copy_with_new_mesh(copy_with_new_mesh) {}
+          copy_with_new_mesh(copy_with_new_mesh),
+          copy_with_new_dist_attr(copy_with_new_dist_attr) {}
     pir::Type (*local_type)(pir::Type);
     ProcessMeshAttribute (*process_mesh_attr)(pir::Type);
     TensorDistAttribute (*tensor_dist_attr)(pir::Type);
     pir::Type (*copy_with_new_mesh)(pir::Type, ProcessMeshAttribute mesh);
+    pir::Type (*copy_with_new_dist_attr)(pir::Type,
+                                         TensorDistAttribute dist_attr);
   };
 
   template <class ConcreteType>
@@ -58,11 +63,16 @@ class IR_API DistTypeInterface
       return pir::cast<ConcreteType>(type).CopyWithNewMesh(mesh);
     }
 
+    static Type CopyWithNewDistAttr(Type type, TensorDistAttribute dist_attr) {
+      return pir::cast<ConcreteType>(type).CopyWithNewDistAttr(dist_attr);
+    }
+
     Model()
         : Concept(local_type,
                   process_mesh_attr,
                   tensor_dist_attr,
-                  CopyWithNewMesh) {}
+                  CopyWithNewMesh,
+                  CopyWithNewDistAttr) {}
   };
 
   DistTypeInterface(pir::Type type, Concept *impl)
@@ -82,6 +92,11 @@ class IR_API DistTypeInterface
     return DistTypeInterface(impl_->copy_with_new_mesh(*this, mesh), impl_);
   }
 
+  DistTypeInterface CopyWithNewDistAttr(TensorDistAttribute dist_attr) {
+    return DistTypeInterface(impl_->copy_with_new_dist_attr(*this, dist_attr),
+                             impl_);
+  }
+
  private:
   Concept *impl_;
 };
 
@@ -37,11 +37,10 @@ void ShardTensorOp::VerifySig() {
   VLOG(4) << "Verifying inputs:";
   {
     auto input_size = num_operands();
-    PADDLE_ENFORCE_EQ(
-        input_size,
-        1u,
-        common::errors::PreconditionNotMet(
-            "The size %d of inputs must be equal to 1.", input_size));
+    PADDLE_ENFORCE_EQ(input_size,
+                      1u,
+                      common::errors::PreconditionNotMet(
+                          "The size of inputs must be equal to 1."));
     PADDLE_ENFORCE_EQ((*this)
                           ->operand_source(0)
                           .type()
@@ -80,19 +79,18 @@ void ShardTensorOp::VerifySig() {
     auto op_dist_attr =
         this->attribute<paddle::dialect::OperationDistAttribute>(
             "op_dist_attr");
-    PADDLE_ENFORCE_EQ(op_dist_attr.num_operand_dist_attrs(),
+    PADDLE_ENFORCE_EQ(op_dist_attr.num_operands(),
                       0u,
-                      common::errors::PreconditionNotMet(
-                          "The op_dist_attr input size %d must be equal to 0.",
-                          op_dist_attr.num_operand_dist_attrs()));
+                      phi::errors::PreconditionNotMet(
+                          "The op_dist_attr input size must be equal to 0."));
 
-    PADDLE_ENFORCE_EQ(op_dist_attr.num_result_dist_attrs(),
-                      num_results(),
-                      common::errors::PreconditionNotMet(
-                          "The op_dist_attr output size %d must "
-                          "be equal to op output size %d.",
-                          op_dist_attr.num_result_dist_attrs(),
-                          num_results()));
+    PADDLE_ENFORCE_EQ(
+        op_dist_attr.num_results(),
+        num_results(),
+        phi::errors::PreconditionNotMet("The op_dist_attr output size %d must "
+                                        "be equal to op output size %d.",
+                                        op_dist_attr.num_results(),
+                                        num_results()));
   }
   VLOG(4) << "End Verifying for: ShardTensorOp.";
 }
@@ -137,8 +135,8 @@ void ShardTensorOp::Build(pir::Builder& builder,
   pir::Attribute op_dist_attr = OperationDistAttribute::get(
       pir::IrContext::Instance(),
       process_mesh_attr,
-      std::vector<TensorDistAttribute>(),
-      std::vector<TensorDistAttribute>{tensor_dist_attr});
+      std::vector<pir::Attribute>(),
+      std::vector<pir::Attribute>{tensor_dist_attr});
   argument.AddAttribute("op_dist_attr", op_dist_attr);
 
   VLOG(4) << "Builder construction outputs";
@@ -254,19 +252,17 @@ void ReshardOp::VerifySig() {
     auto op_dist_attr =
         this->attribute<paddle::dialect::OperationDistAttribute>(
             "op_dist_attr");
-    PADDLE_ENFORCE_EQ(op_dist_attr.num_operand_dist_attrs(),
-                      1u,
-                      common::errors::PreconditionNotMet(
-                          "The op_dist_attr input size %d must be equal to 1.",
-                          op_dist_attr.num_operand_dist_attrs()));
+    PADDLE_ENFORCE_EQ(
+        op_dist_attr.num_operands(),
+        1u,
+        common::errors::PreconditionNotMet(
+            "The op_dist_attr input size of reshard op must be equal to 1."));
 
-    PADDLE_ENFORCE_EQ(op_dist_attr.num_result_dist_attrs(),
+    PADDLE_ENFORCE_EQ(op_dist_attr.num_results(),
                       num_results(),
-                      common::errors::PreconditionNotMet(
-                          "The op_dist_attr output size %d must "
-                          "be equal to op output size %d.",
-                          op_dist_attr.num_result_dist_attrs(),
-                          num_results()));
+                      phi::errors::PreconditionNotMet(
+                          "The op_dist_attr output size of reshard op must be "
+                          "equal to op output size."));
   }
   VLOG(4) << "End Verifying for: ShardTensorOp.";
 }
@@ -293,8 +289,8 @@ void ReshardOp::Build(pir::Builder& builder,
   pir::Attribute op_dist_attr = OperationDistAttribute::get(
       pir::IrContext::Instance(),
       input_tensor_type.tensor_dist_attr().process_mesh_attr(),
-      std::vector<TensorDistAttribute>{input_tensor_type.tensor_dist_attr()},
-      std::vector<TensorDistAttribute>{tensor_dist_attr});
+      std::vector<pir::Attribute>{input_tensor_type.tensor_dist_attr()},
+      std::vector<pir::Attribute>{tensor_dist_attr});
   argument.AddAttribute("op_dist_attr", op_dist_attr);
 
   VLOG(4) << "Builder construction outputs";
 
@@ -64,8 +64,11 @@ class DistDenseTensorType
   DistDenseTensorType CopyWithNewMesh(ProcessMeshAttribute mesh) {
     return get(ir_context(),
                dense_tensor_type(),
-               tensor_dist_attr().CopyWithNewMesh(mesh),
-               local_ddim());
+               tensor_dist_attr().CopyWithNewMesh(mesh));
+  }
+
+  DistDenseTensorType CopyWithNewDistAttr(TensorDistAttribute dist_attr) {
+    return get(ir_context(), dense_tensor_type(), dist_attr);
   }
 
   static DistDenseTensorType get(pir::IrContext* ctx,