[XPU] use non fusion logsumexp impl to normalize input for numeric st… (#68157)

lj970926 · web-flow · commit fa59b7b9c53d · 2024-09-13T13:57:43.000+08:00
* [XPU] use non fusion logsumexp impl to normalize input for numeric stability

* bug fix

* fix bugs in test
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -387,7 +387,10 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::FLOAT32,
                      phi::DataType::BOOL})},
       {"exp_grad", XPUKernelSet({phi::DataType::FLOAT32})},
-      {"exp", XPUKernelSet({phi::DataType::FLOAT32})},
+      {"exp",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"expand_as_v2",
        XPUKernelSet({phi::DataType::INT32,
                      phi::DataType::INT64,
@@ -664,7 +667,10 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::INT8,
                      phi::DataType::INT32,
                      phi::DataType::INT64})},
-      {"log", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+      {"log",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"log_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"log_softmax", XPUKernelSet({phi::DataType::FLOAT32})},
       {"log_softmax_grad", XPUKernelSet({phi::DataType::FLOAT32})},
diff --git a/paddle/phi/kernels/xpu/activation_kernel.cc b/paddle/phi/kernels/xpu/activation_kernel.cc
@@ -644,8 +644,13 @@ PD_REGISTER_KERNEL(square,
                    phi::dtype::float16,
                    phi::dtype::bfloat16) {}
 
-PD_REGISTER_KERNEL(
-    log, XPU, ALL_LAYOUT, phi::LogKernel, float, phi::dtype::float16) {}
+PD_REGISTER_KERNEL(log,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::LogKernel,
+                   float,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
 
 PD_REGISTER_KERNEL(
     relu6, XPU, ALL_LAYOUT, phi::Relu6Kernel, float, phi::dtype::float16) {}
@@ -681,10 +686,17 @@ PD_REGISTER_KERNEL(rsqrt,
                    phi::dtype::float16,
                    phi::dtype::bfloat16) {}
 
+PD_REGISTER_KERNEL(exp,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::ExpKernel,
+                   float,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
+
 #define PD_REGISTER_ACTIVATION_KERNEL(name, func) \
   PD_REGISTER_KERNEL(name, XPU, ALL_LAYOUT, phi::func, float) {}
 
-PD_REGISTER_ACTIVATION_KERNEL(exp, ExpKernel)  // no grad
 PD_REGISTER_ACTIVATION_KERNEL(floor, FloorKernel)
 PD_REGISTER_ACTIVATION_KERNEL(mish, MishKernel)
 PD_REGISTER_ACTIVATION_KERNEL(reciprocal, ReciprocalKernel)
diff --git a/paddle/phi/kernels/xpu/logsumexp_kernel.cc b/paddle/phi/kernels/xpu/logsumexp_kernel.cc
@@ -15,6 +15,11 @@
 
 #include "paddle/phi/backends/xpu/enforce_xpu.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/activation_kernel.h"
+#include "paddle/phi/kernels/elementwise_add_kernel.h"
+#include "paddle/phi/kernels/elementwise_subtract_kernel.h"
+#include "paddle/phi/kernels/reduce_max_kernel.h"
+#include "paddle/phi/kernels/reduce_sum_kernel.h"
 
 namespace phi {
 
@@ -25,44 +30,61 @@ void LogsumexpKernel(const Context& dev_ctx,
                      bool keepdim,
                      bool reduce_all,
                      DenseTensor* out) {
-  auto* output = out;
-  using XPUType = typename XPUTypeTrait<T>::Type;
-  int input_dim_size = x.dims().size();
-  // The dims has full dim, set the reduce_all is True
-  reduce_all |= (static_cast<int>(axis.size()) == input_dim_size);
-
-  auto input_data = reinterpret_cast<const XPUType*>(x.data<T>());
-  auto output_data =
-      reinterpret_cast<XPUType*>(dev_ctx.template Alloc<T>(output));
+  auto xdim = x.dims();
+  for (int i = 0; i < xdim.size(); i++)
+    PADDLE_ENFORCE_LT(0,
+                      xdim[i],
+                      errors::InvalidArgument(
+                          "The dims of Input(X) should be greater than 0."));
 
-  std::vector<int> axis_shape;
-  std::vector<int> xdims = common::vectorize<int>(x.dims());
-  if (input_dim_size == 0 && x.numel() != 0) {
-    // 0-d Tensor.
-    xdims = {1};
-    input_dim_size = 1;
-    reduce_all = true;
+  reduce_all = recompute_reduce_all(x, axis, reduce_all);
+  std::vector<int64_t> outdim_vec, keeped_outdim_vec;
+  std::vector<int> axis_vec;
+  int64_t compute_size = 1, other_size = 1;
+  for (auto i : axis) {
+    auto v = i >= 0 ? i : i + xdim.size();
+    axis_vec.push_back(v);
   }
-  if (reduce_all) {
-    for (int i = 0; i < input_dim_size; ++i) {
-      axis_shape.push_back(i);
-    }
-  } else {
-    for (size_t i = 0; i < axis.size(); ++i) {
-      int rdim = axis[i] < 0 ? axis[i] + input_dim_size : axis[i];
-      axis_shape.push_back(rdim);
+  if (axis.size() == 0 || reduce_all) {
+    axis_vec.clear();
+    for (int i = 0; i < xdim.size(); i++) {
+      axis_vec.push_back(i);
     }
   }
-  for (size_t i = 0; i < xdims.size(); ++i) {
-    PADDLE_ENFORCE_LT(0,
-                      xdims[i],
-                      errors::InvalidArgument(
-                          "The dims of Input(X) should be greater than 0."));
+  for (int i = 0; i < xdim.size(); i++) {
+    bool flag = false;
+    for (auto v : axis_vec) {
+      if (v == i) {
+        flag = true;
+        break;
+      }
+    }
+    if (flag) {
+      compute_size *= xdim[i];
+      keeped_outdim_vec.push_back(1);
+      if (keepdim) outdim_vec.push_back(1);
+    } else {
+      other_size *= xdim[i];
+      outdim_vec.push_back(xdim[i]);
+      keeped_outdim_vec.push_back(xdim[i]);
+    }
   }
+  auto outdim = common::make_ddim(outdim_vec);
+  auto keeped_outdim = common::make_ddim(keeped_outdim_vec);
+
+  // The XPU logsumexp api does not use xmax to normalize its input, so we
+  // fallback to the non fusion impl currently.
+  DenseTensor max_x;
+  max_x.Resize(keeped_outdim);
+  MaxKernel<T, Context>(dev_ctx, x, axis_vec, true, &max_x);
 
-  int r = xpu::logsumexp<XPUType>(
-      dev_ctx.x_context(), input_data, output_data, xdims, axis_shape);
-  PADDLE_ENFORCE_XDNN_SUCCESS(r, "logsumexp");
+  DenseTensor temp_x = Subtract<T, Context>(dev_ctx, x, max_x);
+  ExpKernel<T, Context>(dev_ctx, temp_x, &temp_x);
+  SumKernel<T, Context>(dev_ctx, temp_x, axis_vec, x.dtype(), keepdim, out);
+  LogKernel<T, Context>(dev_ctx, *out, out);
+  max_x.Resize(outdim);
+  out->Resize(outdim);
+  AddKernel<T, Context>(dev_ctx, *out, max_x, out);
 }
 }  // namespace phi
 
diff --git a/test/xpu/test_activation_op_xpu.py b/test/xpu/test_activation_op_xpu.py
@@ -70,10 +70,14 @@ def set_case(self):
             self.op_type = 'exp'
             self.dtype = self.in_type
 
-            x = np.random.uniform(-1, 1, [11, 17]).astype(self.dtype)
+            x = np.random.uniform(-1, 1, [11, 17])
+            if self.dtype == np.uint16:
+                new_x = convert_float_to_uint16(x)
+            else:
+                new_x = x.astype(self.dtype)
             out = np.exp(x)
             self.attrs = {'use_xpu': True}
-            self.inputs = {'X': OpTest.np_dtype_to_base_dtype(x)}
+            self.inputs = {'X': new_x}
             self.outputs = {'Out': out}
 
     class XPUTestExp_ZeroDIm(TestActivationOPBase):
@@ -492,11 +496,15 @@ class XPUTestLog(TestActivationOPBase):
         def set_case(self):
             self.op_type = "log"
             self.dtype = self.in_type
-            x = np.random.uniform(0.1, 1, self.shape).astype(self.dtype)
+            x = np.random.uniform(0.1, 1, self.shape)
+            if self.dtype == np.int16:
+                new_x = convert_float_to_uint16(x)
+            else:
+                new_x = x.astype(self.dtype)
             out = np.log(x)
 
             self.attrs = {'use_xpu': True}
-            self.inputs = {'X': OpTest.np_dtype_to_base_dtype(x)}
+            self.inputs = {'X': new_x}
             self.outputs = {'Out': out}
 
     class TestLogCase_ZeroDim(XPUTestLog):
