[XPU] bf16 for multinomial op (#69898)

houj04 · web-flow · commit d70704410e43 · 2024-12-04T10:47:02.000+08:00
diff --git a/paddle/phi/backends/xpu/xpu2_op_list.cc b/paddle/phi/backends/xpu/xpu2_op_list.cc
@@ -713,7 +713,9 @@ XPUOpMap& get_kl2_ops() {
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"multiclass_nms3", XPUKernelSet({phi::DataType::FLOAT32})},
       {"multinomial",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"nearest_interp_v2",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -813,7 +813,9 @@ XPUOpMap& get_kl3_ops() {
        XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"multiclass_nms3", XPUKernelSet({phi::DataType::FLOAT32})},
       {"multinomial",
-       XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
+       XPUKernelSet({phi::DataType::FLOAT32,
+                     phi::DataType::FLOAT16,
+                     phi::DataType::BFLOAT16})},
       {"nearest_interp_v2",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
diff --git a/paddle/phi/kernels/xpu/multinomial_kernel.cc b/paddle/phi/kernels/xpu/multinomial_kernel.cc
@@ -79,6 +79,7 @@ PD_REGISTER_KERNEL(multinomial,
                    ALL_LAYOUT,
                    phi::MultinomialKernel,
                    float,
-                   phi::dtype::float16) {
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16) {
   kernel->OutputAt(0).SetDataType(phi::DataType::INT64);
 }
diff --git a/test/xpu/test_multinomial_op_xpu.py b/test/xpu/test_multinomial_op_xpu.py
@@ -21,6 +21,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
@@ -61,11 +62,14 @@ def setUp(self):
             self.op_type = "multinomial"
             self.python_api = paddle.multinomial
             self.init_data()
-            self.inputs = {"X": self.input_np}
+            if self.in_type == np.uint16:
+                self.inputs = {"X": convert_float_to_uint16(self.input_np)}
+            else:
+                self.inputs = {"X": self.input_np.astype(self.dtype)}
 
         def init_data(self):
             # input probability is a vector, and replacement is True
-            self.input_np = np.random.rand(4).astype(self.dtype)
+            self.input_np = np.random.rand(4).astype(np.float32)
             self.outputs = {"Out": np.zeros(100000).astype("int64")}
             self.attrs = {"num_samples": 100000, "replacement": True}
 
@@ -95,7 +99,7 @@ def verify_output(self, outs):
     class TestMultinomialOp2(TestMultinomialOp):
         def init_data(self):
             # input probability is a matrix
-            self.input_np = np.random.rand(3, 4).astype(self.dtype)
+            self.input_np = np.random.rand(3, 4).astype(np.float32)
             self.outputs = {"Out": np.zeros((3, 100000)).astype("int64")}
             self.attrs = {"num_samples": 100000, "replacement": True}
 
@@ -105,7 +109,7 @@ def sample_output(self, out):
     class TestMultinomialOp3(TestMultinomialOp):
         def init_data(self):
             # replacement is False. number of samples must be less than number of categories.
-            self.input_np = np.random.rand(1000).astype(self.dtype)
+            self.input_np = np.random.rand(1000).astype(np.float32)
             self.outputs = {"Out": np.zeros(100).astype("int64")}
             self.attrs = {"num_samples": 100, "replacement": False}
 

Original file line number	Diff line number	Diff line change
`@@ -79,6 +79,7 @@ PD_REGISTER_KERNEL(multinomial,`
`79`	`79`	`ALL_LAYOUT,`
`80`	`80`	`phi::MultinomialKernel,`
`81`	`81`	`float,`
`82`		`- phi::dtype::float16) {`
	`82`	`+ phi::dtype::float16,`
	`83`	`+ phi::dtype::bfloat16) {`
`83`	`84`	`kernel->OutputAt(0).SetDataType(phi::DataType::INT64);`
`84`	`85`	`}`