 
 #include <thrust/device_ptr.h>
 #include <thrust/device_vector.h>
+#include <thrust/gather.h>
 #include <thrust/reverse.h>
 #include <thrust/scan.h>
+#include "cub/cub.cuh"
 #include "paddle/fluid/operators/cum_op.h"
 #include "paddle/fluid/platform/gpu_launch_config.h"
 

namespace paddle {
namespace operators {

template <typename T>
__global__ void OuterScan(const T* in, T* out, int inner_dim_size,
                          int outer_dim_size, int scan_dim_size, bool exclusive,
                          bool reverse) {
  int id = blockIdx.y * blockDim.x + threadIdx.x;

  for (int outer_index = blockIdx.x; outer_index < outer_dim_size;
       outer_index += gridDim.x) {
    for (int inner_index = blockIdx.y * blockDim.x + threadIdx.x;
         inner_index < inner_dim_size; inner_index += gridDim.y * blockDim.x) {
      int scan_index_init = 0;
      int forward_direction = 1;
      int src_index =
          outer_index * scan_dim_size * inner_dim_size + inner_index;
      int dst_index =
          outer_index * scan_dim_size * inner_dim_size + inner_index;
      if (reverse) {
        src_index = src_index + (scan_dim_size - 1) * inner_dim_size;
        dst_index = dst_index + (scan_dim_size - 1) * inner_dim_size;
        forward_direction = -1;
      }
      if (exclusive) {
        scan_index_init = 1;
        out[dst_index] = 0;
        dst_index = dst_index + (forward_direction * inner_dim_size);
      }
      T acc = 0;

      for (int scan_index = scan_index_init; scan_index < scan_dim_size;
           ++scan_index) {
        acc = in[src_index] + acc;
        out[dst_index] = acc;
        src_index += (forward_direction * inner_dim_size);
        dst_index += (forward_direction * inner_dim_size);
      }
    }
  }
}

// inclusive scan
template <typename T, int num_threads_x, int num_threads_y>
__global__ void InnerMostDimInclusiveScan(const T* in, T* out,
                                          int inner_dim_size,
                                          int outer_dim_size, int scan_dim_size,
                                          bool reverse) {
  __shared__ T share_data[num_threads_y][num_threads_x * 2];
  T* share_row = share_data[threadIdx.y];
  int forward_direction = 1;
  if (reverse) forward_direction = -1;

  for (int block_row = blockIdx.x * blockDim.y; block_row < outer_dim_size;
       block_row += blockDim.y * gridDim.x) {
    int row = block_row + threadIdx.y;
    T acc = 0;
    const T* row_src = in + row * scan_dim_size;
    T* row_dst = out + row * scan_dim_size;
    int block_col = 0;
    bool loop_condition = (block_col < scan_dim_size);
    if (reverse) {
      loop_condition = (block_col >= 0);
      block_col = scan_dim_size - 1;
    }
    while (loop_condition) {
      // Load data into share memory(two value per thread)
      int col1 = block_col + threadIdx.x * forward_direction;
      int col2 = block_col + (num_threads_x + threadIdx.x) * forward_direction;
      if (row < outer_dim_size) {
        if (col1 < scan_dim_size && col1 >= 0) {
          share_row[threadIdx.x] = row_src[col1];
        } else {
          share_row[threadIdx.x] = 0;
        }

        if (col2 < scan_dim_size && col2 >= 0) {
          share_row[num_threads_x + threadIdx.x] = row_src[col2];
        } else {
          share_row[num_threads_x + threadIdx.x] = 0;
        }

        // Add the previous block acc to the result
        if (threadIdx.x == 0) {
          share_row[0] = share_row[0] + acc;
        }
      }
      __syncthreads();

      // Up-Sweep
      for (unsigned s = num_threads_x, d = 1; s >= 1; s >>= 1, d <<= 1) {
        if (row < outer_dim_size && threadIdx.x < s) {
          unsigned offset = (2 * threadIdx.x + 1) * d - 1;
          share_row[offset + d] = share_row[offset] + share_row[offset + d];
        }
        __syncthreads();
      }
      // Down-Sweep
      for (unsigned s = 2, d = blockDim.x / 2; d >= 1; s <<= 1, d >>= 1) {
        if (row < outer_dim_size && threadIdx.x < s - 1) {
          unsigned offset = 2 * (threadIdx.x + 1) * d - 1;
          share_row[offset + d] = share_row[offset] + share_row[offset + d];
        }
        __syncthreads();
      }

      // Write to the output
      if (row < outer_dim_size) {
        if (col1 < scan_dim_size && col1 >= 0)
          row_dst[col1] = share_row[threadIdx.x];
        if (col2 < scan_dim_size && col2 >= 0)
          row_dst[col2] = share_row[num_threads_x + threadIdx.x];
      }
      acc = share_row[2 * num_threads_x - 1];
      __syncthreads();
      block_col += 2 * num_threads_x * forward_direction;
      if (reverse)
        loop_condition = (block_col >= 0);
      else
        loop_condition = (block_col < scan_dim_size);
    }
  }
}

// exclusive block scan and store block sum for large scan
template <typename T>
__global__ void InnerMostDimExclusiveScan(const T* in, T* out, T* sum_data,
                                          int inner_dim_size,
                                          int outer_dim_size, int scan_dim_size,
                                          int two_power, bool reverse) {
  // https://stackoverflow.com/questions/27570552/templated-cuda-kernel-with-dynamic-shared-memory
  extern __shared__ __align__(sizeof(T)) unsigned char raw_tmp[];
  T* share_tmp = reinterpret_cast<T*>(raw_tmp);
  int thread_id = threadIdx.x;
  int block_id = blockIdx.x;
  int block_scan_size = blockDim.x * 2;
  int remain = scan_dim_size % (2 * blockDim.x);
  if (block_id == gridDim.x - 1 && remain != 0) block_scan_size = remain;
  int col1 = thread_id;
  int col2 = thread_id + (block_scan_size) / 2;
  int index1 = blockIdx.y * (scan_dim_size) + block_id * blockDim.x * 2 + col1;
  int index2 = blockIdx.y * (scan_dim_size) + block_id * blockDim.x * 2 + col2;
  if (reverse) {
    index1 = blockIdx.y * (scan_dim_size) + scan_dim_size - 1 -
             (block_id * blockDim.x * 2 + col1);
    index2 = blockIdx.y * (scan_dim_size) + scan_dim_size - 1 -
             (block_id * blockDim.x * 2 + col2);
  }
  int sum_index = blockIdx.y * gridDim.x + block_id;
  if (thread_id < block_scan_size) {
    share_tmp[col1 + (col1 >> 5)] = in[index1];
    share_tmp[col2 + (col2 >> 5)] = in[index2];
  } else {
    share_tmp[col1 + (col1 >> 5)] = 0;
    share_tmp[col2 + (col2 >> 5)] = 0;
  }

  // Up-Sweep
  int offset = 1;
  for (int d = (two_power / 2); d > 0; d >>= 1) {
    __syncthreads();
    if (thread_id < d) {
      int tmp_index1 = offset * (2 * thread_id + 1) - 1;
      int tmp_index2 = offset * (2 * thread_id + 2) - 1;
      tmp_index1 = tmp_index1 + (tmp_index1 >> 5);
      tmp_index2 = tmp_index2 + (tmp_index2 >> 5);

      share_tmp[tmp_index2] += share_tmp[tmp_index1];
    }
    offset *= 2;
  }
  __syncthreads();

  if (thread_id == 0) {
    int tmp_index = (two_power - 1) + ((two_power - 1) >> 5);
    sum_data[sum_index] = share_tmp[tmp_index];
    share_tmp[tmp_index] = 0;
  }

  // Down Sweep
  for (int d = 1; d < two_power; d *= 2) {
    offset >>= 1;
    __syncthreads();
    if (thread_id < d) {
      int tmp_index1 = offset * (2 * thread_id + 1) - 1;
      int tmp_index2 = offset * (2 * thread_id + 2) - 1;
      tmp_index1 = tmp_index1 + (tmp_index1 >> 5);
      tmp_index2 = tmp_index2 + (tmp_index2 >> 5);

      T tmp = share_tmp[tmp_index1];
      share_tmp[tmp_index1] = share_tmp[tmp_index2];
      share_tmp[tmp_index2] += tmp;
    }
  }

  __syncthreads();

  if (col1 < block_scan_size) out[index1] = share_tmp[col1 + (col1 >> 5)];
  if (col2 < block_scan_size) out[index2] = share_tmp[col2 + (col2 >> 5)];
}

// for large scan_dim_size array we need to add for correct result
template <typename T>
__global__ void AddBlockScan(T* result, T* sum, int size, int scan_dim_size,
                             int sum_size, bool reverse) {
  int idx = threadIdx.x + blockDim.x * (blockIdx.x + blockIdx.y * gridDim.x);
  int block_id_start = blockIdx.y * sum_size;
  int block_id_end = blockIdx.x + blockIdx.y * sum_size;
  int block_id = blockIdx.x;
  int thread_id = threadIdx.x;

  int col = block_id * blockDim.x + thread_id + size;
  int index = blockIdx.y * (scan_dim_size) + col;
  if (reverse) {
    index = blockIdx.y * (scan_dim_size) + scan_dim_size - 1 - col;
  }

  if (col >= scan_dim_size || col < 0) return;
  for (int i = block_id_start; i <= block_id_end; i++) {
    result[index] += sum[i];
  }
}


