[auto parallel] CrossMeshReshard for: p2r, p2s, r2p, r2s, s2p, s2r, s2s.

paddle/phi/core/distributed/auto_parallel/reshard/p_to_r_reshard_function.cc

@@ -20,6 +20,7 @@
 #include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
 #include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h"
 #include "paddle/phi/core/distributed/auto_parallel/reshard/same_status_reshard_function.h"
+#include "paddle/phi/core/distributed/store/store_utils.h"
 #include "paddle/phi/kernels/all_reduce_kernel.h"
 #include "paddle/phi/kernels/elementwise_divide_kernel.h"
 #include "paddle/phi/kernels/full_kernel.h"
@@ -123,14 +124,20 @@ void PToRReshardFunctionCrossMesh::Eval(phi::DeviceContext* dev_ctx,
   tmp_dist_attr.set_process_mesh(out_process_mesh);
   same_status_func.Eval(dev_ctx, in, tmp_dist_attr, &tmp_result);
 
-  PToRReshardFunction p_to_r_func;
-  PADDLE_ENFORCE(
-      p_to_r_func.IsSuitable(tmp_result, out_dist_attr),
-      phi::errors::InvalidArgument(
-          "Invoke the p to r reshard function is not valid from %s to %s.",
-          tmp_result.dist_attr(),
-          out_dist_attr));
-  p_to_r_func.Eval(dev_ctx, tmp_result, out_dist_attr, out);
+  int64_t cur_global_rank = GetCurGlobalRank();
+  if (out_process_mesh.contains(cur_global_rank)) {
+    PToRReshardFunction p_to_r_func;
+    PADDLE_ENFORCE(
+        p_to_r_func.IsSuitable(tmp_result, out_dist_attr),
+        phi::errors::InvalidArgument(
+            "Invoke the p to r reshard function is not valid from %s to %s.",
+            tmp_result.dist_attr(),
+            out_dist_attr));
+    p_to_r_func.Eval(dev_ctx, tmp_result, out_dist_attr, out);
+  } else {
+    SetDistProps(out, in.dims(), out_dist_attr);
+    SetValue(out, tmp_result.value());
+  }
 }
 
 }  // namespace distributed

paddle/phi/core/distributed/auto_parallel/reshard/p_to_s_reshard_function.cc

@@ -18,6 +18,8 @@
 #include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
 #include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
 #include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/same_status_reshard_function.h"
+#include "paddle/phi/core/distributed/store/store_utils.h"
 #include "paddle/phi/kernels/reduce_scatter_kernel.h"
 #include "paddle/phi/kernels/transpose_kernel.h"
 
@@ -89,5 +91,58 @@ void PToSReshardFunction::Eval(DeviceContext* dev_ctx,
   SetDistProps(out, in.dims(), out_dist_attr);
 }
 
+bool PToSReshardFunctionCrossMesh::IsSuitable(
+    const DistTensor& in, const TensorDistAttr& out_dist_attr) {
+  const auto& in_dist_attr = in.dist_attr();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_dist_attr.is_partial());
+  RESHARD_SHORTCUT_IF_FALSE(out_dist_attr.is_shard());
+
+  const auto& in_process_mesh = in_dist_attr.process_mesh();
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(out_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh != out_process_mesh);
+
+  return true;
+}
+
+void PToSReshardFunctionCrossMesh::Eval(DeviceContext* dev_ctx,
+                                        const DistTensor& in,
+                                        const TensorDistAttr& out_dist_attr,
+                                        DistTensor* out) {
+  VLOG(3) << "Call PToSReshardFunctionCrossMesh Eval";
+  const auto& in_dist_attr = in.dist_attr();
+
+  DistTensor tmp_result;
+  TensorDistAttr in_dist_attr_shard = in_dist_attr;
+  in_dist_attr_shard.clean_partial_status();
+  in_dist_attr_shard.set_dims_mapping(out_dist_attr.dims_mapping());
+
+  int64_t cur_global_rank = GetCurGlobalRank();
+  if (in_dist_attr.process_mesh().contains(cur_global_rank)) {
+    PToSReshardFunction p_to_s_func;
+    PADDLE_ENFORCE(
+        p_to_s_func.IsSuitable(in, in_dist_attr_shard),
+        phi::errors::InvalidArgument(
+            "Invoke the p to s reshard function is not valid from %s to %s.",
+            in_dist_attr,
+            in_dist_attr_shard));
+    p_to_s_func.Eval(dev_ctx, in, in_dist_attr_shard, &tmp_result);
+  } else {
+    SetDistProps(&tmp_result, in.dims(), in_dist_attr_shard);
+  }
+
+  SameStatusReshardFunction same_status_func;
+  PADDLE_ENFORCE(
+      same_status_func.IsSuitable(tmp_result, out_dist_attr),
+      phi::errors::InvalidArgument("Invoke the same status reshard function "
+                                   "is not valid from %s to %s.",
+                                   tmp_result.dist_attr(),
+                                   out_dist_attr));
+  same_status_func.Eval(dev_ctx, tmp_result, out_dist_attr, out);
+}
+
 }  // namespace distributed
 }  // namespace phi

