[Prim][PIR] Support dynamic shape for concat_grad (#65148)

zeroRains · web-flow · commit 53c78b1f6b9f · 2024-06-18T10:50:41.000+08:00
* support dynamic shape for cocnat_grad and add the dynamic shape test for split_grad

* fix the bug in detail.h and the test case

* add the allow_dynamic_shape
diff --git a/paddle/fluid/primitive/rule/vjp/details.h b/paddle/fluid/primitive/rule/vjp/details.h
@@ -510,13 +510,31 @@ void concat_grad(const std::vector<Tensor>& x,
     axis_value = axis_value + rank;
   }
   axis_value = axis_value > 0 ? axis_value : 0;
-  std::vector<int> sections;
+
   int x_num = x.size();
-  for (int i = 0; i < x_num; ++i) {
-    sections.push_back(x[i].dims()[axis_value]);
+  std::vector<Tensor> x_grad_tmp;
+  if (has_dynamic_shape(x[0].shape())) {
+    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},
+                                  {}));
+    }
+    Tensor sections_tensor = concat<T>(sections);
+    x_grad_tmp =
+        backend::split<T>(out_grad,
+                          sections_tensor,
+                          full<T>({1}, axis_value, sections_tensor.dtype()));
+  } else {
+    std::vector<int> sections;
+    for (int i = 0; i < x_num; ++i) {
+      sections.push_back(x[i].dims()[axis_value]);
+    }
+    x_grad_tmp = split<T>(out_grad, IntArray(sections), axis_value);
   }
-  std::vector<Tensor> x_grad_tmp =
-      split<T>(out_grad, IntArray(sections), axis_value);
   for (int i = 0; i < x_num; ++i) {
     if (x_grad[i]) {
       set_output<T>(x_grad_tmp.at(i), x_grad.at(i));
diff --git a/python/paddle/autograd/backward_utils.py b/python/paddle/autograd/backward_utils.py
@@ -122,6 +122,8 @@
     "pd_op.mean",
     "pd_op.add",
     "pd_op.subtract",
+    "pd_op.concat",
+    "pd_op.split",
 ]
 
 
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
@@ -61,6 +61,32 @@ def subtract_net(x, y):
     return x - y
 
 
+def concat_net1(x):
+    y = x + 1
+    return paddle.concat([x, y], axis=-1)
+
+
+def concat_net2(x):
+    y = x + 1
+    return paddle.concat([x, y], axis=1)
+
+
+def split_net1(x):
+    res = paddle.split(x, num_or_sections=10, axis=-1)
+    tmp_res = res[0]
+    for i in range(1, len(res)):
+        tmp_res = tmp_res + res[i] * i
+    return tmp_res / len(res)
+
+
+def split_net2(x):
+    res = paddle.split(x, num_or_sections=10, axis=1)
+    tmp_res = res[0]
+    for i in range(1, len(res)):
+        tmp_res = tmp_res + res[i] * i
+    return tmp_res / len(res)
+
+
 def apply_to_static(net, use_cinn, input_spec=None):
     build_strategy = paddle.static.BuildStrategy()
     build_strategy.build_cinn_pass = use_cinn
@@ -521,5 +547,101 @@ def setUp(self):
         self.tol = 1e-6
 
 
+class TestPrimConcatWithGrad1(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = concat_net1
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimConcatWithGrad2(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = concat_net1
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimConcatWithGrad3(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = concat_net2
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimConcatWithGrad4(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, 200, None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = concat_net2
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimSplitWithGrad1(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, 200, None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = split_net1
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimSplitWithGrad2(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = split_net1
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimSplitWithGrad3(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = split_net2
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimSplitWithGrad4(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, 200, None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = split_net2
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
 if __name__ == "__main__":
     unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -122,6 +122,8 @@`
`122`	`122`	`"pd_op.mean",`
`123`	`123`	`"pd_op.add",`
`124`	`124`	`"pd_op.subtract",`
	`125`	`+ "pd_op.concat",`
	`126`	`+ "pd_op.split",`
`125`	`127`	`]`
`126`	`128`
`127`	`129`