[Prim][PIR] Backward decomp (#67988)

cubehan3 · web-flow · commit efe6ef7fb497 · 2024-09-05T19:07:38.000+08:00
* decomp where_grad

* decomp expm1_grad, atan2_grad

* decomp atan_grad
diff --git a/paddle/fluid/primitive/codegen/gen.py b/paddle/fluid/primitive/codegen/gen.py
@@ -93,6 +93,8 @@
 
 OTHER_PRIM_VJP_OPS = [
     'assign_grad',
+    'atan_grad',
+    'atan2_grad',
     'cumsum_grad',
     'sum_grad',
     'cast_grad',
@@ -102,6 +104,7 @@
     'transpose_grad',
     'concat_grad',
     'expand_grad',
+    'expm1_grad',
     'gather_grad',
     'gather_nd_grad',
     'pad_grad',
@@ -117,6 +120,7 @@
     'tile_grad',
     'topk_grad',
     'unsqueeze_grad',
+    'where_grad',
 ]
 
 # whole vjp list of primitive op vjp
diff --git a/paddle/fluid/primitive/rule/vjp/details.h b/paddle/fluid/primitive/rule/vjp/details.h
@@ -2519,6 +2519,57 @@ void softsign_grad(const Tensor& x, const Tensor& out_grad, Tensor* x_grad) {
   }
 }
 
+template <typename T>
+void where_grad(const Tensor& condition,
+                const Tensor& x,
+                const Tensor& y,
+                const Tensor& out_grad,
+                Tensor* x_grad,
+                Tensor* y_grad) {
+  Tensor zero;
+  if (has_dynamic_shape(out_grad.shape())) {
+    zero =
+        backend::full_with_tensor<T>(shape<T>(out_grad), 0.0, out_grad.dtype());
+  } else {
+    zero = full<T>(common::vectorize(out_grad.dims()), 0.0, out_grad.dtype());
+  }
+
+  if (x_grad) {
+    Tensor x_grad_tmp = where<T>(condition, out_grad, zero);
+    set_output<T>(x_grad_tmp, x_grad);
+  }
+  if (y_grad) {
+    Tensor y_grad_tmp = where<T>(condition, zero, out_grad);
+    set_output<T>(y_grad_tmp, y_grad);
+  }
+}
+
+template <typename T>
+void expm1_grad(const Tensor& out, const Tensor& out_grad, Tensor* x_grad) {
+  if (x_grad) {
+    Tensor x_grad_tmp = out_grad * out + out_grad;
+    set_output<T>(x_grad_tmp, x_grad);
+  }
+}
+
+template <typename T>
+void atan2_grad(const Tensor& x,
+                const Tensor& y,
+                const Tensor& out_grad,
+                Tensor* x_grad,
+                Tensor* y_grad) {
+  Tensor tmp = x * x + y * y;
+  if (x_grad) {
+    Tensor x_grad_tmp = (out_grad * y) / tmp;
+    set_output<T>(x_grad_tmp, x_grad);
+  }
+
+  if (y_grad) {
+    Tensor y_grad_tmp = (-out_grad * x) / tmp;
+    set_output<T>(y_grad_tmp, y_grad);
+  }
+}
+
 template <typename T>
 void put_along_axis_grad(const Tensor& x,
                          const Tensor& index,
@@ -2634,6 +2685,15 @@ void put_along_axis_grad(const Tensor& x,
   }
 }
 
+template <typename T>
+void atan_grad(const Tensor& x, const Tensor& out_grad, Tensor* x_grad) {
+  if (x_grad) {
+    const Tensor one = full_scalar<T>(1.0, x.dtype());
+    Tensor x_grad_tmp = out_grad / (one + x * x);
+    set_output<T>(x_grad_tmp, x_grad);
+  }
+}
+
 }  // namespace details
 }  // namespace primitive
 }  // namespace paddle
diff --git a/python/paddle/autograd/backward_utils.py b/python/paddle/autograd/backward_utils.py
@@ -75,6 +75,7 @@
     "pd_op.swiglu",
     "pd_op.transpose",
     "pd_op.tanh",
+    "pd_op.where",
 ]
 
 
diff --git a/test/legacy_test/test_activation_op.py b/test/legacy_test/test_activation_op.py
@@ -250,7 +250,9 @@ def test_api_int(self):
 class TestExpm1(TestActivation):
     def setUp(self):
         self.op_type = "expm1"
+        self.prim_op_type = "prim"
         self.python_api = paddle.expm1
+        self.public_python_api = paddle.expm1
         self.init_dtype()
         self.init_shape()
 
