[PRIM][PIR]Migrate dropout prim rules (#57722)

lxd-cumt · wanghuancoder · web-flow · commit 34210affad6c · 2023-09-26T14:47:10.000+08:00
* fix uniform force gpu

* add composite for dropout_grad op

* add dropout fwd prim rules

* add prim test for dropout

* migrate dropout prim rules to pir

* test dropout prim rules

* fix bugs

* change gpu0 to cpu for testpircompositedropout

* fix bugs

* change uniform place

* add bypass support

* fix bugs, change set_output to by_pass

* support more testcase for dropout op

* fix bugs

* fix bugs

---------

Co-authored-by: Wang Huan &lt;wanghuan29@baidu.com&gt;
diff --git a/paddle/fluid/primitive/codegen/gen.py b/paddle/fluid/primitive/codegen/gen.py
@@ -149,7 +149,11 @@
     'transpose_grad',
     'concat_grad',
 ]  # vjp list of primitive op
-CUSTOM_VJP = ['gelu_grad', 'layer_norm_grad']  # custom vjp list of composite op
+CUSTOM_VJP = [
+    'gelu_grad',
+    'layer_norm_grad',
+    'dropout_grad',
+]  # custom vjp list of composite op
 VJP_COMPS = PRIM_VJP + CUSTOM_VJP
 
 BACKENDS = [
diff --git a/paddle/fluid/primitive/rule/vjp/details.h b/paddle/fluid/primitive/rule/vjp/details.h
@@ -288,7 +288,7 @@ void add_grad(const Tensor& x,
       // Maybe need reduce here
       phi::DDim reduce_dim = get_reduce_dims(y.dims(), x.dims());
       if (!reduce_dim.size()) {
-        set_output<T>(out_grad, dy);
+        by_pass<T>(out_grad, dy);
       } else {
         auto dy_reduce_res =
             out_grad.sum(phi::vectorize(reduce_dim), y.dtype(), false);
@@ -297,23 +297,23 @@ void add_grad(const Tensor& x,
       }
 
     } else {
-      set_output<T>(out_grad, dy);
+      by_pass<T>(out_grad, dy);
     }
   }
   if (dx) {
     if (y.dims() != x.dims()) {
       // Maybe need reduce here
       auto reduce_dim = get_reduce_dims(x.dims(), y.dims());
       if (!reduce_dim.size()) {
-        set_output<T>(out_grad, dx);
+        by_pass<T>(out_grad, dx);
       } else {
         auto dx_reduce_res =
             out_grad.sum(phi::vectorize(reduce_dim), x.dtype(), false);
         auto dx_tmp = reshape<T>(dx_reduce_res, phi::vectorize(x.dims()));
         set_output<T>(dx_tmp, dx);
       }
     } else {
-      set_output<T>(out_grad, dx);
+      by_pass<T>(out_grad, dx);
     }
   }
 }
@@ -502,6 +502,35 @@ void layer_norm_grad(const Tensor& x,
   }
 }
 
+template <typename T>
+void dropout_grad(const Tensor& mask,
+                  const Tensor& out_grad,
+                  const Scalar& p,
+                  bool is_test,
+                  const std::string& mode,
+                  Tensor* x_grad) {
+  if (!x_grad) return;
+  if (is_test) {
+    if (mode == "upscale_in_train") {
+      by_pass<T>(out_grad, x_grad);
+    } else {
+      set_output<T>(out_grad * (1.0 - p.to<float>()), x_grad);
+    }
+  } else {
+    if (mode == "upscale_in_train") {
+      if (p.to<float>() == 1.0f) {
+        set_output<T>(scale<T>(out_grad, 0.0), x_grad);
+      } else {
+        set_output<T>(scale<T>(out_grad * cast<T>(mask, out_grad.dtype()),
+                               1.0 / (1.0 - p.to<float>())),
+                      x_grad);
+      }
+    } else {
+      set_output<T>(out_grad * cast<T>(mask, out_grad.dtype()), x_grad);
+    }
+  }
+}
+
 }  // namespace details
 }  // namespace primitive
 }  // namespace paddle