paddle/phi/core/distributed/auto_parallel/reshard/p_to_s_reshard_function.h

@@ -31,5 +31,18 @@ class PToSReshardFunction final : public ReshardFunction {
   std::string Name() override { return "PToSReshard"; }
 };
 
+class PToSReshardFunctionCrossMesh final : public ReshardFunction {
+ public:
+  bool IsSuitable(const DistTensor& in,
+                  const TensorDistAttr& out_dist_attr) override;
+
+  void Eval(DeviceContext* dev_ctx,
+            const DistTensor& in,
+            const TensorDistAttr& out_dist_attr,
+            DistTensor* out) override;
+
+  std::string Name() override { return "PToSReshardCrossMesh"; }
+};
+
 }  // namespace distributed
 }  // namespace phi

paddle/phi/core/distributed/auto_parallel/reshard/r_to_p_reshard_function.cc

@@ -20,6 +20,7 @@
 #include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
 #include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h"
 #include "paddle/phi/core/distributed/auto_parallel/reshard/same_status_reshard_function.h"
+#include "paddle/phi/core/distributed/store/store_utils.h"
 #include "paddle/phi/kernels/assign_kernel.h"
 #include "paddle/phi/kernels/full_kernel.h"
 
@@ -104,14 +105,20 @@ void RToPReshardFunctionCrossMesh::Eval(phi::DeviceContext* dev_ctx,
   tmp_dist_attr.set_process_mesh(out_process_mesh);
   same_status_func.Eval(dev_ctx, in, tmp_dist_attr, &tmp_result);
 
-  RToPReshardFunction r_to_p_func;
-  PADDLE_ENFORCE(
-      r_to_p_func.IsSuitable(tmp_result, out_dist_attr),
-      phi::errors::InvalidArgument(
-          "Invoke the r to p reshard function is not valid from %s to %s.",
-          tmp_result.dist_attr(),
-          out_dist_attr));
-  r_to_p_func.Eval(dev_ctx, tmp_result, out_dist_attr, out);
+  int64_t cur_global_rank = GetCurGlobalRank();
+  if (out_process_mesh.contains(cur_global_rank)) {
+    RToPReshardFunction r_to_p_func;
+    PADDLE_ENFORCE(
+        r_to_p_func.IsSuitable(tmp_result, out_dist_attr),
+        phi::errors::InvalidArgument(
+            "Invoke the r to p reshard function is not valid from %s to %s.",
+            tmp_result.dist_attr(),
+            out_dist_attr));
+    r_to_p_func.Eval(dev_ctx, tmp_result, out_dist_attr, out);
+  } else {
+    SetDistProps(out, in.dims(), out_dist_attr);
+    SetValue(out, tmp_result.value());
+  }
 }
 
 }  // namespace distributed

paddle/phi/core/distributed/auto_parallel/reshard/r_to_s_reshard_function.cc

@@ -19,6 +19,7 @@
 #include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
 #include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h"
 #include "paddle/phi/core/distributed/auto_parallel/reshard/same_status_reshard_function.h"
