From 1f4d4cd242a13cd93d5986a83a40915d25a1a61c Mon Sep 17 00:00:00 2001
From: Ruibin Cheung <beinggod@foxmail.com>
Date: Tue, 26 Sep 2023 10:47:09 +0800
Subject: [PATCH] [PIR] No.46 Migrate paddle.nn.functional.pad into pir
 (#57348)

---
 .../op_generator/vjp_interface_gen_op_list.py      |  4 ++++
 python/paddle/nn/functional/common.py              |  6 +++---
 test/legacy_test/test_pad3d_op.py                  | 14 ++++++++------
 test/legacy_test/test_pad_op.py                    | 12 +++++++-----
 4 files changed, 22 insertions(+), 14 deletions(-)

diff --git a/paddle/fluid/pir/dialect/op_generator/vjp_interface_gen_op_list.py b/paddle/fluid/pir/dialect/op_generator/vjp_interface_gen_op_list.py
index a42ee2d7ba14a..f6885d7371513 100644
--- a/paddle/fluid/pir/dialect/op_generator/vjp_interface_gen_op_list.py
+++ b/paddle/fluid/pir/dialect/op_generator/vjp_interface_gen_op_list.py
@@ -61,6 +61,8 @@
     'stack',
     'poisson',
     'gumbel_softmax',
+    'pad',
+    'pad3d',
     'squeeze',
     'unsqueeze',
     'tril',
@@ -104,6 +106,8 @@
     'stack',
     'poisson',
     'gumbel_softmax',
+    'pad',
+    'pad3d',
     'squeeze',
     'unsqueeze',
     'tril',
diff --git a/python/paddle/nn/functional/common.py b/python/paddle/nn/functional/common.py
index ea327e17e838d..7918b9223b71e 100644
--- a/python/paddle/nn/functional/common.py
+++ b/python/paddle/nn/functional/common.py
@@ -1658,7 +1658,7 @@ def pad(x, pad, mode='constant', value=0.0, data_format="NCHW", name=None):
         paddings = pad
         pad_value = value
 
-        if in_dynamic_mode():
+        if in_dynamic_or_pir_mode():
             out = _C_ops.pad(x, paddings, float(pad_value))
             return out
 
@@ -1712,7 +1712,7 @@ def pad(x, pad, mode='constant', value=0.0, data_format="NCHW", name=None):
 
     unsqueezed_dim = []
 
-    if isinstance(pad, Variable):
+    if isinstance(pad, (Variable, pir.OpResult)):
         if data_format in ["NCL", "NCHW", "NCDHW"]:
             data_format = "NCDHW"
             if x_dim == 3:
@@ -1756,7 +1756,7 @@ def pad(x, pad, mode='constant', value=0.0, data_format="NCHW", name=None):
                 unsqueezed_dim = [1]
                 x = unsqueeze(x, axis=unsqueezed_dim)
 
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         if isinstance(pad, Variable):
             pad = pad.tolist()
         out = _C_ops.pad3d(x, pad, mode, value, data_format)
diff --git a/test/legacy_test/test_pad3d_op.py b/test/legacy_test/test_pad3d_op.py
index a2fed2b033d85..42efb91a166d1 100644
--- a/test/legacy_test/test_pad3d_op.py
+++ b/test/legacy_test/test_pad3d_op.py
@@ -91,10 +91,10 @@ def setUp(self):
             self.outputs['Out'] = convert_float_to_uint16(self.outputs['Out'])
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_new_ir=True)
 
     def test_check_grad_normal(self):
-        self.check_grad(['X'], 'Out')
+        self.check_grad(['X'], 'Out', check_new_ir=True)
 
     def get_dtype(self):
         return np.float64
@@ -214,10 +214,12 @@ def get_dtype(self):
             return np.float16
 
         def test_check_output(self):
-            self.check_output(atol=1e-3)
+            self.check_output(atol=1e-3, check_new_ir=True)
 
         def test_check_grad_normal(self):
-            self.check_grad(['X'], 'Out', max_relative_error=1.5e-3)
+            self.check_grad(
+                ['X'], 'Out', max_relative_error=1.5e-3, check_new_ir=True
+            )
 
     cls_name = "{}_{}".format(parent.__name__, "FP16OP")
     TestPad3dFp16.__name__ = cls_name
@@ -251,12 +253,12 @@ def get_dtype(self):
 
         def test_check_output(self):
             place = core.CUDAPlace(0)
-            self.check_output_with_place(place, atol=1e-2)
+            self.check_output_with_place(place, atol=1e-2, check_new_ir=True)
 
         def test_check_grad_normal(self):
             place = core.CUDAPlace(0)
             self.check_grad_with_place(
-                place, ['X'], 'Out', max_relative_error=1e-2
+                place, ['X'], 'Out', max_relative_error=1e-2, check_new_ir=True
             )
 
     cls_name = "{}_{}".format(parent.__name__, "BF16OP")
diff --git a/test/legacy_test/test_pad_op.py b/test/legacy_test/test_pad_op.py
index 9e6aea80fd1bd..8054d7c75ffb1 100644
--- a/test/legacy_test/test_pad_op.py
+++ b/test/legacy_test/test_pad_op.py
@@ -57,10 +57,10 @@ def get_dtype(self):
         return np.float64
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_new_ir=True)
 
     def test_check_grad_normal(self):
-        self.check_grad(['X'], 'Out', check_prim=True)
+        self.check_grad(['X'], 'Out', check_prim=True, check_new_ir=True)
 
     def initTestCase(self):
         self.shape = (16, 16)
@@ -101,7 +101,7 @@ def get_dtype(self):
             return np.float16
 
         def test_check_grad_normal(self):
-            self.check_grad(['X'], 'Out', check_prim=True)
+            self.check_grad(['X'], 'Out', check_prim=True, check_new_ir=True)
 
     cls_name = "{}_{}".format(parent.__name__, "Fp16")
     TestPadFp16.__name__ = cls_name
@@ -253,11 +253,13 @@ def initTestCase(self):
 
     def test_check_output(self):
         place = core.CUDAPlace(0)
-        self.check_output_with_place(place)
+        self.check_output_with_place(place, check_new_ir=True)
 
     def test_check_grad(self):
         place = core.CUDAPlace(0)
-        self.check_grad_with_place(place, ['X'], 'Out', check_prim=True)
+        self.check_grad_with_place(
+            place, ['X'], 'Out', check_prim=True, check_new_ir=True
+        )
 
 
 if __name__ == '__main__':