Pad grad (#53374)

* add pad op

* add_some_code

* modify some code

* add some code

* add some code

* modify some code

* add some code

* modify some code

* Update composite_backward_api.h

* modify some code

* add some code

* add some code

* add some code

paddle/fluid/operators/pad_op.cc

@@ -17,6 +17,9 @@ limitations under the License. */
 #include "paddle/fluid/framework/infershape_utils.h"
 #include "paddle/fluid/framework/op_registry.h"
 #include "paddle/fluid/platform/complex.h"
+#include "paddle/fluid/prim/api/composite_backward/composite_backward_api.h"
+#include "paddle/fluid/prim/utils/static/composite_grad_desc_maker.h"
+#include "paddle/fluid/prim/utils/static/desc_tensor.h"
 #include "paddle/phi/infermeta/unary.h"
 
 namespace paddle {
@@ -129,6 +132,27 @@ class PadOpGradMaker : public framework::SingleGradOpMaker<T> {
   }
 };
 
+class PadCompositeGradOpMaker : public prim::CompositeGradOpMakerBase {
+  using prim::CompositeGradOpMakerBase::CompositeGradOpMakerBase;
+
+ public:
+  void Apply() override {
+    paddle::Tensor x = this->GetSingleForwardInput("X");
+    paddle::Tensor out_grad = this->GetSingleOutputGrad("Out");
+    paddle::Tensor x_grad = this->GetSingleInputGrad("X");
+    auto* dx_ptr = this->GetOutputPtr(&x_grad);
+    std::string dx_name = this->GetOutputName(x_grad);
+
+    std::vector<int> paddings =
+        static_cast<std::vector<int>>(this->Attr<std::vector<int>>("paddings"));
+    float pad_value = static_cast<float>(this->Attr<float>("pad_value"));
+    VLOG(6) << "Runing add_grad composite func";
+
+    prim::pad_grad<prim::DescTensor>(x, out_grad, paddings, pad_value, dx_ptr);
+    this->RecoverOutputName(x_grad, dx_name);
+  }
+};
+
 template <typename T>
 class PadOpDoubleGradMaker : public framework::SingleGradOpMaker<T> {
  public:
@@ -155,6 +179,7 @@ REGISTER_OPERATOR(pad,
                   ops::PadOpMaker,
                   ops::PadOpGradMaker<paddle::framework::OpDesc>,
                   ops::PadOpGradMaker<paddle::imperative::OpBase>,
+                  ops::PadCompositeGradOpMaker,
                   PadInferShapeFunctor);
 REGISTER_OPERATOR(pad_grad,
                   ops::PadOpGrad,

paddle/fluid/prim/api/composite_backward/composite_backward_api.h

@@ -1806,6 +1806,32 @@ void roll_grad(const Tensor& x,
   }
 }
 
+template <typename T>
+void pad_grad(const Tensor& input,
+              const Tensor& out_grad,
+              const std::vector<int>& paddings,
+              const Scalar& pad_value,
+              Tensor* input_grad) {
+  if (input_grad) {
+    size_t rank = input.dims().size();
+    auto out_dims = out_grad.dims();
+
+    std::vector<int> starts(rank, 0);
+    std::vector<int64_t> ends(rank, 0);
+    std::vector<int64_t> axes(rank, 0);
+    std::vector<int64_t> infer_flags(rank, 1);
+    std::vector<int64_t> decrease_axis({});
+    for (size_t i = 0; i < rank; ++i) {
+      starts.push_back(static_cast<int>(paddings[2 * i]));
+      ends.push_back(static_cast<int64_t>(out_dims[i] - paddings[2 * i + 1]));
+      axes.push_back(i);
+    }
+    auto out_tmp =
+        slice<T>(out_grad, axes, starts, ends, infer_flags, decrease_axis);
+    set_output<T>(out_tmp, input_grad);
+  }
+}
+
 template <typename T>
 void scatter_nd_add_grad(const Tensor& index,
                          const Tensor& updates,
@@ -1821,5 +1847,6 @@ void scatter_nd_add_grad(const Tensor& index,
     set_output<T>(tmp_updates_grad, updates_grad);
   }
 }
+
 }  // namespace prim
 }  // namespace paddle

