[XPU] use xdnn dropout_v3 (#62726)

houj04 · web-flow · commit b67004fab2a3 · 2024-03-19T17:48:31.000+08:00
* [XPU] use xdnn dropout_v3

* use count_nonzero to check results

* refine ut
diff --git a/paddle/phi/kernels/cpu/dropout_grad_kernel.cc b/paddle/phi/kernels/cpu/dropout_grad_kernel.cc
@@ -89,6 +89,7 @@ PD_REGISTER_KERNEL(dropout_grad,
                    phi::DropoutGradRawKernel,
                    float,
                    double,
+                   phi::dtype::float16,
                    phi::dtype::bfloat16) {}
 
 PD_REGISTER_KERNEL(
diff --git a/paddle/phi/kernels/cpu/dropout_kernel.cc b/paddle/phi/kernels/cpu/dropout_kernel.cc
@@ -209,6 +209,7 @@ PD_REGISTER_KERNEL(dropout,
                    phi::DropoutRawKernel,
                    float,
                    double,
+                   phi::dtype::float16,
                    phi::dtype::bfloat16) {
   kernel->OutputAt(1).SetDataType(phi::DataType::UINT8);
 }
diff --git a/paddle/phi/kernels/cpu/uniform_kernel.cc b/paddle/phi/kernels/cpu/uniform_kernel.cc
@@ -49,4 +49,5 @@ PD_REGISTER_KERNEL(uniform,
                    phi::UniformKernel,
                    float,
                    double,
+                   phi::dtype::float16,
                    phi::dtype::bfloat16) {}
diff --git a/paddle/phi/kernels/xpu/dropout_kernel.cc b/paddle/phi/kernels/xpu/dropout_kernel.cc
@@ -34,15 +34,18 @@ void DropoutRawKernel(const Context& dev_ctx,
                       bool fix_seed,
                       DenseTensor* out,
                       DenseTensor* mask) {
+  bool is_upscale = (mode == "upscale_in_train");
+  dev_ctx.template Alloc<T>(out);
+  if (mask) {
+    dev_ctx.template Alloc<uint8_t>(mask);
+  }
+
   using XPUType = typename XPUTypeTrait<T>::Type;
-  auto* y = out;
   const auto* x_data = x.data<T>();
-  auto* y_data = dev_ctx.template Alloc<T>(y);
+  auto* y_data = out->data<T>();
   float dropout_prob = p.to<float>();
 
-  int is_upscale = (mode == "upscale_in_train");
-
-  if (!is_test) {
+  if (!is_test && mask) {
     int seed_data = 0;
     if (seed_tensor.get_ptr() != nullptr) {
       if ((seed_tensor->place()).GetType() == phi::AllocationType::XPU) {
@@ -54,23 +57,22 @@ void DropoutRawKernel(const Context& dev_ctx,
       } else {
         seed_data = *(seed_tensor->data<int>());
       }
-
     } else {
       seed_data = fix_seed ? seed : 0;
     }
     if (seed_data == 0) {
       seed_data = dev_ctx.GetGenerator()->Random64();
     }
 
-    auto* mask_data = dev_ctx.template Alloc<uint8_t>(mask);
+    auto* mask_data = mask->data<uint8_t>();
     xpu::ctx_guard RAII_GUARD(dev_ctx.x_context());
     auto dev_version =
         phi::backends::xpu::get_xpu_version(dev_ctx.GetPlace().GetDeviceId());
     // Special case when dropout_prob is 1.0
     if (dropout_prob == 1.0f) {
       int r = xpu::constant(dev_ctx.x_context(),
                             reinterpret_cast<XPUType*>(y_data),
-                            y->numel(),
+                            out->numel(),
                             XPUType(0));
       PADDLE_ENFORCE_XDNN_SUCCESS(r, "constant");
       r = xpu::constant(
@@ -79,21 +81,25 @@ void DropoutRawKernel(const Context& dev_ctx,
       return;
     }
     if (dev_version == phi::backends::xpu::XPUVersion::XPU3) {
-      int r = xpu::dropout_v2(dev_ctx.x_context(),
-                              reinterpret_cast<const XPUType*>(x.data<T>()),
-                              reinterpret_cast<XPUType*>(y->data<T>()),
-                              mask->data<uint8_t>(),
+      // int dropout_v3(Context* ctx, const T* input, T* res, uint8_t* mask,
+      // unsigned int seed, int64_t n, bool is_upscale, float dropout_prob);
+      int r = xpu::dropout_v3(dev_ctx.x_context(),
+                              reinterpret_cast<const XPUType*>(x_data),
+                              reinterpret_cast<XPUType*>(y_data),
+                              mask_data,
                               seed_data,
                               mask->numel(),
                               is_upscale,
                               dropout_prob);
-      PADDLE_ENFORCE_XDNN_SUCCESS(r, "dropout_v2");
+      PADDLE_ENFORCE_XDNN_SUCCESS(r, "dropout_v3");
     } else {
       XPUType* mask_tmp_data =
           RAII_GUARD.alloc_l3_or_gm<XPUType>(mask->numel());
+      // int dropout(Context* ctx, const T* input, T* res, T* mask, unsigned int
+      // seed, int64_t n, bool is_upscale, float dropout_prob);
       int r = xpu::dropout(dev_ctx.x_context(),
-                           reinterpret_cast<const XPUType*>(x.data<T>()),
-                           reinterpret_cast<XPUType*>(y->data<T>()),
+                           reinterpret_cast<const XPUType*>(x_data),
+                           reinterpret_cast<XPUType*>(y_data),
                            mask_tmp_data,
                            seed_data,
                            mask->numel(),
@@ -105,16 +111,23 @@ void DropoutRawKernel(const Context& dev_ctx,
       PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast");
     }
   } else {
-    float scale =
-        (is_upscale) ? (1.0) : (static_cast<float>(1.0f - dropout_prob));
-    int r = xpu::scale(dev_ctx.x_context(),
-                       reinterpret_cast<const XPUType*>(x_data),
-                       reinterpret_cast<XPUType*>(y_data),
-                       x.numel(),
-                       false,
-                       scale,
-                       0.0f);
-    PADDLE_ENFORCE_XDNN_SUCCESS(r, "scale");
+    if (is_upscale) {
+      // y = x
+      int ret = xpu::copy(dev_ctx.x_context(),
+                          reinterpret_cast<const int8_t*>(x_data),
+                          reinterpret_cast<int8_t*>(y_data),
+                          x.numel() * phi::SizeOf(x.dtype()));
+      PADDLE_ENFORCE_XDNN_SUCCESS(ret, "copy");
+    } else {
+      int r = xpu::scale(dev_ctx.x_context(),
+                         reinterpret_cast<const XPUType*>(x_data),
+                         reinterpret_cast<XPUType*>(y_data),
+                         x.numel(),
+                         false,
+                         1.0f - dropout_prob,
+                         0.0f);
+      PADDLE_ENFORCE_XDNN_SUCCESS(r, "scale");
+    }
   }
 }
 
@@ -126,5 +139,6 @@ PD_REGISTER_KERNEL(dropout,
                    phi::DropoutRawKernel,
                    float,
                    phi::dtype::float16) {
+  kernel->InputAt(1).SetBackend(phi::Backend::ALL_BACKEND);
   kernel->OutputAt(1).SetDataType(phi::DataType::UINT8);
 }
diff --git a/test/xpu/get_test_cover_info.py b/test/xpu/get_test_cover_info.py
@@ -84,8 +84,6 @@
 xpu_test_op_white_list = []
 xpu_test_device_type_white_list = ['xpu1_float64']
 xpu_test_op_type_white_list = [
-    'dropout_float16',
-    'dropout_grad_float16',
     "grad_add_float32",  # no api for grad_add, skip
     "lamb_float16",
     "lars_momentum_float32",
diff --git a/test/xpu/test_dropout_op_xpu.py b/test/xpu/test_dropout_op_xpu.py
@@ -176,10 +176,15 @@ def cal_grad_downscale_in_infer(self, mask):
         def test_backward_downscale_in_infer(self):
             for place in self.places:
                 with base.dygraph.guard(place):
-                    input = paddle.uniform([40, 40], dtype=self.in_type)
+                    prob = 0.1
+                    input = paddle.uniform([100, 40], dtype=self.in_type)
                     input.stop_gradient = False
                     out, mask = _legacy_C_ops.dropout(
-                        input, 'dropout_prob', 0.5
+                        input, 'dropout_prob', prob
+                    )
+                    nonzero = paddle.count_nonzero(out)
+                    np.testing.assert_allclose(
+                        prob, 1 - nonzero / 4000, atol=0.02
                     )
                     out.backward()
 
@@ -192,7 +197,7 @@ def test_backward_upscale_train(self):
             for place in self.places:
                 with base.dygraph.guard(place):
                     prob = 0.5
-                    input = paddle.uniform([40, 40], dtype=self.in_type)
+                    input = paddle.uniform([100, 40], dtype=self.in_type)
                     input.stop_gradient = False
                     out, mask = _legacy_C_ops.dropout(
                         input,
@@ -201,6 +206,10 @@ def test_backward_upscale_train(self):
                         "dropout_implementation",
                         "upscale_in_train",
                     )
+                    nonzero = paddle.count_nonzero(out)
+                    np.testing.assert_allclose(
+                        prob, 1 - nonzero / 4000, atol=0.02
+                    )
                     out.backward()
 
                     np.testing.assert_allclose(
@@ -211,8 +220,8 @@ def test_backward_upscale_train(self):
         def test_backward_upscale_train_2(self):
             for place in self.places:
                 with base.dygraph.guard(place):
-                    prob = 0.3
-                    input = paddle.uniform([40, 40], dtype=self.in_type)
+                    prob = 0.2
+                    input = paddle.uniform([100, 40], dtype=self.in_type)
                     input.stop_gradient = False
                     out, mask = _legacy_C_ops.dropout(
                         input,
@@ -221,6 +230,10 @@ def test_backward_upscale_train_2(self):
                         "dropout_implementation",
                         "upscale_in_train",
                     )
+                    nonzero = paddle.count_nonzero(out)
+                    np.testing.assert_allclose(
+                        prob, 1 - nonzero / 4000, atol=0.02
+                    )
                     out.backward()
 
                     np.testing.assert_allclose(

Original file line number	Diff line number	Diff line change
`@@ -209,6 +209,7 @@ PD_REGISTER_KERNEL(dropout,`
`209`	`209`	`phi::DropoutRawKernel,`
`210`	`210`	`float,`
`211`	`211`	`double,`
	`212`	`+ phi::dtype::float16,`
`212`	`213`	`phi::dtype::bfloat16) {`
`213`	`214`	`kernel->OutputAt(1).SetDataType(phi::DataType::UINT8);`
`214`	`215`	`}`