[DistDialect] add reshard op and api (#62718)

wentaoyu · web-flow · commit 0718ae37a9af · 2024-03-19T17:32:23.000+08:00
* add reshard op and api

* update

* fix bug

* update ut

* update ut and check logic

* fix by comments

* fix code style

* fix dist_attr print format, local_shape compute

* update

* fix PADDLE_ENFORCE usage

* fix code style

* fix code style
diff --git a/paddle/fluid/pir/dialect/distributed/ir/dist_api.cc b/paddle/fluid/pir/dialect/distributed/ir/dist_api.cc
@@ -45,5 +45,19 @@ pir::Value shard_tensor(const pir::Value& x,
   return shard_tensor_op.out();
 }
 
+pir::Value reshard(const pir::Value& x,
+                   const phi::distributed::ProcessMesh& process_mesh,
+                   const std::vector<int64_t>& dims_mapping) {
+  pir::IrContext* ctx = pir::IrContext::Instance();
+  // TODO(ywt01) get partial_status by func parameter
+  paddle::flat_hash_map<int64_t, phi::ReduceType> partial_status;
+  TensorDistAttribute tensor_dist_attr =
+      TensorDistAttribute::get(ctx, process_mesh, dims_mapping, partial_status);
+
+  auto reshard_op = ApiBuilder::Instance().GetBuilder()->Build<ReShardOp>(
+      x, tensor_dist_attr);
+  return reshard_op.result(0);
+}
+
 }  // namespace dialect
 }  // namespace paddle
diff --git a/paddle/fluid/pir/dialect/distributed/ir/dist_api.h b/paddle/fluid/pir/dialect/distributed/ir/dist_api.h
@@ -27,5 +27,9 @@ namespace dialect {
 pir::Value shard_tensor(const pir::Value& x,
                         const phi::distributed::ProcessMesh& process_mesh,
                         const std::vector<int64_t>& dims_mapping);
+
+pir::Value reshard(const pir::Value& x,
+                   const phi::distributed::ProcessMesh& process_mesh,
+                   const std::vector<int64_t>& dims_mapping);
 }  // namespace dialect
 }  // namespace paddle
diff --git a/paddle/fluid/pir/dialect/distributed/ir/dist_attribute.cc b/paddle/fluid/pir/dialect/distributed/ir/dist_attribute.cc
@@ -111,15 +111,6 @@ OperationDistAttribute OperationDistAttribute::get(
             iter.process_mesh_attr(),
             mesh));
   }
-  for (const auto& iter : result_dist_attrs) {
-    PADDLE_ENFORCE_EQ(
-        mesh,
-        iter.process_mesh_attr(),
-        phi::errors::PreconditionNotMet(
-            "operand_dist_attrs element's mesh(%s) not equal to input mesh(%s)",
-            iter.process_mesh_attr(),
-            mesh));
-  }
   return Base::get(ctx, mesh, operand_dist_attrs, result_dist_attrs);
 }
 
diff --git a/paddle/fluid/pir/dialect/distributed/ir/dist_dialect.cc b/paddle/fluid/pir/dialect/distributed/ir/dist_dialect.cc
@@ -35,7 +35,7 @@ void DistDialect::initialize() {
                      TensorDistAttribute,
                      OperationDistAttribute>();
   RegisterTypes<DistDenseTensorType>();
-  RegisterOps<ShardTensorOp>();
+  RegisterOps<ShardTensorOp, ReShardOp>();
 }
 
 void DistDialect::PrintType(pir::Type type, std::ostream &os) const {
@@ -70,7 +70,6 @@ void DistDialect::PrintAttribute(pir::Attribute attr, std::ostream &os) const {
                   process_mesh_attr.process_ids()) +
               "]";
   } else if (auto tensor_dist_attr = attr.dyn_cast<TensorDistAttribute>()) {
-    // Todo: Design the tensor dist attr print format.
     os << "mesh_shape:[" +
               phi::distributed::auto_parallel::str_join(
                   tensor_dist_attr.process_mesh_attr().shape()) +
@@ -91,14 +90,14 @@ void DistDialect::PrintAttribute(pir::Attribute attr, std::ostream &os) const {
          << phi::distributed::auto_parallel::str_join(partial_status_strs);
     }
   } else if (auto op_dist_attr = attr.dyn_cast<OperationDistAttribute>()) {
-    os << "mesh_shape:[" +
+    os << "{mesh:{shape:[" +
               phi::distributed::auto_parallel::str_join(
                   op_dist_attr.process_mesh_attr().shape()) +
               "]";
     os << ",process_ids:[" +
               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();
     for (uint32_t i = 0; i < num_operand_dist_attrs; ++i) {
       auto dist_attr = op_dist_attr.operand_dist_attr(i);
@@ -159,6 +158,7 @@ void DistDialect::PrintAttribute(pir::Attribute attr, std::ostream &os) const {
         os << "}";
       }
     }
+    os << "}";
   } else {
     os << "error_attribute_type";
   }
