[Prim][PIR] full_like forward sink (#59534)

kevincheng2 · web-flow · commit 2d09f53e2e4b · 2023-12-04T09:37:54.000+08:00
* prim full_like sink

* merge code

* update full_like

* remove code in rules.py

* adjust softmax code
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
@@ -22,6 +22,7 @@
     "add_n",
     "batch_norm",
     "dropout",
+    "full_like",
     "gelu",
     "layer_norm",
     "mean",
@@ -39,6 +40,7 @@
 decomp_interface_implementation_gen_op_list = [
     "add_n",
     "dropout",
+    "full_like",
     "gelu",
     "layer_norm",
     "mean",
diff --git a/paddle/fluid/primitive/composite/composite.h b/paddle/fluid/primitive/composite/composite.h
@@ -85,8 +85,8 @@ template <typename T>
 Tensor pow_decomp(const Tensor& x, const paddle::Scalar& y) {
   auto org_dtype = x.dtype();
   auto x_cast = x;
-  bool need_cast = org_dtype == phi::DataType::FLOAT16 ||
-                   org_dtype == phi::DataType::BFLOAT16;
+
+  bool need_cast = is_half_dtype(org_dtype);
   if (need_cast) {
     x_cast = cast<T>(x, phi::DataType::FLOAT32);
   }
@@ -233,7 +233,7 @@ Tensor softmax_decomp(const Tensor& x, const int& axis) {
   }
 
   auto max_tmp = max<T>(x_tmp, axis_tmp, true);
-  auto molecular = exp<T>(subtract<T>(x_tmp, max_tmp));
+  auto molecular = exp<T>(x_tmp - max_tmp);
   auto res = molecular / sum<T>(molecular, axis_tmp, molecular.dtype(), true);
 
   if (need_cast) {
@@ -384,6 +384,14 @@ std::tuple<Tensor, Tensor, Tensor> layer_norm_decomp(
   return std::make_tuple(out, mean_, variance);
 }
 
+template <typename T>
+Tensor full_like_decomp(const Tensor& x,
+                        const paddle::Scalar& value,
+                        const DataType& dtype,
+                        const Place& place) {
+  return full<T>(phi::vectorize(x.dims()), value, dtype, place);
+}
+
 template <typename T>
 std::tuple<Tensor, Tensor> dropout_decomp(
     const Tensor& x,
@@ -395,7 +403,7 @@ std::tuple<Tensor, Tensor> dropout_decomp(
     bool fix_seed) {
   auto org_dtype = x.dtype();
   bool upscale_in_train = false;
-  if (mode.compare("upscale_in_train") == 0) {
+  if (mode == std::string("upscale_in_train")) {
     upscale_in_train = true;
   }
 
@@ -428,7 +436,7 @@ std::tuple<Tensor, Tensor> dropout_decomp(
         auto zero = full<T>(phi::vectorize(x.dims()), 0.0, org_dtype);
         return std::make_tuple(x * zero, cast<T>(zero, phi::DataType::UINT8));
       } else {
-        auto ans = divide<T>(x * mask, ones_p);
+        auto ans = (x * mask) / ones_p;
         return std::make_tuple(ans, cast<T>(mask, phi::DataType::UINT8));
       }
     }
diff --git a/python/paddle/decomposition/rules.py b/python/paddle/decomposition/rules.py
@@ -17,24 +17,6 @@
 from .register import register_decomp
 
 
-@register_decomp('pd_op.add_n')
-def add_n(x):
-    ans = x[0]
-    for xi in x[1:]:
-        ans = xi + ans
-    return ans
-
-
-@register_decomp('pd_op.full_like')
-def full_like(x, fill_value, dtype, place=None):
-    """define composite rule of op full_like."""
-    """op name: full_like  op type name: fill_any_like."""
-    """arg place is not used, add it here to keep same as python api."""
-    fill_value = fill_value.get_defining_op().attrs()["value"]
-    val = full(x.shape, fill_value, dtype)
-    return val
-
-
 @register_decomp('pd_op.stack')
 def stack(x, axis):
     """
@@ -49,32 +31,6 @@ def stack(x, axis):
     return out
 
 
-@register_decomp('pd_op.squeeze')
-def squeeze(x, axis):
-    """define composite rule of squeeze"""
-    """
-    canonicalize dim within range 0 to rank and
-    determine new shape after squeeze op
-    if axis not specified, remove all dims equal to 1
-    otherwise, remove dims equal to 1 in axis
-    axis can only be list, not int
-    """
-    axis = axis.get_defining_op().attrs()["value"]
-    rank = len(x.shape)
-    if rank == 0:
-        return [assign(x), None]
-    if len(axis) == 0:
-        dims = set(range(rank))
-    else:
-        dims = {ax % rank for ax in axis}
-    new_shape = []
-    for d, s in enumerate(x.shape):
-        if not (s == 1 and (d in dims)):
-            new_shape.append(s)
-    out = reshape(x, new_shape)
-    return [out, None]
-
-
 @register_decomp('pd_op.unsqueeze')
 def unsqueeze(x, axis):
     """define composite rule of op unsqueeze"""