@@ -269,7 +271,11 @@ def setUp(self):
 
     def test_check_grad(self):
         self.check_grad(
-            ['X'], 'Out', check_pir=True, check_pir_onednn=self.check_pir_onednn
+            ['X'],
+            'Out',
+            check_pir=True,
+            check_pir_onednn=self.check_pir_onednn,
+            check_prim_pir=True,
         )
 
     def test_check_output(self):
@@ -929,7 +935,9 @@ def executed_api(self):
 class TestAtan(TestActivation, TestParameter):
     def setUp(self):
         self.op_type = "atan"
+        self.prim_op_type = "prim"
         self.python_api = paddle.atan
+        self.public_python_api = paddle.atan
         self.init_dtype()
         self.init_shape()
 
@@ -954,9 +962,21 @@ def test_check_output(self):
     def test_check_grad(self):
         if self.dtype == np.float16:
             return
-        self.check_grad(
-            ['X'], 'Out', check_pir=True, check_pir_onednn=self.check_pir_onednn
-        )
+        if self.dtype == np.complex64 or self.dtype == np.complex128:
+            self.check_grad(
+                ['X'],
+                'Out',
+                check_pir=True,
+                check_pir_onednn=self.check_pir_onednn,
+            )
+        else:
+            self.check_grad(
+                ['X'],
+                'Out',
+                check_pir=True,
+                check_pir_onednn=self.check_pir_onednn,
+                check_prim_pir=True,
+            )
 
     @test_with_pir_api
     def test_out_name(self):
@@ -5543,7 +5563,7 @@ def test_check_grad(self):
 create_test_act_fp16_class(
     TestExpFp32_Prim, check_prim=True, enable_cinn=True, check_prim_pir=True
 )
-create_test_act_fp16_class(TestExpm1)
+create_test_act_fp16_class(TestExpm1, check_prim_pir=True)
 create_test_act_fp16_class(
     TestSigmoid,
     check_prim=True,
@@ -5598,7 +5618,7 @@ def test_check_grad(self):
 create_test_act_fp16_class(TestSin, check_pir=True, check_prim_pir=True)
 create_test_act_fp16_class(TestSinh, check_pir=True)
 create_test_act_fp16_class(TestAsin, check_pir=True)
-create_test_act_fp16_class(TestAtan, check_pir=True)
+create_test_act_fp16_class(TestAtan, check_pir=True, check_prim_pir=True)
 create_test_act_fp16_class(TestAcosh, check_pir=True)
 create_test_act_fp16_class(TestAsinh, check_pir=True)
 create_test_act_fp16_class(TestAtanh, check_pir=True)
@@ -5739,7 +5759,7 @@ def test_check_grad(self):
 create_test_act_bf16_class(
     TestExpFp32_Prim, check_prim=True, check_prim_pir=True
 )
-create_test_act_bf16_class(TestExpm1)
+create_test_act_bf16_class(TestExpm1, check_prim_pir=True)
 create_test_act_bf16_class(
     TestSigmoid, check_prim=True, check_pir=True, check_prim_pir=True
 )
@@ -5773,7 +5793,7 @@ def test_check_grad(self):
 create_test_act_bf16_class(TestSin, check_pir=True, check_prim_pir=True)
 create_test_act_bf16_class(TestSinh, check_pir=True)
 create_test_act_bf16_class(TestAsin, check_pir=True)
-create_test_act_bf16_class(TestAtan, check_pir=True)
+create_test_act_bf16_class(TestAtan, check_pir=True, check_prim_pir=True)
 create_test_act_bf16_class(TestAcosh, check_pir=True)
 create_test_act_bf16_class(TestAsinh, check_pir=True)
 create_test_act_bf16_class(TestAtanh, check_pir=True)
diff --git a/test/legacy_test/test_atan2_op.py b/test/legacy_test/test_atan2_op.py
@@ -35,7 +35,9 @@ def atan2_grad(x1, x2, dout):
 class TestAtan2(OpTest):
     def setUp(self):
         self.op_type = "atan2"
+        self.prim_op_type = "prim"
         self.python_api = paddle.atan2
+        self.public_python_api = paddle.atan2
         self.check_cinn = True
         self.init_dtype()
 
@@ -48,7 +50,11 @@ def setUp(self):
 
     def test_check_grad(self):
         self.check_grad(
-            ['X1', 'X2'], 'Out', check_cinn=self.check_cinn, check_pir=True
+            ['X1', 'X2'],
+            'Out',
+            check_cinn=self.check_cinn,
+            check_pir=True,
+            check_prim_pir=True,
         )
 
     def test_check_output(self):
@@ -74,6 +80,7 @@ def test_check_grad(self):
                 ),
                 check_cinn=self.check_cinn,
                 check_pir=True,
+                check_prim_pir=True,
             )
 
 