diff --git a/paddle/fluid/pir/dialect/distributed/ir/dist_op.cc b/paddle/fluid/pir/dialect/distributed/ir/dist_op.cc
@@ -27,6 +27,7 @@ namespace paddle {
 namespace dialect {
 
 const char* ShardTensorOp::attributes_name[1] = {"op_dist_attr"};
+const char* ReShardOp::attributes_name[1] = {"op_dist_attr"};
 
 void ShardTensorOp::VerifySig() {
   VLOG(4)
@@ -37,35 +38,38 @@ void ShardTensorOp::VerifySig() {
     PADDLE_ENFORCE_EQ(
         input_size,
         1u,
-        phi::errors::PreconditionNotMet(
+        common::errors::PreconditionNotMet(
             "The size %d of inputs must be equal to 1.", input_size));
-    PADDLE_ENFORCE((*this)
-                       ->operand_source(0)
-                       .type()
-                       .isa<paddle::dialect::DenseTensorType>(),
-                   phi::errors::PreconditionNotMet(
-                       "Type validation failed for the 0th input."));
+    PADDLE_ENFORCE_EQ((*this)
+                          ->operand_source(0)
+                          .type()
+                          .isa<paddle::dialect::DenseTensorType>(),
+                      true,
+                      common::errors::PreconditionNotMet(
+                          "Type validation failed for the 0th input."));
   }
   VLOG(4) << "Verifying attributes:";
   {
     auto& attributes = this->attributes();
-    PADDLE_ENFORCE(attributes.count("op_dist_attr") > 0 &&
+    PADDLE_ENFORCE_EQ((attributes.count("op_dist_attr") > 0 &&
                        attributes.at("op_dist_attr")
-                           .isa<paddle::dialect::OperationDistAttribute>(),
-                   phi::errors::PreconditionNotMet(
-                       "Type of attribute: op_dist_attr is not right."));
+                           .isa<paddle::dialect::OperationDistAttribute>()),
+                      true,
+                      common::errors::PreconditionNotMet(
+                          "Type of attribute: op_dist_attr is not right."));
   }
   VLOG(4) << "Verifying outputs:";
   {
     auto output_size = num_results();
     PADDLE_ENFORCE_EQ(
         output_size,
         1u,
-        phi::errors::PreconditionNotMet(
+        common::errors::PreconditionNotMet(
             "The size %d of outputs must be equal to 1.", output_size));
-    PADDLE_ENFORCE(
+    PADDLE_ENFORCE_EQ(
         (*this)->result(0).type().isa<paddle::dialect::DistDenseTensorType>(),
-        phi::errors::PreconditionNotMet(
+        true,
+        common::errors::PreconditionNotMet(
             "Type validation failed for the 0th output."));
   }
 
@@ -76,17 +80,17 @@ void ShardTensorOp::VerifySig() {
             "op_dist_attr");
     PADDLE_ENFORCE_EQ(op_dist_attr.num_operand_dist_attrs(),
                       0u,
-                      phi::errors::PreconditionNotMet(
+                      common::errors::PreconditionNotMet(
                           "The op_dist_attr input size %d must be equal to 0.",
                           op_dist_attr.num_operand_dist_attrs()));
 
-    PADDLE_ENFORCE_EQ(
-        op_dist_attr.num_result_dist_attrs(),
-        num_results(),
-        phi::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_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()));
   }
   VLOG(4) << "End Verifying for: ShardTensorOp.";
 }
@@ -101,20 +105,22 @@ void ShardTensorOp::Build(pir::Builder& builder,
   PADDLE_ENFORCE_EQ(
       input.use_empty(),
       true,
-      phi::errors::PreconditionNotMet("'input' use_empty is not true"));
+      common::errors::PreconditionNotMet("'input' use_empty is not true"));
 
   paddle::dialect::DenseTensorType input_tensor_type;
   if (input.type().isa<paddle::dialect::DenseTensorType>()) {
     input_tensor_type =
         input.type().dyn_cast<paddle::dialect::DenseTensorType>();
   } else {
-    PADDLE_THROW(phi::errors::Unimplemented(
+    PADDLE_THROW(common::errors::Unimplemented(
         "Only support paddle::dialect::DenseTensorType"));
   }
 
-  PADDLE_ENFORCE(attributes.find("tensor_dist_attr") != attributes.end(),
-                 phi::errors::NotFound(
-                     "'tensor_dist_attr' Attribute is expected for ShardOp"));
+  PADDLE_ENFORCE_NE(
+      attributes.find("tensor_dist_attr"),
+      attributes.end(),
+      common::errors::NotFound(
+          "'tensor_dist_attr' Attribute is expected for ShardOp"));
   paddle::dialect::TensorDistAttribute tensor_dist_attr =
       attributes.at("tensor_dist_attr")
           .dyn_cast<paddle::dialect::TensorDistAttribute>();
@@ -136,36 +142,136 @@ void ShardTensorOp::Build(pir::Builder& builder,
   VLOG(4) << "Builder construction outputs";
   auto global_dims = input_tensor_type.dims();
   auto process_mesh_shape = process_mesh_attr.shape();
-  PADDLE_ENFORCE(static_cast<int>(dims_mapping.size()) == global_dims.size(),
-                 phi::errors::PreconditionNotMet(
-                     "dims_mapping size %d does not match input size %d",
-                     dims_mapping.size(),
-                     global_dims.size()));
-  std::vector<int> local_shape(global_dims.size());
-  for (int i = 0; i < global_dims.size(); ++i) {
-    if (dims_mapping[i] == -1) {
-      local_shape[i] = global_dims[i];
-    } else {
-      auto shard_size = process_mesh_shape[dims_mapping[i]];
-      PADDLE_ENFORCE(
-          global_dims[i] % shard_size == 0,
-          phi::errors::PreconditionNotMet(
-              "global_dims size %d can't be evenly divided by shard_size %d",
-              global_dims[i],
-              shard_size));
-      local_shape[i] = global_dims[i] / shard_size;
-    }
-  }
-
+  PADDLE_ENFORCE_EQ(static_cast<int>(dims_mapping.size()),
+                    global_dims.size(),
+                    common::errors::PreconditionNotMet(
+                        "dims_mapping size %d does not match input size %d",
+                        dims_mapping.size(),
+                        global_dims.size()));
+  auto local_shape = InferLocalDDim(global_dims, tensor_dist_attr);
   pir::Type out_dist_tensor_type =
       paddle::dialect::DistDenseTensorType::get(pir::IrContext::Instance(),
                                                 input_tensor_type,
                                                 tensor_dist_attr,
-                                                phi::make_ddim(local_shape));
+                                                local_shape);
+  argument.AddOutput(out_dist_tensor_type);
+}
+
+void ReShardOp::VerifySig() {
+  VLOG(4) << "Start Verifying inputs, outputs and attributes for: ReShardOp.";
+  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((*this)
+                          ->operand_source(0)
+                          .type()
+                          .isa<paddle::dialect::DistDenseTensorType>(),
+                      true,
+                      common::errors::PreconditionNotMet(
+                          "Type validation failed for the 0th input."));
+  }
+  VLOG(4) << "Verifying attributes:";
+  {
+    auto& attributes = this->attributes();
+    PADDLE_ENFORCE_EQ((attributes.count("op_dist_attr") > 0 &&
+                       attributes.at("op_dist_attr")
+                           .isa<paddle::dialect::OperationDistAttribute>()),
+                      true,
+                      common::errors::PreconditionNotMet(
+                          "Type of attribute: op_dist_attr is not right."));
+  }
+  VLOG(4) << "Verifying outputs:";
+  {
+    auto output_size = num_results();
+    PADDLE_ENFORCE_EQ(
+        output_size,
+        1u,
+        common::errors::PreconditionNotMet(
+            "The size %d of outputs must be equal to 1.", output_size));
+    PADDLE_ENFORCE_EQ(
+        (*this)->result(0).type().isa<paddle::dialect::DistDenseTensorType>(),
+        true,
+        common::errors::PreconditionNotMet(
+            "Type validation failed for the 0th output."));
+  }
+
+  VLOG(4) << "Verifying op dist attrs:";
+  {
+    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_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()));
+  }
+  VLOG(4) << "End Verifying for: ShardTensorOp.";
+}
+
+void ReShardOp::Build(pir::Builder& builder,
+                      pir::OperationArgument& argument,
+                      pir::Value input,
+                      TensorDistAttribute tensor_dist_attr) {
+  VLOG(4) << "Start build ReShardOp";
+
+  paddle::dialect::DistDenseTensorType input_tensor_type;
+  if (input.type().isa<paddle::dialect::DistDenseTensorType>()) {
+    input_tensor_type =
+        input.type().dyn_cast<paddle::dialect::DistDenseTensorType>();
+  } else {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "Only support paddle::dialect::DistDenseTensorType"));
+  }
+
+  VLOG(4) << "Builder construction inputs";
+  argument.AddInput(input);
+
+  VLOG(4) << "Builder construction attributes";
+  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});
+  argument.AddAttribute("op_dist_attr", op_dist_attr);
+
+  VLOG(4) << "Builder construction outputs";
+  auto global_dims = input_tensor_type.global_ddim();
+  auto process_mesh_attr = tensor_dist_attr.process_mesh_attr();
+  auto dims_mapping = tensor_dist_attr.dims_mapping();
+
+  auto process_mesh_shape = process_mesh_attr.shape();
+  PADDLE_ENFORCE_EQ(static_cast<int>(dims_mapping.size()),
+                    global_dims.size(),
+                    common::errors::PreconditionNotMet(
+                        "dst dims_mapping size %d does not match input size %d",
+                        dims_mapping.size(),
+                        global_dims.size()));
+
+  auto local_shape = InferLocalDDim(global_dims, tensor_dist_attr);
+  pir::Type out_dist_tensor_type = paddle::dialect::DistDenseTensorType::get(
+      pir::IrContext::Instance(),
+      input_tensor_type.dense_tensor_type(),
+      tensor_dist_attr,
+      local_shape);
   argument.AddOutput(out_dist_tensor_type);
 }
 
 }  // namespace dialect
 }  // namespace paddle
 
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::ShardTensorOp)
+IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::ReShardOp)
diff --git a/paddle/fluid/pir/dialect/distributed/ir/dist_op.h b/paddle/fluid/pir/dialect/distributed/ir/dist_op.h
@@ -22,6 +22,8 @@
 
 namespace paddle {
 namespace dialect {
+class TensorDistAttribute;
+
 class ShardTensorOp : public pir::Op<ShardTensorOp> {
  public:
   using Op::Op;
@@ -36,7 +38,21 @@ class ShardTensorOp : public pir::Op<ShardTensorOp> {
   pir::Value out() { return result(0); }
   void VerifySig();
 };
+
+class ReShardOp : public pir::Op<ReShardOp> {
+ public:
+  using Op::Op;
+  static const char* name() { return "dist_op.reshard"; }
+  static const char* attributes_name[1];
+  static constexpr uint32_t attributes_num = 1;
+  TEST_API static void Build(pir::Builder& builder,             // NOLINT
+                             pir::OperationArgument& argument,  // NOLINT
+                             pir::Value input,
+                             TensorDistAttribute tensor_dist_attr);
+  void VerifySig();
+};
 }  // namespace dialect
 }  // namespace paddle
 
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::ShardTensorOp)
+IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::ReShardOp)
diff --git a/paddle/fluid/pybind/dist_static_op_function.h b/paddle/fluid/pybind/dist_static_op_function.h
diff --git a/test/cpp/pir/distributed/dist_dialect_test.cc b/test/cpp/pir/distributed/dist_dialect_test.cc
