PaddlePaddle
diff --git a/‎paddle/phi/api/lib/api_gen_utils.cc
Copy file name to clipboard
+24Lines changed: 24 additions & 0 deletions b/‎paddle/phi/api/lib/api_gen_utils.cc
Copy file name to clipboard
+24Lines changed: 24 additions & 0 deletions
diff --git a/‎paddle/phi/api/lib/api_gen_utils.h
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions b/‎paddle/phi/api/lib/api_gen_utils.h
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions
diff --git a/‎paddle/phi/api/lib/data_transform.cc
Copy file name to clipboard
+2-1Lines changed: 2 additions & 1 deletion b/‎paddle/phi/api/lib/data_transform.cc
Copy file name to clipboard
+2-1Lines changed: 2 additions & 1 deletion
diff --git a/‎paddle/phi/api/yaml/backward.yaml
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions b/‎paddle/phi/api/yaml/backward.yaml
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions
diff --git a/‎paddle/phi/api/yaml/generator/dist_api_gen.py
Copy file name to clipboard
+18-3Lines changed: 18 additions & 3 deletions b/‎paddle/phi/api/yaml/generator/dist_api_gen.py
Copy file name to clipboard
+18-3Lines changed: 18 additions & 3 deletions
diff --git a/‎paddle/phi/api/yaml/legacy_ops.yaml
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/phi/api/yaml/legacy_ops.yaml
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/phi/api/yaml/ops.yaml
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/phi/api/yaml/ops.yaml
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/phi/infermeta/spmd_rules/concat.cc
Copy file name to clipboard
+92-114Lines changed: 92 additions & 114 deletions b/‎paddle/phi/infermeta/spmd_rules/concat.cc
Copy file name to clipboard
+92-114Lines changed: 92 additions & 114 deletions
diff --git a/‎paddle/phi/infermeta/spmd_rules/concat.h
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions b/‎paddle/phi/infermeta/spmd_rules/concat.h
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions
diff --git a/‎paddle/phi/infermeta/spmd_rules/rules.h
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/phi/infermeta/spmd_rules/rules.h
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
@@ -628,6 +628,30 @@ std::vector<phi::distributed::DistTensor*> SetKernelDistOutput(
   return result;
 }
 
+std::vector<phi::distributed::DistTensor*> SetKernelDistOutput(
+    const phi::distributed::ArgDistAttr& dist_attr, std::vector<Tensor>* out) {
+  PADDLE_ENFORCE_EQ(
+      paddle::holds_alternative<std::vector<phi::distributed::TensorDistAttr>>(
+          dist_attr),
+      true,
+      phi::errors::PreconditionNotMet(
+          "Arg must be a vector of  TensorDistAttr"));
+  const std::vector<phi::distributed::TensorDistAttr>& dist_attrs =
+      PADDLE_GET_CONST(std::vector<phi::distributed::TensorDistAttr>,
+                       dist_attr);
+  auto out_size = dist_attrs.size();
+  out->reserve(out_size);
+  std::vector<phi::distributed::DistTensor*> results(out_size);
+  for (size_t i = 0; i < out_size; ++i) {
+    auto dist_t = std::make_shared<phi::distributed::DistTensor>(phi::DDim(),
+                                                                 dist_attrs[i]);
+    results[i] = dist_t.get();
+    out->emplace_back();
+    out->back().set_impl(dist_t);
+  }
+  return results;
+}
+
 std::vector<phi::distributed::DistTensor*> SetKernelDistOutput(
     size_t out_size, std::vector<Tensor>* out) {
   out->reserve(out_size);
 
@@ -166,6 +166,9 @@ std::vector<phi::distributed::DistTensor*> SetKernelDistOutput(
 std::vector<phi::distributed::DistTensor*> SetKernelDistOutput(
     size_t out_size, std::vector<Tensor>* out);
 
+std::vector<phi::distributed::DistTensor*> SetKernelDistOutput(
+    const phi::distributed::ArgDistAttr& dist_attr, std::vector<Tensor>* out);
+
 std::vector<phi::distributed::DistTensor*> SetKernelDistInplaceOutput(
     size_t out_size, std::vector<Tensor>* out);
 
 
@@ -656,7 +656,8 @@ ReshardApiInputToKernelInput(phi::DeviceContext* dev_ctx,
       true,
       phi::errors::PreconditionNotMet(
           "Arg must be a vector of TensorDistAttr"));
-  const auto& tensor_dist_attrs = paddle::get<1>(dist_attrs);
+  const auto& tensor_dist_attrs = PADDLE_GET_CONST(
+      std::vector<phi::distributed::TensorDistAttr>, dist_attrs);
 
   PADDLE_ENFORCE_EQ(tensors.size(),
                     tensor_dist_attrs.size(),
 
@@ -341,6 +341,7 @@
   infer_meta :
     func : UnchangedMultiInferMeta
     param : [x]
+    spmd_rule: ConcatGradInferSpmdDynamic
   kernel :
     func : concat_grad
     data_type : out_grad
@@ -2289,6 +2290,7 @@
   infer_meta :
     func : StackGradInferMeta
     param: [out_grad, axis]
+    spmd_rule : StackGradInferSpmd
   kernel :
     func : stack_grad
     param : [out_grad, axis]
 
@@ -147,6 +147,8 @@
             phi::DenseTensorMeta());
     }}
 """
+
+
 VECTOR_OUT_CREATION_TEMPLATE = """
     auto dist_out = SetKernelDistOutput({}, &api_output);
     std::vector<phi::DenseTensor*> dense_out(dist_out.size());
