[Prim][Pir] Only set primitive backward for dynamic shape tests. (#70270)

cubehan3 · web-flow · commit 7203a13c7c25 · 2024-12-18T14:16:27.000+08:00
* only prim backward

* polish more python test files

* add relu6_grad and hardsigmoid_grad interface
diff --git a/paddle/fluid/pir/dialect/op_generator/decomp_interface_gen_op_list.py b/paddle/fluid/pir/dialect/op_generator/decomp_interface_gen_op_list.py
@@ -103,6 +103,7 @@
     'floor_grad',
     'gelu_grad',
     'group_norm_grad',
+    'hardsigmoid_grad',
     'hardswish_grad',
     'leaky_relu_grad',
     'layer_norm_grad',
@@ -116,6 +117,7 @@
     'pow_grad',
     'reduce_as_grad',
     'relu_grad',
+    'relu6_grad',
     'reshape_grad',
     'roll_grad',
     'rsqrt_grad',
diff --git a/paddle/fluid/primitive/codegen/decomp_vjp_gen.py b/paddle/fluid/primitive/codegen/decomp_vjp_gen.py
@@ -143,12 +143,14 @@
     'dropout_grad',
     'gelu_grad',
     'group_norm_grad',
+    'hardsigmoid_grad',
     'hardswish_grad',
     'instance_norm_grad',
     'layer_norm_grad',
     'leaky_relu_grad',
     'mean_grad',
     'relu_grad',
+    'relu6_grad',
     'silu_grad',
     'softmax_grad',
     'softsign_grad',
diff --git a/test/prim/pir_prim/test_prim_sub_graph_abcde_backward_dynamic_shape.py b/test/prim/pir_prim/test_prim_sub_graph_abcde_backward_dynamic_shape.py
@@ -829,7 +829,7 @@ def setUp(self):
 
     def base_net(self, flag=None):
         if flag == "prim":
-            core._set_prim_all_enabled(True)
+            core._set_prim_backward_enabled(True)
         x = [paddle.to_tensor(self.x[i], stop_gradient=False) for i in range(4)]
         if flag == "prim":
             fn = apply_to_static(
@@ -861,7 +861,7 @@ def base_net(self, flag=None):
                 .ops
             ]
             assert self.op_name not in ops
-            core._set_prim_all_enabled(False)
+            core._set_prim_backward_enabled(False)
         return res, [x_grad1, x_grad2, x_grad3, x_grad4]
 
     def test_prim_all_dynamic(self):
diff --git a/test/prim/pir_prim/test_prim_sub_graph_backward_dynamic_shape.py b/test/prim/pir_prim/test_prim_sub_graph_backward_dynamic_shape.py
@@ -64,7 +64,7 @@ def setUp(self):
 
     def base_net(self, flag=None):
         if flag == "prim":
-            core._set_prim_all_enabled(True)
+            core._set_prim_backward_enabled(True)
         x = paddle.to_tensor(self.x, stop_gradient=False)
         if flag == "prim":
             fn = apply_to_static(
@@ -88,7 +88,7 @@ def base_net(self, flag=None):
                 .ops
             ]
             assert self.op_name not in ops
-            core._set_prim_all_enabled(False)
+            core._set_prim_backward_enabled(False)
         return res, x_grad
 
     def test_prim_all_dynamic(self):
@@ -121,7 +121,7 @@ def setUp(self):
 
     def base_net(self, flag=None):
         if flag == "prim":
-            core._set_prim_all_enabled(True)
+            core._set_prim_backward_enabled(True)
         x = paddle.to_tensor(self.x, stop_gradient=False)
         y = paddle.to_tensor(self.y, stop_gradient=False)
         if flag == "prim":
@@ -148,7 +148,7 @@ def base_net(self, flag=None):
                 .ops
             ]
             assert self.op_name not in ops
-            core._set_prim_all_enabled(False)
+            core._set_prim_backward_enabled(False)
         return res, [x_grad, y_grad]
 
     def test_prim_all_dynamic(self):
@@ -182,7 +182,7 @@ def setUp(self):
 
     def base_net(self, flag=None):
         if flag == "prim":
-            core._set_prim_all_enabled(True)
+            core._set_prim_backward_enabled(True)
         x = paddle.to_tensor(self.x, stop_gradient=False)
         y = paddle.to_tensor(self.y, stop_gradient=False)
         if flag == "prim":
@@ -211,7 +211,7 @@ def base_net(self, flag=None):
                 .ops
             ]
             assert self.op_name not in ops
-            core._set_prim_all_enabled(False)
+            core._set_prim_backward_enabled(False)
         return res, [grad]
 
     def test_prim_all_dynamic(self):
