[XPU] update xhpc to support bf16 gather_element (#67914)

runzhech · web-flow · commit 1e0c8fb165fc · 2024-09-03T15:09:12.000+08:00
diff --git a/paddle/phi/backends/xpu/xpu3_op_list.cc b/paddle/phi/backends/xpu/xpu3_op_list.cc
@@ -604,6 +604,7 @@ XPUOpMap& get_kl3_ops() {
        XPUKernelSet({phi::DataType::INT8,
                      phi::DataType::INT16,
                      phi::DataType::INT32,
+                     phi::DataType::BFLOAT16,
                      phi::DataType::FLOAT16,
                      phi::DataType::FLOAT32,
                      phi::DataType::BOOL})},
diff --git a/paddle/phi/kernels/xpu/index_sample_kernel.cc b/paddle/phi/kernels/xpu/index_sample_kernel.cc
@@ -82,6 +82,7 @@ PD_REGISTER_KERNEL(index_sample,
                    XPU,
                    ALL_LAYOUT,
                    phi::IndexSampleKernel,
+                   phi::dtype::bfloat16,
                    phi::dtype::float16,
                    float,
                    int8_t,
diff --git a/paddle/phi/kernels/xpu/take_along_axis_kernel.cc b/paddle/phi/kernels/xpu/take_along_axis_kernel.cc
@@ -61,11 +61,7 @@ void TakeAlongAxisKernel(const Context& dev_ctx,
     }
   }
 
-  using XPUType = typename std::conditional<
-      std::is_same<T, phi::dtype::bfloat16>::value ||
-          std::is_same<T, phi::dtype::float16>::value,
-      typename XPUTypeTrait<phi::dtype::float16>::Type,
-      typename XPUTypeTrait<T>::Type>::type;
+  using XPUType = typename XPUTypeTrait<T>::Type;
   int r = XPU_SUCCESS;
 #ifndef PADDLE_WITH_XPU_PLUGIN
   if (index_dtype == DataType::INT32) {

Original file line number	Diff line number	Diff line change
`@@ -61,11 +61,7 @@ void TakeAlongAxisKernel(const Context& dev_ctx,`
`61`	`61`	`}`
`62`	`62`	`}`
`63`	`63`
`64`		`- using XPUType = typename std::conditional<`
`65`		`- std::is_same<T, phi::dtype::bfloat16>::value \|\|`
`66`		`- std::is_same<T, phi::dtype::float16>::value,`
`67`		`- typename XPUTypeTrait<phi::dtype::float16>::Type,`
`68`		`- typename XPUTypeTrait<T>::Type>::type;`
	`64`	`+ using XPUType = typename XPUTypeTrait<T>::Type;`
`69`	`65`	`int r = XPU_SUCCESS;`
`70`	`66`	`#ifndef PADDLE_WITH_XPU_PLUGIN`
`71`	`67`	`if (index_dtype == DataType::INT32) {`