[Prim][PIR] stack prim sink (#59713)

kevincheng2 · web-flow · commit 6e704556d6f2 · 2023-12-07T10:45:17.000+08:00
* stack sink

* prim stack sink

* stack sink
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
@@ -33,6 +33,7 @@
     "softmax",
     "sqrt",
     "squeeze",
+    "stack",
 ]
 
 # come into effect in generated file op_decomp.cc
@@ -51,6 +52,7 @@
     "softmax",
     "sqrt",
     "squeeze",
+    "stack",
 ]
 
 
diff --git a/paddle/fluid/primitive/composite/composite.h b/paddle/fluid/primitive/composite/composite.h
@@ -244,6 +244,24 @@ 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<Tensor> concat_x;
+  for (size_t i = 0; i < x.size(); ++i) {
+    concat_x.push_back(reshape<T>(x[i], out_shape));
+  }
+  return concat<T>(concat_x, axis);
+}
+
 template <typename T>
 Tensor silu_decomp(const Tensor& x) {
   auto org_dtype = x.dtype();
diff --git a/python/paddle/decomposition/rules.py b/python/paddle/decomposition/rules.py
@@ -17,20 +17,6 @@
 from .register import register_decomp
 
 
-@register_decomp('pd_op.stack')
-def stack(x, axis):
-    """
-    define composite rule of op stack
-    unsqueeze each dimension of the input (use reshape), and then concat
-    """
-    x_shape = x[0].shape
-    if axis < 0:
-        axis += len(x_shape) + 1
-    out_shape = x_shape[:axis] + [1] + x_shape[axis:]
-    out = concat([reshape(item, out_shape) for item in x], axis)
-    return out
-
-
 @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
`@@ -33,6 +33,7 @@`
`33`	`33`	`"softmax",`
`34`	`34`	`"sqrt",`
`35`	`35`	`"squeeze",`
	`36`	`+ "stack",`
`36`	`37`	`]`
`37`	`38`
`38`	`39`	`# come into effect in generated file op_decomp.cc`
`@@ -51,6 +52,7 @@`
`51`	`52`	`"softmax",`
`52`	`53`	`"sqrt",`
`53`	`54`	`"squeeze",`
	`55`	`+ "stack",`
`54`	`56`	`]`
`55`	`57`
`56`	`58`