+#include "paddle/phi/core/distributed/store/store_utils.h"
 #include "paddle/phi/kernels/split_kernel.h"
 
 namespace phi {
@@ -115,16 +116,21 @@ void RToSReshardFunctionCrossMesh::Eval(phi::DeviceContext* dev_ctx,
   DistTensor tmp_result;
   TensorDistAttr in_dist_attr_shard = in_dist_attr;
   in_dist_attr_shard.set_dims_mapping(out_dist_attr.dims_mapping());
-  RToSReshardFunction r_to_s_func;
-  PADDLE_ENFORCE(
-      r_to_s_func.IsSuitable(in, in_dist_attr_shard),
-      phi::errors::InvalidArgument(
-          "Invoke the r to s reshard function is not valid from %s to %s.",
-          tmp_result.dist_attr(),
-          out_dist_attr));
-  r_to_s_func.Eval(dev_ctx, in, in_dist_attr_shard, &tmp_result);
-
-  // Step 2: Same status from the input mesh to output mesh
+
+  int64_t cur_global_rank = GetCurGlobalRank();
+  if (in_dist_attr.process_mesh().contains(cur_global_rank)) {
+    RToSReshardFunction r_to_s_func;
+    PADDLE_ENFORCE(
+        r_to_s_func.IsSuitable(in, in_dist_attr_shard),
+        phi::errors::InvalidArgument(
+            "Invoke the r to s reshard function is not valid from %s to %s.",
+            in_dist_attr,
+            in_dist_attr_shard));
+    r_to_s_func.Eval(dev_ctx, in, in_dist_attr_shard, &tmp_result);
+  } else {
+    SetDistProps(&tmp_result, in.dims(), in_dist_attr_shard);
+  }
+
   SameStatusReshardFunction same_status_func;
   PADDLE_ENFORCE(
       same_status_func.IsSuitable(tmp_result, out_dist_attr),

paddle/phi/core/distributed/auto_parallel/reshard/reshard_function_registry.cc

@@ -52,16 +52,19 @@ std::vector<std::unique_ptr<ReshardFunction>>& GetReshardFunctionList() {
 // Reshard function with higher priority will be evoked
 // when more than one reshard function satisfy the request.
 REGISTER_RESHARD_FUNC(SToRReshardFunction);
-REGISTER_RESHARD_FUNC(SToPReshardFunction);
 REGISTER_RESHARD_FUNC(SToRReshardFunctionCrossMesh);
+REGISTER_RESHARD_FUNC(SToPReshardFunction);
+REGISTER_RESHARD_FUNC(SToPReshardFunctionCrossMesh);
 REGISTER_RESHARD_FUNC(RToSReshardFunction);
 REGISTER_RESHARD_FUNC(RToSReshardFunctionCrossMesh);
 REGISTER_RESHARD_FUNC(RToPReshardFunction);
 REGISTER_RESHARD_FUNC(RToPReshardFunctionCrossMesh);
 REGISTER_RESHARD_FUNC(PToRReshardFunction);
 REGISTER_RESHARD_FUNC(PToRReshardFunctionCrossMesh);
 REGISTER_RESHARD_FUNC(PToSReshardFunction);
+REGISTER_RESHARD_FUNC(PToSReshardFunctionCrossMesh);
 REGISTER_RESHARD_FUNC(SToSReshardFunction);
+REGISTER_RESHARD_FUNC(SToSReshardFunctionCrossMesh);
 REGISTER_RESHARD_FUNC(SameStatusReshardFunction);
 REGISTER_RESHARD_FUNC(SameNdMeshReshardFunction);
 

paddle/phi/core/distributed/auto_parallel/reshard/s_to_p_reshard_function.cc

@@ -20,6 +20,7 @@
 #include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
 #include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
 #include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h"
+#include "paddle/phi/core/distributed/store/store_utils.h"
 #include "paddle/phi/kernels/reduce_scatter_kernel.h"
 #include "paddle/phi/kernels/transpose_kernel.h"
 
@@ -63,5 +64,52 @@ void SToPReshardFunction::Eval(DeviceContext* dev_ctx,
   r_to_p.Eval(dev_ctx, tmp_tensor, out_dist_attr, out);
 }
 
+bool SToPReshardFunctionCrossMesh::IsSuitable(
+    const DistTensor& in, const TensorDistAttr& out_dist_attr) {
+  const auto& in_dist_attr = in.dist_attr();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_dist_attr.is_shard());
+  RESHARD_SHORTCUT_IF_FALSE(out_dist_attr.is_partial());
+
+  const auto& in_process_mesh = in_dist_attr.process_mesh();
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(out_process_mesh.ndim() == 1);
+  RESHARD_SHORTCUT_IF_FALSE(in_process_mesh != out_process_mesh);
+
+  return true;
+}
+
+void SToPReshardFunctionCrossMesh::Eval(DeviceContext* dev_ctx,
+                                        const DistTensor& in,
+                                        const TensorDistAttr& out_dist_attr,
+                                        DistTensor* out) {
+  VLOG(3) << "Call SToPReshardFunctionCrossMesh Eval";
+  const auto& out_process_mesh = out_dist_attr.process_mesh();
+
+  DistTensor tmp_result;
+
+  SameStatusReshardFunction same_status_func;
+  TensorDistAttr tmp_dist_attr = in.dist_attr();
+  tmp_dist_attr.set_process_mesh(out_process_mesh);
+  same_status_func.Eval(dev_ctx, in, tmp_dist_attr, &tmp_result);
+
+  int64_t cur_global_rank = GetCurGlobalRank();
+  if (out_process_mesh.contains(cur_global_rank)) {
+    SToPReshardFunction s_to_p_func;
+    PADDLE_ENFORCE(
+        s_to_p_func.IsSuitable(tmp_result, out_dist_attr),
+        phi::errors::InvalidArgument(
+            "Invoke the s to p reshard function is not valid from %s to %s.",
+            tmp_result.dist_attr(),
+            out_dist_attr));
+    s_to_p_func.Eval(dev_ctx, tmp_result, out_dist_attr, out);
+  } else {
+    SetDistProps(out, in.dims(), out_dist_attr);
+    SetValue(out, tmp_result.value());
+  }
+}
+
 }  // namespace distributed
 }  // namespace phi

paddle/phi/core/distributed/auto_parallel/reshard/s_to_p_reshard_function.h

@@ -31,5 +31,18 @@ class SToPReshardFunction final : public ReshardFunction {
   std::string Name() override { return "SToPReshard"; }
 };
 
+class SToPReshardFunctionCrossMesh final : public ReshardFunction {
+ public:
+  bool IsSuitable(const DistTensor& in,
+                  const TensorDistAttr& out_dist_attr) override;
+
+  void Eval(DeviceContext* dev_ctx,
+            const DistTensor& in,
+            const TensorDistAttr& out_dist_attr,
+            DistTensor* out) override;
+
+  std::string Name() override { return "SToPReshardCrossMesh"; }
+};
+
 }  // namespace distributed
 }  // namespace phi

