Skip to content

Commit

Permalink
Addition of bf16 type support for Compare OP (PaddlePaddle#46413)
Browse files Browse the repository at this point in the history 
* first commit

* clarify the quotes

* change code style format

* support bfloat16
  • Loading branch information
@JamesLim-sy @Xreki
JamesLim-sy authored and Xreki committed Apr 6, 2023
1 parent 6c6d290 commit f62a96e
Showing 1 changed file with 12 additions and 6 deletions.
18 changes: 12 additions & 6 deletions paddle/phi/kernels/kps/compare_kernel.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -114,7 +114,8 @@ PD_REGISTER_KERNEL(less_than,
int64_t,
float,
double,
phi::dtype::float16) {}
phi::dtype::float16,
phi::dtype::bfloat16) {}
PD_REGISTER_KERNEL(less_equal,
KPS,
ALL_LAYOUT,
Expand All @@ -125,7 +126,8 @@ PD_REGISTER_KERNEL(less_equal,
int64_t,
float,
double,
phi::dtype::float16) {}
phi::dtype::float16,
phi::dtype::bfloat16) {}
PD_REGISTER_KERNEL(greater_than,
KPS,
ALL_LAYOUT,
Expand All @@ -136,7 +138,8 @@ PD_REGISTER_KERNEL(greater_than,
int64_t,
float,
double,
phi::dtype::float16) {}
phi::dtype::float16,
phi::dtype::bfloat16) {}
PD_REGISTER_KERNEL(greater_equal,
KPS,
ALL_LAYOUT,
Expand All @@ -147,7 +150,8 @@ PD_REGISTER_KERNEL(greater_equal,
int64_t,
float,
double,
phi::dtype::float16) {}
phi::dtype::float16,
phi::dtype::bfloat16) {}
PD_REGISTER_KERNEL(equal,
KPS,
ALL_LAYOUT,
Expand All @@ -158,7 +162,8 @@ PD_REGISTER_KERNEL(equal,
int64_t,
float,
double,
phi::dtype::float16) {}
phi::dtype::float16,
phi::dtype::bfloat16) {}
PD_REGISTER_KERNEL(not_equal,
KPS,
ALL_LAYOUT,
Expand All @@ -169,7 +174,8 @@ PD_REGISTER_KERNEL(not_equal,
int64_t,
float,
double,
phi::dtype::float16) {}
phi::dtype::float16,
phi::dtype::bfloat16) {}

PD_REGISTER_KERNEL(equal_all,
KPS,
Expand Down

0 comments on commit f62a96e

Please sign in to comment.