[XPU] update dtype for numel/concat/slice (#63715)

houj04 · web-flow · commit ffcca4dcaad0 · 2024-04-22T19:06:51.000+08:00
* [XPU] update dtype for numel op

* dtype for concat and slice.
diff --git a/cmake/external/xpu.cmake b/cmake/external/xpu.cmake
@@ -32,7 +32,7 @@ if(NOT DEFINED XPU_XDNN_BASE_DATE)
   set(XPU_XDNN_BASE_DATE "20240327")
 endif()
 if(NOT DEFINED XPU_XHPC_BASE_DATE)
-  set(XPU_XHPC_BASE_DATE "20240413")
+  set(XPU_XHPC_BASE_DATE "20240422")
 endif()
 set(XPU_XCCL_BASE_VERSION "1.2.0.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
@@ -202,9 +202,11 @@ XPUOpMap& get_kl3_ops() {
                      phi::DataType::BFLOAT16,
                      phi::DataType::FLOAT64,
                      phi::DataType::BOOL,
+                     phi::DataType::UINT8,
                      phi::DataType::INT8,
-                     phi::DataType::INT64,
-                     phi::DataType::INT32})},
+                     phi::DataType::INT16,
+                     phi::DataType::INT32,
+                     phi::DataType::INT64})},
       {"conv2d_grad",
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"conv2d",
@@ -855,8 +857,11 @@ XPUOpMap& get_kl3_ops() {
        XPUKernelSet({phi::DataType::INT64,
                      phi::DataType::INT32,
                      phi::DataType::INT16,
+                     phi::DataType::UINT8,
                      phi::DataType::BOOL,
                      phi::DataType::FLOAT16,
+                     phi::DataType::FLOAT64,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::FLOAT32})},
       {"sigmoid_cross_entropy_with_logits_grad",
        XPUKernelSet({phi::DataType::FLOAT32})},
@@ -881,6 +886,10 @@ XPUOpMap& get_kl3_ops() {
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
                      phi::DataType::BFLOAT16,
+                     phi::DataType::FLOAT64,
+                     phi::DataType::UINT8,
+                     phi::DataType::INT8,
+                     phi::DataType::INT16,
                      phi::DataType::INT32,
                      phi::DataType::INT64})},
       {"softmax",
diff --git a/paddle/phi/kernels/xpu/concat_kernel.cc b/paddle/phi/kernels/xpu/concat_kernel.cc
@@ -116,10 +116,12 @@ PD_REGISTER_KERNEL(concat,
                    ALL_LAYOUT,
                    phi::ConcatKernel,
                    float,
-                   double,
                    phi::dtype::float16,
                    phi::dtype::bfloat16,
-                   int64_t,
-                   int,
+                   double,
+                   bool,
+                   uint8_t,
                    int8_t,
-                   bool) {}
+                   int16_t,
+                   int32_t,
+                   int64_t) {}
diff --git a/paddle/phi/kernels/xpu/numel_kernel.cc b/paddle/phi/kernels/xpu/numel_kernel.cc
@@ -22,11 +22,14 @@ PD_REGISTER_KERNEL(numel,
                    XPU,
                    ALL_LAYOUT,
                    phi::NumelKernel,
+                   uint8_t,
                    int16_t,
                    int,
                    int64_t,
                    phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    float,
+                   double,
                    bool,
                    phi::dtype::complex<float>,
                    phi::dtype::complex<double>) {
diff --git a/paddle/phi/kernels/xpu/slice_kernel.cc b/paddle/phi/kernels/xpu/slice_kernel.cc
@@ -118,7 +118,11 @@ PD_REGISTER_KERNEL(slice,
                    ALL_LAYOUT,
                    phi::SliceKernel,
                    float,
-                   int,
                    phi::dtype::float16,
                    phi::dtype::bfloat16,
+                   double,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int32_t,
                    int64_t) {}
