[XPU] add some bf16 ops and update xdnn (#59653)

houj04 · web-flow · commit e3a535ac4e52 · 2023-12-05T17:07:08.000+08:00
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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)