diff --git a/paddle/fluid/primitive/utils/CMakeLists.txt b/paddle/fluid/primitive/utils/CMakeLists.txt
@@ -7,4 +7,4 @@ endif()
 cc_library(
   primitive_static_utils_experimental
   SRCS static_utils.cc
-  DEPS phi common_infer_shape_functions)
+  DEPS phi common_infer_shape_functions pd_op_dialect_api)
diff --git a/paddle/fluid/primitive/utils/eager_utils.cc b/paddle/fluid/primitive/utils/eager_utils.cc
@@ -22,5 +22,10 @@ void set_output<Tensor>(const paddle::Tensor& x_tmp, paddle::Tensor* x) {
   x->set_autograd_meta(x_tmp.mutable_autograd_meta());
 }
 
+template <>
+void by_pass<Tensor>(const paddle::Tensor& x, Tensor* out) {
+  set_output<Tensor>(x, out);
+}
+
 }  // namespace primitive
 }  // namespace paddle
diff --git a/paddle/fluid/primitive/utils/static_utils.cc b/paddle/fluid/primitive/utils/static_utils.cc
@@ -11,6 +11,7 @@
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
+#include "paddle/fluid/pir/dialect/operator/ir/pd_api.h"
 #include "paddle/fluid/primitive/type/lazy_tensor.h"
 #include "paddle/fluid/primitive/utils/utils.h"
 
@@ -21,6 +22,14 @@ void set_output<LazyTensor>(const paddle::Tensor& x_tmp, paddle::Tensor* x) {
   x->set_impl(x_tmp.impl());
 }
 
+template <>
+void by_pass<LazyTensor>(const paddle::Tensor& x, paddle::Tensor* real_out) {
+  pir::Value x_res = std::static_pointer_cast<LazyTensor>(x.impl())->value();
+  auto op_res = paddle::dialect::assign(x_res);
+  Tensor out(std::make_shared<LazyTensor>(op_res));
+  set_output<LazyTensor>(out, real_out);
+}
+
 /**
  * @brief set output with empty grads in pir.
  *
diff --git a/paddle/fluid/primitive/utils/utils.h b/paddle/fluid/primitive/utils/utils.h
@@ -26,6 +26,9 @@ namespace primitive {
 template <typename T>
 void set_output(const Tensor& x_tmp, Tensor* x);
 
+template <typename T>
+void by_pass(const Tensor& x_tmp, Tensor* x);
+
 // This fucction compute unsqueeze dims for reshape to replace unsqueeze.
 static std::vector<int64_t> get_unsqueeze_dims(
     const Tensor& origin, const std::vector<int64_t>& axis) {
diff --git a/paddle/phi/api/yaml/legacy_backward.yaml b/paddle/phi/api/yaml/legacy_backward.yaml
@@ -206,6 +206,7 @@
     param : [out_grad]
   kernel :
     func : dropout_grad
+  composite : dropout_grad(mask, out_grad, p, is_test, mode, x_grad)
 
 - backward_op : einsum_grad
   forward : einsum (Tensor[] x, str equation) -> Tensor(out), Tensor[](inner_cache), Tensor[](x_shape)
diff --git a/python/paddle/decomposition/rules.py b/python/paddle/decomposition/rules.py
@@ -145,6 +145,67 @@ def layernorm_composite(x, scale, bias, epsilon, begin_norm_axis):
     return out, mean_, variance
 
 
+@register_decomp('pd_op.dropout')
+def dropout_composite(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.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 x, cast(
+                mask, uint8_type
+            )  # assign(x), cast(mask, mask, core.VarDesc.VarType.UINT8)
+    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 sum_composite(x):
     ans = x[0]
diff --git a/test/cpp/pir/kernel_dialect/CMakeLists.txt b/test/cpp/pir/kernel_dialect/CMakeLists.txt
@@ -9,5 +9,4 @@ cc_test_old(
   pd_trait
   pir
   phi
-  standalone_executor
   gtest)
diff --git a/test/legacy_test/test_dropout_op.py b/test/legacy_test/test_dropout_op.py
