[XPU] add more dtypes to logsumexp and repeat_interleave (#66585)

lj970926 · web-flow · commit be7e429c3af7 · 2024-08-01T10:33:42.000+08:00
* [XPU] add more dtypes to logsumexp and repeat_interleave

* buf fix

* bugfix
diff --git a/cmake/external/xpu.cmake b/cmake/external/xpu.cmake
@@ -30,7 +30,7 @@ if(NOT DEFINED XPU_XRE_BASE_VERSION)
   set(XPU_XRE_BASE_VERSION "4.32.0.1")
 endif()
 if(NOT DEFINED XPU_XHPC_BASE_DATE)
-  set(XPU_XHPC_BASE_DATE "20240712")
+  set(XPU_XHPC_BASE_DATE "20240730")
 endif()
 set(XPU_XCCL_BASE_VERSION "1.2.5")
 if(NOT DEFINED XPU_XFT_BASE_VERSION)
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -625,6 +625,10 @@ XPUOpMap& get_kl3_ops() {
       {"logical_not", XPUKernelSet({phi::DataType::BOOL})},
       {"logical_or", XPUKernelSet({phi::DataType::BOOL})},
       {"logical_xor", XPUKernelSet({phi::DataType::BOOL})},
+      {"logsumexp",
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"lookup_table_v2_grad",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
@@ -824,7 +828,8 @@ XPUOpMap& get_kl3_ops() {
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT32,
                      phi::DataType::INT64,
-                     phi::DataType::BFLOAT16})},
+                     phi::DataType::BFLOAT16,
+                     phi::DataType::FLOAT16})},
       {"repeat_interleave_with_tensor_index",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::INT32,
diff --git a/paddle/phi/kernels/xpu/logsumexp_kernel.cc b/paddle/phi/kernels/xpu/logsumexp_kernel.cc
@@ -11,33 +11,37 @@
 // 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/phi/kernels/logsumexp_kernel.h"
 
-#include "paddle/phi/backends/xpu/xpu_header.h"
+#include "paddle/phi/backends/xpu/enforce_xpu.h"
 #include "paddle/phi/core/kernel_registry.h"
 
 namespace phi {
 
 template <typename T, typename Context>
-void XPULogsumexpKernel(const Context& dev_ctx,
-                        const DenseTensor& x,
-                        const std::vector<int>& axis,
-                        bool keepdim,
-                        bool reduce_all,
-                        DenseTensor* out) {
-  auto* input = &x;
+void LogsumexpKernel(const Context& dev_ctx,
+                     const DenseTensor& x,
+                     const std::vector<int>& axis,
+                     bool keepdim,
+                     bool reduce_all,
+                     DenseTensor* out) {
   auto* output = out;
-
-  const auto& input_dim_size = input->dims().size();
+  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);
 
-  const T* input_data = input->data<T>();
-  T* output_data = dev_ctx.template Alloc<T>(output);
+  auto input_data = reinterpret_cast<const XPUType*>(x.data<T>());
+  auto output_data =
+      reinterpret_cast<XPUType*>(dev_ctx.template Alloc<T>(output));
 
   std::vector<int> axis_shape;
-  std::vector<int> xdims(input_dim_size);
-  for (int i = 0; i < input_dim_size; ++i) {
-    xdims[i] = input->dims()[i];
+  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;
   }
   if (reduce_all) {
     for (int i = 0; i < input_dim_size; ++i) {
@@ -49,17 +53,23 @@ void XPULogsumexpKernel(const Context& dev_ctx,
       axis_shape.push_back(rdim);
     }
   }
+  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."));
+  }
 
-  int r = xpu::logsumexp<T>(
+  int r = xpu::logsumexp<XPUType>(
       dev_ctx.x_context(), input_data, output_data, xdims, axis_shape);
-  PADDLE_ENFORCE_EQ(
-      r,
-      xpu::Error_t::SUCCESS,
-      phi::errors::External("XPU logsumexp kernel error! error value[%d %]",
-                            r,
-                            XPUAPIErrorMsg[r]));
+  PADDLE_ENFORCE_XDNN_SUCCESS(r, "logsumexp");
 }
 }  // namespace phi
 
-PD_REGISTER_KERNEL(logsumexp, XPU, ALL_LAYOUT, phi::XPULogsumexpKernel, float) {
-}
+PD_REGISTER_KERNEL(logsumexp,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::LogsumexpKernel,
+                   float,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
diff --git a/paddle/phi/kernels/xpu/repeat_interleave_kernel.cc b/paddle/phi/kernels/xpu/repeat_interleave_kernel.cc
@@ -136,7 +136,8 @@ PD_REGISTER_KERNEL(repeat_interleave,
                    float,
                    int,
                    int64_t,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16,
+                   phi::dtype::float16) {}
 
 PD_REGISTER_KERNEL(repeat_interleave_with_tensor_index,
                    XPU,
@@ -145,4 +146,5 @@ PD_REGISTER_KERNEL(repeat_interleave_with_tensor_index,
                    float,
                    int,
                    int64_t,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16,
+                   phi::dtype::float16) {}
diff --git a/test/xpu/test_logsumexp_op_xpu.py b/test/xpu/test_logsumexp_op_xpu.py
@@ -15,10 +15,15 @@
 import unittest
 
 import numpy as np
+from get_test_cover_info import (
+    XPUOpTestWrapper,
+    create_test_class,
+    get_xpu_op_support_types,
+)
+from op_test import convert_float_to_uint16, convert_uint16_to_float
 from op_test_xpu import XPUOpTest
 
 import paddle
-from paddle.base import core
 
 paddle.enable_static()
 
@@ -34,69 +39,80 @@ def ref_logsumexp(x, axis=None, keepdim=False, reduce_all=False):
     return out
 
 
-@unittest.skipIf(
-    core.is_compiled_with_xpu()
-    and core.get_xpu_device_version(0) == core.XPUVersion.XPU3,
-    "logsumexp do not support XPU3 currently.",
-)
-class XPUTestLogsumexp(XPUOpTest):
-    def setUp(self):
-        self.op_type = 'logsumexp'
-        self.shape = [2, 3, 4, 5]
-        self.dtype = 'float32'
-        self.axis = [-1]
-        self.keepdim = False
-        self.reduce_all = False
-        self.set_attrs()
-
-        np.random.seed(10)
-        x = np.random.uniform(-1, 1, self.shape).astype(self.dtype)
-        out = ref_logsumexp(x, self.axis, self.keepdim, self.reduce_all)
-
-        self.inputs = {'X': x}
-        self.outputs = {'Out': out}
-        self.attrs = {
-            'axis': self.axis,
-            'keepdim': self.keepdim,
-            'reduce_all': self.reduce_all,
-        }
-
-    def set_attrs(self):
-        pass
-
-    def test_check_output(self):
-        if paddle.is_compiled_with_xpu():
-            place = paddle.XPUPlace(0)
-            self.check_output_with_place(place)
-
-    def test_check_grad(self):
-        pass
-
-
-class TestLogsumexp_shape(XPUTestLogsumexp):
-    def set_attrs(self):
-        self.shape = [4, 5, 6]
-
-
-class TestLogsumexp_axis(XPUTestLogsumexp):
-    def set_attrs(self):
-        self.axis = [0, -1]
-
-
-class TestLogsumexp_axis_all(XPUTestLogsumexp):
-    def set_attrs(self):
-        self.axis = [0, 1, 2, 3]
-
-
-class TestLogsumexp_keepdim(XPUTestLogsumexp):
-    def set_attrs(self):
-        self.keepdim = True
-
-
-class TestLogsumexp_reduce_all(XPUTestLogsumexp):
-    def set_attrs(self):
-        self.reduce_all = True
-
+class XPUTestLogsumexpOp(XPUOpTestWrapper):
+    def __init__(self):
+        self.op_name = "logsumexp"
+
+    class XPUTestLogsumexp(XPUOpTest):
+        def setUp(self):
+            self.op_type = 'logsumexp'
+            self.shape = [2, 3, 4, 5]
+            self.dtype = self.in_type
+            self.axis = [-1]
+            self.keepdim = False
+            self.reduce_all = False
+            self.set_attrs()
+
+            np.random.seed(10)
+
+            x = np.random.uniform(-1, 1, self.shape).astype("float32")
+            if self.dtype == np.uint16:
+                x = convert_float_to_uint16(x)
+                out = ref_logsumexp(
+                    convert_uint16_to_float(x),
+                    self.axis,
+                    self.keepdim,
+                    self.reduce_all,
+                )
+            else:
+                x = x.astype(self.dtype)
+                out = ref_logsumexp(x, self.axis, self.keepdim, self.reduce_all)
+
+            self.inputs = {'X': x}
+            self.outputs = {'Out': out}
+            self.attrs = {
+                'axis': self.axis,
+                'keepdim': self.keepdim,
+                'reduce_all': self.reduce_all,
+            }
+
+        def set_attrs(self):
+            pass
+
+        def test_check_output(self):
+            if paddle.is_compiled_with_xpu():
+                place = paddle.XPUPlace(0)
+                self.check_output_with_place(place)
+
+        def test_check_grad(self):
+            if paddle.is_compiled_with_xpu():
+                place = paddle.XPUPlace(0)
+                self.check_grad_with_place(place, ['X'], 'Out')
+
+    class TestLogsumexp_shape(XPUTestLogsumexp):
+        def set_attrs(self):
+            self.shape = [4, 5, 6]
+
+    class TestLogsumexp_axis(XPUTestLogsumexp):
+        def set_attrs(self):
+            self.axis = [0, -1]
+
+    class TestLogsumexp_axis_all(XPUTestLogsumexp):
+        def set_attrs(self):
+            self.axis = [0, 1, 2, 3]
+
+    class TestLogsumexp_keepdim(XPUTestLogsumexp):
+        def set_attrs(self):
+            self.keepdim = True
+
+    class TestLogsumexp_reduce_all(XPUTestLogsumexp):
+        def set_attrs(self):
+            self.reduce_all = True
+
+
+support_types = get_xpu_op_support_types("logsumexp")
+for stype in support_types:
+    create_test_class(globals(), XPUTestLogsumexpOp, stype)
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/test/xpu/test_repeat_interleave_op_xpu.py b/test/xpu/test_repeat_interleave_op_xpu.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.`
	`1`	`+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.`
`2`	`2`	`#`
`3`	`3`	`# Licensed under the Apache License, Version 2.0 (the "License");`
`4`	`4`	`# you may not use this file except in compliance with the License.`