[PHI bf16] add bf16 support for where op (#58616)

RuohengMa · web-flow · commit 81782d6c78d1 · 2023-11-06T10:26:01.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 "20231025")
+  set(XPU_BASE_DATE "20231101")
 endif()
 set(XPU_XCCL_BASE_VERSION "1.0.53.6")
 if(NOT DEFINED XPU_XFT_BASE_VERSION)
diff --git a/paddle/phi/backends/xpu/xpu2_op_list.cc b/paddle/phi/backends/xpu/xpu2_op_list.cc
@@ -1048,7 +1048,8 @@ XPUOpMap& get_kl2_ops() {
        XPUKernelSet({phi::DataType::INT32,
                      phi::DataType::INT64,
                      phi::DataType::FLOAT32,
-                     phi::DataType::FLOAT16})},
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"sin", XPUKernelSet({phi::DataType::FLOAT32})},
       {"sin_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"cos", XPUKernelSet({phi::DataType::FLOAT32})},
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -1013,7 +1013,8 @@ XPUOpMap& get_kl3_ops() {
        XPUKernelSet({phi::DataType::INT32,
                      phi::DataType::INT64,
                      phi::DataType::FLOAT32,
-                     phi::DataType::FLOAT16})},
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"sin", XPUKernelSet({phi::DataType::FLOAT32})},
       {"sin_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"cos", XPUKernelSet({phi::DataType::FLOAT32})},
diff --git a/paddle/phi/kernels/xpu/where_kernel.cc b/paddle/phi/kernels/xpu/where_kernel.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+// Copyright (c) 2023 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.
@@ -57,4 +57,5 @@ PD_REGISTER_KERNEL(where,
                    float,
                    int,
                    int64_t,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {}
diff --git a/test/xpu/test_where_op_xpu.py b/test/xpu/test_where_op_xpu.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+# Copyright (c) 2023 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.
@@ -20,6 +20,7 @@
     create_test_class,
     get_xpu_op_support_types,
 )
+from op_test import convert_float_to_uint16
 from op_test_xpu import XPUOpTest
 
 import paddle
@@ -37,14 +38,23 @@ class TestXPUWhereOp(XPUOpTest):
         def setUp(self):
             self.init_config()
             self.init_data()
+            self.convert_data_if_bf16()
             self.inputs = {'Condition': self.cond, 'X': self.x, 'Y': self.y}
             self.outputs = {'Out': np.where(self.cond, self.x, self.y)}
 
         def init_data(self):
-            self.x = np.random.uniform(-3, 5, (100)).astype(self.dtype)
-            self.y = np.random.uniform(-3, 5, (100)).astype(self.dtype)
+            self.x = np.random.uniform(-3, 5, (100))
+            self.y = np.random.uniform(-3, 5, (100))
             self.cond = np.zeros(100).astype("bool")
 
+        def convert_data_if_bf16(self):
+            if self.dtype == np.uint16:
+                self.x = convert_float_to_uint16(self.x)
+                self.y = convert_float_to_uint16(self.y)
+            else:
+                self.x = self.x.astype(self.dtype)
+                self.y = self.y.astype(self.dtype)
+
         def init_config(self):
             self.op_type = "where"
             self.dtype = self.in_type
@@ -59,14 +69,14 @@ def test_check_grad(self):
 
     class TestXPUWhereOp2(TestXPUWhereOp):
         def init_data(self):
-            self.x = np.random.uniform(-5, 5, (60, 2)).astype(self.dtype)
-            self.y = np.random.uniform(-5, 5, (60, 2)).astype(self.dtype)
+            self.x = np.random.uniform(-5, 5, (60, 2))
+            self.y = np.random.uniform(-5, 5, (60, 2))
             self.cond = np.ones((60, 2)).astype("bool")
 
     class TestXPUWhereOp3(TestXPUWhereOp):
         def init_data(self):
-            self.x = np.random.uniform(-3, 5, (20, 2, 4)).astype(self.dtype)
-            self.y = np.random.uniform(-3, 5, (20, 2, 4)).astype(self.dtype)
+            self.x = np.random.uniform(-3, 5, (20, 2, 4))
+            self.y = np.random.uniform(-3, 5, (20, 2, 4))
             self.cond = np.array(
                 np.random.randint(2, size=(20, 2, 4)), dtype=bool
             )