paddle/fluid/prim/api/manual_prim/eager_prim_api.cc

@@ -33,5 +33,16 @@ Tensor cast<Tensor>(const Tensor& x, DataType dtype) {
   return ::cast_ad_func(x, dtype);
 }
 
+template <>
+Tensor slice<Tensor>(const Tensor& input,
+                     const std::vector<int64_t>& axes,
+                     const IntArray& starts,
+                     const IntArray& ends,
+                     const std::vector<int64_t>& infer_flags,
+                     const std::vector<int64_t>& decrease_axis) {
+  VLOG(4) << "Eager Prim API slice_ad_func call";
+  return ::slice_ad_func(input, axes, starts, ends, infer_flags, decrease_axis);
+}
+
 }  // namespace prim
 }  // namespace paddle

paddle/fluid/prim/api/manual_prim/prim_manual_api.h

@@ -38,5 +38,13 @@ Tensor full(const IntArray& shape,
 template <typename T>
 Tensor cast(const Tensor& x, DataType dtype);
 
+template <typename T>
+Tensor slice(const Tensor& input,
+             const std::vector<int64_t>& axes,
+             const IntArray& starts,
+             const IntArray& ends,
+             const std::vector<int64_t>& infer_flags,
+             const std::vector<int64_t>& decrease_axis);
+
 }  // namespace prim
 }  // namespace paddle

paddle/fluid/prim/api/manual_prim/static_prim_api.cc

@@ -127,5 +127,32 @@ Tensor cast<DescTensor>(const Tensor& x, DataType dtype) {
   return out;
 }
 
+template <>
+Tensor slice<DescTensor>(const Tensor& input,
+                         const std::vector<int64_t>& axes,
+                         const IntArray& starts,
+                         const IntArray& ends,
+                         const std::vector<int64_t>& infer_flags,
+                         const std::vector<int64_t>& decrease_axis) {
+  framework::BlockDesc* block = StaticCompositeContext::Instance().GetBlock();
+  framework::OpDesc* op = block->AppendOp();
+  op->SetType("slice");
+  op->SetInput(
+      "Input",
+      {std::static_pointer_cast<prim::DescTensor>(input.impl())->Name()});
+  auto out = empty<DescTensor>({}, phi::DataType::FLOAT32, paddle::Place());
+  op->SetOutput(
+      "Out", {std::static_pointer_cast<prim::DescTensor>(out.impl())->Name()});
+  op->SetAttr("axes", unsafe_vector_cast<int64_t, int>(axes));
+  op->SetAttr("starts", unsafe_vector_cast<int64_t, int>(starts.GetData()));
+  op->SetAttr("ends", unsafe_vector_cast<int64_t, int>(ends.GetData()));
+  op->SetAttr("infer_flags", unsafe_vector_cast<int64_t, int>(infer_flags));
+  op->SetAttr("decrease_axis", unsafe_vector_cast<int64_t, int>(decrease_axis));
+  op->CheckAttrs();
+  op->InferVarType(block);
+  op->InferShape(*block);
+  return out;
+}
+
 }  // namespace prim
 }  // namespace paddle

paddle/phi/api/yaml/legacy_backward.yaml

@@ -722,6 +722,7 @@
     func : pad_grad
     param: [out_grad, paddings, pad_value]
   no_need_buffer : x
+  composite : pad_grad(x, out_grad, paddings, pad_value, x_grad)
   backward : pad_double_grad
 
 - backward_op : pool2d_double_grad

python/paddle/fluid/tests/unittests/test_pad_op.py

@@ -24,7 +24,9 @@ from paddle.fluid import Program, core, program_guard
 
 
 def pad_wrapper(x, paddings, pad_value):
-    return paddle._C_ops.pad(x, paddings, float(pad_value))
+    return paddle.nn.functional.pad(
+        x, pad=list(paddings), mode='constant', value=pad_value
+    )
 
 
 class TestPadOp(OpTest):
