[Prim][PIR] rsqrt forward sink (#59202)

kevincheng2 · web-flow · commit dbc296c27dda · 2023-11-28T13:57:49.000+08:00
* prim rsqrt sink c++

* update 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
@@ -26,6 +26,7 @@
     "mean",
     "pow",
     "relu",
+    "rsqrt",
     "silu",
     "softmax",
     "sqrt",
@@ -41,6 +42,7 @@
     "mean",
     "pow",
     "relu",
+    "rsqrt",
     "silu",
     "softmax",
     "sqrt",
diff --git a/paddle/fluid/primitive/composite/composite.h b/paddle/fluid/primitive/composite/composite.h
@@ -270,6 +270,25 @@ Tensor relu_decomp(const Tensor& x) {
   return maximum<T>(x, full<T>(phi::vectorize(x.dims()), 0.0, x.dtype()));
 }
 
+template <typename T>
+Tensor rsqrt_decomp(const Tensor& x) {
+  auto org_dtype = x.dtype();
+  Tensor x_cast = x;
+
+  bool need_cast = is_half_dtype(org_dtype);
+  if (need_cast) {
+    x_cast = cast<T>(x, phi::DataType::FLOAT32);
+  }
+
+  auto ans = elementwise_pow<T>(
+      x_cast, full<T>(phi::vectorize(x_cast.dims()), -0.5, x_cast.dtype()));
+  if (need_cast) {
+    return cast<T>(ans, org_dtype);
+  } else {
+    return ans;
+  }
+}
+
 template <typename T>
 std::tuple<Tensor, Tensor> squeeze_decomp(const Tensor& x,
                                           const IntArray& axis) {
diff --git a/python/paddle/decomposition/rules.py b/python/paddle/decomposition/rules.py
@@ -17,63 +17,6 @@
 from .register import register_decomp
 
 
-def mean(x, axis, keepdim):
-    """define composite rule of op mean"""
-    x_shape = x.shape
-    if axis in (None, []):
-        axis = tuple(range(0, len(x_shape)))
-    axes = (axis,) if isinstance(axis, int) else axis
-    sum_x = sum(x, axis=axes, keepdim=keepdim)
-    value_to_fill = 1
-    for axis in axes:
-        value_to_fill *= x_shape[axis]
-    norm = fill_constant(
-        shape=[],
-        value=value_to_fill,
-        dtype=sum_x.dtype,
-    )
-    res = sum_x / norm
-    return res
-
-
-@register_decomp('pd_op.rsqrt')
-def rsqrt(x):
-    """define composite rule of op rsqrt."""
-    # rsqrt(x) = x^(-0.5)
-    is_amp = False
-    from paddle.base.data_feeder import convert_dtype
-
-    dtype = convert_dtype(x.dtype)
-    if dtype in ["float16", "uint16"]:
-        is_amp = True
-        x = cast(x, "float32")
-    y = full(x.shape if len(x.shape) == 0 else [1], -0.5, x.dtype)
-    res = pow_composite(x, y)
-    return res if not is_amp else cast(res, dtype)
-
-
-# @register_decomp('pd_op.pow')
-def pow_composite(x, y):
-    """
-    define composite rule of op pow
-    res = x^y
-    """
-    is_amp = False
-    from paddle.base.data_feeder import convert_dtype
-
-    dtype = convert_dtype(x.dtype)
-    if dtype in ["float16", "uint16"]:
-        is_amp = True
-        x = cast(x, "float32")
-
-    if isinstance(y, (int, float)):
-        y = full(x.shape if len(x.shape) == 0 else [1], y, x.dtype)
-    res = pow(x, y)
-    if is_amp:
-        res = cast(res, dtype)
-    return res
-
-
 @register_decomp('pd_op.dropout')
 def dropout(x, seed_tensor, p, is_test, mode, seed, fix_seed):
     """define composite rule of op dropout.