      return;
    }

    size_t height = 1;
    size_t width = 1;
    for (size_t i = 0; i <= axis; i++) {
      height *= out_dims[i];
    }

    for (size_t i = axis + 1; i < out_dims.size(); i++) {
      width *= out_dims[i];
    }
    int scan_size = out_dims[axis];
    bool transpose = (axis != out_dims.size() - 1);

    int tile_size = 32;
    dim3 blocks(32, 8);
    dim3 transpose_grids((width + tile_size - 1) / tile_size,
                         (height + tile_size - 1) / tile_size);
    auto& dev_ctx = context.template device_context<DeviceContext>();
    Tensor tmp;
    tmp.Resize(out_dims);
    auto* tmp_data = tmp.mutable_data<T>(context.GetPlace());
    T* next_in_data = out_data;
    T* next_out_data = tmp_data;
    if (transpose) {
      MatrixTranspose<T, 32,
                      8><<<transpose_grids, blocks, 0, dev_ctx.stream()>>>(
          out_data, in_data, height, width);
      next_in_data = out_data;
      next_out_data = tmp_data;
    }
    auto swap_ptr = [](T*& ptr1, T*& ptr2) {
      T* tmp = ptr2;
      ptr2 = ptr1;
      ptr1 = tmp;
    };
    int outer_size = height / scan_size;
    int inner_size = width;
    // Consider the size of shared memory, here block size is 128
    dim3 scan_grid(outer_size, inner_size);
    dim3 reverse_grid = scan_grid;
    if (reverse) {
      if (transpose) {
        reverse_grid.x = scan_grid.y;
        reverse_grid.y = scan_grid.x;
        MatrixRowReverse<T><<<reverse_grid, 1024, 0, dev_ctx.stream()>>>(
            next_in_data, next_out_data, scan_size, outer_size, inner_size);
        if (!transpose) next_in_data = tmp_data;
        swap_ptr(next_in_data, next_out_data);
      } else {
        MatrixRowReverse<T><<<reverse_grid, 1024, 0, dev_ctx.stream()>>>(
            in_data, out_data, scan_size, outer_size, inner_size);
      }
    }
    if (!transpose && !reverse) {
      BlockScanKernel<T, 128, 4><<<scan_grid, 128, 0, dev_ctx.stream()>>>(
          out_data, in_data, outer_size, inner_size, scan_size, exclusive);

    } else {
      BlockScanKernel<T, 128, 4><<<scan_grid, 128, 0, dev_ctx.stream()>>>(
          next_out_data, next_in_data, outer_size, inner_size, scan_size,
          exclusive);
    }
    swap_ptr(next_in_data, next_out_data);
    if (reverse) {
      MatrixRowReverse<T><<<reverse_grid, 1024, 0, dev_ctx.stream()>>>(
          next_in_data, next_out_data, scan_size, outer_size, inner_size);
      swap_ptr(next_in_data, next_out_data);
    }
    if (transpose) {
      transpose_grids.x = (height + tile_size - 1) / tile_size;
      transpose_grids.y = (width + tile_size - 1) / tile_size;
      MatrixTranspose<T, 32,
                      8><<<transpose_grids, blocks, 0, dev_ctx.stream()>>>(
          next_out_data, next_in_data, width, height);
    }
  }
};
