[XPU] update xhpc date, support bf16 arange, clip, clip_grad, norm and reduce_max_grad (#68411)

runzhech · web-flow · commit ebb194d30614 · 2024-10-18T16:02:50.000+08:00
* [XPU] support bf16 arange, clip, clip_grad, norm, reduce_max_grad

* [XPU] update xhpc date to eb35/20241015
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -50,6 +50,8 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::BFLOAT16})},
       {"arange_tensor",
        XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::INT32,
                      phi::DataType::INT64})},
       {"arg_max",
@@ -215,11 +217,13 @@ XPUOpMap& get_kl3_ops() {
       {"clip",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::INT64,
                      phi::DataType::INT32})},
       {"clip_by_norm", XPUKernelSet({phi::DataType::FLOAT32})},
       {"clip_grad",
        XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
                      phi::DataType::INT64,
                      phi::DataType::INT32})},
       {"coalesce_tensor",
@@ -808,7 +812,12 @@ XPUOpMap& get_kl3_ops() {
       {"nearest_interp_v2_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"nll_loss", XPUKernelSet({phi::DataType::FLOAT32})},
       {"nll_loss_grad", XPUKernelSet({phi::DataType::FLOAT32})},
-      {"norm", XPUKernelSet({phi::DataType::FLOAT32})},
+      {"norm",
+       XPUKernelSet({
+           phi::DataType::FLOAT32,
+           phi::DataType::FLOAT16,
+           phi::DataType::BFLOAT16,
+       })},
       {"not_equal",
        XPUKernelSet({phi::DataType::INT64,
                      phi::DataType::INT32,
@@ -867,6 +876,8 @@ XPUOpMap& get_kl3_ops() {
        })},
       {"range",
        XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::INT64,
                      phi::DataType::INT32})},
       {"randperm",
@@ -879,7 +890,12 @@ XPUOpMap& get_kl3_ops() {
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"reduce_all", XPUKernelSet({phi::DataType::BOOL})},
       {"reduce_any", XPUKernelSet({phi::DataType::BOOL})},
-      {"reduce_max_grad", XPUKernelSet({phi::DataType::FLOAT32})},
+      {"reduce_max_grad",
+       XPUKernelSet({
+           phi::DataType::FLOAT32,
+           phi::DataType::FLOAT16,
+           phi::DataType::BFLOAT16,
+       })},
       {"reduce_max",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
diff --git a/paddle/phi/kernels/xpu/arange_kernel.cc b/paddle/phi/kernels/xpu/arange_kernel.cc
@@ -14,7 +14,7 @@ limitations under the License. */
 
 #include "paddle/phi/kernels/arange_kernel.h"
 #include "paddle/phi/backends/xpu/enforce_xpu.h"
-
+#include "paddle/phi/common/amp_type_traits.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/funcs/range_function.h"
 
@@ -26,17 +26,24 @@ void ArangeTensorKernel(const Context& dev_ctx,
                         const DenseTensor& end,
                         const DenseTensor& step,
                         DenseTensor* out) {
-  T start_value = GetValue<T, Context>(dev_ctx, start);
-  T end_value = GetValue<T, Context>(dev_ctx, end);
-  T step_value = GetValue<T, Context>(dev_ctx, step);
+  using MPType = typename phi::dtype::MPTypeTrait<T>::Type;
+  using XPUType = typename XPUTypeTrait<T>::Type;
+  MPType start_value =
+      static_cast<MPType>(GetValue<T, Context>(dev_ctx, start));
+  MPType end_value = static_cast<MPType>(GetValue<T, Context>(dev_ctx, end));
+  MPType step_value = static_cast<MPType>(GetValue<T, Context>(dev_ctx, step));
 
   int64_t size = 0;
   phi::funcs::GetSize(start_value, end_value, step_value, &size);
   out->Resize(common::make_ddim({size}));
-  auto* out_data = dev_ctx.template Alloc<T>(out);
+  XPUType* out_data =
+      reinterpret_cast<XPUType*>(dev_ctx.template Alloc<T>(out));
 
-  int ret = xpu::range<T>(
-      dev_ctx.x_context(), out_data, start_value, step_value, size);
+  int ret = xpu::range(dev_ctx.x_context(),
+                       out_data,
+                       static_cast<XPUType>(start_value),
+                       static_cast<XPUType>(step_value),
+                       size);
   PADDLE_ENFORCE_XDNN_SUCCESS(ret, "range");
 }
 
@@ -47,6 +54,8 @@ PD_REGISTER_KERNEL(arange_tensor,
                    ALL_LAYOUT,
                    phi::ArangeTensorKernel,
                    float,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    int,
                    int64_t) {
   kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
diff --git a/paddle/phi/kernels/xpu/clip_grad_kernel.cc b/paddle/phi/kernels/xpu/clip_grad_kernel.cc
@@ -34,11 +34,17 @@ void ClipGradKernel(const Context& ctx,
                      reinterpret_cast<const XPUDataType*>(out_grad.data<T>()),
                      reinterpret_cast<XPUDataType*>(x_grad->data<T>()),
                      x.numel(),
-                     min.to<XPUDataType>(),
-                     max.to<XPUDataType>());
+                     static_cast<XPUDataType>(min.to<T>()),
+                     static_cast<XPUDataType>(max.to<T>()));
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "clip_grad");
 }
 }  // namespace phi
 
-PD_REGISTER_KERNEL(
-    clip_grad, XPU, ALL_LAYOUT, phi::ClipGradKernel, float, int64_t, int) {}
+PD_REGISTER_KERNEL(clip_grad,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::ClipGradKernel,
+                   float,
+                   phi::dtype::float16,
+                   int64_t,
+                   int) {}
diff --git a/paddle/phi/kernels/xpu/clip_kernel.cc b/paddle/phi/kernels/xpu/clip_kernel.cc
@@ -55,5 +55,6 @@ PD_REGISTER_KERNEL(clip,
                    phi::ClipKernel,
                    float,
                    phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    int64_t,
                    int) {}
diff --git a/paddle/phi/kernels/xpu/norm_kernel.cc b/paddle/phi/kernels/xpu/norm_kernel.cc
@@ -29,6 +29,7 @@ void NormKernel(const Context& ctx,
                 bool is_test,
                 DenseTensor* out,
                 DenseTensor* norm) {
+  using XPUType = typename XPUTypeTrait<T>::Type;
   ctx.template Alloc<T>(out);
   ctx.template Alloc<T>(norm);
 
@@ -59,16 +60,22 @@ void NormKernel(const Context& ctx,
     xshape[i] = static_cast<int>(x_dims[i]);
   }
 
-  int r = xpu::l2_norm<T>(ctx.x_context(),
-                          x.data<T>(),
-                          out->data<T>(),
-                          norm->data<T>(),
-                          xshape,
-                          axis,
-                          epsilon);
+  int r = xpu::l2_norm(ctx.x_context(),
+                       reinterpret_cast<const XPUType*>(x.data<T>()),
+                       reinterpret_cast<XPUType*>(out->data<T>()),
+                       reinterpret_cast<XPUType*>(norm->data<T>()),
+                       xshape,
+                       axis,
+                       epsilon);
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "l2_norm");
 }
 
 }  // namespace phi
 
