slice infer spmd support decrease dim slice (#59986)

liuzhenhai93 · web-flow · commit fcc8760c7b0a · 2023-12-15T15:04:24.000+08:00
* decrease_dim_slice

* decrease_dim_slice

* decrease_dim_slice
diff --git a/paddle/phi/infermeta/spmd_rules/slice.cc b/paddle/phi/infermeta/spmd_rules/slice.cc
@@ -26,11 +26,28 @@ namespace distributed {
 
 using phi::distributed::auto_parallel::str_join;
 
+std::vector<int64_t> BuildOutputAxisToInputAxisMap(
+    const std::vector<int64_t>& decrease_axis, int input_ndim) {
+  std::vector<int64_t> output_axis_to_input_axis(input_ndim -
+                                                 decrease_axis.size());
+  int index = 0;
+  for (int i = 0; i < input_ndim; ++i) {
+    if (std::find(decrease_axis.begin(), decrease_axis.end(), i) ==
+        decrease_axis.end()) {
+      output_axis_to_input_axis[index] = i;
+      ++index;
+    }
+  }
+  return output_axis_to_input_axis;
+}
+
 SpmdInfo SliceInferSpmdBase(const DistMetaTensor& input,
-                            const std::vector<int64_t>& axes) {
+                            const std::vector<int64_t>& axes,
+                            const std::vector<int64_t>& decrease_axis) {
   // Step0: Verify input args based on slice logic
   auto input_shape = common::vectorize(input.dims());
   int input_ndim = input_shape.size();
+  int output_ndim = input_ndim - static_cast<int>(decrease_axis.size());
   auto input_dist_attr_src = input.dist_attr();
   std::vector<int64_t> input_dims_mapping = input_dist_attr_src.dims_mapping();
   PADDLE_ENFORCE_EQ(
@@ -47,31 +64,15 @@ SpmdInfo SliceInferSpmdBase(const DistMetaTensor& input,
   // get einsum notation for input
   std::string input_axes = alphabet.substr(0, input_ndim);
 
+  auto output_input_axis_mapping =
+      BuildOutputAxisToInputAxisMap(decrease_axis, input_ndim);
   // get einsum notation for output
-  std::string out_axes(input_axes);
-
-  for (int i = 0; i < static_cast<int>(axes.size()); i++) {
-    int axis = axes[i] < 0 ? axes[i] + input_ndim : axes[i];
-    // the sliced axis cannot be sharded, set its notation
-    // with the special '1' to set its dim mapping to -1.
-    out_axes[axis] = '1';
+  std::string out_axes = alphabet.substr(0, output_ndim);
+  for (int i = 0; i < output_ndim; i++) {
+    auto input_axis = output_input_axis_mapping[i];
+    out_axes[i] = input_axes[input_axis];
   }
 
-  // Step2: Sharding Propogation
-  // Step2.1: merge input shardings
-  std::unordered_map<std::string, int64_t> axis_to_dim_map =
-      ShardingMergeForTensors({{input_axes, input_dims_mapping}});
-
-  // Step2.2: infer output dims mapping from merged input dims mapping
-  std::vector<int64_t> out_dims_mapping =
-      GetDimsMappingForAxes(out_axes, axis_to_dim_map);
-
-  // get the dist attributes for output. the sliced
-  // cannot be sharded, if it is sharded, set it to replicated.
-  TensorDistAttr out_dist_attr =
-      CopyTensorDistAttrForOutput(input_dist_attr_src);
-  out_dist_attr.set_dims_mapping(out_dims_mapping);
-
   // Step2.3 get new dist attribute for input. the sliced
   // cannot be sharded, if it is sharded, set it to replicated.
   TensorDistAttr input_dist_attr_dst(input_dist_attr_src);
@@ -81,6 +82,15 @@ SpmdInfo SliceInferSpmdBase(const DistMetaTensor& input,
   }
   input_dist_attr_dst.set_dims_mapping(input_dims_mapping);
 
+  std::vector<int64_t> out_dims_mapping(output_ndim);
+  for (int i = 0; i < output_ndim; i++) {
+    auto input_axis = output_input_axis_mapping[i];
+    out_dims_mapping[i] = input_dims_mapping[input_axis];
+  }
+  TensorDistAttr out_dist_attr =
+      CopyTensorDistAttrForOutput(input_dist_attr_src);
+  out_dist_attr.set_dims_mapping(out_dims_mapping);
+
   VLOG(4) << "SliceInferSpmd:";
   VLOG(4) << "Einsum Notation: " << input_axes << "-->" << out_axes;
   VLOG(4) << "Input shape: [" << str_join(input_shape) << "] "
@@ -103,12 +113,13 @@ SpmdInfo SliceInferSpmd(const DistMetaTensor& input,
                         const std::vector<int64_t>& decrease_axis) {
   // starts, ends, infer_flags and decrease_axis have no impact on the
   // derivation, only to align with the definition in phi api
-  return SliceInferSpmdBase(input, axes);
+  return SliceInferSpmdBase(input, axes, {});
 }
 
 SpmdInfo SliceInferSpmdReverseBase(const DistMetaTensor& input,
                                    const DistMetaTensor& output,
-                                   const std::vector<int64_t>& axes) {
+                                   const std::vector<int64_t>& axes,
+                                   const std::vector<int64_t>& decrease_axis) {
   auto output_shape = common::vectorize(output.dims());
   int out_ndim = output_shape.size();
   auto out_dist_attr = output.dist_attr();
@@ -118,9 +129,11 @@ SpmdInfo SliceInferSpmdReverseBase(const DistMetaTensor& input,
   auto input_dist_attr = input.dist_attr();
   std::vector<int64_t> input_dims_mapping = input_dist_attr.dims_mapping();
 
+  int decrease_axis_num = decrease_axis.size();
+
   PADDLE_ENFORCE_EQ(
       input_ndim,
-      out_ndim,
+      out_ndim + decrease_axis_num,
       phi::errors::InvalidArgument("The Tensor Input's rank [%d] is not equal "
                                    "to the Tensor Output's rank [%d]",
                                    input_ndim,
@@ -140,8 +153,14 @@ SpmdInfo SliceInferSpmdReverseBase(const DistMetaTensor& input,
   // get einsum notation for input
   std::string input_axes = alphabet.substr(0, input_ndim);
 
+  auto output_input_axis_mapping =
+      BuildOutputAxisToInputAxisMap(decrease_axis, input_ndim);
   // get einsum notation for output
-  std::string out_axes(input_axes);
+  std::string out_axes = alphabet.substr(0, out_ndim);
+  for (int i = 0; i < out_ndim; i++) {
+    auto input_axis = output_input_axis_mapping[i];
+    out_axes[i] = input_axes[input_axis];
+  }
 
   for (int i = 0; i < static_cast<int>(axes.size()); i++) {
     int axis = axes[i] < 0 ? axes[i] + input_ndim : axes[i];
@@ -197,7 +216,7 @@ SpmdInfo SliceInferSpmdReverse(const DistMetaTensor& input,
                                const std::vector<int64_t>& decrease_axis) {
   // starts, ends, infer_flags and decrease_axis have no impact on the
   // derivation, only to align with the definition in phi api
-  return SliceInferSpmdReverseBase(input, output, axes);
+  return SliceInferSpmdReverseBase(input, output, axes, {});
 }
 
 SpmdInfo SliceInferSpmdDynamic(const DistMetaTensor& input,
@@ -211,27 +230,45 @@ SpmdInfo SliceInferSpmdDynamic(const DistMetaTensor& input,
   std::vector<int> start_indexes(starts.GetData().begin(),
                                  starts.GetData().end());
   std::vector<int> end_indexes(ends.GetData().begin(), ends.GetData().end());
-  return SliceInferSpmdBase(input, axes);
+  return SliceInferSpmdBase(input, axes, decrease_axis);
 }
 
 SpmdInfo SliceGradInferBase(const DistMetaTensor& input,
                             const DistMetaTensor& out_grad,
-                            const std::vector<int64_t>& axes) {
+                            const std::vector<int64_t>& axes,
+                            const std::vector<int64_t>& decrease_axis) {
   // Step0: Verify input args based on slice logic
+  auto input_shape = common::vectorize(input.dims());
+  int input_ndim = input_shape.size();
   auto input_dist_attr = input.dist_attr();
-  auto out_dist_attr = out_grad.dist_attr();
+
   input_dist_attr = UnShardTensorDims(input_dist_attr, axes);
-  out_dist_attr = UnShardTensorDims(out_dist_attr, axes);
+  std::vector<int64_t> input_dims_mapping = input_dist_attr.dims_mapping();
+
+  auto output_axis_to_input_axis_mapping =
+      BuildOutputAxisToInputAxisMap(decrease_axis, input_ndim);
+  std::unordered_map<int, int> reverse_output_axis_to_input_axis_mapping;
+  for (size_t i = 0; i < output_axis_to_input_axis_mapping.size(); ++i) {
+    reverse_output_axis_to_input_axis_mapping
+        [output_axis_to_input_axis_mapping[i]] = i;
+  }
+  std::vector<int64_t> mapped_axes;
+  for (size_t i = 0; i < axes.size(); ++i) {
+    int axis = axes[i] < 0 ? axes[i] + input_ndim : axes[i];
+    if (reverse_output_axis_to_input_axis_mapping.count(axis) > 0) {
+      mapped_axes.push_back(reverse_output_axis_to_input_axis_mapping[axis]);
+    }
+  }
+  auto out_dist_attr = out_grad.dist_attr();
+  out_dist_attr = UnShardTensorDims(out_dist_attr, mapped_axes);
   auto output_shape = common::vectorize(out_grad.dims());
   int out_ndim = output_shape.size();
   int out_dims_mapping_size = out_dist_attr.dims_mapping().size();
-  auto input_shape = common::vectorize(input.dims());
-  int input_ndim = input_shape.size();
-  std::vector<int64_t> input_dims_mapping = input_dist_attr.dims_mapping();
+  int decrease_axis_num = decrease_axis.size();
 
   PADDLE_ENFORCE_EQ(
       input_ndim,
-      out_ndim,
+      out_ndim + decrease_axis_num,
       phi::errors::InvalidArgument("The Tensor Input's rank [%d] is not equal "
                                    "to the Tensor Output's rank [%d]",
                                    input_ndim,
@@ -249,11 +286,13 @@ SpmdInfo SliceGradInferBase(const DistMetaTensor& input,
   std::string alphabet = "abcdefghijklmnopqrstuvwxyz";
 
   // get einsum notation for input
-  std::string align_axes = alphabet.substr(0, input_ndim);
-  std::string input_axes = align_axes;
+  std::string input_axes = alphabet.substr(0, input_ndim);
 
   // get einsum notation for output
-  std::string out_axes(input_axes);
+  std::string out_axes(out_ndim, ' ');
+  for (int i = 0; i < out_ndim; ++i) {
+    out_axes[i] = input_axes[output_axis_to_input_axis_mapping[i]];
+  }
 
   // Step2: Sharding Propogation
   // Step2.1: merge input shardings
@@ -266,14 +305,17 @@ SpmdInfo SliceGradInferBase(const DistMetaTensor& input,
       ShardingMergeForTensors(axes_sharding_info);
 
   // Step2.2: infer output dims mapping from merged input dims mapping
-  auto aligned_dim_mapping =
-      GetDimsMappingForAxes(align_axes, axis_to_dim_map, true);
+  auto input_dim_mapping_dst =
+      GetDimsMappingForAxes(input_axes, axis_to_dim_map, true);
 
   // get the dist attributes for output
-  TensorDistAttr aligned_dist_attr = CopyTensorDistAttrForOutput(out_dist_attr);
-  input_dist_attr.set_dims_mapping(aligned_dim_mapping);
-  out_dist_attr.set_dims_mapping(aligned_dim_mapping);
-  aligned_dist_attr.set_dims_mapping(aligned_dim_mapping);
+  TensorDistAttr input_grad = CopyTensorDistAttrForOutput(input_dist_attr);
+  input_dist_attr.set_dims_mapping(input_dim_mapping_dst);
+  input_grad.set_dims_mapping(input_dim_mapping_dst);
+
+  auto out_grad_dims_mapping =
+      GetDimsMappingForAxes(out_axes, axis_to_dim_map, true);
+  out_dist_attr.set_dims_mapping(out_grad_dims_mapping);
 
   VLOG(4) << "SliceGradInfer:";
 
@@ -293,10 +335,9 @@ SpmdInfo SliceGradInferBase(const DistMetaTensor& input,
 
   VLOG(4) << "input Grad"
           << " shape: [" << str_join(output_shape) << "] "
-          << "dims_mapping: [" << str_join(aligned_dist_attr.dims_mapping())
-          << "] ";
+          << "dims_mapping: [" << str_join(input_grad.dims_mapping()) << "] ";
 
-  return {{input_dist_attr, out_dist_attr}, {aligned_dist_attr}};
+  return {{input_dist_attr, out_dist_attr}, {input_grad}};
 }
 
 SpmdInfo SliceGradInferSpmdDynamic(const DistMetaTensor& input,
@@ -308,7 +349,7 @@ SpmdInfo SliceGradInferSpmdDynamic(const DistMetaTensor& input,
                                    const std::vector<int64_t>& decrease_axis) {
   // starts, ends, infer_flags and decrease_axis have no impact on the
   // derivation, only to align with the definition in phi api
-  return SliceGradInferBase(input, out_grad, axes);
+  return SliceGradInferBase(input, out_grad, axes, decrease_axis);
 }
 
 SpmdInfo StridedSliceInferSpmdDynamic(const DistMetaTensor& input,
@@ -319,7 +360,7 @@ SpmdInfo StridedSliceInferSpmdDynamic(const DistMetaTensor& input,
   // starts, ends and strides have no impact on the derivation,
   // only to align with the definition in phi api
   std::vector<int64_t> axes_bridge(axes.begin(), axes.end());
-  return SliceInferSpmdBase(input, axes_bridge);
+  return SliceInferSpmdBase(input, axes_bridge, {});
 }
 
 SpmdInfo StridedSliceGradInferSpmdDynamic(const DistMetaTensor& input,
@@ -331,7 +372,7 @@ SpmdInfo StridedSliceGradInferSpmdDynamic(const DistMetaTensor& input,
   // starts, ends and strides have no impact on the derivation,
   // only to align with the definition in phi api
   std::vector<int64_t> axes_bridge(axes.begin(), axes.end());
-  return SliceGradInferBase(input, out_grad, axes_bridge);
+  return SliceGradInferBase(input, out_grad, axes_bridge, {});
 }
 
 }  // namespace distributed
diff --git a/paddle/phi/infermeta/spmd_rules/utils.cc b/paddle/phi/infermeta/spmd_rules/utils.cc
@@ -170,6 +170,8 @@ TensorDistAttr GetReplicatedDistAttr(const TensorDistAttr& dist_attr) {
 TensorDistAttr ReplicateTensorDim(const TensorDistAttr& dist_attr, int dim) {
   TensorDistAttr dst_dist_attr = CopyTensorDistAttrForOutput(dist_attr);
   std::vector<int64_t> dims_mapping = dist_attr.dims_mapping();
+  int64_t n_dim = dims_mapping.size();
+  dim = dim < 0 ? n_dim + dim : dim;
   dims_mapping[dim] = kReplicateDim;
   dst_dist_attr.set_dims_mapping(dims_mapping);
   return dst_dist_attr;
@@ -178,6 +180,8 @@ TensorDistAttr ReplicateTensorDim(const TensorDistAttr& dist_attr, int dim) {
 TensorDistAttr UnShardTensorDim(const TensorDistAttr& dist_attr, int dim) {
   TensorDistAttr dst_dist_attr = CopyTensorDistAttrForOutput(dist_attr);
   std::vector<int64_t> dims_mapping = dist_attr.dims_mapping();
+  int64_t n_dim = dims_mapping.size();
+  dim = dim < 0 ? n_dim + dim : dim;
   dims_mapping[dim] = kReplicateDim;
   dst_dist_attr.set_dims_mapping(dims_mapping);
   return dst_dist_attr;
diff --git a/test/auto_parallel/semi_auto_parallel_for_concat.py b/test/auto_parallel/semi_auto_parallel_for_concat.py
@@ -104,6 +104,34 @@ def test_slice(self):
             ends=[3, 3],
         )
 
+    def test_decrease_axes_slice_shard(self):
+        def decrease_axes_slice(x):
+            return x[0]
+
+        shapes = [64, 4, 4]
+        specs = ["x", None, None]
+        inputs, outputs = self.runfunc_and_check(
+            inputs_shape=shapes,
+            inputs_specs=specs,
+            op_func=decrease_axes_slice,
+            with_backward=True,
+        )
+        self.check_dim_mapping(outputs, [-1, -1])
+
+    def test_decrease_axes_slice(self):
+        def decrease_axes_slice(x):
+            return x[0]
+
+        shapes = [64, 4, 4]
+        specs = [None, None, 'x']
+        inputs, outputs = self.runfunc_and_check(
+            inputs_shape=shapes,
+            inputs_specs=specs,
+            op_func=decrease_axes_slice,
+            with_backward=True,
+        )
+        self.check_dim_mapping(outputs, [-1, 0])
+
     def test_slice_reshard(self):
         shapes = [64, 4, 4]
         specs = [None, 'x', None]
@@ -157,12 +185,14 @@ def run_test_case(self):
         self.test_stack_forward()
         self.test_slice()
         self.test_stride_slice()
+        self.test_decrease_axes_slice()
         # all to all is not supported yet for cpu
         if self._backend == "gpu":
             self.test_concat_forward_reshard()
             self.test_slice_reshard()
             self.test_stride_slice_reshard()
             self.test_stack_forward_reshard()
+            self.test_decrease_axes_slice_shard()
 
 
 if __name__ == '__main__':
