use concat instead of padding when axes.size() is 1 (#68928)

cubehan3 · web-flow · commit edf6a63d9794 · 2024-10-28T14:11:54.000+08:00
diff --git a/paddle/fluid/primitive/rule/vjp/details.h b/paddle/fluid/primitive/rule/vjp/details.h
@@ -1623,14 +1623,42 @@ void slice_grad(const Tensor& input,
       reshape_out_grad = out_grad;
     }
 
-    if (decrease_size > 0 &&
-        (decrease_size != static_cast<size_t>(in_dims.size()))) {
-      auto out_tmp =
-          pad<T>(reshape<T>(reshape_out_grad, origin_out_shape), paddings, 0.0);
-      set_output<T>(out_tmp, input_grad);
+    // If axes.size() is 1, we can attempt to use concatenation instead of
+    // padding.
+    if (axes.size() == 1) {
+      const int64_t axis = axes[0];
+      const std::vector<int64_t> input_shape = input.shape();
+      if (decrease_size > 0 &&
+          (decrease_size != static_cast<size_t>(in_dims.size()))) {
+        reshape_out_grad = reshape<T>(reshape_out_grad, origin_out_shape);
+      }
+
+      std::vector<Tensor> concat_tensors;
+      // if concat axis has a shape of 0, concatenation may lead to errors.
+      if (paddings[2 * axis] != 0) {
+        std::vector<int64_t> left_shape(input_shape);
+        left_shape[axis] = paddings[2 * axis];
+        concat_tensors.push_back(full<T>(left_shape, 0.0, out_grad.dtype()));
+      }
+      concat_tensors.push_back(reshape_out_grad);
+
+      if (paddings[2 * axis + 1] != 0) {
+        std::vector<int64_t> right_shape(input_shape);
+        right_shape[axis] = paddings[2 * axis + 1];
+        concat_tensors.push_back(full<T>(right_shape, 0.0, out_grad.dtype()));
+      }
+
+      set_output<T>(concat<T>(concat_tensors, axis), input_grad);
     } else {
-      auto out_tmp = pad<T>(reshape_out_grad, paddings, 0.0);
-      set_output<T>(out_tmp, input_grad);
+      if (decrease_size > 0 &&
+          (decrease_size != static_cast<size_t>(in_dims.size()))) {
+        auto out_tmp = pad<T>(
+            reshape<T>(reshape_out_grad, origin_out_shape), paddings, 0.0);
+        set_output<T>(out_tmp, input_grad);
+      } else {
+        auto out_tmp = pad<T>(reshape_out_grad, paddings, 0.0);
+        set_output<T>(out_tmp, input_grad);
+      }
     }
   }
 }
diff --git a/test/legacy_test/test_slice_op.py b/test/legacy_test/test_slice_op.py
@@ -101,6 +101,46 @@ def config(self):
         self.out = self.input[-3:3, 0:100, :, 2:-1]
 
 
+class TestCase3(TestSliceOp):
+    def config(self):
+        self.input = np.random.random([4, 4, 5, 6]).astype("float64")
+        self.starts = [-3]
+        self.ends = [3]
+        self.axes = [0]
+        self.infer_flags = [1, 1, 1]
+        self.out = self.input[-3:3, :, :, :]
+
+
+class TestCase4(TestSliceOp):
+    def config(self):
+        self.input = np.random.random([3, 4, 5, 6]).astype("float64")
+        self.starts = [0]
+        self.ends = [4]
+        self.axes = [1]
+        self.infer_flags = [1, 1, 1]
+        self.out = self.input[:, :, :, :]
+
+
+class TestCase5(TestSliceOp):
+    def config(self):
+        self.input = np.random.random([3, 4, 5, 6]).astype("float64")
+        self.starts = [0]
+        self.ends = [2]
+        self.axes = [1]
+        self.infer_flags = [1, 1, 1]
+        self.out = self.input[:, 0:2, :, :]
+
+
+class TestCase6(TestSliceOp):
+    def config(self):
+        self.input = np.random.random([3, 4, 5, 6]).astype("float64")
+        self.starts = [2]
+        self.ends = [4]
+        self.axes = [1]
+        self.infer_flags = [1, 1, 1]
+        self.out = self.input[:, 2:4, :, :]
+
+
 class TestSliceZerosShapeTensor(OpTest):
     def setUp(self):
         self.op_type = "slice"
