Modify reduce_op.op.h for xpu2 with kernel primitive api

paddle/fluid/operators/kernel_primitives/datamover_primitives.h

@@ -360,12 +360,12 @@ __device__ __forceinline__ void ReadDataBc(
  * reduce_last_dim: Used to indicate whether the dimension of reduce contains
  * the lowest dimension.
  */
-template <typename T, int NX, int NY, int BlockSize, int Rank,
- typename IndexCal, bool IsBoundary = false>
+template <typename Tx, typename Ty, int NX, int NY, int BlockSize, int Rank,
+ typename IndexCal, typename Functor, bool IsBoundary = false>
 __device__ __forceinline__ void ReadDataReduce(
- T* dst, const T* __restrict__ src, int block_offset,
+ Ty* dst, const Tx* __restrict__ src, int block_offset,
  const IndexCal& index_cal, int size_nx, int size_ny, int stride_nx,
- int stride_ny, bool reduce_last_dim) {
+ int stride_ny, Functor func, bool reduce_last_dim) {
  int thread_offset = 0;
  int left_idx = 0;
  if (reduce_last_dim) {
@@ -385,7 +385,7 @@ __device__ __forceinline__ void ReadDataReduce(
  }
  }
  uint32_t index_src = index_cal(thread_offset + block_offset);
- dst[ny] = src[index_src];
+ dst[ny] = static_cast<Ty>(func(src[index_src]));
  thread_offset += stride_ny;
  }
  } else {
@@ -400,7 +400,7 @@ __device__ __forceinline__ void ReadDataReduce(
  }
  }
  uint32_t index_src = index_cal(thread_offset + block_offset);
- dst[nx + ny * NX] = src[index_src];
+ dst[nx + ny * NX] = static_cast<Ty>(func(src[index_src]));
  thread_offset += stride_ny;
  }
  }

paddle/fluid/operators/kernel_primitives/helper_primitives.h