@@ -247,7 +247,7 @@ def setUp(self):
 
     def base_net(self, flag=None):
         if flag == "prim":
-            core._set_prim_all_enabled(True)
+            core._set_prim_backward_enabled(True)
         x = paddle.to_tensor(self.x, stop_gradient=False)
         y = paddle.to_tensor(self.y, stop_gradient=False)
         z = paddle.to_tensor(self.z, stop_gradient=False)
@@ -277,7 +277,7 @@ def base_net(self, flag=None):
                 .ops
             ]
             assert self.op_name not in ops
-            core._set_prim_all_enabled(False)
+            core._set_prim_backward_enabled(False)
         return res, [x_grad, y_grad, z_grad]
 
     def test_prim_all_dynamic(self):
diff --git a/test/prim/pir_prim/test_prim_sub_graph_fghij_backward_dynamic_shape.py b/test/prim/pir_prim/test_prim_sub_graph_fghij_backward_dynamic_shape.py
@@ -231,7 +231,7 @@ def setUp(self):
 
     def base_net(self, flag=None):
         if flag == "prim":
-            core._set_prim_all_enabled(True)
+            core._set_prim_backward_enabled(True)
         x = paddle.to_tensor(self.x, stop_gradient=False)
         y = paddle.to_tensor(self.y)
         if flag == "prim":
@@ -250,7 +250,7 @@ def base_net(self, flag=None):
         res.backward()
         x_grad = x.gradient()
         if flag == "prim":
-            core._set_prim_all_enabled(False)
+            core._set_prim_backward_enabled(False)
         return res, [x_grad]
 
 
diff --git a/test/prim/pir_prim/test_prim_sub_graph_pqrst_backward_dynamic_shape.py b/test/prim/pir_prim/test_prim_sub_graph_pqrst_backward_dynamic_shape.py
@@ -454,7 +454,7 @@ def setUp(self):
 
     def base_net(self, flag=None):
         if flag == "prim":
-            core._set_prim_all_enabled(True)
+            core._set_prim_backward_enabled(True)
         x = paddle.to_tensor(self.x, stop_gradient=False)
         y = paddle.to_tensor(self.y)
         z = paddle.to_tensor(self.z, stop_gradient=False)
@@ -483,7 +483,7 @@ def base_net(self, flag=None):
                 .ops
             ]
             assert self.op_name not in ops
-            core._set_prim_all_enabled(False)
+            core._set_prim_backward_enabled(False)
         return res, [x_grad, z_grad]
 
 
@@ -742,7 +742,7 @@ def setUp(self):
 
     def base_net(self, flag=None):
         if flag == "prim":
-            core._set_prim_all_enabled(True)
+            core._set_prim_backward_enabled(True)
         x = [paddle.to_tensor(self.x[i], stop_gradient=False) for i in range(4)]
         if flag == "prim":
             fn = apply_to_static(
@@ -774,7 +774,7 @@ def base_net(self, flag=None):
                 .ops
             ]
             assert self.op_name not in ops
-            core._set_prim_all_enabled(False)
+            core._set_prim_backward_enabled(False)
         return res, [x_grad1, x_grad2, x_grad3, x_grad4]
 
     def test_prim_all_dynamic(self):
@@ -1105,7 +1105,7 @@ def setUp(self):
 
     def base_net(self, flag=None):
         if flag == "prim":
-            core._set_prim_all_enabled(True)
+            core._set_prim_backward_enabled(True)
         x = paddle.to_tensor(self.x, stop_gradient=False)
         y = paddle.to_tensor(self.y)
         y = paddle.broadcast_to(y, [1, 200, 40])
@@ -1132,7 +1132,7 @@ def base_net(self, flag=None):
                 .ops
             ]
             assert self.op_name not in ops
-            core._set_prim_all_enabled(False)
+            core._set_prim_backward_enabled(False)
         return res, [x_grad]
 
 
@@ -1153,7 +1153,7 @@ def setUp(self):
 
     def base_net(self, flag=None):
         if flag == "prim":
-            core._set_prim_all_enabled(True)
+            core._set_prim_backward_enabled(True)
         x = paddle.to_tensor(self.x, stop_gradient=False)
         y = paddle.to_tensor(self.y)
         y = paddle.broadcast_to(y, [2, 1, 200])
@@ -1180,7 +1180,7 @@ def base_net(self, flag=None):
                 .ops
             ]
             assert self.op_name not in ops
-            core._set_prim_all_enabled(False)
+            core._set_prim_backward_enabled(False)
         return res, [x_grad]
 
 
