[PHI] Refine softmax kernel support for big tensor (#73460)

Author: lshpku

paddle/phi/kernels/funcs/axis_utils.h

@@ -36,7 +36,7 @@ static inline T SizeToAxis(const int axis, DDim dims) {
 }
 
 template <typename T = int>
-static inline int SizeFromAxis(const int axis, DDim dims) {
+static inline T SizeFromAxis(const int axis, DDim dims) {
   T size = 1;
   for (int i = axis; i < dims.size(); i++) {
     size *= dims[i];
@@ -45,7 +45,7 @@ static inline int SizeFromAxis(const int axis, DDim dims) {
 }
 
 template <typename T = int>
-static inline int SizeOutAxis(const int axis, DDim dims) {
+static inline T SizeOutAxis(const int axis, DDim dims) {
   T size = 1;
   for (int i = axis + 1; i < dims.size(); i++) {
     size *= dims[i];

paddle/phi/kernels/gpudnn/softmax_gpudnn.h

@@ -82,9 +82,9 @@ class VecT2<phi::dtype::bfloat16> {
   using Type = int;
 };
 
-static inline int Log2Ceil(int value) {
+static inline int Log2Ceil(int64_t value) {
   int log2_value = 0;
-  while ((1 << log2_value) < value) ++log2_value;
+  while ((int64_t(1) << log2_value) < value) ++log2_value;
   return log2_value;
 }
 
@@ -836,37 +836,42 @@ void SwitchWarpSoftmaxBackward(const IndexType blocks,
  * Better performance when axis != -1
  */
 
-static void GetGridDim(
-    int high_dim, int mid_dim, int low_dim, const dim3& block, dim3* grid) {
+static void GetGridDim(int64_t high_dim,
+                       int64_t low_dim,
+                       const dim3& block,
+                       dim3* grid) {
   int device_id = phi::backends::gpu::GetCurrentDeviceId();
   int max_mp = phi::backends::gpu::GetGPUMultiProcessors(device_id);
   int max_threads_per_mp =
       phi::backends::gpu::GetGPUMaxThreadsPerMultiProcessor(device_id);
   int max_threads = max_threads_per_mp * max_mp;
   int num_threads = block.x * block.y;
-  int max_num_blocks = max_threads / num_threads;
+  int64_t max_num_blocks = max_threads / num_threads;
 
-  int grid_x = (low_dim + block.x - 1) / block.x;
+  int64_t grid_x = (low_dim + block.x - 1) / block.x;
   grid_x = std::min(grid_x, max_num_blocks);
-  int grid_y = (max_num_blocks + grid_x - 1) / grid_x;
+  int64_t grid_y = (max_num_blocks + grid_x - 1) / grid_x;
   grid_y = std::min(grid_y, high_dim);
   grid->x = grid_x;
   grid->y = grid_y;
 }
 
-static void GetBlockDim(int mid_dim, int low_dim, dim3* block) {
+static void GetBlockDim(int64_t mid_dim, int64_t low_dim, dim3* block) {
   constexpr int max_num_threads = 1024;
-  int block_x = 1 << Log2Ceil(low_dim);
-  int block_y = 1 << Log2Ceil(mid_dim);
-  block->x = std::min(block_x, 32);
-  block->y = std::min(block_y, static_cast<int>(max_num_threads / block->x));
-  block->x = std::min(block_x, static_cast<int>(max_num_threads / block->y));
+  int64_t block_x = int64_t(1) << Log2Ceil(low_dim);
+  int64_t block_y = int64_t(1) << Log2Ceil(mid_dim);
+  block->x = std::min<int64_t>(block_x, 32);
+  block->y = std::min<int64_t>(block_y, max_num_threads / block->x);
+  block->x = std::min<int64_t>(block_x, max_num_threads / block->y);
 }
 
-static void GetLaunchConfig(
-    int high_dim, int mid_dim, int low_dim, dim3* grid, dim3* block) {
+static void GetLaunchConfig(int64_t high_dim,
+                            int64_t mid_dim,
+                            int64_t low_dim,
+                            dim3* grid,
+                            dim3* block) {
   GetBlockDim(mid_dim, low_dim, block);
-  GetGridDim(high_dim, mid_dim, low_dim, *block, grid);
+  GetGridDim(high_dim, low_dim, *block, grid);
 }
 
 template <typename T,