【PIR API adaptor No.71、186、203、215】 Migrate paddle.fill_diagonal_tensor，roll，shard_index，strided_slice into pir (#59091)

enkilee · web-flow · commit 2499332524fd · 2023-11-23T14:59:15.000+08:00
diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py
@@ -660,7 +660,7 @@ def shard_index(input, index_num, nshards, shard_id, ignore_value=-1):
             [[-1]
              [ 1]]
     """
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         return _C_ops.shard_index(
             input, index_num, nshards, shard_id, ignore_value
         )
@@ -1021,11 +1021,11 @@ def _fill_diagonal_tensor_impl(x, y, offset=0, dim1=0, dim2=1, inplace=False):
     if len(y.shape) == 1:
         y = y.reshape([1, -1])
 
-    if inplace:
-        return _C_ops.fill_diagonal_tensor_(x, y, offset, dim1, dim2)
-
-    if in_dynamic_mode():
-        return _C_ops.fill_diagonal_tensor(x, y, offset, dim1, dim2)
+    if in_dynamic_or_pir_mode():
+        if inplace:
+            return _C_ops.fill_diagonal_tensor_(x, y, offset, dim1, dim2)
+        else:
+            return _C_ops.fill_diagonal_tensor(x, y, offset, dim1, dim2)
     else:
         check_variable_and_dtype(
             x,
@@ -1843,7 +1843,7 @@ def roll(x, shifts, axis=None, name=None):
     else:
         axis = []
 
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         return _C_ops.roll(x, shifts, axis)
     else:
         check_variable_and_dtype(
@@ -4411,7 +4411,7 @@ def strided_slice(x, axes, starts, ends, strides, name=None):
             >>> sliced_2 = paddle.strided_slice(x, axes=axes, starts=[minus_3, 0, 2], ends=ends, strides=strides_2)
             >>> # sliced_2 is x[:, 1:3:1, 0:2:1, 2:4:2].
     """
-    if in_dynamic_mode():
+    if in_dynamic_or_pir_mode():
         return _C_ops.strided_slice(x, axes, starts, ends, strides)
     else:
         helper = LayerHelper('strided_slice', **locals())
diff --git a/test/legacy_test/test_fill_diagonal_tensor_op.py b/test/legacy_test/test_fill_diagonal_tensor_op.py
@@ -103,10 +103,10 @@ def init_kernel_type(self):
         self.dtype = np.float64
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_pir=True)
 
     def test_check_grad(self):
-        self.check_grad(['X'], 'Out')
+        self.check_grad(['X'], 'Out', check_pir=True)
 
 
 class TensorFillDiagTensor_Test2(TensorFillDiagTensor_Test):
@@ -193,11 +193,11 @@ def init_input_output(self):
 
     def test_check_output(self):
         place = core.CUDAPlace(0)
-        self.check_output_with_place(place)
+        self.check_output_with_place(place, check_pir=True)
 
     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_pir=True)
 
 
 if __name__ == '__main__':
diff --git a/test/legacy_test/test_roll_op.py b/test/legacy_test/test_roll_op.py
@@ -19,7 +19,8 @@
 
 import paddle
 from paddle import base
-from paddle.base import Program, core, program_guard
+from paddle.base import core
+from paddle.pir_utils import test_with_pir_api
 
 
 class TestRollOp(OpTest):
@@ -48,10 +49,10 @@ def init_dtype_type(self):
         self.axis = [0, -2]
 
     def test_check_output(self):
-        self.check_output(check_prim=True)
+        self.check_output(check_prim=True, check_pir=True)
 
     def test_check_grad_normal(self):
-        self.check_grad(['X'], 'Out', check_prim=True)
+        self.check_grad(['X'], 'Out', check_prim=True, check_pir=True)
 
 
 class TestRollOpCase2(TestRollOp):
@@ -108,10 +109,14 @@ def init_dtype_type(self):
         self.place = core.CUDAPlace(0)
 
     def test_check_output(self):
-        self.check_output_with_place(self.place, check_prim=True)
+        self.check_output_with_place(
+            self.place, check_prim=True, check_pir=True
+        )
 
     def test_check_grad_normal(self):