-PD_REGISTER_KERNEL(norm, XPU, ALL_LAYOUT, phi::NormKernel, float) {}
+PD_REGISTER_KERNEL(norm,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::NormKernel,
+                   float,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
diff --git a/paddle/phi/kernels/xpu/reduce_max_grad_kernel.cc b/paddle/phi/kernels/xpu/reduce_max_grad_kernel.cc
@@ -31,14 +31,17 @@ void ReduceMaxGradKernel(const Context& dev_ctx,
                          bool keep_dim,
                          bool reduce_all,
                          DenseTensor* x_grad) {
+  using XPUDataType = typename XPUTypeTrait<T>::Type;
   reduce_all = recompute_reduce_all(x, dims_arr, reduce_all);
   auto dims = dims_arr.GetData();
 
   dev_ctx.template Alloc<T>(x_grad);
-  const T* x_data = x.data<T>();
-  const T* out_data = out.data<T>();
-  const T* out_grad_data = out_grad.data<T>();
-  auto* x_grad_data = x_grad->data<T>();
+  const XPUDataType* x_data = reinterpret_cast<const XPUDataType*>(x.data<T>());
+  const XPUDataType* out_data =
+      reinterpret_cast<const XPUDataType*>(out.data<T>());
+  const XPUDataType* out_grad_data =
+      reinterpret_cast<const XPUDataType*>(out_grad.data<T>());
+  XPUDataType* x_grad_data = reinterpret_cast<XPUDataType*>(x_grad->data<T>());
   const auto& input_dim_size = x.dims().size();
   std::vector<int> true_dims;
   for (size_t i = 0; i < dims.size(); ++i) {
@@ -60,21 +63,21 @@ void ReduceMaxGradKernel(const Context& dev_ctx,
     }
   }
 
-  T* broadcast1 = nullptr;
-  T* broadcast2 = nullptr;
+  XPUDataType* broadcast1 = nullptr;
+  XPUDataType* broadcast2 = nullptr;
   bool* equal = nullptr;
 
   xpu::ctx_guard RAII_GUARD(dev_ctx.x_context());
 
-  broadcast1 = RAII_GUARD.alloc_l3_or_gm<T>(x.numel());
+  broadcast1 = RAII_GUARD.alloc_l3_or_gm<XPUDataType>(x.numel());
   PADDLE_ENFORCE_NOT_NULL(
       broadcast1, errors::ResourceExhausted("XPU has no enough memory"));
 
   equal = RAII_GUARD.alloc_l3_or_gm<bool>(x.numel());
   PADDLE_ENFORCE_NOT_NULL(
       equal, errors::ResourceExhausted("XPU has no enough memory"));
 
-  broadcast2 = RAII_GUARD.alloc_l3_or_gm<T>(x.numel());
+  broadcast2 = RAII_GUARD.alloc_l3_or_gm<XPUDataType>(x.numel());
   PADDLE_ENFORCE_NOT_NULL(
       broadcast2, errors::ResourceExhausted("XPU has no enough memory"));
 
@@ -87,31 +90,37 @@ void ReduceMaxGradKernel(const Context& dev_ctx,
   }
 
   // step 1. broadcast out and out_grad
-  int r = xpu::broadcast<T>(
-      dev_ctx.x_context(), out_data, broadcast1, ydims, xdims);
+  int r =
+      xpu::broadcast(dev_ctx.x_context(), out_data, broadcast1, ydims, xdims);
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "broadcast");
 
-  r = xpu::broadcast<T>(
+  r = xpu::broadcast(
       dev_ctx.x_context(), out_grad_data, broadcast2, ydims, xdims);
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "broadcast");
 
   // step 2. compare out_broadcast and x
-  r = xpu::equal<T>(dev_ctx.x_context(), x_data, broadcast1, equal, x.numel());
+  r = xpu::equal(dev_ctx.x_context(), x_data, broadcast1, equal, x.numel());
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "equal");
   // step 3. get x_grad
