 #include <cfloat>
+#include <cub/block/block_reduce.cuh>
 
 #include "caffe2/core/context_gpu.h"
 #include "softmax_op.h"

  }
}

#define REDUCTION_KERNEL_THREADS_X 128
#define REDUCTION_KERNEL_THREADS_Y 4
#define REDUCTION_THREADS \
  (REDUCTION_KERNEL_THREADS_X * REDUCTION_KERNEL_THREADS_Y)

__global__ void
RowMaxKernelLargeD(const int num, const int D, const float* data, float* out) {
  __shared__ float
      max_buffer[REDUCTION_KERNEL_THREADS_Y * REDUCTION_KERNEL_THREADS_X];
  const int threadId = threadIdx.y * REDUCTION_KERNEL_THREADS_X + threadIdx.x;

  for (int index = blockIdx.y * blockDim.y + threadIdx.y; index < num;
       index += blockDim.y * gridDim.y) {
    float maxval = -FLT_MAX;
    for (int x = blockIdx.x * blockDim.x + threadIdx.x; x < D;
         x += blockDim.x * gridDim.x) {
      maxval = fmaxf(data[index * D + x], maxval);
    }
    max_buffer[threadId] = maxval;

    __syncthreads();

    if (threadIdx.x < 32) {
      maxval = fmaxf(
          fmaxf(
              fmaxf(maxval, max_buffer[threadId + 32]),
              max_buffer[threadId + 64]),
          max_buffer[threadId + 96]);
      max_buffer[threadId] = maxval;
    }

    __syncthreads();

    if (threadIdx.x == 0) {
#pragma unroll
      for (int j = 1; j < 32; j++) {
        maxval = max(max_buffer[threadId + j], maxval);
      }
      out[index] = maxval;
    }
    __syncthreads();
  }
}

__global__ void RowMaxKernel(const int num, const int D, const float* data,
    float* out) {
  CUDA_1D_KERNEL_LOOP(index, num) {
    float maxval = -FLT_MAX;
    for (int d = 0; d < D; ++d) {
      maxval = max(data[index * D + d], maxval);
    }
    out[index] = maxval;
  }
}


    CUDAContext* context) {
  const int size = N * D;

  RowMaxKernel<<<
      std::min(N, CAFFE_MAXIMUM_NUM_BLOCKS),
      CAFFE_CUDA_NUM_THREADS,
      0,
      context->cuda_stream()>>>(N, D, logits, rowmax);
  // Put the intermediate result X - max(X) into Y
  context->Copy<float, CUDAContext, CUDAContext>(size, logits, probs);
  // Subtract the scale
