[auto parallel] add squeeze/unsqueeze backward spmd rules

paddle/phi/api/yaml/backward.yaml

@@ -2282,6 +2282,7 @@
   infer_meta :
     func : KernelWithXShapeInferMeta
     param: [xshape, out_grad]
+    spmd_rule : SqueezeGradInferSpmd
   kernel :
     func : squeeze_grad
     data_type : out_grad
@@ -2527,6 +2528,7 @@
   infer_meta :
     func : KernelWithXShapeInferMeta
     param: [xshape, out_grad]
+    spmd_rule : UnsqueezeGradInferSpmd
   kernel :
     func : unsqueeze_grad
     param : [xshape, out_grad]

paddle/phi/api/yaml/ops.yaml

@@ -2489,6 +2489,7 @@
   output : Tensor(out), Tensor(xshape)
   infer_meta :
     func : SqueezeWithXShapeInferMeta
+    spmd_rule : SqueezeInferSpmd
   kernel :
     func : squeeze
     data_type : x
@@ -2714,6 +2715,7 @@
   output : Tensor(out), Tensor(xshape)
   infer_meta :
     func : UnsqueezeWithXShapeInferMeta
+    spmd_rule : UnsqueezeInferSpmd
   kernel :
     func : unsqueeze
     data_type : x

paddle/phi/infermeta/spmd_rules/squeeze.cc

@@ -1,16 +1,16 @@
-/* Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License. */
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
 
 #include "paddle/phi/infermeta/spmd_rules/squeeze.h"
 #include <algorithm>
@@ -22,13 +22,22 @@ limitations under the License. */
 #include "paddle/phi/core/distributed/auto_parallel/inferspmd_utils.h"
 #include "paddle/phi/core/distributed/auto_parallel/utils.h"
 #include "paddle/phi/infermeta/spmd_rules/dim_trans.h"