@@ -17,64 +17,49 @@
 namespace paddle {
 namespace operators {
 namespace kernel_primitives {
-namespace details {
 
-static __device__ __forceinline__ platform::float16 ExpFunctor(
- platform::float16 x) {
- return ::Eigen::numext::exp(x);
-}
-static __device__ __forceinline__ float ExpFunctor(float x) { return expf(x); }
-static __device__ __forceinline__ double ExpFunctor(double x) { return exp(x); }
-static __device__ __forceinline__ platform::float16 LogFunctor(
- platform::float16 x) {
- return ::Eigen::numext::log(x);
-}
-static __device__ __forceinline__ float LogFunctor(float x) { return logf(x); }
-static __device__ __forceinline__ double LogFunctor(double x) { return log(x); }
+#ifdef PADDLE_WITH_XPU2
+struct dim3 {
+ int x;
+ int y;
+ int z;
 
-/*************************** Compute Functor****************************/
-// for margin_cross_entropy
-template <typename Tx, typename Ty = Tx>
-struct ExpLogitTransformer {
- HOSTDEVICE explicit inline ExpLogitTransformer(int n) {}
-
- HOSTDEVICE inline Ty operator()(const Tx* x) const {
- return static_cast<Ty>(details::ExpFunctor(x[0]));
- }
-
- HOSTDEVICE inline Ty operator()(const Tx& x) const {
- return static_cast<Ty>(details::ExpFunctor(x));
+ explicit inline dim3(int split_x, int split_y = 1, int split_z = 1) {
+ x = split_x;
+ y = split_y;
+ z = split_z;
  }
 };
+#endif
 
-// Post processing function for sum, max, min, prod, any
-template <typename Tx, typename Ty = Tx>
-struct IdentityFunctor {
- HOSTDEVICE explicit inline IdentityFunctor(int n) {}
+struct DimConfig {
+ int split_num_x;
+ int split_num_y;
+ int split_num_z;
+ int deal_size_x;
+ int deal_size_y;
+ int deal_size_z;
+ int rem_x;
+ int rem_y;
+ int rem_z;
 
- HOSTDEVICE inline Ty operator()(const Tx* x) const {
- return static_cast<Ty>(x[0]);
+ HOSTDEVICE explicit inline DimConfig(int split_x, int split_y, int split_z,
+ int size_x, int size_y, int size_z) {
+ split_num_x = split_x;
+ split_num_y = split_y;
+ split_num_z = split_z;
+ deal_size_x = size_x;
+ deal_size_y = size_y;
+ deal_size_z = size_z;
  }
 
- HOSTDEVICE inline Ty operator()(const Tx& x) const {
- return static_cast<Ty>(x);
+ HOSTDEVICE void SetRem(int rem_nx, int rem_ny, int rem_nz) {
+ rem_x = rem_nx;
+ rem_y = rem_ny;
+ rem_z = rem_nz;
  }
 };
 
-// Post processing function for mean
-template <typename T>
-struct DivideFunctor {
- HOSTDEVICE explicit inline DivideFunctor(int n) : n_inv((T)(1.0 / n)) {}
-
- HOSTDEVICE inline T operator()(const T* x) const { return x[0] * n_inv; }
-
- HOSTDEVICE inline T operator()(const T& x) const { return x * n_inv; }
-
- private:
- T n_inv;
-};
-
-} // namespace details
 } // namespace kernel_primitives
 } // namespace operators
 } // namespace paddle

paddle/fluid/operators/kernel_primitives/kernel_primitives.h

@@ -13,11 +13,45 @@
 // limitations under the License.
 
 #pragma once
+#include "paddle/fluid/operators/kernel_primitives/functor_primitives.h"
+#include "paddle/fluid/operators/kernel_primitives/helper_primitives.h"
+#ifdef PADDLE_WITH_XPU2
+#include "paddle/fluid/operators/kernel_primitives/compute_primitives_xpu2.h"
+#include "paddle/fluid/operators/kernel_primitives/datamover_primitives_xpu2.h"
+#define THREAD_ID_X core_id()
+#define THREAD_ID_Y 0
+#define THREAD_ID_Z 0
+
+#define BLOCK_NUM_X core_num()
+#define BLOCK_NUM_Y 0
+#define BLOCK_NUM_Z 0
 
+#define BLOCK_ID_X cluster_id()
+#define BLOCK_ID_Y 0
+#define BLOCK_ID_Z 0
+
+#define GRID_NUM_X cluster_num()
+#define GRID_NUM_Y 0
+#define GRID_NUM_Z 0
+#else
 #include "paddle/fluid/operators/kernel_primitives/compute_primitives.h"
 #include "paddle/fluid/operators/kernel_primitives/datamover_primitives.h"
-#include "paddle/fluid/operators/kernel_primitives/functor_primitives.h"
-#include "paddle/fluid/operators/kernel_primitives/helper_primitives.h"
+#define THREAD_ID_X threadIdx.x
+#define THREAD_ID_Y threadIdx.y
+#define THREAD_ID_Z threadIdx.z
+
+#define BLOCK_NUM_X blockDim.x
+#define BLOCK_NUM_Y blockDim.y
+#define BLOCK_NUM_Z blockDim.z
+
+#define BLOCK_ID_X blockIdx.x
+#define BLOCK_ID_Y blockIdx.y
+#define BLOCK_ID_Z blockIdx.z
+
+#define GRID_NUM_X gridDim.x
+#define GRID_NUM_Y gridDim.y
+#define GRID_NUM_Z gridDim.z
+#endif
 
 namespace paddle {
 namespace operators {

paddle/fluid/operators/margin_cross_entropy_op.cu

@@ -130,7 +130,7 @@ __global__ void AddMarginToPositiveLogitsKernel(
 
 template <typename Tx, typename Ty = Tx>
 struct ExpAndSum {
- using Transformer = kpds::ExpLogitTransformer<Tx>;
+ using Transformer = kps::ExpFunctor<Tx>;
 
  inline Ty initial() { return static_cast<Ty>(0.0f); }
 
@@ -159,7 +159,7 @@ __global__ void LogitsMinusLogSumKernel(T* logits, const T* logits_sum_per_row,
  const int64_t N, const int64_t D) {
  CUDA_KERNEL_LOOP(i, N * D) {
  auto row = i / D;
- logits[i] -= kpds::LogFunctor(logits_sum_per_row[row]);
+ logits[i] -= kps::details::Log(logits_sum_per_row[row]);
  }
 }
 
@@ -174,9 +174,9 @@ __global__ void HardLabelSoftmaxWithCrossEntropyKernel(
  if ((col + start_index) == labels[row]) {
  auto softmax = log_softmax[i];
  loss[row] = -softmax;
- log_softmax[i] = kpds::ExpFunctor(softmax);
+ log_softmax[i] = kps::details::Exp(softmax);
  } else {
- log_softmax[i] = kpds::ExpFunctor(log_softmax[i]);
+ log_softmax[i] = kps::details::Exp(log_softmax[i]);
  }
  }
 }

paddle/fluid/operators/reduce_ops/reduce_functor_op.h

@@ -24,11 +24,11 @@ limitations under the License. */
 namespace paddle {
 namespace operators {
 
-namespace kpds = paddle::operators::kernel_primitives::details;
+namespace kps = paddle::operators::kernel_primitives;
 
 template <typename Tx, typename Ty = Tx>
 struct CustomMin {
- using Transformer = kpds::IdentityFunctor<Tx>;
+ using Transformer = kps::IdentityFunctor<Tx>;
 
  inline Ty initial() {
  return static_cast<Ty>(std::numeric_limits<Ty>::max());
@@ -41,7 +41,7 @@ struct CustomMin {
 
 template <typename Tx, typename Ty = Tx>
 struct CustomMax {
- using Transformer = kpds::IdentityFunctor<Tx>;
+ using Transformer = kps::IdentityFunctor<Tx>;
 
  inline Ty initial() {
  return static_cast<Ty>(std::numeric_limits<Ty>::lowest());
@@ -55,7 +55,7 @@ struct CustomMax {
 // for cub::Reduce
 template <typename Tx, typename Ty = Tx>
 struct CustomSum {
- using Transformer = kpds::IdentityFunctor<Tx, Ty>;
+ using Transformer = kps::IdentityFunctor<Tx, Ty>;
 
  inline Ty initial() { return static_cast<Ty>(0.0f); }
 
@@ -66,7 +66,7 @@ struct CustomSum {
 
 template <typename Tx, typename Ty = Tx>
 struct CustomMean {
- using Transformer = kpds::DivideFunctor<Tx>;
+ using Transformer = kps::DivideFunctor<Tx>;
 
  inline Ty initial() { return static_cast<Ty>(0.0f); }
 
@@ -77,7 +77,7 @@ struct CustomMean {
 
 template <typename Tx, typename Ty = Tx>
 struct CustomMul {
- using Transformer = kpds::IdentityFunctor<Tx>;
+ using Transformer = kps::IdentityFunctor<Tx>;
 
  inline Ty initial() { return static_cast<Ty>(1.0f); }
 
@@ -88,7 +88,7 @@ struct CustomMul {
 
 template <typename Tx, typename Ty = Tx>
 struct CustomLogicalOr {
- using Transformer = kpds::IdentityFunctor<Tx>;
+ using Transformer = kps::IdentityFunctor<Tx>;
 
  inline Ty initial() { return static_cast<Ty>(false); }
 
@@ -99,7 +99,7 @@ struct CustomLogicalOr {
 
 template <typename Tx, typename Ty = Tx>
 struct CustomLogicalAnd {
- using Transformer = kpds::IdentityFunctor<Tx>;
+ using Transformer = kps::IdentityFunctor<Tx>;
 
  inline Ty initial() { return static_cast<Ty>(true); }