[Prim][PIR] add the test case for concat_grad backward decomposite (#65252)

zeroRains · web-flow · commit 1eb6f6a30267 · 2024-07-01T10:47:36.000+08:00
* move to utils func and add some test for concat_grad

* fix bug

* fix code style

* polish
diff --git a/paddle/fluid/primitive/composite/composite.h b/paddle/fluid/primitive/composite/composite.h
@@ -15,26 +15,13 @@
 #pragma once
 
 #include <numeric>
-#include "paddle/fluid/primitive/primitive/primitive.h"
 #include "paddle/fluid/primitive/type/lazy_tensor.h"
 #include "paddle/fluid/primitive/utils/utils.h"
 
 namespace paddle {
 namespace primitive {
 namespace details {
 
-template <typename T>
-static Tensor get_slice(const Tensor& x, int64_t idx) {
-  return slice<T>(x, {0}, {idx}, {idx + 1}, {1}, {});
-}
-
-template <typename T>
-static Tensor get_slice_vec(const Tensor& x,
-                            int64_t start_idx,
-                            int64_t end_idx) {
-  return slice<T>(x, {0}, {start_idx}, {end_idx}, {1}, {});
-}
-
 template <typename T>
 Tensor any_decomp(const Tensor& x, const IntArray& axis, bool keepdim) {
   auto org_dtype = x.dtype();
diff --git a/paddle/fluid/primitive/rule/vjp/details.h b/paddle/fluid/primitive/rule/vjp/details.h
@@ -20,8 +20,6 @@
 
 #include <math.h>
 #include <vector>
-
-#include "paddle/fluid/primitive/primitive/primitive.h"
 #include "paddle/fluid/primitive/type/lazy_tensor.h"
 #include "paddle/fluid/primitive/utils/utils.h"
 
@@ -193,8 +191,7 @@ void sum_grad(const Tensor& x,
             result_shape.push_back(ones);
             j++;
           } else {
-            result_shape.push_back(slice<T>(
-                out_grad_shape, {0}, {int64_t(k)}, {int64_t(k) + 1}, {1}, {}));
+            result_shape.push_back(get_slice<T>(out_grad_shape, int64_t(k)));
             k++;
           }
         }
@@ -513,15 +510,17 @@ void concat_grad(const std::vector<Tensor>& x,
 
   int x_num = x.size();
   std::vector<Tensor> x_grad_tmp;
-  if (has_dynamic_shape(x[0].shape())) {
+  bool has_dynamic = false;
+  for (size_t i = 0; i < x.size(); i++) {
+    if (has_dynamic_shape(x[i].shape())) {
+      has_dynamic = true;
+      break;
+    }
+  }
+  if (has_dynamic) {
     std::vector<Tensor> sections;
     for (int i = 0; i < x_num; i++) {
-      sections.push_back(slice<T>(shape<T>(x[i]),
-                                  {0},
-                                  {int64_t(axis_value)},
-                                  {int64_t(axis_value) + 1},
-                                  {1},
-                                  {}));
+      sections.push_back(get_slice<T>(shape<T>(x[i]), int64_t(axis_value)));
     }
     Tensor sections_tensor = concat<T>(sections);
     x_grad_tmp =
diff --git a/paddle/fluid/primitive/utils/utils.h b/paddle/fluid/primitive/utils/utils.h
@@ -18,13 +18,26 @@
 #include "paddle/common/ddim.h"
 #include "paddle/fluid/framework/details/op_registry.h"
 #include "paddle/fluid/framework/operator.h"
+#include "paddle/fluid/primitive/primitive/primitive.h"
 #include "paddle/fluid/primitive/type/lazy_tensor.h"
 #include "paddle/phi/api/include/tensor.h"
 #include "paddle/phi/kernels/funcs/common_infer_shape_functions.h"
 
 namespace paddle {
 class Tensor;
 namespace primitive {
+template <typename T>
+static Tensor get_slice(const Tensor& x, int64_t idx) {
+  return slice<T>(x, {0}, {idx}, {idx + 1}, {1}, {});
+}
+
+template <typename T>
+static Tensor get_slice_vec(const Tensor& x,
+                            int64_t start_idx,
+                            int64_t end_idx) {
+  return slice<T>(x, {0}, {start_idx}, {end_idx}, {1}, {});
+}
+
 template <typename T>
 void set_output(const Tensor& x_tmp, Tensor* x);
 
diff --git a/test/prim/pir_prim/test_prim_sub_graph_backward_dynamic_shape.py b/test/prim/pir_prim/test_prim_sub_graph_backward_dynamic_shape.py
@@ -71,6 +71,10 @@ def concat_net2(x):
     return paddle.concat([x, y], axis=1)
 
 
+def concat_net3(x):
+    return paddle.concat(x, axis=0)
+
+
 def split_net1(x):
     res = paddle.split(x, num_or_sections=10, axis=-1)
     tmp_res = res[0]
@@ -607,6 +611,74 @@ def setUp(self):
         self.tol = 1e-6
 
 
+class TestPrimConcatWithGrad5(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, 40]
+        x = np.random.random(self.x_shape).astype(self.dtype)
+        self.x = [x + i for i in range(4)]
+        self.net = concat_net3
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+    def base_net(self, flag=None):
+        if flag == "prim":
+            core._set_prim_all_enabled(True)
+        x = [paddle.to_tensor(self.x[i], stop_gradient=False) for i in range(4)]
+        if flag == "prim":
+            fn = apply_to_static(
+                self.net,
+                use_cinn=self.enable_cinn,
+                input_spec=[
+                    [
+                        InputSpec(shape=self.x_shape, dtype='float32'),
+                        InputSpec(shape=self.init_x_shape, dtype='float32'),
+                        InputSpec(shape=self.init_x_shape, dtype='float32'),
+                        InputSpec(shape=self.x_shape, dtype='float32'),
+                    ]
+                ],
+            )
+            fn.train()
+        else:
+            fn = self.net
+        res = fn(x)
+        res.backward()
+        x_grad1 = x[0].gradient()
+        x_grad2 = x[1].gradient()
+        x_grad3 = x[2].gradient()
+        x_grad4 = x[3].gradient()
+        if flag == "prim":
+            core._set_prim_all_enabled(False)
+        return res, [x_grad1, x_grad2, x_grad3, x_grad4]
+
+    def test_prim_all_dynamic(self):
+        res_ref, grad_ref = self.base_net()
+        res, grad = self.base_net("prim")
+
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_allclose(
+                ref, actual, rtol=self.tol, atol=self.tol
+            )
+
+        for dr, d in zip(grad_ref, grad):
+            np.testing.assert_allclose(dr, d, rtol=self.tol, atol=self.tol)
+
+
+class TestPrimConcatWithGrad6(TestPrimConcatWithGrad5):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, 200, None]
+        x = np.random.random(self.x_shape).astype(self.dtype)
+        self.x = [x + i for i in range(4)]
+        self.net = concat_net3
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
 class TestPrimSplitWithGrad1(TestPrimBaseWithGrad):
     def setUp(self):
         np.random.seed(2023)
