[Prim][PIR] relu forward sink

paddle/fluid/pir/dialect/op_generator/decomp_interface_gen_op_list.py

@@ -19,8 +19,13 @@
 
 # come into effect in generated file pd_op.h
 # manual decomp interface declare are located in manual_op.h
-decomp_interface_declare_gen_op_list = ["mean", "squeeze", "add_n"]
+decomp_interface_declare_gen_op_list = ["mean", "squeeze", "add_n", "relu"]
 
 # come into effect in generated file op_decomp.cc
 # manual decomp interface implementation are located in manual_op_decomp.cc
-decomp_interface_implementation_gen_op_list = ["mean", "squeeze", "add_n"]
+decomp_interface_implementation_gen_op_list = [
+    "mean",
+    "squeeze",
+    "add_n",
+    "relu",
+]

paddle/fluid/primitive/composite/composite.h

@@ -62,6 +62,11 @@ Tensor mean_decomp(const Tensor& x, const IntArray& axis, bool keepdim) {
   }
 }
 
+template <typename T>
+Tensor relu_decomp(const Tensor& x) {
+  return maximum<T>(x, full<T>(phi::vectorize(x.dims()), 0.0, x.dtype()));
+}
+
 template <typename T>
 std::tuple<Tensor, Tensor> squeeze_decomp(const Tensor& x,
                                           const IntArray& axis) {

test/legacy_test/test_activation_op.py

@@ -2399,7 +2399,7 @@ def test_check_grad(self):
         self.check_grad(['X'], 'Out', check_prim=True, check_pir=True)
 
     def test_check_output(self):
-        self.check_output(check_prim=True, check_pir=True)
+        self.check_output(check_prim=True, check_pir=True, check_prim_pir=True)
 
     def if_enable_cinn(self):
         pass
@@ -4767,7 +4767,11 @@ def test_check_grad(self):
 create_test_act_fp16_class(TestAtanh, check_pir=True)
 create_test_act_fp16_class(TestRound, grad_check=False, check_pir=True)
 create_test_act_fp16_class(
-    TestRelu, check_prim=True, enable_cinn=True, check_pir=True
+    TestRelu,
+    check_prim=True,
+    enable_cinn=True,
+    check_pir=True,
+    check_prim_pir=True,
 )
 create_test_act_fp16_class(
     TestGelu,
@@ -4834,6 +4838,7 @@ def create_test_act_bf16_class(
     check_prim=False,
     enable_cinn=False,
     check_pir=False,
+    check_prim_pir=False,
     grad_atol=1e-2,
     **kwargs
 ):
@@ -4866,6 +4871,7 @@ def test_check_output(self):
                 atol=atol,
                 check_prim=check_prim,
                 check_pir=check_pir,
+                check_prim_pir=check_prim_pir,
             )
 
         def test_check_grad(self):
@@ -4878,6 +4884,7 @@ def test_check_grad(self):
                     max_relative_error=grad_atol,
                     check_prim=check_prim,
                     check_pir=check_pir,
+                    check_prim_pir=check_prim_pir,
                 )
 
     cls_name = "{}_{}".format(parent.__name__, "BF16OP")
@@ -4920,7 +4927,9 @@ def test_check_grad(self):
 create_test_act_bf16_class(TestAsinh, check_pir=True)
 create_test_act_bf16_class(TestAtanh, check_pir=True)
 create_test_act_bf16_class(TestRound, grad_check=False, check_pir=True)
-create_test_act_bf16_class(TestRelu, check_prim=True, check_pir=True)
+create_test_act_bf16_class(
+    TestRelu, check_prim=True, check_pir=True, check_prim_pir=True
+)
 create_test_act_bf16_class(
     TestGelu,
     check_prim=True,

test/prim/pir_prim/test_sink_decomp.py

@@ -17,6 +17,7 @@
 import numpy as np
 
 import paddle
+import paddle.nn.functional as F
 from paddle.autograd.ir_backward import grad
 from paddle.base import core
 from paddle.decomposition import decompose
@@ -109,5 +110,44 @@ def test_has_decomp(self):
                 self.assertEqual(core.has_decomp(op), True)
 
 
+class TestReluSink(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [8, 16, 32, 64]
+        self.x = np.random.random(self.shape_x).astype("float32")
+        self.prog = None
+
+    def base_net(self, flag=None):
+        if flag == "forward":
+            core._set_prim_forward_enabled(True)
+        main_program = paddle.static.Program()
+        with paddle.static.program_guard(main_program):
+            x = paddle.static.data('x', self.shape_x, dtype='float32')
+            x.stop_gradient = False
+            sum_out = F.relu(x)
+            [new_out] = decompose(main_program, [sum_out])
+            gradients = grad(new_out, x)
+
+            exe = paddle.static.Executor()
+            [fwd, dx] = exe.run(
+                feed={'x': self.x}, fetch_list=[new_out, gradients]
+            )
+
+        whole_ops = [op.name() for op in main_program.global_block().ops]
+        self.prog = main_program
+        if flag == "forward":
+            core._set_prim_forward_enabled(False)
+            assert 'pd_op.relu' not in whole_ops
+        else:
+            assert 'pd_op.relu' in whole_ops
+        return fwd, dx
+
+    def test_relu_forward(self):
+        res_ref = self.base_net()
+        res = self.base_net("forward")
+        for ref, actual in zip(res_ref, res):
+            np.testing.assert_equal(ref, actual)
+
+
 if __name__ == "__main__":
     unittest.main()