[NPU] Add bf16 register (#1403)

Author: Birdylx

backends/npu/kernels/activation_kernel.cc

@@ -1916,14 +1916,16 @@ PD_REGISTER_PLUGIN_KERNEL(gelu,
                           ALL_LAYOUT,
                           custom_kernel::GeluKernel,
                           float,
-                          phi::dtype::float16) {}
+                          phi::dtype::float16,
+                          phi::dtype::bfloat16) {}
 
 PD_REGISTER_PLUGIN_KERNEL(gelu_grad,
                           npu,
                           ALL_LAYOUT,
                           custom_kernel::GeluGradKernel,
                           float,
-                          phi::dtype::float16) {}
+                          phi::dtype::float16,
+                          phi::dtype::bfloat16) {}
 
 PD_REGISTER_PLUGIN_KERNEL(tanh,
                           npu,
@@ -1945,6 +1947,7 @@ PD_REGISTER_PLUGIN_KERNEL(sigmoid,
                           custom_kernel::SigmoidKernel,
                           float,
                           phi::dtype::float16,
+                          phi::dtype::bfloat16,
                           double) {}
 
 PD_REGISTER_PLUGIN_KERNEL(sigmoid_grad,
@@ -1953,6 +1956,7 @@ PD_REGISTER_PLUGIN_KERNEL(sigmoid_grad,
                           custom_kernel::SigmoidGradKernel,
                           float,
                           phi::dtype::float16,
+                          phi::dtype::bfloat16,
                           double) {}
 
 PD_REGISTER_PLUGIN_KERNEL(sqrt,

backends/npu/kernels/conv2d_kernel.cc

@@ -492,11 +492,13 @@ PD_REGISTER_PLUGIN_KERNEL(conv2d,
                           ALL_LAYOUT,
                           custom_kernel::Conv2dKernel,
                           float,
-                          phi::dtype::float16) {}
+                          phi::dtype::float16,
+                          phi::dtype::bfloat16) {}
 
 PD_REGISTER_PLUGIN_KERNEL(conv2d_grad,
                           npu,
                           ALL_LAYOUT,
                           custom_kernel::Conv2DGradKernel,
                           float,
-                          phi::dtype::float16) {}
+                          phi::dtype::float16,
+                          phi::dtype::bfloat16) {}

backends/npu/kernels/grid_sample_kernel.cc

@@ -301,7 +301,7 @@ void GridSampleGradKernel(const Context& dev_ctx,
       if (find_it != paddlingModeMap.end()) {
         paddle_mode_to_int = find_it->second;
       }
-      EXEC_NPU_CMD(aclnnGridSamplerd2DBackward,
+      EXEC_NPU_CMD(aclnnGridSampler2DBackward,
                    dev_ctx,
                    out_grad,
                    x,

backends/npu/kernels/scatter_kernel.cc

@@ -150,7 +150,8 @@ PD_REGISTER_PLUGIN_KERNEL(scatter,
                           float,
                           int64_t,
                           int,
-                          phi::dtype::float16) {}
+                          phi::dtype::float16,
+                          phi::dtype::bfloat16) {}
 
 PD_REGISTER_PLUGIN_KERNEL(scatter_grad,
                           npu,
@@ -159,4 +160,5 @@ PD_REGISTER_PLUGIN_KERNEL(scatter_grad,
                           float,
                           int64_t,
                           int,
-                          phi::dtype::float16) {}
+                          phi::dtype::float16,
+                          phi::dtype::bfloat16) {}