@@ -151,7 +158,9 @@ def run(place):
 class TestAtan2BF16OP(OpTest):
     def setUp(self):
         self.op_type = 'atan2'
+        self.prim_op_type = 'prim'
         self.python_api = paddle.atan2
+        self.public_python_api = paddle.atan2
         self.dtype = np.uint16
         self.check_cinn = True
         x1 = np.random.uniform(-1, -0.1, [15, 17]).astype('float32')
@@ -178,6 +187,7 @@ def test_check_grad(self):
             'Out',
             check_cinn=self.check_cinn,
             check_pir=True,
+            check_prim_pir=True,
         )
 
 
diff --git a/test/legacy_test/test_where_op.py b/test/legacy_test/test_where_op.py
@@ -28,7 +28,9 @@
 class TestWhereOp(OpTest):
     def setUp(self):
         self.op_type = 'where'
+        self.prim_op_type = 'prim'
         self.python_api = paddle.where
+        self.public_python_api = paddle.where
         self.check_cinn = True
         self.init_config()
         self.inputs = {'Condition': self.cond, 'X': self.x, 'Y': self.y}
@@ -39,7 +41,11 @@ def test_check_output(self):
 
     def test_check_grad(self):
         self.check_grad(
-            ['X', 'Y'], 'Out', check_cinn=self.check_cinn, check_pir=True
+            ['X', 'Y'],
+            'Out',
+            check_cinn=self.check_cinn,
+            check_pir=True,
+            check_prim_pir=True,
         )
 
     def init_config(self):
@@ -71,8 +77,10 @@ def init_config(self):
 class TestWhereBF16OP(OpTest):
     def setUp(self):
         self.op_type = 'where'
+        self.prim_op_type = 'prim'
         self.dtype = np.uint16
         self.python_api = paddle.where
+        self.public_python_api = paddle.where
         self.check_cinn = True
         self.init_config()
         self.inputs = {
@@ -99,6 +107,7 @@ def test_check_grad(self):
             numeric_grad_delta=0.05,
             check_cinn=self.check_cinn,
             check_pir=True,
+            check_prim_pir=True,
         )
 
     def init_config(self):
diff --git a/test/prim/pir_prim/CMakeLists.txt b/test/prim/pir_prim/CMakeLists.txt
@@ -23,6 +23,7 @@ set(TEST_PRIM_PURE_PIR_CASES
     test_prim_sub_graph_fghij_backward_dynamic_shape
     test_prim_sub_graph_klmno_backward_dynamic_shape
     test_prim_sub_graph_pqrst_backward_dynamic_shape
+    test_prim_sub_graph_uvwxyz_backward_dynamic_shape
     test_decompose_control_flow
     test_decomp_whole_program
     test_dynamic_combine1
diff --git a/test/prim/pir_prim/test_prim_sub_graph_uvwxyz_backward_dynamic_shape.py b/test/prim/pir_prim/test_prim_sub_graph_uvwxyz_backward_dynamic_shape.py
@@ -0,0 +1,75 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+from test_prim_sub_graph_backward_dynamic_shape import (
+    TestPrimTwoWithGrad,
+)
+
+import paddle
+
+
+def where_net(x, y):
+    return paddle.where(x > y, x, y)
+
+
+class TestPrimWhereWithGrad1(TestPrimTwoWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 30, 40]
+        self.init_x_shape = [None, None, 40]
+        self.y_shape = [30, 30, 40]
+        self.init_y_shape = [None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.net = where_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimWhereWithGrad2(TestPrimTwoWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 30, 40]
+        self.init_x_shape = [None, None, 40]
+        self.y_shape = [30, 30, 40]
+        self.init_y_shape = [30, 30, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.net = where_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+class TestPrimWhereWithGrad3(TestPrimTwoWithGrad):
+    def setUp(self):
+        np.random.seed(2023)
+        self.dtype = "float32"
+        self.x_shape = [30, 30, 40]
+        self.init_x_shape = [30, 30, 40]
+        self.y_shape = [30, 30, 40]
+        self.init_y_shape = [None, None, 40]
+        self.x = np.random.random(self.x_shape).astype(self.dtype)
+        self.y = np.random.random(self.y_shape).astype(self.dtype)
+        self.net = where_net
+        self.enable_cinn = False
+        self.tol = 1e-6
+
+
+if __name__ == "__main__":
+    unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -75,6 +75,7 @@`
`75`	`75`	`"pd_op.swiglu",`
`76`	`76`	`"pd_op.transpose",`
`77`	`77`	`"pd_op.tanh",`
	`78`	`+ "pd_op.where",`
`78`	`79`	`]`
`79`	`80`
`80`	`81`