   }
 }
 
+/*
+  Cross entropy soft label with dynamic size on axis (log2_elements is
+  varibale).
+  - if the input is softmax，compute loss with softmax
+  - if the input is log_softmax, compute loss with log_softmax and update
+  softmax
+*/
+template <typename T, typename VecT, bool InLogMode = false>
+__global__ void CrossEntropySoftLabel(T* loss, T* softmaxwrt, const T* softmax,
+                                      const T* labels, const int n,
+                                      const int dim, const int d,
+                                      int log2_elements) {
+  const int kDimCeil = 1 << log2_elements;
+  const int kVSize = sizeof(VecT) / sizeof(T);
+
+#ifdef __HIPCC__
+  const int kThreadPerBlock = 256;
+#else
+  const int kThreadPerBlock = 512;
+#endif
+  const int kBatchPerBlock = 1;
+  const int kWarpSize = 32;  // (dim < 32) ? dim : 32;
+  const int kBatchSize = 1;
+  const int kThreadPerBatch = kThreadPerBlock / kBatchPerBlock;
+  const int kWarpPerBatch = kThreadPerBatch / kWarpSize;
+
+  const int kIterations = (dim + kThreadPerBatch - 1) / kThreadPerBatch;
+  const int kIterationsV = (kIterations >= kVSize) ? (kIterations / kVSize) : 1;
+
+  const int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * kBatchSize;
+
+  T sum[kBatchSize]{static_cast<T>(0.0)};
+#pragma unroll
+  for (int i = 0; i < kBatchSize; ++i) {
+    int ids = first_batch + i;
+    if (ids >= n * d) break;
+    int idx_n = ids / d;
+    int idx_d = ids % d;
+#pragma unroll
+    for (int it = 0; it < kIterations; ++it) {
+      int idx_dim = it * kThreadPerBatch + threadIdx.x;
+      int idx = idx_n * dim * d + idx_dim * d + idx_d;
+
+      if (idx_n < n && idx_dim < dim) {
+        VecT softmaxdata;
+        if (InLogMode) {
+          softmaxdata = reinterpret_cast<VecT*>(&softmaxwrt[idx])[0];
+        } else {
+          softmaxdata = reinterpret_cast<const VecT*>(&softmax[idx])[0];
+        }
+        VecT labelsdata = reinterpret_cast<const VecT*>(&labels[idx])[0];
+        T* softmaxptr = reinterpret_cast<T*>(&softmaxdata);
+        T* labelsptr = reinterpret_cast<T*>(&labelsdata);
+#pragma unroll
+        for (int s = 0; s < kVSize; s++) {
+          if (InLogMode) {
+            sum[i] -= softmaxptr[s] * labelsptr[s];
+            softmaxptr[s] = Exp(softmaxptr[s]);
+          } else {
+            sum[i] -= Log(softmaxptr[s]) * labelsptr[s];
+          }
+        }
+        if (InLogMode) {
+          reinterpret_cast<VecT*>(&softmaxwrt[idx])[0] = softmaxdata;
+        }
+      }
+    }
+  }
+  WarpReduceSum<T, kBatchSize, kWarpSize>(sum);
+  __syncthreads();
+
+  __shared__ T sumshare[kWarpPerBatch][kBatchPerBlock][kBatchSize];
+  if (threadIdx.x % kWarpSize == 0) {
+#pragma unroll
+    for (int i = 0; i < kBatchSize; i++) {
+      sumshare[threadIdx.x / kWarpSize][threadIdx.y][i] = sum[i];
+    }
+  }
+  __syncthreads();
+
+  // write
+  if (threadIdx.x == 0) {
+    for (int i = 0; i < kBatchSize; i++) {
+      int ids = first_batch + i;
+      if (ids < n * d) {
+        loss[ids] = sumshare[0][threadIdx.y][i];
+        for (int s = 1; s < kWarpPerBatch; s++) {
+          loss[ids] += sumshare[s][threadIdx.y][i];
+        }
+      }
+    }
+  }
+}
+
+/*
+Core function of softmax with cross entropy forward soft label.
+The computation includes
+  - Compute maximum of batch: maxvalue_{i} = max_j src_{i,j}
+  - Compute sum of exp batch: s_{i} = sum_{j}{ exp(src_{i,j} - maxvalue_{i} }
+  - Compute: sum of - sum_{j}{ label_{i,j} * (src_{i,j} - maxvalue_{i} -
+log(sum[i]))}
+One warp (32 threads) is used to compute 1 or 2 batch (kBatchSize).
+For reduction max (sum), firstly compute max (sum) to one warp, then use shuffle
+api to compute max (sum) in one warp.
+*/
+template <typename T, typename VecT, typename AccT, int Log2Elements>
+__global__ void WarpSoftmaxForwardSoftLabel(T* loss, T* softmax, const T* src,
+                                            const T* label,
+                                            const int batch_size,
+                                            const int stride,
+                                            const int element_count) {
+  const bool LogMode = true;
+
+  constexpr int kDimCeil = 1 << Log2Elements;
+  constexpr int kWarpSize = (kDimCeil < 32) ? kDimCeil : 32;
+  constexpr int kVSize = sizeof(VecT) / sizeof(T);
+  constexpr int kIterations = kDimCeil / kWarpSize;
+  constexpr int kIterationsV =
+      (kIterations >= kVSize) ? (kIterations / kVSize) : 1;
+  constexpr int kBatchSize = (kDimCeil <= 128) ? 2 : 1;
+
+  int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * kBatchSize;
+  int local_batches = batch_size - first_batch;
+  if (local_batches > kBatchSize) {
+    local_batches = kBatchSize;
+  }
+
+  // read data from global memory
+  VecT srcdata[kBatchSize][kIterationsV];
+  VecT labeldata[kBatchSize][kIterationsV];
+
+  for (int i = 0; i < kBatchSize; ++i) {
+    const VecT* src_v =
+        reinterpret_cast<const VecT*>(&src[(first_batch + i) * stride]);
+    const VecT* label_v =
+        reinterpret_cast<const VecT*>(&label[(first_batch + i) * stride]);
+
+    // max index to read
+    int idx_max = (i < local_batches) ? element_count : 0;
+    int idx_max_v = idx_max / kVSize;
+
+    // read data
+    for (int it = 0; it < kIterationsV; ++it) {
+      int src_idx = threadIdx.x + it * kWarpSize;
+      if (src_idx < idx_max_v) {
+        srcdata[i][it] = src_v[src_idx];
+        labeldata[i][it] = label_v[src_idx];
+      } else {
+#pragma unroll
+        for (int s = 0; s < kVSize; s++) {
+          reinterpret_cast<T*>(&srcdata[i][it])[s] =
+              -std::numeric_limits<AccT>::max();
+          reinterpret_cast<T*>(&labeldata[i][it])[s] = 0.0;
+        }
+      }
+    }
+  }
+
+  // compute max value
+  AccT max_value[kBatchSize];
+#pragma unroll
+  for (int i = 0; i < kBatchSize; ++i) {
+    max_value[i] = -std::numeric_limits<AccT>::infinity();
+#pragma unroll
+    for (int it = 0; it < kIterationsV; ++it) {
+      T* srcptr_v = reinterpret_cast<T*>(&srcdata[i][it]);
+      T valmax = srcptr_v[0];
+#pragma unroll
+      for (int s = 1; s < kVSize; ++s) {
+        valmax = (valmax > srcptr_v[s]) ? valmax : srcptr_v[s];
+      }
+      max_value[i] = (max_value[i] > static_cast<AccT>(valmax))
+                         ? max_value[i]
+                         : static_cast<AccT>(valmax);
+    }
+  }
+  WarpReduceMax<AccT, kBatchSize, kWarpSize>(max_value);
+
+  // compute sum
+  AccT sum[kBatchSize]{0.0};
+#pragma unroll
+  for (int i = 0; i < kBatchSize; ++i) {
+#pragma unroll
+    for (int it = 0; it < kIterationsV; ++it) {
+      T* srcptr_v = reinterpret_cast<T*>(&srcdata[i][it]);
+#pragma unroll
+      for (int s = 0; s < kVSize; ++s) {
+        if (LogMode) {
+          sum[i] += std::exp(static_cast<AccT>(srcptr_v[s]) - max_value[i]);
+        } else {
+          srcptr_v[s] = std::exp(static_cast<AccT>(srcptr_v[s]) - max_value[i]);
+          sum[i] += static_cast<AccT>(srcptr_v[s]);
+        }
+      }
+    }
+  }
+  WarpReduceSum<AccT, kBatchSize, kWarpSize>(sum);
+
+  // log_softmax and loss
+  AccT sumloss[kBatchSize]{0.0};
+#pragma unroll
+  for (int i = 0; i < kBatchSize; ++i) {
+    if (i >= local_batches) break;
+
+    VecT* softmax_v =
+        reinterpret_cast<VecT*>(&softmax[(first_batch + i) * stride]);
+
+    // max index to write
+    int idx_max = (i < local_batches) ? element_count : 0;
+    int idx_max_v = idx_max / kVSize;
+
+    if (LogMode) {
+      sum[i] = std::log(sum[i]);
+    }
+#pragma unroll
+    for (int it = 0; it < kIterationsV; ++it) {
+      T* srcvp = reinterpret_cast<T*>(&srcdata[i][it]);
+      T* labelvp = reinterpret_cast<T*>(&labeldata[i][it]);
+      VecT tmpv;
+      T* tmpvp = reinterpret_cast<T*>(&tmpv);
+#pragma unroll
+      for (int s = 0; s < kVSize; ++s) {
+        if (LogMode) {
+          AccT logsoftmax = static_cast<AccT>(srcvp[s]) - max_value[i] - sum[i];
+          sumloss[i] -= logsoftmax * static_cast<AccT>(labelvp[s]);
+          tmpvp[s] = std::exp(logsoftmax);
+        } else {
+          tmpvp[s] = static_cast<AccT>(srcvp[s]) / sum[i];
+        }
+      }
+
+      int idx = threadIdx.x + it * kWarpSize;
+      if (idx < idx_max_v) {
+        softmax_v[idx] = tmpv;
+      }
+    }
+  }
+
+  // loss
+  WarpReduceSum<AccT, kBatchSize, kWarpSize>(sumloss);
+
+  for (int i = 0; i < kBatchSize; i++) {
+    if (i >= local_batches) break;
+    loss[first_batch + i] = sumloss[i];
+  }
+}
+
+#define SOFTMAX_WARP_FORWARD_SOFT_CASE(Log2Elements, VecT, AccT)               \
+  case Log2Elements:                                                           \
+    WarpSoftmaxForwardSoftLabel<T, VecT, AccT,                                 \
+                                Log2Elements><<<blocks, threads, 0, stream>>>( \
+        loss, softmax, src, label, batch_size, stride, element_count);         \
+    break;
+
+/*
+  Wrapper of softmax with cross entropy forward soft label.
+*/
+template <typename T>
+void SwitchWarpSoftmaxForwardSoftLabel(const int blocks, const dim3 threads,
+                                       gpuStream_t stream, T* loss, T* softmax,
+                                       const T* src, const T* label,
+                                       const int batch_size, const int stride,
+                                       const int element_count,
+                                       const int log2_elements) {
+  using AccT = typename details::MPTypeTrait<T>::Type;
+  switch (log2_elements) {
+    SOFTMAX_WARP_FORWARD_SOFT_CASE(0, T, AccT);
+    SOFTMAX_WARP_FORWARD_SOFT_CASE(1, T, AccT);
+    SOFTMAX_WARP_FORWARD_SOFT_CASE(2, T, AccT);
+    SOFTMAX_WARP_FORWARD_SOFT_CASE(3, T, AccT);
+    SOFTMAX_WARP_FORWARD_SOFT_CASE(4, T, AccT);
+    SOFTMAX_WARP_FORWARD_SOFT_CASE(5, T, AccT);
+    SOFTMAX_WARP_FORWARD_SOFT_CASE(6, T, AccT);
+    SOFTMAX_WARP_FORWARD_SOFT_CASE(7, T, AccT);
+    SOFTMAX_WARP_FORWARD_SOFT_CASE(8, T, AccT);
+    SOFTMAX_WARP_FORWARD_SOFT_CASE(9, T, AccT);
+    default:
+      break;
+  }
+}
+
+template <typename T>
+static void SoftmaxWithCrossEntropySoftLabel(
+    const platform::CUDADeviceContext& ctx, const int rank, const int axis,
+    const T* logits_data, const T* labels_data, T* softmax_data, T* loss_data,
+    int N, int dim, int D) {
+#ifdef __HIPCC__
+  constexpr int kMaxBlockDim = 256;
+#else
+  constexpr int kMaxBlockDim = 512;
+#endif
+  int64_t block_dim = dim >= kMaxBlockDim
+                          ? kMaxBlockDim
+                          : (1 << static_cast<int>(std::log2(dim)));
+
+  int64_t grid_dim = N * D;
+  constexpr int max_dim = 320;
+
+  const int kDimLog2 = static_cast<int>(Log2Ceil(dim));
+  const int kDimCeil = 1 << kDimLog2;
+  auto stream = ctx.stream();
+
+  if (D == 1 && dim <= max_dim) {
+    int kWarpSize = (kDimCeil < 32) ? kDimCeil : 32;
+    int batches_per_warp = (kDimCeil <= 128) ? 2 : 1;
+
+    // use 128 threads per block to maximimize gpu utilization
+    constexpr int threads_per_block = 128;
+    int warps_per_block = (threads_per_block / kWarpSize);
+    int batches_per_block = warps_per_block * batches_per_warp;
+    int blocks = (N + batches_per_block - 1) / batches_per_block;
+    dim3 threads(kWarpSize, warps_per_block, 1);
+
+    SwitchWarpSoftmaxForwardSoftLabel<T>(blocks, threads, stream, loss_data,
+                                         softmax_data, logits_data, labels_data,
+                                         N, dim, dim, kDimLog2);
+
+  } else {
+    ScopedTensorDescriptor desc;
+    std::vector<int> tensor_dims = {N, dim, D, 1};
+    DataLayout layout = DataLayout::kNCHW;
+#ifdef PADDLE_WITH_HIP
+    miopenTensorDescriptor_t descp = desc.descriptor<T>(layout, tensor_dims);
+#else
+    cudnnTensorDescriptor_t descp = desc.descriptor<T>(layout, tensor_dims);
+#endif
+
+    auto handle = ctx.cudnn_handle();
+
+#ifdef PADDLE_WITH_HIP
+    auto mode = axis == rank - 1 ? MIOPEN_SOFTMAX_MODE_INSTANCE
+                                 : MIOPEN_SOFTMAX_MODE_CHANNEL;
+    PADDLE_ENFORCE_GPU_SUCCESS(platform::dynload::miopenSoftmaxForward_V2(
+        handle, platform::CudnnDataType<T>::kOne(), descp, logits_data,
+        platform::CudnnDataType<T>::kZero(), descp, softmax_data,
+        MIOPEN_SOFTMAX_LOG, mode));
+#else
+    auto mode = axis == rank - 1 ? CUDNN_SOFTMAX_MODE_INSTANCE
+                                 : CUDNN_SOFTMAX_MODE_CHANNEL;
+    PADDLE_ENFORCE_GPU_SUCCESS(platform::dynload::cudnnSoftmaxForward(
+        handle, CUDNN_SOFTMAX_LOG, mode, platform::CudnnDataType<T>::kOne(),
+        descp, logits_data, platform::CudnnDataType<T>::kZero(), descp,
+        softmax_data));
+#endif
+
+    const int kDimLog2 = static_cast<int>(Log2Ceil(dim));
+    const int kDimCeil = 1 << kDimLog2;
+#ifdef __HIPCC__
+    int kThreadPerBlock = 256;
+#else
+    int kThreadPerBlock = 512;
+#endif
+
+    int kBatchPerBlock = 1;
+    int blocks = (N * D + kBatchPerBlock - 1) / kBatchPerBlock;
+    dim3 threads(kThreadPerBlock / kBatchPerBlock, kBatchPerBlock, 1);
+
+    CrossEntropySoftLabel<T, T, true><<<blocks, threads, 0, stream>>>(
+        loss_data, softmax_data, NULL, labels_data, N, dim, D, kDimLog2);
+  }
+}
+
 template <typename T>
 __global__ void SoftCrossEntropyGradientKernel(T* logit_grad,
                                                const T* loss_grad,

#include "paddle/fluid/operators/math/math_function.h"
#include "paddle/fluid/operators/softmax_cudnn_op.cu.h"
#include "paddle/fluid/operators/softmax_with_cross_entropy_op.h"
#include "paddle/fluid/platform/device/gpu/gpu_dnn.h"
#include "paddle/fluid/platform/for_range.h"


  }
}

template <typename T>
class SoftmaxWithCrossEntropyCUDAKernel : public framework::OpKernel<T> {
 public:
