[Prim][PIR] sqrt forward sink and fix Decomp bug (#59135)

kevincheng2 · web-flow · commit 8321febbd6a3 · 2023-11-20T18:56:24.000+08:00
* sqrt prim sink c++

* prim sqrt sink c++

* remove sqrt op with python
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
@@ -27,6 +27,7 @@
     "softmax",
     "layer_norm",
     "gelu",
+    "sqrt",
 ]
 
 # come into effect in generated file op_decomp.cc
@@ -39,6 +40,7 @@
     "softmax",
     "layer_norm",
     "gelu",
+    "sqrt",
 ]
 
 
diff --git a/paddle/fluid/pir/dialect/op_generator/op_gen.py b/paddle/fluid/pir/dialect/op_generator/op_gen.py
@@ -1185,22 +1185,9 @@ def OpGenerator(
 
         # If op has inplace info, we will generate inplace op and non-inplace op.
         for op_name in op_info.op_phi_name:
-            if op_name in decomp_interface_declare_gen_op_list:
-                op_interfaces = op_interfaces + [
-                    "paddle::dialect::DecompInterface"
-                ]
-                exclusive_interface_str += "\n  static std::vector<std::vector<pir::OpResult>> Decomp(pir::Operation* op);"
-            else:
-                op_interfaces = op_interfaces_tmp
-                exclusive_interface_str = exclusive_interface_str_tmp
-
             # =================================== #
-            #    gen interface/trait list str     #
+            #        gen trait list str           #
             # =================================== #
-            op_interfaces_str = ""
-            if len(op_interfaces) > 0:
-                op_interfaces_str = "," + ",".join(op_interfaces)
-
             if op_name[-1] == "_":
                 op_traits += ["paddle::dialect::InplaceTrait"]
 
@@ -1216,6 +1203,25 @@ def OpGenerator(
                 func_list = op_kernel_map['func']
 
             for kernel_func_name in func_list:
+                if (
+                    op_name in decomp_interface_declare_gen_op_list
+                    and kernel_func_name in decomp_interface_declare_gen_op_list
+                ):
+                    op_interfaces = op_interfaces + [
+                        "paddle::dialect::DecompInterface"
+                    ]
+                    exclusive_interface_str += "\n  static std::vector<std::vector<pir::OpResult>> Decomp(pir::Operation* op);"
+                else:
+                    op_interfaces = op_interfaces_tmp
+                    exclusive_interface_str = exclusive_interface_str_tmp
+
+                # =================================== #
+                #      gen interface list str         #
+                # =================================== #
+                op_interfaces_str = ""
+                if len(op_interfaces) > 0:
+                    op_interfaces_str = "," + ",".join(op_interfaces)
+
                 if len(func_list) == 1:
                     op_class_name = to_pascal_case(op_name) + "Op"
                     op_dialect_name = dialect_name + "." + op_name
diff --git a/paddle/fluid/primitive/composite/composite.h b/paddle/fluid/primitive/composite/composite.h
@@ -304,6 +304,28 @@ std::tuple<Tensor, Tensor, Tensor> layer_norm_decomp(
   return std::make_tuple(out, mean_, variance);
 }
 
+template <typename T>
+Tensor sqrt_decomp(const Tensor& x) {
+  auto org_dtype = x.dtype();
+  bool need_cast =
+      org_dtype == phi::DataType::FLOAT16 || org_dtype == phi::DataType::UINT16;
+
+  Tensor x_cast;
+  if (need_cast) {
+    x_cast = cast<T>(x, phi::DataType::FLOAT32);
+  } else {
+    x_cast = x;
+  }
+
+  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>
 Tensor gelu_decomp(const Tensor& x, bool approximate) {
   const double PM_2_SQRTPI = 1.12837916709551257390; /* 2/sqrt(pi) */
diff --git a/python/paddle/decomposition/rules.py b/python/paddle/decomposition/rules.py
@@ -36,25 +36,6 @@ def mean(x, axis, keepdim):
     return res
 
 
-@register_decomp('pd_op.sqrt')
-def sqrt(x):
-    """
-    define composite rule of op sqrt
-    res = pow(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.rsqrt')
 def rsqrt(x):
     """define composite rule of op rsqrt."""

Original file line number	Diff line number	Diff line change
`@@ -27,6 +27,7 @@`
`27`	`27`	`"softmax",`
`28`	`28`	`"layer_norm",`
`29`	`29`	`"gelu",`
	`30`	`+ "sqrt",`
`30`	`31`	`]`
`31`	`32`
`32`	`33`	`# come into effect in generated file op_decomp.cc`
`@@ -39,6 +40,7 @@`
`39`	`40`	`"softmax",`
`40`	`41`	`"layer_norm",`
`41`	`42`	`"gelu",`
	`43`	`+ "sqrt",`
`42`	`44`	`]`
`43`	`45`
`44`	`46`