[Prim][PIR] Support for rule operator backward under new ir (#58210)

kevincheng2 · web-flow · commit 3b83b14e9c1a · 2023-11-06T11:27:09.000+08:00
* [Prim][PIR] Sink relu backward Prim

* relu op registration

* codestyle

* codestyle
diff --git a/paddle/fluid/primitive/codegen/gen.py b/paddle/fluid/primitive/codegen/gen.py
@@ -83,6 +83,7 @@
     'silu_grad',
     'softmax_grad',
     'sqrt_grad',
+    'relu_grad',
 ]  # custom vjp list of composite op
 VJP_COMPS = PRIM_VJP + CUSTOM_VJP
 
diff --git a/paddle/fluid/primitive/rule/vjp/details.h b/paddle/fluid/primitive/rule/vjp/details.h
@@ -674,6 +674,18 @@ void softmax_grad(const Tensor& out,
   }
 }
 
+template <typename T>
+void relu_grad(const Tensor& out, const Tensor& out_grad, Tensor* x_grad) {
+  if (x_grad) {
+    auto condition = greater_than<T>(
+        out, full<T>(phi::vectorize(out.dims()), 0.0, out.dtype()));
+    auto res = where<T>(condition,
+                        out_grad,
+                        full<T>(phi::vectorize(out.dims()), 0.0, out.dtype()));
+    set_output<T>(res, x_grad);
+  }
+}
+
 template <typename T>
 void gather_nd_grad(const Tensor& x,
                     const Tensor& index,
diff --git a/test/legacy_test/test_activation_op.py b/test/legacy_test/test_activation_op.py
@@ -2405,7 +2405,9 @@ def setUp(self):
     def test_check_grad(self):
         if self.dtype == np.float16:
             return
-        self.check_grad(['X'], 'Out', check_prim=True, check_pir=True)
+        self.check_grad(
+            ['X'], 'Out', check_prim=True, check_pir=True, check_prim_pir=True
+        )
 
     def test_check_output(self):
         self.check_output(check_prim=True, check_pir=True)