@@ -37,7 +39,7 @@ class TestPadOp(OpTest):
             'X': np.random.random(self.shape).astype(self.dtype),
         }
         self.attrs = {}
-        self.attrs['paddings'] = np.array(self.paddings).flatten()
+        self.attrs['paddings'] = list(np.array(self.paddings).flatten())
         self.attrs['pad_value'] = self.pad_value
         self.outputs = {
             'Out': np.pad(
@@ -47,6 +49,9 @@ class TestPadOp(OpTest):
                 constant_values=self.pad_value,
             )
         }
+        self.prim_op_type = "prim"
+        self.public_python_api = pad_wrapper
+        self.enable_cinn = False
 
     def get_dtype(self):
         return np.float64
@@ -55,7 +60,7 @@ class TestPadOp(OpTest):
         self.check_output()
 
     def test_check_grad_normal(self):
-        self.check_grad(['X'], 'Out')
+        self.check_grad(['X'], 'Out', check_prim=True)
 
     def initTestCase(self):
         self.shape = (16, 16)
@@ -111,16 +116,19 @@ create_test_fp16(TestCase3)
 
 class TestPadOpError(unittest.TestCase):
     def test_errors(self):
-        with program_guard(Program(), Program()):
-            input_data = np.random.random((2, 2)).astype("float32")
+        with paddle.fluid.framework._static_guard():
+            with program_guard(Program(), Program()):
+                input_data = np.random.random((2, 2)).astype("float32")
 
-            def test_Variable():
-                paddle.nn.functional.pad(x=input_data, pad=[1, 1, 1, 1])
+                def test_Variable():
+                    paddle.nn.functional.pad(x=input_data, pad=[1, 1, 1, 1])
 
-            self.assertRaises(TypeError, test_Variable)
+                self.assertRaises(TypeError, test_Variable)
 
-            data = paddle.static.data(name='data', shape=[4], dtype='float16')
-            paddle.nn.functional.pad(x=data, pad=[0, 1])
+                data = paddle.static.data(
+                    name='data', shape=[4], dtype='float16'
+                )
+                paddle.nn.functional.pad(x=data, pad=[0, 1])
 
 
 class TestPaddingValueTensor(UnittestBase):
@@ -129,34 +137,40 @@ class TestPaddingValueTensor(UnittestBase):
         self.save_path = os.path.join(self.temp_dir.name, self.path_prefix())
 
     def test_static(self):
-        main_prog = Program()
-        starup_prog = Program()
-        with program_guard(main_prog, starup_prog):
-            fc = paddle.nn.Linear(4, 10)
-            x = paddle.randn([2, 4])
-            x.stop_gradient = False
-            feat = fc(x)  # [2,3,10]
-
-            out = self.call_func(feat)
-
-            sgd = paddle.optimizer.SGD()
-            sgd.minimize(paddle.mean(out))
-            self.assertTrue(self.var_prefix() in str(main_prog))
-
-            exe = paddle.static.Executor()
-            exe.run(starup_prog)
-            res = exe.run(fetch_list=[feat, out])
-            gt = np.pad(res[0], [1, 1], 'constant', constant_values=[1.0, 1.0])
-            np.testing.assert_allclose(res[1], gt)
-            paddle.static.save_inference_model(
-                self.save_path, [x], [feat, out], exe
-            )
-            # Test for Inference Predictor
-            infer_outs = self.infer_prog()
-            gt = np.pad(
-                infer_outs[0], [1, 1], 'constant', constant_values=[1.0, 1.0]
-            )
-            np.testing.assert_allclose(infer_outs[1], gt)
+        with paddle.fluid.framework._static_guard():
+            main_prog = Program()
+            starup_prog = Program()
+            with program_guard(main_prog, starup_prog):
+                fc = paddle.nn.Linear(4, 10)
+                x = paddle.randn([2, 4])
+                x.stop_gradient = False
+                feat = fc(x)  # [2,3,10]
+
+                out = self.call_func(feat)
+
+                sgd = paddle.optimizer.SGD()
+                sgd.minimize(paddle.mean(out))
+                self.assertTrue(self.var_prefix() in str(main_prog))
+
+                exe = paddle.static.Executor()
+                exe.run(starup_prog)
+                res = exe.run(fetch_list=[feat, out])
+                gt = np.pad(
+                    res[0], [1, 1], 'constant', constant_values=[1.0, 1.0]
+                )
+                np.testing.assert_allclose(res[1], gt)
+                paddle.static.save_inference_model(
+                    self.save_path, [x], [feat, out], exe
+                )
+                # Test for Inference Predictor
+                infer_outs = self.infer_prog()
+                gt = np.pad(
+                    infer_outs[0],
+                    [1, 1],
+                    'constant',
+                    constant_values=[1.0, 1.0],
+                )
+                np.testing.assert_allclose(infer_outs[1], gt)
 
     def path_prefix(self):
         return 'padding_value'