+#include "paddle/phi/infermeta/spmd_rules/reshape.h"
 #include "paddle/phi/infermeta/spmd_rules/utils.h"
 
 namespace phi {
 namespace distributed {
 
 using phi::distributed::auto_parallel::str_join;
 
+TensorDistAttr CreateSqueezeXshape(const TensorDistAttr& x) {
+  TensorDistAttr out(x);
+  auto dims_mapping = x.dims_mapping();
+  dims_mapping.insert(dims_mapping.begin(), -1);
+  out.set_dims_mapping(dims_mapping);
+  return out;
+}
+
 void MakeSqueezeDimTransWithoutAxis(
     const std::vector<int64_t>& x_shape,
     std::vector<int64_t>* out_shape,
@@ -137,9 +146,18 @@ SpmdInfo SqueezeInferSpmd(const DistMetaTensor& x,
   // and output with the inferred dims mapping.
   TensorDistAttr x_dist_attr_dst(x_dist_attr_src);
   x_dist_attr_dst.set_dims_mapping(dims_mapping_vec[0]);
+  if (x_dist_attr_dst.dynamic_dims().size() !=
+      x_dist_attr_dst.dims_mapping().size()) {
+    VLOG(4) << "SqueezeInferSPMD change x dist attr dynamic dims";
+    x_dist_attr_dst.set_default_dynamic_dims(x_dist_attr_dst.dims_mapping());
+  }
   TensorDistAttr out_dist_attr(x_dist_attr_src);
   out_dist_attr.set_dims_mapping(dims_mapping_vec[1]);
-
+  if (out_dist_attr.dynamic_dims().size() !=
+      out_dist_attr.dims_mapping().size()) {
+    VLOG(4) << "SqueezeInferSPMD change output dist attr dynamic dims";
+    out_dist_attr.set_default_dynamic_dims(out_dist_attr.dims_mapping());
+  }
   VLOG(4) << "SqueezeInferSpmd: X shape: [" << str_join(x_shape)
           << "] Out shape: [" << str_join(out_shape) << "]";
   VLOG(4) << "Transformation from input to output:";
@@ -151,7 +169,8 @@ SpmdInfo SqueezeInferSpmd(const DistMetaTensor& x,
           << "]\n Out dims_mapping: [" << str_join(dims_mapping_vec[1])
           << "]\n\n";
 
-  return {{x_dist_attr_dst}, {out_dist_attr}};
+  return {{x_dist_attr_dst},
+          {out_dist_attr, CreateSqueezeXshape(x_dist_attr_dst)}};
 }
 
 SpmdInfo SqueezeInferSpmdReverse(const DistMetaTensor& x,
@@ -202,8 +221,18 @@ SpmdInfo SqueezeInferSpmdReverse(const DistMetaTensor& x,
   // and output with the inferred dims mapping
   TensorDistAttr out_dist_attr_dst(out_dist_attr_src);
   out_dist_attr_dst.set_dims_mapping(dims_mapping_vec[0]);
+  if (out_dist_attr_dst.dynamic_dims().size() !=
+      out_dist_attr_dst.dims_mapping().size()) {
+    VLOG(4) << "SqueezeInferSPMD change output dist attr dynamic dims";
+    out_dist_attr_dst.set_default_dynamic_dims(
+        out_dist_attr_dst.dims_mapping());
+  }
   TensorDistAttr x_dist_attr(x.dist_attr());
   x_dist_attr.set_dims_mapping(dims_mapping_vec[1]);
+  if (x_dist_attr.dynamic_dims().size() != x_dist_attr.dims_mapping().size()) {
+    VLOG(4) << "SqueezeInferSPMD change x dist attr dynamic dims";
+    x_dist_attr.set_default_dynamic_dims(x_dist_attr.dims_mapping());
+  }
 
   VLOG(4) << "SqueezeInferSpmdReverse: Out shape: [" << str_join(out_shape)
           << "] X shape: [" << str_join(x_shape) << "]";
@@ -218,5 +247,14 @@ SpmdInfo SqueezeInferSpmdReverse(const DistMetaTensor& x,
   return {{x_dist_attr}, {out_dist_attr_dst}};
 }
 
+SpmdInfo SqueezeGradInferSpmd(const DistMetaTensor& xshape,
+                              const DistMetaTensor& out_grad,
+                              const IntArray& axis) {
+  auto shape = phi::vectorize(xshape.dims());
+  shape = std::vector<int64_t>(shape.begin() + 1, shape.end());
+  const auto& spmd = ReshapeInferSpmd(out_grad, shape);
+  return {{xshape.dist_attr(), spmd.first[0]}, {spmd.second[0]}};
+}
+
 }  // namespace distributed
 }  // namespace phi

paddle/phi/infermeta/spmd_rules/squeeze.h

@@ -16,6 +16,7 @@ limitations under the License. */
 
 #include <vector>
 
+#include "paddle/phi/common/int_array.h"
 #include "paddle/phi/core/distributed/auto_parallel/dist_meta_tensor.h"
 #include "paddle/phi/core/distributed/type_defs.h"
 
@@ -28,5 +29,9 @@ SpmdInfo SqueezeInferSpmd(const DistMetaTensor& x,
 SpmdInfo SqueezeInferSpmdReverse(const DistMetaTensor& x,
                                  const DistMetaTensor& out,
                                  const std::vector<int64_t>& axis);
+
+SpmdInfo SqueezeGradInferSpmd(const DistMetaTensor& xshape,
+                              const DistMetaTensor& out_grad,
+                              const IntArray& axis = {});
 }  // namespace distributed
 }  // namespace phi

paddle/phi/infermeta/spmd_rules/unsqueeze.cc

@@ -1,16 +1,16 @@
-/* Copyright (c) 2023 PaddlePaddle Authors. All Rights resized.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License. */
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
 
 #include "paddle/phi/infermeta/spmd_rules/unsqueeze.h"
 #include <algorithm>
@@ -22,13 +22,22 @@ limitations under the License. */
 #include "paddle/phi/core/distributed/auto_parallel/inferspmd_utils.h"
 #include "paddle/phi/core/distributed/auto_parallel/utils.h"
 #include "paddle/phi/infermeta/spmd_rules/dim_trans.h"
+#include "paddle/phi/infermeta/spmd_rules/reshape.h"
 #include "paddle/phi/infermeta/spmd_rules/utils.h"
 
 namespace phi {
 namespace distributed {
 
 using phi::distributed::auto_parallel::str_join;
 
+TensorDistAttr CreateUnsqueezeXshape(const TensorDistAttr& x) {
+  TensorDistAttr out(x);
+  auto dims_mapping = x.dims_mapping();
+  dims_mapping.insert(dims_mapping.begin(), -1);
+  out.set_dims_mapping(dims_mapping);
+  return out;
+}
+
 std::vector<std::shared_ptr<DimTrans>> MakeUnsqueezeDimTrans(
     const std::vector<int64_t>& x_shape,
     std::vector<int64_t>* out_shape,
@@ -119,8 +128,18 @@ SpmdInfo UnsqueezeInferSpmd(const DistMetaTensor& x,
   // and output with the inferred dims mapping.
   TensorDistAttr x_dist_attr_dst(x_dist_attr_src);
   x_dist_attr_dst.set_dims_mapping(dims_mapping_vec[0]);
+  if (x_dist_attr_dst.dynamic_dims().size() !=
+      x_dist_attr_dst.dims_mapping().size()) {
+    VLOG(4) << "UnSqueezeInferSPMD change output dist attr dynamic dims";
+    x_dist_attr_dst.set_default_dynamic_dims(x_dist_attr_dst.dims_mapping());
+  }
   TensorDistAttr out_dist_attr(x_dist_attr_src);
   out_dist_attr.set_dims_mapping(dims_mapping_vec[1]);
+  if (out_dist_attr.dynamic_dims().size() !=
+      out_dist_attr.dims_mapping().size()) {
+    VLOG(4) << "UnSqueezeInferSPMD change output dist attr dynamic dims";
+    out_dist_attr.set_default_dynamic_dims(out_dist_attr.dims_mapping());
+  }
 
   VLOG(4) << "UnsqueezeInferSpmd: X shape: [" << str_join(x_shape)
           << "] Out shape: [" << str_join(out_shape) << "]";
@@ -134,7 +153,8 @@ SpmdInfo UnsqueezeInferSpmd(const DistMetaTensor& x,
           << "]\n Out dims_mapping: [" << str_join(dims_mapping_vec[1])
           << "]\n\n";
 
-  return {{x_dist_attr_dst}, {out_dist_attr}};
+  return {{x_dist_attr_dst},
+          {out_dist_attr, CreateUnsqueezeXshape(x_dist_attr_dst)}};
 }
 
 SpmdInfo UnsqueezeInferSpmdReverse(const DistMetaTensor& x,
@@ -181,9 +201,18 @@ SpmdInfo UnsqueezeInferSpmdReverse(const DistMetaTensor& x,
   // and output with the inferred dims mapping
   TensorDistAttr out_dist_attr_dst(out_dist_attr_src);
   out_dist_attr_dst.set_dims_mapping(dims_mapping_vec[0]);
+  if (out_dist_attr_dst.dynamic_dims().size() !=
+      out_dist_attr_dst.dims_mapping().size()) {
+    VLOG(4) << "UnSqueezeInferSPMDReverse change output dist attr dynamic dims";
+    out_dist_attr_dst.set_default_dynamic_dims(
+        out_dist_attr_dst.dims_mapping());
+  }
   TensorDistAttr x_dist_attr(x.dist_attr());
   x_dist_attr.set_dims_mapping(dims_mapping_vec[1]);
-
+  if (x_dist_attr.dynamic_dims().size() != x_dist_attr.dims_mapping().size()) {
+    VLOG(4) << "UnSqueezeInferSPMDReverse change x dist attr dynamic dims";
+    x_dist_attr.set_default_dynamic_dims(x_dist_attr.dims_mapping());
+  }
   VLOG(4) << "UnsqueezeInferSpmdReverse: Out shape: [" << str_join(out_shape)
           << "] X shape: [" << str_join(x_shape) << "]";
   VLOG(4) << "Transformation from output to input:";
@@ -198,5 +227,14 @@ SpmdInfo UnsqueezeInferSpmdReverse(const DistMetaTensor& x,
   return {{x_dist_attr}, {out_dist_attr_dst}};
 }
 
+SpmdInfo UnsqueezeGradInferSpmd(const DistMetaTensor& xshape,
+                                const DistMetaTensor& out_grad,
+                                const IntArray& axis) {
+  auto shape = phi::vectorize(xshape.dims());
+  shape = std::vector<int64_t>(shape.begin() + 1, shape.end());
+  const auto& spmd = ReshapeInferSpmd(out_grad, shape);
+  return {{xshape.dist_attr(), spmd.first[0]}, {spmd.second[0]}};
+}
+
 }  // namespace distributed
 }  // namespace phi

paddle/phi/infermeta/spmd_rules/unsqueeze.h

@@ -16,6 +16,7 @@ limitations under the License. */
 
 #include <vector>
 
+#include "paddle/phi/common/int_array.h"
 #include "paddle/phi/core/distributed/auto_parallel/dist_meta_tensor.h"
 #include "paddle/phi/core/distributed/type_defs.h"
 
@@ -28,5 +29,8 @@ SpmdInfo UnsqueezeInferSpmd(const DistMetaTensor& x,
 SpmdInfo UnsqueezeInferSpmdReverse(const DistMetaTensor& x,
                                    const DistMetaTensor& out,
                                    const std::vector<int64_t>& axis);
+SpmdInfo UnsqueezeGradInferSpmd(const DistMetaTensor& xshape,
+                                const DistMetaTensor& out_grad,
+                                const IntArray& axis = {});
 }  // namespace distributed
 }  // namespace phi