adapt pir api (#60416)

YuanRisheng · web-flow · commit f02f3d6c4f80 · 2023-12-29T10:53:58.000+08:00
diff --git a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py b/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
@@ -102,6 +102,7 @@
     'print',
     'number_count',
     'assign_value',
+    'share_data',
     'onednn_to_paddle_layout',
 ]
 
diff --git a/paddle/fluid/pir/dialect/operator/ir/ops.yaml b/paddle/fluid/pir/dialect/operator/ir/ops.yaml
@@ -1153,6 +1153,9 @@
 - op : share_data
   args : (Tensor x)
   output : Tensor(out)
+  infer_meta:
+    func: UnchangedInferMeta
+    param: [x]
   kernel:
     func: share_data
     param: [x]
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -777,6 +777,18 @@ void BindValue(py::module *m) {
       .def("apply", &apply)
       .def("is_same", &Value::operator==)
       .def("hash", [](Value self) { return std::hash<pir::Value>{}(self); })
+      .def("detach",
+           [](Value self) {
+             auto share_data_op =
+                 ApiBuilder::Instance()
+                     .GetBuilder()
+                     ->Build<paddle::dialect::ShareDataOp>(self);
+             auto out = share_data_op.out();
+             out.set_attribute(
+                 kAttrStopGradients,
+                 BoolAttribute::get(pir::IrContext::Instance(), false));
+             return out;
+           })
       .def("__repr__", &Value2String);
 }
 
diff --git a/python/paddle/incubate/optimizer/functional/utils.py b/python/paddle/incubate/optimizer/functional/utils.py
@@ -23,7 +23,7 @@ def check_input_type(input, name, op_name):
         if not isinstance(input, paddle.Tensor):
             raise ValueError(f"The input: {input} must be tensor.")
     else:
-        check_type(input, name, Variable, op_name)
+        check_type(input, name, (Variable, paddle.pir.Value), op_name)
 
 
 def check_initial_inverse_hessian_estimate(H0):
diff --git a/test/legacy_test/test_lr_scheduler.py b/test/legacy_test/test_lr_scheduler.py
@@ -1231,6 +1231,38 @@ def test_linear_warmp(self):
                 natural_lr.step()
             natural_lr_warmup.step()
 
+    def test_pir_linear_warmup_lr(self):
+        params = {
+            'learning_rate': 0.5,
+            'warmup_steps': 10,
+            'start_lr': 0,
+            'end_lr': 0.5,
+        }
+        scheduler = paddle.optimizer.lr.LinearWarmup(**params)
+        adam = paddle.optimizer.Adam(learning_rate=scheduler)
+        with paddle.pir_utils.IrGuard():
+            main_prog = paddle.static.Program()
+            start_prog = paddle.static.Program()
+            with paddle.static.program_guard(main_prog, start_prog):
+                x = paddle.static.data(name='x', shape=[3, 4, 5])
+                loss = paddle.mean(x)
+                adam.minimize(loss)
+                lr_var = adam._global_learning_rate()
+
+            exe = paddle.static.Executor()
+            exe.run(start_prog)
+            for epoch in range(5):
+                for batch_id in range(2):
+                    out = exe.run(
+                        main_prog,
+                        feed={'x': np.random.randn(3, 4, 5).astype('float32')},
+                        fetch_list=[lr_var],
+                    )
+                self.assertEqual(
+                    out, np.array(linear_warmup_lr(epoch, **params))
+                )
+                scheduler.step()
+
 
 if __name__ == '__main__':
     paddle.enable_static()
diff --git a/test/legacy_test/test_lrn_op.py b/test/legacy_test/test_lrn_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 TestLRNOp(OpTest):
@@ -115,98 +116,96 @@ def setUp(self):
             self.places.append(base.CUDAPlace(0))
 
     def check_static_3d_input(self, place):