@@ -183,23 +197,26 @@ class TestPaddingValueTensor2(TestPaddingValueTensor):
 
 class TestPaddingValueTensor3(unittest.TestCase):
     def test_static(self):
-        np_x = np.random.random((16, 16)).astype('float32')
-        main_prog = Program()
-        starup_prog = Program()
-        with program_guard(main_prog, starup_prog):
-            x = paddle.assign(np_x).astype('float32')
-            pad_value = paddle.assign([0.0]).astype('float64')
-            y = paddle.nn.functional.pad(x, [0, 1, 2, 3], value=pad_value)
-            loss = y.sum()
-            optimize_ops, params_grads = paddle.optimizer.SGD(0.01).minimize(
-                loss
+        with paddle.fluid.framework._static_guard():
+            np_x = np.random.random((16, 16)).astype('float32')
+            main_prog = Program()
+            starup_prog = Program()
+            with program_guard(main_prog, starup_prog):
+                x = paddle.assign(np_x).astype('float32')
+                pad_value = paddle.assign([0.0]).astype('float64')
+                y = paddle.nn.functional.pad(x, [0, 1, 2, 3], value=pad_value)
+                loss = y.sum()
+                optimize_ops, params_grads = paddle.optimizer.SGD(
+                    0.01
+                ).minimize(loss)
+
+            exe = paddle.static.Executor(paddle.CPUPlace())
+            res = exe.run(
+                main_prog, fetch_list=[y] + [g for p, g in params_grads]
             )
-
-        exe = paddle.static.Executor(paddle.CPUPlace())
-        res = exe.run(main_prog, fetch_list=[y] + [g for p, g in params_grads])
-        pd_out = res[0]
-        np_out = np.pad(np_x, [(0, 1), (2, 3)], constant_values=0.0)
-        np.testing.assert_allclose(pd_out, np_out)
+            pd_out = res[0]
+            np_out = np.pad(np_x, [(0, 1), (2, 3)], constant_values=0.0)
+            np.testing.assert_allclose(pd_out, np_out)
 
 
 @unittest.skipIf(
@@ -215,13 +232,16 @@ class TestPadBP16Op(OpTest):
         self.python_api = pad_wrapper
         x = np.random.random(self.shape).astype(np.float32)
         self.attrs = {}
-        self.attrs['paddings'] = np.array(self.paddings).flatten()
+        self.attrs['paddings'] = list(np.array(self.paddings).flatten())
         self.attrs['pad_value'] = self.pad_value
         out = np.pad(
             x, self.paddings, mode='constant', constant_values=self.pad_value
         )
         self.inputs = {'X': convert_float_to_uint16(x)}
         self.outputs = {'Out': convert_float_to_uint16(out)}
+        self.enable_cinn = False
+        self.prim_op_type = "prim"
+        self.public_python_api = pad_wrapper
 
     def initTestCase(self):
         self.shape = (16, 16)
@@ -234,9 +254,9 @@ class TestPadBP16Op(OpTest):
 
     def test_check_grad(self):
         place = core.CUDAPlace(0)
-        self.check_grad_with_place(place, ['X'], 'Out')
+        self.check_grad_with_place(place, ['X'], 'Out', check_prim=True)
 
 
 if __name__ == '__main__':
-    paddle.enable_static()
+    # paddle.enable_static()
     unittest.main()