[Prim][PIR] softsign decomp (#66450)

cubehan3 · web-flow · commit d9f2901c2a57 · 2024-07-26T10:09:26.000+08:00
* add softsign decomp

* string format

* refactor: remove unnecessary fp16 support
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
@@ -57,6 +57,7 @@
     "silu",
     "swiglu",
     "softmax",
+    "softsign",
     "square",
     "squared_l2_norm",
     "squeeze",
@@ -103,6 +104,7 @@
     "silu",
     "swiglu",
     "softmax",
+    "softsign",
     "square",
     "squared_l2_norm",
     "squeeze",
diff --git a/paddle/fluid/primitive/codegen/gen.py b/paddle/fluid/primitive/codegen/gen.py
@@ -136,6 +136,7 @@
     'sigmoid_grad',
     'silu_grad',
     'softmax_grad',
+    'softsign_grad',
     'sqrt_grad',
     'stack_grad',
     'swiglu',
diff --git a/paddle/fluid/primitive/composite/composite.h b/paddle/fluid/primitive/composite/composite.h
@@ -1555,6 +1555,15 @@ Tensor log_loss_decomp(const Tensor& input,
   return term1 - term2;
 }
 
+template <typename T>
+Tensor softsign_decomp(const Tensor& x) {
+  // softsign = x / (1 + abs(x))
+
+  Tensor x_abs = abs<T>(x);
+  Tensor one = full_scalar<T>(1.0, x.dtype());
+  return x / (one + x_abs);
+}
+
 }  // namespace details
 
 }  // namespace primitive
diff --git a/paddle/fluid/primitive/rule/vjp/details.h b/paddle/fluid/primitive/rule/vjp/details.h
@@ -2251,6 +2251,18 @@ void swiglu_grad(const Tensor& x,
   set_output<T>(x_grad, dx);
 }
 
+template <typename T>
+void softsign_grad(const Tensor& x, const Tensor& out_grad, Tensor* x_grad) {
+  // x_grad = out_grad / ((1 + abs(x))^2)
+
+  if (x_grad) {
+    Tensor x_abs = abs<T>(x);
+    Tensor x_abs_plusone = x_abs + full_scalar<T>(1.0, x.dtype());
+    Tensor x_grad_tmp = out_grad / (x_abs_plusone * x_abs_plusone);
+    set_output<T>(x_grad_tmp, x_grad);
+  }
+}
+
 }  // namespace details
 }  // namespace primitive
 }  // namespace paddle
diff --git a/test/legacy_test/test_activation_op.py b/test/legacy_test/test_activation_op.py
@@ -4875,10 +4875,12 @@ def ref_softsign(x):
 class TestSoftsign(TestActivation):
     def setUp(self):
         self.op_type = "softsign"
+        self.prim_op_type = "comp"
         self.init_dtype()
         self.init_shape()
 
         self.python_api = paddle.nn.functional.softsign
+        self.public_python_api = paddle.nn.functional.softsign
 
         np.random.seed(1024)
         x = np.random.uniform(-1, 1, self.shape).astype(self.dtype)
@@ -4897,16 +4899,35 @@ def init_shape(self):
         self.shape = [10, 12]
 
     def test_check_output(self):
-        self.check_output(
-            check_pir=True, check_pir_onednn=self.check_pir_onednn
-        )
+        if self.dtype == np.complex64 or self.dtype == np.complex128:
+            self.check_output(
+                check_pir=True, check_pir_onednn=self.check_pir_onednn
+            )
+        else:
+            self.check_output(
+                check_pir=True,
+                check_pir_onednn=self.check_pir_onednn,
+                check_prim_pir=True,
+            )
 
     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,
+            )
 
 
 class TestSoftsign_Complex64(TestSoftsign):
