[PHI] Fix cum kernel for big tensor (PaddlePaddle#72562)

Author: lshpku

paddle/phi/kernels/gpu/cum_kernel.cu

@@ -36,53 +36,23 @@ namespace cub = hipcub;
 
 namespace phi {
 
-template <typename T, int BLOCK_SIZE>
-__device__ void BlockReverse(
-    const T* idata, T* odata, int src_base, int dst_base, int valid_item) {
-  __shared__ T sh_mem[BLOCK_SIZE];
-  int tx = threadIdx.x;
-
-  int offset = tx;
-  T src_data = static_cast<T>(0);
-  int src_offset = BLOCK_SIZE - offset - 1;
-  if (src_offset < valid_item) {
-    src_data = idata[src_base + src_offset];
-  }
-  sh_mem[offset] = src_data;
-
-  __syncthreads();
-  int out_index = dst_base - offset;
-  if (offset < valid_item) {
-    int sh_mem_index = BLOCK_SIZE - offset - 1;
-    odata[out_index] = sh_mem[sh_mem_index];
-  }
-}
-
 template <typename T>
 __global__ void MatrixRowReverse(const T* matrix_data,
                                  T* reverse_data,
-                                 int reverse_size,
-                                 int outer_size,
-                                 int inner_size) {
-  int bx = blockIdx.x;
-  int by = blockIdx.y;
+                                 int64_t grid_size,
+                                 int64_t reverse_size) {
   int item_per_block = 1024;
-
-  for (int block_offset = 0; block_offset < reverse_size;
-       block_offset += item_per_block) {
-    int valid_item = (reverse_size - block_offset > item_per_block)
-                         ? item_per_block
-                         : reverse_size - block_offset;
-    int src_offset =
-        bx * reverse_size + block_offset + by * (inner_size * reverse_size);
-    int dst_offset = bx * reverse_size + by * (inner_size * reverse_size) +
-                     reverse_size - 1 - block_offset;
-    if (reverse_size < item_per_block) {
-      valid_item = reverse_size;
+  for (int64_t bx = blockIdx.x; bx < grid_size; bx += gridDim.x) {
+    for (int64_t block_offset = 0; block_offset < reverse_size;
+         block_offset += item_per_block) {
+      int64_t reverse_offset = block_offset + threadIdx.x;
+      int64_t src_offset = bx * reverse_size + reverse_offset;
+      int64_t dst_offset =
+          bx * reverse_size + (reverse_size - reverse_offset - 1);
+      if (reverse_offset < reverse_size) {
+        reverse_data[dst_offset] = matrix_data[src_offset];
+      }
     }
-
-    BlockReverse<T, 1024>(
-        matrix_data, reverse_data, src_offset, dst_offset, valid_item);
   }
 }
 
@@ -112,24 +82,30 @@ __global__ void MatrixTranspose(T* odata,
                                 size_t width) {
   __shared__ T tile[TILE_DIM][TILE_DIM + 1];
 
-  int x = blockIdx.x * TILE_DIM + threadIdx.x;
-  int y = blockIdx.y * TILE_DIM + threadIdx.y;
-  for (int j = 0; j < TILE_DIM; j += BLOCK_ROWS) {
-    if (x < width && (y + j) < height) {
-      tile[threadIdx.y + j][threadIdx.x] = idata[(y + j) * width + x];
-    } else {
-      tile[threadIdx.y + j][threadIdx.x] = 0;
-    }
-  }
+  int64_t wblocks = (width + TILE_DIM - 1) / TILE_DIM;
+  int64_t hblocks = (height + TILE_DIM - 1) / TILE_DIM;
+
+  int64_t block_i = blockIdx.x;
+  for (; block_i < wblocks * hblocks; block_i += gridDim.x) {
+    int64_t block_y = block_i / wblocks;
+    int64_t block_x = block_i % wblocks;
+    int64_t x = block_x * TILE_DIM + threadIdx.x;
+    int64_t y = block_y * TILE_DIM + threadIdx.y;
 
-  __syncthreads();
+    for (int j = 0; j < TILE_DIM; j += BLOCK_ROWS) {
+      if (x < width && (y + j) < height) {
+        tile[threadIdx.y + j][threadIdx.x] = idata[(y + j) * width + x];
+      }
+    }
+    __syncthreads();
 
-  x = blockIdx.y * TILE_DIM + threadIdx.x;  // transpose block offset
-  y = blockIdx.x * TILE_DIM + threadIdx.y;
+    x = block_y * TILE_DIM + threadIdx.x;  // transpose block offset
+    y = block_x * TILE_DIM + threadIdx.y;
 
-  for (int j = 0; j < TILE_DIM; j += BLOCK_ROWS) {
-    if (x < height && (y + j) < width) {
-      odata[(y + j) * height + x] = tile[threadIdx.x][threadIdx.y + j];
+    for (int j = 0; j < TILE_DIM; j += BLOCK_ROWS) {
+      if (x < height && (y + j) < width) {
+        odata[(y + j) * height + x] = tile[threadIdx.x][threadIdx.y + j];
+      }
     }
   }
 }
@@ -172,9 +148,8 @@ struct Identity<T, ComplexSum> {
 template <typename T, int BLOCK_THREADS, int ITEMS_PER_THREAD, typename Op>
 __global__ void BlockScanKernel(T* d_out,
                                 const T* d_in,
-                                int inner_size,
-                                int outer_size,
-                                int scan_size,
+                                int64_t grid_size,
+                                int64_t scan_size,
                                 bool exclusive,
                                 Op op) {
   using MT = typename phi::dtype::MPTypeTrait<T>::Type;
@@ -196,38 +171,40 @@ __global__ void BlockScanKernel(T* d_out,
     typename BlockScanT::TempStorage scan;
   } temp_storage;
 
-  int bx = blockIdx.x;
-  BlockPrefixCallbackOp<MT, Op> prefix_op(Identity<MT, Op>::value, op);
-
   // Obtain this block's segment of consecutive keys (blocked across threads)
-  int item_per_block = BLOCK_THREADS * ITEMS_PER_THREAD;
-  for (int block_offset = 0; block_offset < scan_size;
-       block_offset += BLOCK_THREADS * ITEMS_PER_THREAD) {
-    int valid_item = (scan_size - block_offset > item_per_block)
-                         ? item_per_block
-                         : (scan_size - block_offset);
-    if (scan_size < item_per_block) {
-      valid_item = scan_size;
+  int64_t item_per_block = BLOCK_THREADS * ITEMS_PER_THREAD;
+
+  for (int64_t bx = blockIdx.x; bx < grid_size; bx += gridDim.x) {
+    BlockPrefixCallbackOp<MT, Op> prefix_op(Identity<MT, Op>::value, op);
+
+    for (int64_t block_offset = 0; block_offset < scan_size;
+         block_offset += item_per_block) {
+      int64_t valid_item = (scan_size - block_offset > item_per_block)
+                               ? item_per_block
+                               : (scan_size - block_offset);
+      if (scan_size < item_per_block) {
+        valid_item = scan_size;
+      }
+
+      int64_t offset = bx * scan_size + block_offset;
+
+      MT thread_keys[ITEMS_PER_THREAD];
+      BlockLoadT(temp_storage.load)
+          .Load(d_in + offset, thread_keys, valid_item, 0);
+
+      __syncthreads();
+      if (exclusive) {
+        BlockScanT(temp_storage.scan)
+            .ExclusiveScan(thread_keys, thread_keys, op, prefix_op);
+      } else {
+        BlockScanT(temp_storage.scan)
+            .InclusiveScan(thread_keys, thread_keys, op, prefix_op);
+      }
+      __syncthreads();
+
+      BlockStoreT(temp_storage.store)
+          .Store(d_out + offset, thread_keys, valid_item);
     }
-
-    int offset = block_offset + bx * scan_size;
-
-    MT thread_keys[ITEMS_PER_THREAD];
-    BlockLoadT(temp_storage.load)
-        .Load(d_in + offset, thread_keys, valid_item, 0);
-
-    __syncthreads();
-    if (exclusive) {
-      BlockScanT(temp_storage.scan)
-          .ExclusiveScan(thread_keys, thread_keys, op, prefix_op);
-    } else {
-      BlockScanT(temp_storage.scan)
-          .InclusiveScan(thread_keys, thread_keys, op, prefix_op);
-    }
-    __syncthreads();
-
-    BlockStoreT(temp_storage.store)
-        .Store(d_out + offset, thread_keys, valid_item);
   }
 }
 
@@ -347,14 +324,24 @@ void ScanKernel(const Context& dev_ctx,
   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);
   DenseTensor tmp_tensor;
   tmp_tensor.Resize(out_dims);
   auto* tmp_data = dev_ctx.template Alloc<T>(&tmp_tensor);
 
+  auto swap_ptr = [](T*& ptr1, T*& ptr2) {
+    T* tmp = ptr2;
+    ptr2 = ptr1;
+    ptr1 = tmp;
+  };
+
+  int64_t max_grid_x = dev_ctx.GetCUDAMaxGridDimSize()[0];
+
+  // Do pre-process transpose
+  int tile_size = 32;
+  dim3 blocks(32, 8);
+  int64_t transpose_grids = ((width + tile_size - 1) / tile_size) *
+                            ((height + tile_size - 1) / tile_size);
+  transpose_grids = std::min(transpose_grids, max_grid_x);
   T* next_in_data = out_data;
   T* next_out_data = tmp_data;
   if (transpose) {
@@ -363,53 +350,42 @@ void ScanKernel(const Context& dev_ctx,
     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;
+
+  // Do pre-process reverse
+  int64_t outer_size = height / scan_size;
+  int64_t inner_size = width;
+  int64_t grid_size = outer_size * inner_size;
+  int64_t scan_grid = std::min(grid_size, max_grid_x);
   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);
+      MatrixRowReverse<T><<<scan_grid, 1024, 0, dev_ctx.stream()>>>(
+          next_in_data, next_out_data, grid_size, scan_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);
+      MatrixRowReverse<T><<<scan_grid, 1024, 0, dev_ctx.stream()>>>(
+          in_data, out_data, grid_size, scan_size);
     }
   }
-  int64_t grid_size = outer_size * inner_size;
+
+  // Do scan
   if (!transpose && !reverse) {
-    BlockScanKernel<T, 128, 4, Op><<<grid_size, 128, 0, dev_ctx.stream()>>>(
-        out_data, in_data, outer_size, inner_size, scan_size, exclusive, op);
+    BlockScanKernel<T, 128, 4, Op><<<scan_grid, 128, 0, dev_ctx.stream()>>>(
+        out_data, in_data, grid_size, scan_size, exclusive, op);
 
   } else {
-    BlockScanKernel<T, 128, 4, Op>
-        <<<grid_size, 128, 0, dev_ctx.stream()>>>(next_out_data,
-                                                  next_in_data,
-                                                  outer_size,
-                                                  inner_size,
-                                                  scan_size,
-                                                  exclusive,
-                                                  op);
+    BlockScanKernel<T, 128, 4, Op><<<scan_grid, 128, 0, dev_ctx.stream()>>>(
+        next_out_data, next_in_data, grid_size, scan_size, exclusive, op);
   }
   swap_ptr(next_in_data, next_out_data);
+
+  // Do post-process reverse and transpose
   if (reverse) {
-    MatrixRowReverse<T><<<reverse_grid, 1024, 0, dev_ctx.stream()>>>(
-        next_in_data, next_out_data, scan_size, outer_size, inner_size);
+    MatrixRowReverse<T><<<scan_grid, 1024, 0, dev_ctx.stream()>>>(
+        next_in_data, next_out_data, grid_size, scan_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);
   }