paddle/phi/core/distributed/auto_parallel/reshard/s_to_r_reshard_function.cc

@@ -19,6 +19,8 @@
 #include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
 #include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
 #include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard/same_status_reshard_function.h"
+#include "paddle/phi/core/distributed/store/store_utils.h"
 #include "paddle/phi/kernels/all_gather_kernel.h"
 #include "paddle/phi/kernels/concat_kernel.h"
 #include "paddle/phi/kernels/split_kernel.h"
@@ -126,11 +128,15 @@ bool SToRReshardFunctionCrossMesh::IsSuitable(
   const auto& in_process_mesh = in_dist_attr.process_mesh();
   const auto& out_process_mesh = out_dist_attr.process_mesh();
 
-  int split_axis = GetSplitAxisWithDimsMapping(in_dims_mapping).begin()->first;
-  int64_t num_of_process = in_process_mesh.size();
-  RESHARD_SHORTCUT_IF_FALSE(in.local_dims()[static_cast<int>(split_axis)] *
-                                num_of_process ==
-                            in.dims()[static_cast<int>(split_axis)]);
+  int64_t cur_global_rank = GetCurGlobalRank();
+  if (in_process_mesh.contains(cur_global_rank)) {
+    int split_axis =
+        GetSplitAxisWithDimsMapping(in_dims_mapping).begin()->first;
+    int64_t num_of_process = in_process_mesh.size();
+    RESHARD_SHORTCUT_IF_FALSE(in.local_dims()[static_cast<int>(split_axis)] *
+                                  num_of_process ==
+                              in.dims()[static_cast<int>(split_axis)]);
+  }
 
   RESHARD_SHORTCUT_IF_FALSE(in_process_mesh.ndim() == 1);
   RESHARD_SHORTCUT_IF_FALSE(out_process_mesh.ndim() == 1);
@@ -155,8 +161,20 @@ void SToRReshardFunctionCrossMesh::Eval(DeviceContext* dev_ctx,
   tmp_dist_attr.set_process_mesh(out_process_mesh);
   same_status_func.Eval(dev_ctx, in, tmp_dist_attr, &tmp_result);
 
-  SToRReshardFunction s_to_r_func;
-  s_to_r_func.Eval(dev_ctx, tmp_result, out_dist_attr, out);
+  int64_t cur_global_rank = GetCurGlobalRank();
+  if (out_process_mesh.contains(cur_global_rank)) {
+    SToRReshardFunction s_to_r_func;
+    PADDLE_ENFORCE(
+        s_to_r_func.IsSuitable(tmp_result, out_dist_attr),
+        phi::errors::InvalidArgument(
+            "Invoke the s to r reshard function is not valid from %s to %s.",
+            tmp_result.dist_attr(),
+            out_dist_attr));
+    s_to_r_func.Eval(dev_ctx, tmp_result, out_dist_attr, out);
+  } else {
+    SetDistProps(out, in.dims(), out_dist_attr);
+    SetValue(out, tmp_result.value());
+  }
 }
 
 }  // namespace distributed