-        self.check_grad_with_place(self.place, ['X'], 'Out', check_prim=True)
+        self.check_grad_with_place(
+            self.place, ['X'], 'Out', check_prim=True, check_pir=True
+        )
 
 
 @unittest.skipIf(
@@ -128,10 +133,14 @@ def init_dtype_type(self):
         self.place = core.CUDAPlace(0)
 
     def test_check_output(self):
-        self.check_output_with_place(self.place, check_prim=True)
+        self.check_output_with_place(
+            self.place, check_prim=True, check_pir=True
+        )
 
     def test_check_grad_normal(self):
-        self.check_grad_with_place(self.place, ['X'], 'Out', check_prim=True)
+        self.check_grad_with_place(
+            self.place, ['X'], 'Out', check_prim=True, check_pir=True
+        )
 
 
 @unittest.skipIf(
@@ -148,10 +157,14 @@ def init_dtype_type(self):
         self.place = core.CUDAPlace(0)
 
     def test_check_output(self):
-        self.check_output_with_place(self.place, check_prim=True)
+        self.check_output_with_place(
+            self.place, check_prim=True, check_pir=True
+        )
 
     def test_check_grad_normal(self):
-        self.check_grad_with_place(self.place, ['X'], 'Out', check_prim=True)
+        self.check_grad_with_place(
+            self.place, ['X'], 'Out', check_prim=True, check_pir=True
+        )
 
 
 class TestRollAPI(unittest.TestCase):
@@ -160,37 +173,53 @@ def input_data(self):
             [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]
         )
 
-    def test_roll_op_api(self):
-        self.input_data()
-
+    @test_with_pir_api
+    def test_roll_op_api_case1(self):
         paddle.enable_static()
-        # case 1:
-        with program_guard(Program(), Program()):
+        with paddle.static.program_guard(
+            paddle.static.Program(), paddle.static.Program()
+        ):
             x = paddle.static.data(name='x', shape=[-1, 3], dtype='float32')
-            x.desc.set_need_check_feed(False)
+            data_x = np.array(
+                [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]
+            ).astype('float32')
             z = paddle.roll(x, shifts=1)
-            exe = base.Executor(base.CPUPlace())
+            exe = paddle.static.Executor(paddle.CPUPlace())
             (res,) = exe.run(
-                feed={'x': self.data_x}, fetch_list=[z.name], return_numpy=False
+                paddle.static.default_main_program(),
+                feed={'x': data_x},
+                fetch_list=[z],
+                return_numpy=False,
             )
             expect_out = np.array(
                 [[9.0, 1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0]]
             )
-            np.testing.assert_allclose(expect_out, np.array(res), rtol=1e-05)
+        np.testing.assert_allclose(expect_out, np.array(res), rtol=1e-05)
+        paddle.disable_static()
 
-        # case 2:
-        with program_guard(Program(), Program()):
+    @test_with_pir_api
+    def test_roll_op_api_case2(self):
+        paddle.enable_static()
+        with paddle.static.program_guard(
+            paddle.static.Program(), paddle.static.Program()
+        ):
             x = paddle.static.data(name='x', shape=[-1, 3], dtype='float32')
-            x.desc.set_need_check_feed(False)
+            data_x = np.array(
+                [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]
+            ).astype('float32')
             z = paddle.roll(x, shifts=1, axis=0)
-            exe = base.Executor(base.CPUPlace())
+            exe = paddle.static.Executor(paddle.CPUPlace())
             (res,) = exe.run(
-                feed={'x': self.data_x}, fetch_list=[z.name], return_numpy=False
+                paddle.static.default_main_program(),
+                feed={'x': data_x},
+                fetch_list=[z],
+                return_numpy=False,
+            )
+            expect_out = np.array(
+                [[7.0, 8.0, 9.0], [1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]
             )
-        expect_out = np.array(
-            [[7.0, 8.0, 9.0], [1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]
-        )
         np.testing.assert_allclose(expect_out, np.array(res), rtol=1e-05)
+        paddle.disable_static()
 
     def test_dygraph_api(self):
         self.input_data()
@@ -214,22 +243,27 @@ def test_dygraph_api(self):
         )
         np.testing.assert_allclose(expect_out, np_z, rtol=1e-05)
 
+    @test_with_pir_api
     def test_roll_op_false(self):
-        self.input_data()
-
         def test_axis_out_range():
-            with program_guard(Program(), Program()):
+            paddle.enable_static()
+            with paddle.static.program_guard(
+                paddle.static.Program(), paddle.static.Program()
+            ):
                 x = paddle.static.data(name='x', shape=[-1, 3], dtype='float32')
-                x.desc.set_need_check_feed(False)
+                data_x = np.array(
+                    [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]
+                ).astype('float32')
                 z = paddle.roll(x, shifts=1, axis=10)
                 exe = base.Executor(base.CPUPlace())
                 (res,) = exe.run(
-                    feed={'x': self.data_x},
-                    fetch_list=[z.name],
+                    feed={'x': data_x},
+                    fetch_list=[z],
                     return_numpy=False,
                 )
 
         self.assertRaises(ValueError, test_axis_out_range)
+        paddle.disable_static()
 
     def test_shifts_as_tensor_dygraph(self):
         with base.dygraph.guard():
@@ -241,23 +275,28 @@ def test_shifts_as_tensor_dygraph(self):
             expected_out = np.array([[8, 6, 7], [2, 0, 1], [5, 3, 4]])
             np.testing.assert_allclose(out, expected_out, rtol=1e-05)
 
+    @test_with_pir_api
     def test_shifts_as_tensor_static(self):
-        with program_guard(Program(), Program()):
+        paddle.enable_static()
+        with paddle.static.program_guard(
+            paddle.static.Program(), paddle.static.Program()
+        ):
             x = paddle.arange(9).reshape([3, 3]).astype('float32')
             shape = paddle.shape(x)
             shifts = shape // 2
             axes = [0, 1]
             out = paddle.roll(x, shifts=shifts, axis=axes)
             expected_out = np.array([[8, 6, 7], [2, 0, 1], [5, 3, 4]])
 
-            exe = base.Executor(base.CPUPlace())
+            exe = paddle.static.Executor(paddle.CPUPlace())
             [out_np] = exe.run(fetch_list=[out])
             np.testing.assert_allclose(out_np, expected_out, rtol=1e-05)
 
             if paddle.is_compiled_with_cuda():
                 exe = base.Executor(base.CPUPlace())
                 [out_np] = exe.run(fetch_list=[out])
                 np.testing.assert_allclose(out_np, expected_out, rtol=1e-05)
+        paddle.disable_static()
 
 
 if __name__ == "__main__":
diff --git a/test/legacy_test/test_strided_slice_op.py b/test/legacy_test/test_strided_slice_op.py