@@ -159,8 +161,9 @@
       }}
     }}
 """
+
 MULTI_VECTOR_OUT_CREATION_TEMPLATE = """
-    auto dist_out_{out_name} = SetKernelDistOutput({size}, {in_name});
+    auto dist_out_{out_name} = SetKernelDistOutput({dist_output_arg}, {in_name});
     std::vector<phi::DenseTensor*> dense_out_{out_name}(dist_out_{out_name}.size());
     for (size_t i = 0; i < dist_out_{out_name}.size(); ++i) {{
         dense_out_{out_name}[i] = const_cast<phi::DenseTensor*>(&dist_out_{out_name}[i]->value());
@@ -829,9 +832,16 @@ def generate_output_creation_code(self) -> str:
                 else:
                     output_creation_code += SINGLE_OUT_CREATION_TEMPLATE_NO_SPMD
             elif self.outputs['types'][0] == 'std::vector<Tensor>':
+                # SetKernelDistOutput arg
+                dist_output_arg = (
+                    "spmd_info.second[0]"
+                    if self.infer_meta['spmd_rule'] is not None
+                    else self.outputs['out_size_expr'][0]
+                )
                 output_creation_code += VECTOR_OUT_CREATION_TEMPLATE.format(
-                    self.outputs['out_size_expr'][0]
+                    dist_output_arg
                 )
+
             else:
                 self.vector_output_size_assertion_check()
         elif output_num > 1:
@@ -876,10 +886,15 @@ def generate_output_creation_code(self) -> str:
                             in_name=get_out_code,
                         )
                     else:
+                        dist_output_arg = (
+                            f"spmd_info.second[{i}]"
+                            if self.infer_meta['spmd_rule'] is not None
+                            else self.outputs['out_size_expr'][i]
+                        )
                         output_creation_code += (
                             MULTI_VECTOR_OUT_CREATION_TEMPLATE.format(
                                 out_name=i,
-                                size=self.outputs['out_size_expr'][i],
+                                dist_output_arg=dist_output_arg,
                                 in_name=get_out_code,
                             )
                         )
 
@@ -1017,6 +1017,7 @@
   output : Tensor[]{num}
   infer_meta :
     func : SplitWithNumInferMeta
+    spmd_rule : SplitWithNumInferSpmdDynamic
   kernel :
     func : split_with_num
   backward : split_with_num_grad
 
@@ -2496,6 +2496,7 @@
   output : Tensor (out)
   infer_meta :
     func : StackInferMeta
+    spmd_rule : StackInferSpmd
   kernel :
     func : stack
   backward : stack_grad
 
@@ -17,6 +17,8 @@ limitations under the License. */
 #include <limits>
 #include <set>
 
+#include "glog/logging.h"
+
 #include "paddle/phi/infermeta/spmd_rules/elementwise.h"
 #include "paddle/phi/infermeta/spmd_rules/utils.h"
 
@@ -25,8 +27,25 @@ namespace distributed {
 
 using phi::distributed::auto_parallel::str_join;
 
-static bool IsEmpty(const std::vector<int64_t>& shape) {
-  return shape.empty() || shape.at(0) == 0;
+std::tuple<std::string, std::string> FillConcatNotation(int64_t n_axis,
+                                                        int64_t concat_axis) {
+  PADDLE_ENFORCE_GT(
+      n_axis,
+      concat_axis,
+      phi::errors::InvalidArgument(
+          "n_axis [%d] and concat_axis[%d] not match", n_axis, concat_axis));
+  static const std::string alphabet = "abcdefghijlopqrstuvwxyz";
+  PADDLE_ENFORCE_GT(alphabet.size(),
+                    static_cast<size_t>(n_axis),
+                    phi::errors::InvalidArgument(
+                        "alphabet.size() [%d]; n_axis [%d] is too large",
+                        alphabet.size(),
+                        n_axis));
+  std::string all_axis = alphabet.substr(0, n_axis);
+  std::string align_axis =
+      std::string(all_axis.begin(), all_axis.begin() + concat_axis) +
+      std::string(all_axis.begin() + concat_axis + 1, all_axis.end());
+  return {all_axis, align_axis};
 }
 
 SpmdInfo ConcatInferSpmd(const std::vector<DistMetaTensor>& x, int axis) {
@@ -58,130 +77,89 @@ SpmdInfo ConcatInferSpmd(const std::vector<DistMetaTensor>& x, int axis) {
   auto non_empty_index = non_empty_iter - tensor_shapes.begin();
   int64_t ndim = static_cast<int64_t>(tensor_shapes[non_empty_index].size());
   // normlize dim
-  int64_t dim = axis;
-  dim = dim < 0 ? dim + ndim : dim;
-
+  auto dim = axis < 0 ? ndim + axis : axis;
   std::vector<TensorDistAttr> input_attrs;
-  // 2、make sure all tensors replicated on concat dim
-  auto n_inputs = x.size();
-  for (size_t i = 0; i < n_inputs; ++i) {
-    const auto& dist_attr = x[i].dist_attr();
-    if ((!IsEmpty(tensor_shapes[i])) && IsDimSharded(dist_attr, dim)) {
-      auto sharded_dist_attr = ReplicateTensorDim(dist_attr, dim);
-      input_attrs.emplace_back(sharded_dist_attr);
-    } else {
-      input_attrs.emplace_back(dist_attr);
-    }
-  }
-  // 3、align non-concat dimensions according to cost
-  std::vector<std::vector<std::shared_ptr<PlacementStatus>>> inputs_placements;
   std::transform(
-      input_attrs.begin(),
-      input_attrs.end(),
-      std::back_inserter(inputs_placements),
-      [](const TensorDistAttr& attr) { return attr.to_placement(); });
-  const auto& process_mess = input_attrs[non_empty_index].process_mesh();
-  auto has_mismatch = [&](int32_t mesh_dim) {
-    bool mismatch = false;
-    for (size_t i = 0; i < n_inputs; i++) {
-      if ((!IsEmpty(tensor_shapes[i])) &&
-          !PlacementEqual(inputs_placements[non_empty_index][mesh_dim],
-                          inputs_placements[i][mesh_dim])) {
-        mismatch = true;
-        break;
-      }
-    }
-    return mismatch;
-  };
-  bool need_reshard = false;
-  int32_t n_mesh_dim = process_mess.ndim();
-  std::vector<std::shared_ptr<PlacementStatus>> best_placements(
-      n_mesh_dim, std::make_shared<ReplicatedStatus>());
-  // a dim can not be sharded twice along diffrent mesh_dim
-  std::set<int64_t> sharded_dims = {dim};
-
-  for (int32_t mesh_dim = 0; mesh_dim < process_mess.ndim(); ++mesh_dim) {
-    if (!has_mismatch(mesh_dim)) {
-      // use the old placement
-      auto& best = inputs_placements[non_empty_index][mesh_dim];
-      if (best->is_shard()) {
-        auto shard_placement = std::dynamic_pointer_cast<ShardStatus>(best);
-        sharded_dims.insert(shard_placement->get_axis());
-      }
-      best_placements[mesh_dim] = best;
-    }
-  }
+      x.begin(), x.end(), std::back_inserter(input_attrs), [](auto& meta) {
+        return meta.dist_attr();
+      });
 
-  for (int32_t mesh_dim = 0; mesh_dim < process_mess.ndim(); ++mesh_dim) {
-    if (!has_mismatch(mesh_dim)) {
-      continue;
-    }
-    need_reshard = true;
-    std::vector<double> costs;
-    for (int32_t shard_dim = 0; shard_dim < ndim; shard_dim++) {
-      double cost = std::numeric_limits<double>::infinity();
-      if (!sharded_dims.count(shard_dim)) {
-        cost = 0.0;
-        for (size_t i = 0; i < n_inputs; i++) {
-          auto& tensor_shape = tensor_shapes[i];
-          auto& tensor_dist_attr = input_attrs[i];
-          if (IsEmpty(tensor_shape)) {
-            continue;
-          }
-
-          if (tensor_shape[shard_dim] < process_mess.dim_size(mesh_dim)) {
-            // should not be selected
-            cost += std::numeric_limits<double>::infinity();
-            continue;
-          }
-          if (IsDimSharded(tensor_dist_attr, shard_dim)) {
-            continue;
-          }
-          int64_t num = std::accumulate(tensor_shape.begin(),
-                                        tensor_shape.end(),
-                                        1,
-                                        std::multiplies<int64_t>());
-          if (num == static_cast<int64_t>(0)) {
-            continue;
-          }
-          std::vector<int64_t> local_shape =
-              GetLocalShape(tensor_shape, process_mess, inputs_placements[i]);
-          cost += std::accumulate(local_shape.begin(),
-                                  local_shape.end(),
-                                  1,
-                                  std::multiplies<int64_t>()) *
-                  process_mess.dim_size(mesh_dim);
-        }
-      }
-      costs.push_back(cost);
-    }
-    auto min_itr = std::min_element(costs.begin(), costs.end());
-    auto min_dim = min_itr - costs.begin();
-    if (!sharded_dims.count(min_dim)) {
-      best_placements[mesh_dim] = std::make_shared<ShardStatus>(min_dim);
-      sharded_dims.insert(min_dim);
-    }
-  }
-  // set placement to the best placements
-  if (need_reshard) {
-    std::vector<TensorDistAttr> new_input_attrs;
-    for (auto& e : input_attrs) {
-      new_input_attrs.emplace_back(FromPlacements(e, best_placements));
-    }
-    std::swap(input_attrs, new_input_attrs);
-  }
-  return {{input_attrs}, {input_attrs[non_empty_index]}};
+  std::string all_aixs;
+  std::string align_axis;
+  std::tie(all_aixs, align_axis) = FillConcatNotation(ndim, dim);
+  std::vector<std::string> axis_names(input_attrs.size(), all_aixs);
+  AlignDimsSharding(
+      &input_attrs, tensor_shapes, axis_names, {}, align_axis, true);
+
+  auto out_dist_attr =
+      CopyTensorDistAttrForOutput(input_attrs[non_empty_index]);
+  out_dist_attr.set_dims_mapping(input_attrs[non_empty_index].dims_mapping());
+  VLOG(4) << "concat out " << out_dist_attr.to_string();
+  return {{input_attrs}, {out_dist_attr}};
 }
 
 SpmdInfo ConcatInferSpmdReverse(const std::vector<DistMetaTensor>& x,
                                 const DistMetaTensor& output,
                                 int axis) {
-  // TODO(liuzhenhai): add latter
-  return SpmdInfo();
+  auto out_dist_attr = output.dist_attr();
+  out_dist_attr = UnShardTensorDim(out_dist_attr, axis);
+  auto n_inputs = x.size();
+  TensorDistAttr input_attr = CopyTensorDistAttrForOutput(out_dist_attr);
+  const auto& input_dim_mapping = out_dist_attr.dims_mapping();
+  input_attr.set_dims_mapping(input_dim_mapping);
+  std::vector<TensorDistAttr> input_attrs(n_inputs, input_attr);
+  return {{input_attrs}, {output.dist_attr()}};
 }
+
 SpmdInfo ConcatInferSpmdDynamic(const std::vector<DistMetaTensor>& x,
                                 const Scalar& axis) {
   return ConcatInferSpmd(x, axis.to<int32_t>());
 }
+
+SpmdInfo ConcatGradInferSpmdDynamic(const std::vector<DistMetaTensor>& x,
+                                    const DistMetaTensor& output_grad,
+                                    const Scalar& axis) {
+  // 1、check tensors shapes
+  std::vector<std::vector<int64_t>> tensor_shapes;
+  std::transform(x.begin(),
+                 x.end(),
+                 std::back_inserter(tensor_shapes),
+                 [](const DistMetaTensor& meta) {
+                   return phi::vectorize<int64_t>(meta.dims());
+                 });
+  bool all_empty =
+      std::all_of(tensor_shapes.begin(), tensor_shapes.end(), IsEmpty);
+  if (all_empty) {
+    return SpmdInfo();
+  }
+
+  auto non_empty_iter =
+      std::find_if(tensor_shapes.begin(), tensor_shapes.end(), [](auto& shape) {
+        return !IsEmpty(shape);
+      });
+  auto non_empty_index = non_empty_iter - tensor_shapes.begin();
+  int64_t ndim = static_cast<int64_t>(tensor_shapes[non_empty_index].size());
+  auto dim = axis.to<int64_t>();
+  // normlize dim
+  dim = dim < 0 ? ndim + dim : dim;
+  std::vector<TensorDistAttr> input_attrs;
+  std::transform(
+      x.begin(), x.end(), std::back_inserter(input_attrs), [](auto& meta) {
+        return meta.dist_attr();
+      });
+  input_attrs.push_back(output_grad.dist_attr());
+  tensor_shapes.push_back(phi::vectorize<int64_t>(output_grad.dims()));
+  std::string all_aixs;
+  std::string align_axis;
+  std::tie(all_aixs, align_axis) = FillConcatNotation(ndim, dim);
+  std::vector<std::string> axis_names(input_attrs.size(), all_aixs);
+  AlignDimsSharding(
+      &input_attrs, tensor_shapes, axis_names, {}, align_axis, true);
+  auto output_grad_attr = input_attrs.back();
+  input_attrs.pop_back();
+  std::vector<TensorDistAttr> inputs_grad = input_attrs;
+  return {{input_attrs, output_grad_attr}, {inputs_grad}};
+}
+
 }  // namespace distributed
 }  // namespace phi
@@ -30,5 +30,9 @@ SpmdInfo ConcatInferSpmdReverse(const std::vector<DistMetaTensor>& x,
 SpmdInfo ConcatInferSpmdDynamic(const std::vector<DistMetaTensor>& x,
                                 const Scalar& axis);
 
+SpmdInfo ConcatGradInferSpmdDynamic(const std::vector<DistMetaTensor>& x,
+                                    const DistMetaTensor& output_grad,
+                                    const Scalar& axis);
+
 }  // namespace distributed
 }  // namespace phi
@@ -29,6 +29,7 @@ limitations under the License. */
 #include "paddle/phi/infermeta/spmd_rules/slice.h"
 #include "paddle/phi/infermeta/spmd_rules/softmax.h"
 #include "paddle/phi/infermeta/spmd_rules/split.h"
+#include "paddle/phi/infermeta/spmd_rules/stack.h"
 #include "paddle/phi/infermeta/spmd_rules/transpose.h"
 #include "paddle/phi/infermeta/spmd_rules/unsqueeze.h"