[Prim][PIR] support square op backward in prim pir (PaddlePaddle#64381)

* update square

* update

Author: Eddie-Wang1120

paddle/fluid/primitive/codegen/gen.py

@@ -75,6 +75,7 @@
     'sin_grad',
     'cos_grad',
     'tanh_grad',
+    'square_grad',
 ]
 
 # prim op with two inputs and one output, with no attribute

paddle/fluid/primitive/rule/vjp/details.h

@@ -905,6 +905,14 @@ void log_grad(const Tensor& x, const Tensor& out_grad, Tensor* x_grad) {
   }
 }
 
+template <typename T>
+void square_grad(const Tensor& x, const Tensor& out_grad, Tensor* x_grad) {
+  if (x_grad) {
+    Tensor x_grad_tmp = 2 * x * out_grad;
+    set_output<T>(x_grad_tmp, x_grad);
+  }
+}
+
 template <typename T>
 void exp_grad(const Tensor& out, const Tensor& out_grad, Tensor* x_grad) {
   if (x_grad) {

test/legacy_test/test_activation_op.py

@@ -4305,6 +4305,7 @@ def test_check_grad(self):
             'Out',
             max_relative_error=0.007,
             check_pir=True,
+            check_prim_pir=True,
             check_pir_onednn=self.check_pir_onednn,
         )
 
@@ -4323,6 +4324,17 @@ def init_dtype(self):
     def test_check_output(self):
         self.check_output(check_pir=True)
 
+    def test_check_grad(self):
+        if self.dtype == np.float16:
+            return
+        self.check_grad(
+            ['X'],
+            'Out',
+            max_relative_error=0.007,
+            check_pir=True,
+            check_pir_onednn=self.check_pir_onednn,
+        )
+
 
 class TestSquare_Complex128(TestSquare):
     def init_dtype(self):
@@ -4331,6 +4343,17 @@ def init_dtype(self):
     def test_check_output(self):
         self.check_output(check_pir=True)
 
+    def test_check_grad(self):
+        if self.dtype == np.float16:
+            return
+        self.check_grad(
+            ['X'],
+            'Out',
+            max_relative_error=0.007,
+            check_pir=True,
+            check_pir_onednn=self.check_pir_onednn,
+        )
+
 
 class TestSquare_ZeroDim(TestSquare):
     def init_shape(self):
@@ -4373,7 +4396,12 @@ def test_check_output(self):
     def test_check_grad(self):
         place = core.CUDAPlace(0)
         self.check_grad_with_place(
-            place, ['X'], 'Out', numeric_grad_delta=0.5, check_pir=True
+            place,
+            ['X'],
+            'Out',
+            numeric_grad_delta=0.5,
+            check_pir=True,
+            check_prim_pir=True,
         )