-        with paddle_static_guard():
-            with base.program_guard(base.Program(), base.Program()):
-                in_np1 = np.random.random([3, 40, 40]).astype("float32")
-                in_np2 = np.transpose(in_np1, (0, 2, 1))
-
-                input1 = paddle.static.data(
-                    name="input1", shape=[3, 40, 40], dtype="float32"
-                )
-                input2 = paddle.static.data(
-                    name="input2", shape=[3, 40, 40], dtype="float32"
-                )
-                res1 = paddle.nn.functional.local_response_norm(
-                    x=input1, size=5, data_format='NCL'
-                )
-                res2 = paddle.nn.functional.local_response_norm(
-                    x=input2, size=5, data_format='NLC'
-                )
-                exe = base.Executor(place)
-                fetches = exe.run(
-                    base.default_main_program(),
-                    feed={"input1": in_np1, "input2": in_np2},
-                    fetch_list=[res1, res2],
-                )
-
-                fetches1_tran = np.transpose(fetches[1], (0, 2, 1))
-                np.testing.assert_allclose(
-                    fetches[0], fetches1_tran, rtol=1e-05
-                )
+        with paddle.static.program_guard(
+            paddle.static.Program(), paddle.static.Program()
+        ):
+            in_np1 = np.random.random([3, 40, 40]).astype("float32")
+            in_np2 = np.transpose(in_np1, (0, 2, 1))
+
+            input1 = paddle.static.data(
+                name="input1", shape=[3, 40, 40], dtype="float32"
+            )
+            input2 = paddle.static.data(
+                name="input2", shape=[3, 40, 40], dtype="float32"
+            )
+            res1 = paddle.nn.functional.local_response_norm(
+                x=input1, size=5, data_format='NCL'
+            )
+            res2 = paddle.nn.functional.local_response_norm(
+                x=input2, size=5, data_format='NLC'
+            )
+            exe = base.Executor(place)
+            fetches = exe.run(
+                paddle.static.default_main_program(),
+                feed={"input1": in_np1, "input2": in_np2},
+                fetch_list=[res1, res2],
+            )
+
+            fetches1_tran = np.transpose(fetches[1], (0, 2, 1))
+            np.testing.assert_allclose(fetches[0], fetches1_tran, rtol=1e-05)
 
     def check_static_4d_input(self, place):
-        with paddle_static_guard():
-            with base.program_guard(base.Program(), base.Program()):
-                input1 = paddle.static.data(
-                    name="input1", shape=[3, 3, 40, 40], dtype="float32"
-                )
-                input2 = paddle.static.data(
-                    name="input2", shape=[3, 40, 40, 3], dtype="float32"
-                )
-
-                res1 = paddle.nn.functional.local_response_norm(
-                    x=input1, size=5, data_format='NCHW'
-                )
-                res2 = paddle.nn.functional.local_response_norm(
-                    x=input2, size=5, data_format='NHWC'
-                )
-
-                in_np1 = np.random.random([3, 3, 40, 40]).astype("float32")
-                in_np2 = np.transpose(in_np1, (0, 2, 3, 1))
-
-                exe = base.Executor(place)
-                fetches = exe.run(
-                    base.default_main_program(),
-                    feed={"input1": in_np1, "input2": in_np2},
-                    fetch_list=[res1, res2],
-                )
-
-                fetches1_tran = np.transpose(fetches[1], (0, 3, 1, 2))
-                np.testing.assert_allclose(
-                    fetches[0], fetches1_tran, rtol=1e-05
-                )
+        with paddle.static.program_guard(
+            paddle.static.Program(), paddle.static.Program()
+        ):
+            input1 = paddle.static.data(
+                name="input1", shape=[3, 3, 40, 40], dtype="float32"
+            )
+            input2 = paddle.static.data(
+                name="input2", shape=[3, 40, 40, 3], dtype="float32"
+            )
+
+            res1 = paddle.nn.functional.local_response_norm(
+                x=input1, size=5, data_format='NCHW'
+            )
+            res2 = paddle.nn.functional.local_response_norm(
+                x=input2, size=5, data_format='NHWC'
+            )
+
+            in_np1 = np.random.random([3, 3, 40, 40]).astype("float32")
+            in_np2 = np.transpose(in_np1, (0, 2, 3, 1))
+
+            exe = base.Executor(place)
+            fetches = exe.run(
+                paddle.static.default_main_program(),
+                feed={"input1": in_np1, "input2": in_np2},
+                fetch_list=[res1, res2],
+            )
+
+            fetches1_tran = np.transpose(fetches[1], (0, 3, 1, 2))
+            np.testing.assert_allclose(fetches[0], fetches1_tran, rtol=1e-05)
 
     def check_static_5d_input(self, place):
