[Prim][PIR] softmax forward sink

paddle/fluid/pir/dialect/op_generator/decomp_interface_gen_op_list.py

@@ -19,8 +19,13 @@
 
 # come into effect in generated file pd_op.h
 # manual decomp interface declare are located in manual_op.h
-decomp_interface_declare_gen_op_list = ["mean", "squeeze", "add_n"]
+decomp_interface_declare_gen_op_list = ["mean", "squeeze", "add_n", "softmax"]
 
 # come into effect in generated file op_decomp.cc
 # manual decomp interface implementation are located in manual_op_decomp.cc
-decomp_interface_implementation_gen_op_list = ["mean", "squeeze", "add_n"]
+decomp_interface_implementation_gen_op_list = [
+    "mean",
+    "squeeze",
+    "add_n",
+    "softmax",
+]

paddle/fluid/primitive/composite/composite.h

@@ -62,6 +62,31 @@ Tensor mean_decomp(const Tensor& x, const IntArray& axis, bool keepdim) {
   }
 }
 
+template <typename T>
+Tensor softmax_decomp(const Tensor& x, const int& axis) {
+  auto org_dtype = x.dtype();
+  auto x_tmp = x;
+  auto axis_tmp = IntArray({axis});
+
+  bool need_cast =
+      org_dtype == phi::DataType::FLOAT16 || org_dtype == phi::DataType::UINT16;
+  if (need_cast) {
+    x_tmp = cast<T>(x, phi::DataType::FLOAT32);
+  }
+
+  auto max_tmp = max<T>(x_tmp, axis_tmp, true);
+  auto molecular = exp<T>(subtract<T>(x_tmp, max_tmp));
+
+  auto denominator = sum<T>(molecular, axis_tmp, molecular.dtype(), true);
+  auto res = divide<T>(molecular, denominator);
+
+  if (need_cast) {
+    return cast<T>(res, org_dtype);
+  } else {
+    return res;
+  }
+}
+
 template <typename T>
 std::tuple<Tensor, Tensor> squeeze_decomp(const Tensor& x,
                                           const IntArray& axis) {

python/paddle/decomposition/rules.py

@@ -248,7 +248,7 @@ def silu(x):
     return res if not is_amp else cast(res, dtype)
 
 
-@register_decomp('pd_op.softmax')
+# @register_decomp('pd_op.softmax')
 def softmax(x, axis):
     """define composite rule of op softmax"""
     is_amp = False