[Prim][PIR] Add backward_decomp and support dynamic shape for ceil_grad (#68829)

eggman-1024 · web-flow · commit 1cd4e07902aa · 2024-11-25T18:37:45.000+08:00
* add back_decomp and support dynamic shape for ceil_grad

* add bf16, fp16 test for ceil_grad
diff --git a/paddle/fluid/primitive/codegen/decomp_vjp_gen.py b/paddle/fluid/primitive/codegen/decomp_vjp_gen.py
@@ -68,6 +68,7 @@
 # prim op with one input and one output, with no attribute
 UNARY_PRIM_VJP_OPS = [
     'abs_grad',
+    'ceil_grad',
     'erf_grad',
     'exp_grad',
     'floor_grad',
diff --git a/paddle/fluid/primitive/decomp_rule/decomp_vjp/details.h b/paddle/fluid/primitive/decomp_rule/decomp_vjp/details.h
@@ -3259,6 +3259,21 @@ void take_along_axis_grad(const Tensor& arr,
   }
 }
 
+template <typename T>
+void ceil_grad(const Tensor& out_grad, Tensor* x_grad) {
+  if (x_grad) {
+    Tensor zero_tensor;
+    if (has_dynamic_shape(out_grad.shape())) {
+      zero_tensor = backend::full_with_tensor<T>(
+          shape<T>(out_grad), 0.0, out_grad.dtype());
+    } else {
+      zero_tensor =
+          full<T>(common::vectorize(out_grad.dims()), 0.0, out_grad.dtype());
+    }
+    set_output<T>(zero_tensor, x_grad);
+  }
+}
+
 }  // namespace details
 }  // namespace primitive
 }  // namespace paddle
diff --git a/python/paddle/autograd/backward_utils.py b/python/paddle/autograd/backward_utils.py
@@ -35,6 +35,7 @@
     "pd_op.assign",
     "pd_op.batch_norm_",
     "pd_op.cast",
+    "pd_op.ceil",
     "pd_op.concat",
     "pd_op.cos",
     "pd_op.cumprod",
diff --git a/test/legacy_test/test_activation_op.py b/test/legacy_test/test_activation_op.py
@@ -2006,7 +2006,9 @@ def init_shape(self):
 class TestCeil(TestActivation):
     def setUp(self):
         self.op_type = "ceil"
+        self.prim_op_type = "prim"
         self.python_api = paddle.ceil
+        self.public_python_api = paddle.ceil
         self.init_dtype()
         self.init_shape()
 
@@ -2032,6 +2034,20 @@ def test_check_output(self):
     def test_check_grad(self):
         pass
 
+    def test_check_grad_for_prim(self):
+        # the gradient on floor, ceil, round is undefined.
+        # we return zero as gradient, but the numpy return nan.
+        # for prim, we compare result with eager python api,
+        # so, we use only_prim flag to express we only test prim.
+        if core.is_compiled_with_cuda():
+            self.check_grad_with_place(
+                paddle.CUDAPlace(0),
+                ['X'],
+                'Out',
+                check_pir=True,
+                check_prim_pir=True,
+            )
+
 
 class TestCeil_ZeroDim(TestCeil):
     def init_shape(self):
@@ -5639,7 +5655,12 @@ def test_check_grad(self):
     check_pir=True,
     check_prim_pir=True,
 )
-create_test_act_fp16_class(TestCeil, grad_check=False, check_pir=True)
+create_test_act_fp16_class(
+    TestCeil,
+    grad_check=False,
+    check_pir=True,
+    check_prim_pir=True,
+)
 create_test_act_fp16_class(
     TestFloor,
     check_prim=True,
@@ -5815,7 +5836,12 @@ def test_check_grad(self):
 create_test_act_bf16_class(
     TestAbs, check_prim=True, check_pir=True, check_prim_pir=True
 )
-create_test_act_bf16_class(TestCeil, grad_check=False, check_pir=True)
+create_test_act_bf16_class(
+    TestCeil,
+    grad_check=False,
+    check_pir=True,
+    check_prim_pir=True,
+)
 create_test_act_bf16_class(
     TestFloor,
     grad_check=False,
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
@@ -69,6 +69,10 @@ def batch_norm_net4(x, y, z):
     )
 
 
+def ceil_net(x):
+    return paddle.ceil(x)
+
+
 def concat_net1(x):
     y = x + 1
     return paddle.concat([x, y], axis=-1)
@@ -520,6 +524,19 @@ def setUp(self):
         self.tol = 1e-5
 
 
+class TestPrimCeilWithGrad(TestPrimBaseWithGrad):
+    def setUp(self):
+        np.random.seed(2024)
+        self.op_name = "pd_op.ceil_grad"
+        self.dtype = "float32"
+        self.x_shape = [30, 200, 40]
+        self.init_x_shape = [None, None, None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.net = ceil_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
 class TestPrimConcatWithGrad1(TestPrimBaseWithGrad):
     def setUp(self):
         np.random.seed(2023)
