[PIR]Migrate cumsum into pir (#58318)

0x45f · web-flow · commit 62c97fb32924 · 2023-10-24T19:27:03.000+08:00
diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
@@ -3925,7 +3925,7 @@ def cumsum(x, axis=None, dtype=None, name=None):
     if dtype is not None and x.dtype != convert_np_dtype_to_dtype_(dtype):
         x = cast(x, dtype)
 
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         if axis is None:
             axis = -1
         return _C_ops.cumsum(x, axis, flatten, False, False)
diff --git a/test/legacy_test/test_cumsum_op.py b/test/legacy_test/test_cumsum_op.py
@@ -23,6 +23,7 @@
 import paddle.inference as paddle_infer
 from paddle import base
 from paddle.base import core
+from paddle.pir_utils import test_with_pir_api
 
 
 class TestCumsumOp(unittest.TestCase):
@@ -53,7 +54,7 @@ def run_cases(self):
         np.testing.assert_array_equal(z, y.numpy())
 
     def run_static(self, use_gpu=False):
-        with base.program_guard(base.Program()):
+        with paddle.static.program_guard(paddle.static.Program()):
             data_np = np.random.random((100, 100)).astype(np.float32)
             x = paddle.static.data('X', [100, 100])
             y = paddle.cumsum(x)
@@ -65,16 +66,16 @@ def run_static(self, use_gpu=False):
 
             place = base.CUDAPlace(0) if use_gpu else base.CPUPlace()
             exe = base.Executor(place)
-            exe.run(base.default_startup_program())
+            exe.run(paddle.static.default_startup_program())
             out = exe.run(
                 feed={'X': data_np},
                 fetch_list=[
-                    y.name,
-                    y2.name,
-                    y3.name,
-                    y4.name,
-                    y5.name,
-                    y6.name,
+                    y,
+                    y2,
+                    y3,
+                    y4,
+                    y5,
+                    y6,
                 ],
             )
 
@@ -89,20 +90,26 @@ def run_static(self, use_gpu=False):
             z = np.cumsum(data_np, axis=-2)
             np.testing.assert_allclose(z, out[5], rtol=1e-05)
 
-    def test_cpu(self):
+    def test_cpu_dygraph(self):
         paddle.disable_static(paddle.base.CPUPlace())
         self.run_cases()
         paddle.enable_static()
 
+    @test_with_pir_api
+    def test_cpu_static(self):
         self.run_static()
 
-    def test_gpu(self):
+    def test_gpu_dygraph(self):
         if not base.core.is_compiled_with_cuda():
             return
         paddle.disable_static(paddle.base.CUDAPlace(0))
         self.run_cases()
         paddle.enable_static()
 
+    @test_with_pir_api
+    def test_gpu_static(self):
+        if not base.core.is_compiled_with_cuda():
+            return
         self.run_static(use_gpu=True)
 
     def test_name(self):
@@ -133,10 +140,10 @@ def setUp(self):
             self.outputs = {'Out': self.out}
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_pir=True)
 
     def test_check_grad(self):
-        self.check_grad(['X'], 'Out', check_prim=True)
+        self.check_grad(['X'], 'Out', check_prim=True, check_pir=True)
 
     def init_dtype(self):
         self.dtype = self.dtype_ = np.float64
@@ -242,10 +249,10 @@ def setUp(self):
             self.outputs = {'Out': self.out}
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_pir=True)
 
     def test_check_grad(self):
-        self.check_grad(['X'], 'Out', check_prim=True)
+        self.check_grad(['X'], 'Out', check_prim=True, check_pir=True)
 
     def init_dtype(self):
         self.dtype = self.dtype_ = np.float64
@@ -341,10 +348,10 @@ def setUp(self):
             self.outputs = {'Out': self.out}
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_pir=True)
 
     def test_check_grad(self):
-        self.check_grad(['X'], 'Out', check_prim=True)
+        self.check_grad(['X'], 'Out', check_prim=True, check_pir=True)
 
     def init_dtype(self):
         self.dtype = np.float16
@@ -380,10 +387,10 @@ def setUp(self):
             self.outputs = {'Out': self.out}
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_pir=True)
 
     def test_check_grad(self):
-        self.check_grad(['X'], 'Out', check_prim=True)
+        self.check_grad(['X'], 'Out', check_prim=True, check_pir=True)
 
     def init_dtype(self):
         self.dtype = self.dtype_ = np.float64
@@ -401,14 +408,10 @@ def if_enable_cinn(self):
             pass
 
         def test_check_output(self):
-            self.check_output()
+            self.check_output(check_pir=True)
 
         def test_check_grad(self):
-            self.check_grad(
-                ['X'],
-                'Out',
-                check_prim=True,
-            )
+            self.check_grad(['X'], 'Out', check_prim=True, check_pir=True)
 
     cls_name = "{}_{}".format(parent.__name__, "Fp16")
     TestCumsumFP16Op.__name__ = cls_name
@@ -445,12 +448,17 @@ def if_enable_cinn(self):
 
         def test_check_output(self):
             place = paddle.CUDAPlace(0)
-            self.check_output_with_place(place, check_prim=True)
+            self.check_output_with_place(place, check_prim=True, check_pir=True)
 
         def test_check_grad(self):
             place = paddle.CUDAPlace(0)
             self.check_grad_with_place(
-                place, ["X"], "Out", check_prim=True, numeric_grad_delta=0.05
+                place,
+                ["X"],
+                "Out",
+                check_prim=True,
+                numeric_grad_delta=0.05,
+                check_pir=True,
             )
 
     cls_name = "{}_{}".format(parent.__name__, "BF16")
@@ -552,6 +560,7 @@ def test_static_and_infer(self):
 
 
 class TestCumSumOpFp16(unittest.TestCase):
+    @test_with_pir_api
     def test_fp16(self):
         paddle.enable_static()
         x_np = np.random.random((100, 100)).astype('float16')
