[Prim][PIR] forward decomp unstack op (#67593)

cubehan3 · web-flow · commit 069f9199bf6f · 2024-08-22T09:36:30.000+08:00
* forward decomp unstack op

* fix bugs
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
@@ -36,6 +36,7 @@
     "hardswish",
     "hardsigmoid",
     "heaviside",
+    "huber_loss",
     "group_norm",
     "index_sample",
     "index_select",
@@ -64,9 +65,9 @@
     "squared_l2_norm",
     "squeeze",
     "stack",
-    "unsqueeze",
     "unbind",
-    "huber_loss",
+    "unsqueeze",
+    "unstack",
 ]
 
 # come into effect in generated file op_decomp.cc
@@ -86,6 +87,7 @@
     "hardswish",
     "hardsigmoid",
     "heaviside",
+    "huber_loss",
     "group_norm",
     "index_sample",
     "index_select",
@@ -113,9 +115,9 @@
     "squared_l2_norm",
     "squeeze",
     "stack",
-    "unsqueeze",
     "unbind",
-    "huber_loss",
+    "unsqueeze",
+    "unstack",
 ]
 
 # xshape output will no longer used after decomp, but return none to keep output num the same as origin op
diff --git a/paddle/fluid/primitive/composite/composite.h b/paddle/fluid/primitive/composite/composite.h
@@ -1612,6 +1612,56 @@ Tensor softsign_decomp(const Tensor& x) {
   return x / (one + x_abs);
 }
 
+template <typename T>
+std::vector<Tensor> unstack_decomp(const Tensor& x, int axis, const int num) {
+  if (axis < 0) {
+    axis += x.dims().size();
+  }
+  std::vector<int64_t> x_shape = x.shape();
+  if (x_shape[axis] < 0) {
+    PADDLE_THROW(
+        common::errors::Unimplemented("unstack axis must not be dynamic."));
+  }
+  PADDLE_ENFORCE_EQ(
+      num,
+      x_shape[axis],
+      common::errors::InvalidArgument(
+          "The number of unstacks should be equal to the value of "
+          "x.shape[axis], but received num is %d and x.shape[axis] is %d.",
+          num,
+          x_shape[axis]));
+
+  std::vector<int> sections(num, 1);
+  std::vector<Tensor> res = backend::split<T>(x, sections, axis);
+  if (has_dynamic_shape(x_shape)) {
+    const Tensor x_shape_tensor = shape<T>(x);
+
+    // find new shape of each tensor.
+    std::vector<Tensor> new_shape_vec;
+    for (size_t i = 0; i < x_shape.size(); ++i) {
+      if (static_cast<int>(i) != axis) {
+        new_shape_vec.push_back(get_slice<T>(x_shape_tensor, i));
+      }
+    }
+    const Tensor new_shape = concat<T>(new_shape_vec);
+    std::transform(res.begin(), res.end(), res.begin(), [&](Tensor& x) {
+      return backend::reshape_with_tensor<T>(x, new_shape);
+    });
+  } else {
+    std::vector<int64_t> new_shape;
+    // find new shape of each tensor.
+    for (size_t i = 0; i < x_shape.size(); ++i) {
+      if (static_cast<int>(i) != axis) {
+        new_shape.push_back(x_shape[i]);
+      }
+    }
+    std::transform(res.begin(), res.end(), res.begin(), [&](Tensor& x) {
+      return reshape<T>(x, new_shape);
+    });
+  }
+  return res;
+}
+
 }  // namespace details
 
 }  // namespace primitive
diff --git a/test/legacy_test/test_unstack_op.py b/test/legacy_test/test_unstack_op.py
@@ -42,7 +42,9 @@ def setUp(self):
         self.initDefaultParameters()
         self.initParameters()
         self.op_type = 'unstack'
+        self.prim_op_type = "comp"
         self.python_api = paddle.unstack
+        self.public_python_api = paddle.unstack
         self.x = np.random.random(size=self.input_dim).astype(self.dtype)
 
         outs = np.split(self.x, self.input_dim[self.axis], self.axis)
@@ -61,10 +63,12 @@ def setUp(self):
         self.attrs = {'axis': self.axis, 'num': self.input_dim[self.axis]}
 
     def test_check_output(self):
-        self.check_output(check_pir=True)
+        self.check_output(check_pir=True, check_prim_pir=True)
 
     def test_check_grad(self):
-        self.check_grad(['X'], self.get_y_names(), check_pir=True)
+        self.check_grad(
+            ['X'], self.get_y_names(), check_pir=True, check_prim_pir=True
+        )
 
 
 class TestUnStackFP16Op(TestUnStackOpBase):
