[Prim][PIR] unsequeeze prim sink (#59798)

kevincheng2 · web-flow · commit 50d8240141da · 2023-12-09T09:24:45.000+08:00
* prim unsqueeze sink

* prim unsequeeze sink

* update code

* update code

* update code

* update code
diff --git a/paddle/fluid/pir/dialect/op_generator/decomp_interface_gen_op_list.py b/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
@@ -53,6 +54,7 @@
     "sqrt",
     "squeeze",
     "stack",
+    "unsqueeze",
 ]
 
 
diff --git a/paddle/fluid/primitive/composite/composite.h b/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];
diff --git a/paddle/fluid/primitive/utils/utils.h b/paddle/fluid/primitive/utils/utils.h
@@ -39,6 +39,29 @@ 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;
+    }
+
+    PADDLE_ENFORCE_LE(
+        offset,
+        result.size(),
+        platform::errors::OutOfRange("Your index [%lu] exceeds the number of "
+                                     "elements in origin_dims[%lu].",
+                                     offset,
+                                     result.size()));
+    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) {
diff --git a/python/paddle/decomposition/rules.py b/python/paddle/decomposition/rules.py
@@ -17,6 +17,8 @@
 from .register import register_decomp
 
 
+# TODO(kevincheng2): python implementation of prim feature,
+# now it has been sinked to c++, waiting for further deletion.
 @register_decomp('pd_op.unsqueeze')
 def unsqueeze(x, axis):
     """define composite rule of op unsqueeze"""

Original file line number	Diff line number	Diff line change
`@@ -34,6 +34,7 @@`
`34`	`34`	`"sqrt",`
`35`	`35`	`"squeeze",`
`36`	`36`	`"stack",`
	`37`	`+ "unsqueeze",`
`37`	`38`	`]`
`38`	`39`
`39`	`40`	`# come into effect in generated file op_decomp.cc`
`@@ -53,6 +54,7 @@`
`53`	`54`	`"sqrt",`
`54`	`55`	`"squeeze",`
`55`	`56`	`"stack",`
	`57`	`+ "unsqueeze",`
`56`	`58`	`]`
`57`	`59`
`58`	`60`