[XPU] fix the conflict between xdnn and plugin kernel (#59326)

runzhech · web-flow · commit b715c2e61366 · 2023-11-27T10:16:30.000+08:00
diff --git a/paddle/phi/kernels/xpu/plugin/src/kernel/kunlun2cpp/fast_reduce.xpu b/paddle/phi/kernels/xpu/plugin/src/kernel/kunlun2cpp/fast_reduce.xpu
@@ -239,12 +239,14 @@ __global__ void fast_reduce_min_tiny(const T* x, T* y, int m, int t) {
       const DTYPE* x, DTYPE* y, int m, int t);
 _XPU_DEF__FAST_REDUCE_SUM_TINY_(float);
 _XPU_DEF__FAST_REDUCE_SUM_TINY_(float16);
+_XPU_DEF__FAST_REDUCE_SUM_TINY_(bfloat16);
 
 #define _XPU_DEF__FAST_REDUCE_MEAN_TINY_(DTYPE)          \
   template __global__ void fast_reduce_mean_tiny<DTYPE>( \
       const DTYPE* x, DTYPE* y, int m, int t);
 _XPU_DEF__FAST_REDUCE_MEAN_TINY_(float);
 _XPU_DEF__FAST_REDUCE_MEAN_TINY_(float16);
+_XPU_DEF__FAST_REDUCE_MEAN_TINY_(bfloat16);
 
 #define _XPU_DEF__FAST_REDUCE_MAX_TINY_(DTYPE)          \
   template __global__ void fast_reduce_max_tiny<DTYPE>( \
diff --git a/paddle/phi/kernels/xpu/plugin/src/wrapper/fast_embedding.cpp b/paddle/phi/kernels/xpu/plugin/src/wrapper/fast_embedding.cpp
@@ -124,6 +124,9 @@ int fast_embedding(Context* ctx,
                    int64_t padding_idx,
                    TID start_index) {
   WRAPPER_CHECK_CTX(ctx);
+  if (std::is_same<T, bfloat16>::value) {
+    WRAPPER_UNIMPLEMENTED(ctx);
+  }
   WRAPPER_DUMP_FUNCTION_T2(ctx, "fast_embedding", T, TID);
   WRAPPER_DUMP_PARAM6(ctx, x, indices, y, xm, n, ym);
   WRAPPER_DUMP_PARAM3(ctx, padding_idx, start_index, ctx->_l3_mgr.get_size());
@@ -182,6 +185,24 @@ template int fast_embedding(Context*,
                             int64_t,
                             int64_t,
                             int64_t);
+template int fast_embedding(Context*,
+                            const bfloat16*,
+                            const int*,
+                            bfloat16*,
+                            int64_t,
+                            int64_t,
+                            int64_t,
+                            int64_t,
+                            int);
+template int fast_embedding(Context*,
+                            const bfloat16*,
+                            const int64_t*,
+                            bfloat16*,
+                            int64_t,
+                            int64_t,
+                            int64_t,
+                            int64_t,
+                            int64_t);
 
 }  // namespace plugin
 }  // namespace api
diff --git a/paddle/phi/kernels/xpu/plugin/src/wrapper/fast_gather_nd.cpp b/paddle/phi/kernels/xpu/plugin/src/wrapper/fast_gather_nd.cpp
@@ -190,6 +190,9 @@ int fast_gather_nd(Context* ctx,
                    const VectorParam<int64_t>& x_shape,
                    const std::vector<int64_t>& index_shape) {
   WRAPPER_CHECK_CTX(ctx);
+  if (std::is_same<T, bfloat16>::value) {
+    WRAPPER_UNIMPLEMENTED(ctx);
+  }
   WRAPPER_DUMP_FUNCTION_T2(ctx, "fast_gather_nd", T, TID);
   WRAPPER_DUMP_PARAM6(
       ctx, x, index, y, x_shape, index_shape, ctx->_l3_mgr.get_size());
@@ -274,6 +277,18 @@ template int fast_gather_nd(Context*,
                             float16*,
                             const VectorParam<int64_t>&,
                             const std::vector<int64_t>&);
+template int fast_gather_nd(Context*,
+                            const bfloat16*,
+                            const int*,
+                            bfloat16*,
+                            const VectorParam<int64_t>&,
+                            const std::vector<int64_t>&);
+template int fast_gather_nd(Context*,
+                            const bfloat16*,
+                            const int64_t*,
+                            bfloat16*,
+                            const VectorParam<int64_t>&,
+                            const std::vector<int64_t>&);
 
 }  // namespace plugin
 }  // namespace api
diff --git a/paddle/phi/kernels/xpu/plugin/src/wrapper/fast_reduce.cpp b/paddle/phi/kernels/xpu/plugin/src/wrapper/fast_reduce.cpp
@@ -147,6 +147,30 @@ int xpu2_wrapper<float16>(Context* ctx,
   return SUCCESS;
 }
 
+template <>
+int xpu2_wrapper<bfloat16>(Context* ctx,
+                           const bfloat16* x,
+                           bfloat16* y,
+                           const std::vector<int>& xshape,
+                           int op_type) {
+  int t = xshape[xshape.size() - 1];
+  int xlen = vector_prod(xshape);
+  int m = xlen / t;
+  switch (op_type) {
+    case 0:
+      xpu2::plugin::fast_reduce_sum_tiny<bfloat16>
+          <<<ctx->ncluster(), 64, ctx->xpu_stream>>>(x, y, m, t);
+      break;
+    case 1:
+      xpu2::plugin::fast_reduce_mean_tiny<bfloat16>
+          <<<ctx->ncluster(), 64, ctx->xpu_stream>>>(x, y, m, t);
+      break;
+    default:
+      return NOT_IMPLEMENT;
+  }
+  return SUCCESS;
+}
+
 template <typename T>
 int fast_reduce_tiny(Context* ctx,
                      const T* x,
@@ -239,6 +263,11 @@ template int fast_reduce_sum(Context*,
                              float16*,
                              const std::vector<int>&,
                              const std::vector<int>&);
+template int fast_reduce_sum(Context*,
+                             const bfloat16*,
+                             bfloat16*,
+                             const std::vector<int>&,
+                             const std::vector<int>&);
 template int fast_reduce_sum(Context*,
                              const int*,
                              int*,
@@ -264,6 +293,11 @@ template int fast_reduce_mean(Context*,
                               float16*,
                               const std::vector<int>&,
                               const std::vector<int>&);
+template int fast_reduce_mean(Context*,
+                              const bfloat16*,
+                              bfloat16*,
+                              const std::vector<int>&,
+                              const std::vector<int>&);
 template int fast_reduce_min(Context*,
                              const float*,
                              float*,