@@ -188,7 +192,9 @@ def setUp(self):
         self.initDefaultParameters()
         self.initParameters()
         self.op_type = 'unstack'
+        self.prim_op_type = "comp"
         self.python_api = paddle.unstack
+        self.public_python_api = paddle.unstack
         self.x = np.random.random(size=self.input_dim).astype(np.float32)
         outs = np.split(self.x, self.input_dim[self.axis], self.axis)
         new_shape = list(self.input_dim)
@@ -213,7 +219,7 @@ def setUp(self):
 
     def test_check_output(self):
         place = core.CUDAPlace(0)
-        self.check_output_with_place(place, check_pir=True)
+        self.check_output_with_place(place, check_pir=True, check_prim_pir=True)
 
     def test_check_grad(self):
         with base.dygraph.guard():
diff --git a/test/prim/pir_prim/test_prim_sub_graph_dynamic_shape.py b/test/prim/pir_prim/test_prim_sub_graph_dynamic_shape.py
@@ -276,6 +276,30 @@ def sigmoid_cross_entropy_with_logits_net2(x, y, z):
     return paddle._C_ops.sigmoid_cross_entropy_with_logits(x, y, z, True, -1)
 
 
+def unstack_net1(x):
+    return paddle.unstack(x, axis=0)
+
+
+def unstack_net2(x):
+    return paddle.unstack(x, axis=1)
+
+
+def unstack_net3(x):
+    return paddle.unstack(x, axis=2)
+
+
+def unstack_net4(x):
+    return paddle.unstack(x, axis=-3)
+
+
+def unstack_net5(x):
+    return paddle.unstack(x, axis=-2)
+
+
+def unstack_net6(x):
+    return paddle.unstack(x, axis=-1)
+
+
 class TestPrimBase(unittest.TestCase):
     def setUp(self):
         np.random.seed(2023)
@@ -1320,5 +1344,89 @@ def setUp(self):
         self.tol = 1e-6
 
 
+class TestPrimUnstack1(TestPrimBase):
+    def setUp(self):
+        np.random.seed(2024)
+        paddle.seed(2024)
+        self.shape_x = [5, 10, 15]
+        self.dtype_x = "float32"
+        self.init_x_shape = [5, None, None]
+        self.x = np.random.random(self.shape_x).astype(self.dtype_x)
+        self.net = unstack_net1
+        self.necessary_ops = "pd_op.unstack"
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimUnstack2(TestPrimBase):
+    def setUp(self):
+        np.random.seed(2024)
+        paddle.seed(2024)
+        self.shape_x = [5, 10, 15]
+        self.dtype_x = "float32"
+        self.init_x_shape = [None, 10, None]
+        self.x = np.random.random(self.shape_x).astype(self.dtype_x)
+        self.net = unstack_net2
+        self.necessary_ops = "pd_op.unstack"
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimUnstack3(TestPrimBase):
+    def setUp(self):
+        np.random.seed(2024)
+        paddle.seed(2024)
+        self.shape_x = [5, 10, 15]
+        self.dtype_x = "float32"
+        self.init_x_shape = [None, None, 15]
+        self.x = np.random.random(self.shape_x).astype(self.dtype_x)
+        self.net = unstack_net3
+        self.necessary_ops = "pd_op.unstack"
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimUnstack4(TestPrimBase):
+    def setUp(self):
+        np.random.seed(2024)
+        paddle.seed(2024)
+        self.shape_x = [5, 10, 15]
+        self.dtype_x = "float32"
+        self.init_x_shape = [5, None, None]
+        self.x = np.random.random(self.shape_x).astype(self.dtype_x)
+        self.net = unstack_net4
+        self.necessary_ops = "pd_op.unstack"
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimUnstack5(TestPrimBase):
+    def setUp(self):
+        np.random.seed(2024)
+        paddle.seed(2024)
+        self.shape_x = [5, 10, 15]
+        self.dtype_x = "float32"
+        self.init_x_shape = [None, 10, None]
+        self.x = np.random.random(self.shape_x).astype(self.dtype_x)
+        self.net = unstack_net5
+        self.necessary_ops = "pd_op.unstack"
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimUnstack6(TestPrimBase):
+    def setUp(self):
+        np.random.seed(2024)
+        paddle.seed(2024)
+        self.shape_x = [5, 10, 15]
+        self.dtype_x = "float32"
+        self.init_x_shape = [None, None, 15]
+        self.x = np.random.random(self.shape_x).astype(self.dtype_x)
+        self.net = unstack_net6
+        self.necessary_ops = "pd_op.unstack"
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
 if __name__ == "__main__":
     unittest.main()
