diff --git a/cmake/external/xpu.cmake b/cmake/external/xpu.cmake index 5534a24c8c297..1f804fb8d4d42 100644 --- a/cmake/external/xpu.cmake +++ b/cmake/external/xpu.cmake @@ -24,7 +24,7 @@ set(XPU_XFT_LIB_NAME "libxft.so") set(XPU_XPTI_LIB_NAME "libxpti.so") if(NOT DEFINED XPU_BASE_DATE) - set(XPU_BASE_DATE "20231128") + set(XPU_BASE_DATE "20231203") endif() set(XPU_XCCL_BASE_VERSION "1.1.6.1") 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 index 4e55220f79aec..6d52ef3dfdf35 100644 --- a/paddle/phi/backends/xpu/xpu3_op_list.cc +++ b/paddle/phi/backends/xpu/xpu3_op_list.cc @@ -814,9 +814,13 @@ XPUOpMap& get_kl3_ops() { phi::DataType::INT32, phi::DataType::INT64})}, {"softmax", - XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})}, + XPUKernelSet({phi::DataType::FLOAT32, + phi::DataType::FLOAT16, + phi::DataType::BFLOAT16})}, {"softmax_grad", - XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})}, + XPUKernelSet({phi::DataType::FLOAT32, + phi::DataType::FLOAT16, + phi::DataType::BFLOAT16})}, {"fused_softmax_mask_grad", XPUKernelSet({phi::DataType::FLOAT32})}, {"softmax_with_cross_entropy_grad", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})}, @@ -1094,11 +1098,14 @@ XPUOpMap& get_kl3_ops() { {"where_grad", XPUKernelSet({phi::DataType::INT32, phi::DataType::INT64, + phi::DataType::FLOAT64, + phi::DataType::FLOAT32, phi::DataType::FLOAT16, - phi::DataType::FLOAT32})}, + phi::DataType::BFLOAT16})}, {"where", XPUKernelSet({phi::DataType::INT32, phi::DataType::INT64, + phi::DataType::FLOAT64, phi::DataType::FLOAT32, phi::DataType::FLOAT16, phi::DataType::BFLOAT16})}, diff --git a/paddle/phi/kernels/xpu/softmax_grad_kernel.cc b/paddle/phi/kernels/xpu/softmax_grad_kernel.cc index 809534b8d96f1..933fed1472f25 100644 --- a/paddle/phi/kernels/xpu/softmax_grad_kernel.cc +++ b/paddle/phi/kernels/xpu/softmax_grad_kernel.cc @@ -64,4 +64,5 @@ PD_REGISTER_KERNEL(softmax_grad, ALL_LAYOUT, phi::SoftmaxGradKernel, float, - phi::dtype::float16) {} + phi::dtype::float16, + phi::dtype::bfloat16) {} diff --git a/paddle/phi/kernels/xpu/softmax_kernel.cc b/paddle/phi/kernels/xpu/softmax_kernel.cc index e8d8cd3cc77a3..09ca288fe9e46 100644 --- a/paddle/phi/kernels/xpu/softmax_kernel.cc +++ b/paddle/phi/kernels/xpu/softmax_kernel.cc @@ -77,5 +77,10 @@ void SoftmaxKernel(const Context& dev_ctx, } // namespace phi -PD_REGISTER_KERNEL( - softmax, XPU, ALL_LAYOUT, phi::SoftmaxKernel, float, phi::dtype::float16) {} +PD_REGISTER_KERNEL(softmax, + XPU, + ALL_LAYOUT, + phi::SoftmaxKernel, + float, + phi::dtype::float16, + phi::dtype::bfloat16) {} diff --git a/paddle/phi/kernels/xpu/where_grad_kernel.cc b/paddle/phi/kernels/xpu/where_grad_kernel.cc index 49a5a1b22685d..f3bd915fb945b 100644 --- a/paddle/phi/kernels/xpu/where_grad_kernel.cc +++ b/paddle/phi/kernels/xpu/where_grad_kernel.cc @@ -68,6 +68,8 @@ PD_REGISTER_KERNEL(where_grad, ALL_LAYOUT, phi::WhereGradKernel, float, - phi::dtype::float16, + double, int, - int64_t) {} + int64_t, + phi::dtype::float16, + phi::dtype::bfloat16) {} diff --git a/paddle/phi/kernels/xpu/where_kernel.cc b/paddle/phi/kernels/xpu/where_kernel.cc index 1edfc693cff86..4c76367d221d3 100644 --- a/paddle/phi/kernels/xpu/where_kernel.cc +++ b/paddle/phi/kernels/xpu/where_kernel.cc @@ -55,6 +55,7 @@ PD_REGISTER_KERNEL(where, ALL_LAYOUT, phi::WhereKernel, float, + double, int, int64_t, phi::dtype::float16, diff --git a/test/xpu/test_softmax_op_xpu.py b/test/xpu/test_softmax_op_xpu.py index 9b849832bd984..e42bf9411b34f 100644 --- a/test/xpu/test_softmax_op_xpu.py +++ b/test/xpu/test_softmax_op_xpu.py @@ -54,7 +54,7 @@ def __init__(self): def dynamic_create_class(self): base_class = self.TestSoftmaxOp classes = [] - shapes = [[2, 3, 4, 5], [7, 1], [63, 18], [2, 38512], [3, 4095]] + shapes = [[2, 3, 4, 5], [63, 18], [2, 38512], [3, 4095]] axis = [-1, 0, 1] for shape in shapes: for axi in axis: @@ -67,9 +67,9 @@ class TestSoftmaxOp(XPUOpTest): def setUp(self): self.op_type = "softmax" if not hasattr(self, 'shape'): - self.shape = [1, 7] + self.shape = [2, 3, 4, 5] self.axis = -1 - self.dtype = np.float32 + self.dtype = self.in_type x = np.random.uniform(-1, 1, self.shape).astype(self.dtype) out = np.apply_along_axis(stable_softmax, self.axis, x)