Merge commit 'refs/pull/59798/head' of github.com:PaddlePaddle/Paddle into sink

Author: cyber-pioneer

paddle/fluid/pir/dialect/op_generator/decomp_interface_gen_op_list.py

@@ -34,6 +34,7 @@
     "sqrt",
     "squeeze",
     "stack",
+    "unsqueeze",
 ]
 
 # come into effect in generated file op_decomp.cc
@@ -52,7 +53,9 @@
     "softmax",
     "sqrt",
     "squeeze",
+    "unsqueeze",
     "stack",
+    "unsqueeze",
 ]
 
 

paddle/fluid/primitive/composite/composite.h

@@ -246,14 +246,8 @@ Tensor softmax_decomp(const Tensor& x, const int& axis) {
 
 template <typename T>
 Tensor stack_decomp(const std::vector<Tensor>& x, const int& axis) {
-  auto tensor_dims = x[0].dims();
-  int tmp_axis = axis;
-  if (tmp_axis < 0) {
-    tmp_axis += tensor_dims.size() + 1;
-  }
-
-  auto out_shape = phi::vectorize(tensor_dims);
-  out_shape.insert(out_shape.begin() + tmp_axis, 1);
+  std::vector<int64_t> axis_tmp = {axis};
+  auto out_shape = get_expand_dims(x[0], axis_tmp);
 
   std::vector<Tensor> concat_x;
   for (size_t i = 0; i < x.size(); ++i) {
@@ -318,6 +312,15 @@ std::tuple<Tensor, Tensor> squeeze_decomp(const Tensor& x,
   return std::make_tuple(out, xshape);
 }
 
+template <typename T>
+std::tuple<Tensor, Tensor> unsqueeze_decomp(const Tensor& x,
+                                            const IntArray& axis) {
+  auto out_shape = get_expand_dims(x, axis.GetData());
+  Tensor out = reshape<T>(x, out_shape);
+  Tensor xshape;
+  return std::make_tuple(out, xshape);
+}
+
 template <typename T>
 Tensor add_n_decomp(const std::vector<Tensor>& x) {
   Tensor res = x[0];

paddle/fluid/primitive/utils/utils.h

@@ -39,6 +39,21 @@ static bool is_half_dtype(const DataType& dtype) {
   }
 }
 
+// This function expands the dimension of origin Tensor based on the value of
+// axis
+static std::vector<int64_t> get_expand_dims(const Tensor& origin,
+                                            const std::vector<int64_t>& axis) {
+  std::vector<int64_t> result(origin.shape());
+  for (size_t i = 0; i < axis.size(); ++i) {
+    int64_t offset = axis[i];
+    if (offset < 0) {
+      offset += result.size() + 1;
+    }
+    result.insert(result.begin() + offset, 1);
+  }
+  return result;
+}
+
 // This fucction compute unsqueeze dims for reshape to replace unsqueeze.
 static std::vector<int64_t> get_unsqueeze_dims(
     const Tensor& origin, const std::vector<int64_t>& axis) {

python/paddle/decomposition/rules.py

@@ -11,28 +11,3 @@
 # 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.
-
-
-from .primitives import *  # noqa: F403
-from .register import register_decomp
-
-
-@register_decomp('pd_op.unsqueeze')
-def unsqueeze(x, axis):
-    """define composite rule of op unsqueeze"""
-    """using reshape to implement unsqueeze op"""
-    axis = axis.get_defining_op().attrs()["value"]
-    x_shape = list(x.shape)
-    axis_list = list(axis)
-    for i in axis_list:
-        if i < 0:
-            i += len(x_shape) + 1
-        x_shape = (
-            x_shape[:i]
-            + [
-                1,
-            ]
-            + x_shape[i:]
-        )
-    out = reshape(x, x_shape)
-    return [out, None]