                               const int64_t output_rows,
                               const int64_t output_cols,
                               T* output) {
-  int tid_x = blockIdx.x * blockDim.x + threadIdx.x;
-  int curr_segment = 0;
-  int curr_offset = input_cols[0];
-  for (; tid_x < output_cols; tid_x += blockDim.x * gridDim.x) {
-    int curr_col_offset = input_cols[curr_segment + 1];
+  int64_t curr_segment = 0;
+  int64_t curr_offset = input_cols[0];
+  CUDA_KERNEL_LOOP_TYPE(tid_x, output_cols, int64_t) {
+    int64_t curr_col_offset = input_cols[curr_segment + 1];
     while (curr_col_offset <= tid_x) {
       curr_offset = curr_col_offset;
       ++curr_segment;
       curr_col_offset = input_cols[curr_segment + 1];
     }
 
-    int local_col = tid_x - curr_offset;
-    int segment_width = curr_col_offset - curr_offset;
+    int64_t local_col = tid_x - curr_offset;
+    int64_t segment_width = curr_col_offset - curr_offset;
 
     const T* input_ptr = inputs[curr_segment];
-    int tid_y = blockIdx.y * blockDim.y + threadIdx.y;
+    int64_t tid_y = blockIdx.y * blockDim.y + threadIdx.y;
     for (; tid_y < output_rows; tid_y += blockDim.y * gridDim.y)
       output[tid_y * output_cols + tid_x] =
           input_ptr[tid_y * segment_width + local_col];

    int64_t in_col = input[0].numel() / in_row;
    int64_t out_row = in_row, out_col = 0;

    int inputs_col_num = in_num + 1;
    std::vector<const T*> inputs_data_vec(in_num);
    std::vector<int64_t> inputs_col_vec(inputs_col_num);
    const T** inputs_data = inputs_data_vec.data();

                  int axis,
                  phi::DenseTensor* output) {
    // TODO(zcd): Add input data validity checking
    int64_t in_num = input.size();
    int64_t in_row = 1;
    auto dim_0 = input[0].dims();
    for (int i = 0; i < axis; ++i) {
