[Prim][PIR] dropout forward sink

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

@@ -21,6 +21,7 @@
 decomp_interface_declare_gen_op_list = [
     "add_n",
     "batch_norm",
+    "dropout",
     "gelu",
     "layer_norm",
     "mean",
@@ -37,6 +38,7 @@
 # manual decomp interface implementation are located in manual_op_decomp.cc
 decomp_interface_implementation_gen_op_list = [
     "add_n",
+    "dropout",
     "gelu",
     "layer_norm",
     "mean",

paddle/fluid/primitive/composite/composite.h

@@ -384,6 +384,65 @@ std::tuple<Tensor, Tensor, Tensor> layer_norm_decomp(
   return std::make_tuple(out, mean_, variance);
 }
 
+template <typename T>
+std::tuple<Tensor, Tensor> dropout_decomp(
+    const Tensor& x,
+    const paddle::optional<Tensor>& seed_tensor,
+    const paddle::Scalar& p,
+    bool is_test,
+    const std::string& mode,
+    const int seed,
+    bool fix_seed) {
+  auto org_dtype = x.dtype();
+  bool upscale_in_train = false;
+  if (mode.compare("upscale_in_train") == 0) {
+    upscale_in_train = true;
+  }
+
+  int seed_tmp = 0;
+  if (fix_seed) {
+    seed_tmp = seed;
+  }
+
+  auto dtype_tmp = org_dtype;
+  if (is_half_dtype(org_dtype)) {
+    dtype_tmp = phi::DataType::FLOAT32;
+  }
+
+  auto uniform_tensor =
+      uniform<T>(phi::vectorize(x.dims()), dtype_tmp, 0.0, 1.0, seed_tmp);
+  auto mask =
+      cast<T>(greater_equal<T>(uniform_tensor,
+                               full<T>(phi::vectorize(x.dims()), p, dtype_tmp)),
+              org_dtype);
+  auto ones_p =
+      full<T>(phi::vectorize(x.dims()), 1.0 - p.to<float>(), org_dtype);
+  if (upscale_in_train) {
+    if (is_test) {
+      // inference: out = input
+      return std::make_tuple(x, cast<T>(mask, phi::DataType::UINT8));
+    } else {
+      // train: out = input * mask / ( 1.0 - p )
+      if (p.to<float>() == 1.0) {
+        // Process p=1. for avoid devide zero error (x*mask/(1.0-p))
+        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);
+        return std::make_tuple(ans, cast<T>(mask, phi::DataType::UINT8));
+      }
+    }
+  } else {
+    if (is_test) {
+      // inference: out = input * (1.0 - p)
+      return std::make_tuple(x * ones_p, cast<T>(mask, phi::DataType::UINT8));
+    } else {
+      // train: out = input * mask
+      return std::make_tuple(x * mask, cast<T>(mask, phi::DataType::UINT8));
+    }
+  }
+}
+
 template <typename T>
 Tensor sqrt_decomp(const Tensor& x) {
   auto org_dtype = x.dtype();

paddle/fluid/primitive/primitive.yaml

@@ -52,3 +52,4 @@
 - cast
 - sign
 - slice
+- uniform

python/paddle/decomposition/rules.py

@@ -17,66 +17,6 @@
 from .register import register_decomp
 
 
-@register_decomp('pd_op.dropout')
-def dropout(x, seed_tensor, p, is_test, mode, seed, fix_seed):
-    """define composite rule of op dropout.
-    upscale_in_train:
-        train: out = input * mask / ( 1.0 - p )
-        inference: out = input
-    downscale_in_infer
-        train: out = input * mask
-        inference: out = input * (1.0 - p)
-    """
-    from paddle import assign
-    from paddle.base import core
-    from paddle.base.data_feeder import convert_dtype
-
-    fix_seed = True if fix_seed is None else fix_seed
-    seed = seed if fix_seed else 0
-    upscale_in_train = mode == "upscale_in_train"
-
-    x_dtype = convert_dtype(x.dtype)
-    mask = bernoulli(shape=x.shape, dtype=x_dtype, p=p, seed=seed)
-
-    uint8_type = convert_dtype(core.VarDesc.VarType.UINT8)
-    if upscale_in_train:
-        if not is_test:
-            # Process p=1.0 for avoid devide zero error (x*mask/(1.0-p))
-            if p == 1.0:
-                return fill_constant(
-                    shape=x.shape, value=0.0, dtype=x.dtype
-                ) * x, zeros(x.shape, uint8_type)
-            else:
-                return x * mask / fill_constant(
-                    shape=x.shape, value=(1.0 - p), dtype=x.dtype
-                ), cast(mask, uint8_type)
-        else:
-            return assign(x), cast(mask, uint8_type)
-    else:
-        if not is_test:
-            return x * mask, cast(mask, uint8_type)
-        else:
-            return x * fill_constant(
-                shape=x.shape, value=(1.0 - p), dtype=x.dtype
-            ), cast(mask, uint8_type)
-
-
-def bernoulli(shape, dtype, p, seed=0):
-    from paddle.base.data_feeder import convert_dtype
-
-    # TODO(jiabin) Fix uniform doesn't support float16 error in CINN
-    new_dtype = (
-        "float32" if convert_dtype(dtype) in ["float16", "uint16"] else dtype
-    )
-    return cast(
-        greater_equal(
-            uniform(shape, new_dtype, min=0.0, max=1.0, seed=seed),
-            fill_constant(shape if len(shape) == 0 else [1], new_dtype, p),
-        ),
-        dtype,
-    )
-
-
 @register_decomp('pd_op.add_n')
 def add_n(x):
     ans = x[0]