-        with paddle_static_guard():
-            with base.program_guard(base.Program(), base.Program()):
-                input1 = paddle.static.data(
-                    name="input1", shape=[3, 3, 3, 40, 40], dtype="float32"
-                )
-                input2 = paddle.static.data(
-                    name="input2", shape=[3, 3, 40, 40, 3], dtype="float32"
-                )
-                res1 = paddle.nn.functional.local_response_norm(
-                    x=input1, size=5, data_format='NCDHW'
-                )
-                res2 = paddle.nn.functional.local_response_norm(
-                    x=input2, size=5, data_format='NDHWC'
-                )
-
-                in_np1 = np.random.random([3, 3, 3, 40, 40]).astype("float32")
-                in_np2 = np.transpose(in_np1, (0, 2, 3, 4, 1))
-
-                exe = base.Executor(place)
-                fetches = exe.run(
-                    base.default_main_program(),
-                    feed={"input1": in_np1, "input2": in_np2},
-                    fetch_list=[res1, res2],
-                )
-
-                fetches1_tran = np.transpose(fetches[1], (0, 4, 1, 2, 3))
-                np.testing.assert_allclose(
-                    fetches[0], fetches1_tran, rtol=1e-05
-                )
+        with paddle.static.program_guard(
+            paddle.static.Program(), paddle.static.Program()
+        ):
+            input1 = paddle.static.data(
+                name="input1", shape=[3, 3, 3, 40, 40], dtype="float32"
+            )
+            input2 = paddle.static.data(
+                name="input2", shape=[3, 3, 40, 40, 3], dtype="float32"
+            )
+            res1 = paddle.nn.functional.local_response_norm(
+                x=input1, size=5, data_format='NCDHW'
+            )
+            res2 = paddle.nn.functional.local_response_norm(
+                x=input2, size=5, data_format='NDHWC'
+            )
+
+            in_np1 = np.random.random([3, 3, 3, 40, 40]).astype("float32")
+            in_np2 = np.transpose(in_np1, (0, 2, 3, 4, 1))
+
+            exe = base.Executor(place)
+            fetches = exe.run(
+                paddle.static.default_main_program(),
+                feed={"input1": in_np1, "input2": in_np2},
+                fetch_list=[res1, res2],
+            )
+
+            fetches1_tran = np.transpose(fetches[1], (0, 4, 1, 2, 3))
+            np.testing.assert_allclose(fetches[0], fetches1_tran, rtol=1e-05)
 
+    @test_with_pir_api
     def test_static(self):
         with paddle_static_guard():
             for place in self.places:
@@ -276,9 +275,12 @@ def test_dygraph(self):
 
 
 class TestLocalResponseNormFAPIError(unittest.TestCase):
+    @test_with_pir_api
     def test_errors(self):
         with paddle_static_guard():
-            with program_guard(Program(), Program()):
+            with paddle.static.program_guard(
+                paddle.static.Program(), paddle.static.Program()
+            ):
 
                 def test_Variable():
                     # the input of lrn must be Variable.
@@ -346,6 +348,7 @@ def test_dygraph(self):
                 res2_tran = np.transpose(res2.numpy(), (0, 3, 1, 2))
                 np.testing.assert_allclose(res1.numpy(), res2_tran, rtol=1e-05)
 
+    @test_with_pir_api
     def test_static_fp16_gpu(self):
         if paddle.base.core.is_compiled_with_cuda():
             place = paddle.CUDAPlace(0)

Original file line number	Diff line number	Diff line change
`@@ -102,6 +102,7 @@`
`102`	`102`	`'print',`
`103`	`103`	`'number_count',`
`104`	`104`	`'assign_value',`
	`105`	`+ 'share_data',`
`105`	`106`	`'onednn_to_paddle_layout',`
`106`	`107`	`]`
`107`	`108`