-  r = xpu::constant<T>(dev_ctx.x_context(), broadcast1, x.numel(), 0);
+  r = xpu::constant(
+      dev_ctx.x_context(), broadcast1, x.numel(), static_cast<XPUDataType>(0));
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "constant");
-  r = xpu::select<T>(dev_ctx.x_context(),
-                     equal,
-                     broadcast2,
-                     broadcast1,
-                     x_grad_data,
-                     xdims,
-                     xdims);
+  r = xpu::select(dev_ctx.x_context(),
+                  equal,
+                  broadcast2,
+                  broadcast1,
+                  x_grad_data,
+                  xdims,
+                  xdims);
   PADDLE_ENFORCE_XDNN_SUCCESS(r, "select");
 }
 
 }  // namespace phi
 
-PD_REGISTER_KERNEL(max_grad, XPU, ALL_LAYOUT, phi::ReduceMaxGradKernel, float) {
-}
+PD_REGISTER_KERNEL(max_grad,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::ReduceMaxGradKernel,
+                   float,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
diff --git a/test/xpu/test_clip_op_xpu.py b/test/xpu/test_clip_op_xpu.py
@@ -34,6 +34,7 @@ def __init__(self):
 
     class TestClipOp(XPUOpTest):
         def setUp(self):
+            self.python_api = paddle.clip
             self.init_dtype()
             self.set_xpu()
             self.op_type = "clip"
diff --git a/test/xpu/test_norm_op_xpu.py b/test/xpu/test_norm_op_xpu.py
@@ -20,6 +20,7 @@
     create_test_class,
     get_xpu_op_support_types,
 )
+from op_test import convert_float_to_uint16
 from op_test_xpu import XPUOpTest
 
 import paddle
@@ -43,11 +44,15 @@ def __init__(self):
     class TestXPUNormOp(XPUOpTest):
         def setUp(self):
             self.op_type = "norm"
-            self.dtype = self.in_type
+            self.dtype = (
+                self.in_type if self.in_type != np.uint16 else np.float32
+            )
             self.place = paddle.XPUPlace(0)
             self.init_test_case()
             x = np.random.random(self.shape).astype(self.dtype)
             y, norm = l2_norm(x, self.axis, self.epsilon)
+            if self.in_type == np.uint16:
+                x = convert_float_to_uint16(x)
             self.inputs = {'X': x}
             self.attrs = {'epsilon': self.epsilon, 'axis': self.axis}
             self.outputs = {'Out': y, 'Norm': norm}
diff --git a/test/xpu/test_range_xpu.py b/test/xpu/test_range_xpu.py
@@ -20,6 +20,7 @@
     create_test_class,
     get_xpu_op_support_types,
 )
+from op_test import convert_float_to_uint16
 from op_test_xpu import XPUOpTest
 
 import paddle
@@ -38,11 +39,20 @@ def setUp(self):
             self.op_type = "range"
             self.init_dtype()
             self.init_config()
-            self.inputs = {
-                'Start': np.array([self.case[0]]).astype(self.dtype),
-                'End': np.array([self.case[1]]).astype(self.dtype),
-                'Step': np.array([self.case[2]]).astype(self.dtype),
-            }
+            start = np.array([self.case[0]]).astype(
+                self.dtype if self.dtype != np.uint16 else np.float32
+            )
+            end = np.array([self.case[1]]).astype(
+                self.dtype if self.dtype != np.uint16 else np.float32
+            )
+            step = np.array([self.case[2]]).astype(
+                self.dtype if self.dtype != np.uint16 else np.float32
+            )
+            if self.dtype == np.uint16:
+                start = convert_float_to_uint16(start)
+                end = convert_float_to_uint16(end)
+                step = convert_float_to_uint16(step)
+            self.inputs = {'Start': start, 'End': end, 'Step': step}
 
             self.outputs = {
                 'Out': np.arange(
