[PHI] Fix transpose kernel dim3 index type (#73203)

lshpku · web-flow · commit bef3a0b8a2e5 · 2025-06-10T10:44:47.000+08:00
diff --git a/paddle/phi/backends/gpu/gpu_utils.h b/paddle/phi/backends/gpu/gpu_utils.h
@@ -66,25 +66,18 @@ struct DeviceArray {
   T data[Size];
 };
 
-struct Dim3 : DeviceArray<int, 3, 1> {
-  typedef DeviceArray<int, 3, 1> Base;
+template <typename T>
+struct Dim3 : DeviceArray<T, 3, 1> {
+  typedef DeviceArray<T, 3, 1> Base;
   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Dim3() : Base() {}
-  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Dim3(int a0, int a1, int a2)
-      : Base(a0, a1, a2) {}
-  EIGEN_STRONG_INLINE Dim3(const std::array<int, 3>& array) : Base(array) {}
-};
-
-struct Index3 : DeviceArray<int, 3, 0> {
-  typedef DeviceArray<int, 3, 0> Base;
-  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Index3() : Base() {}
-  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Index3(int a0, int a1, int a2)
+  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Dim3(T a0, T a1, T a2)
       : Base(a0, a1, a2) {}
 };
 
 // Flat index with real dimension
 template <typename IndexType = int>
 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE IndexType
-FlatTensorIndex(const Index3& index, const Dim3& dims) {
+FlatTensorIndex(const Dim3<IndexType>& index, const Dim3<IndexType>& dims) {
   IndexType flat_index = index[0];
 #pragma unroll
   for (int i = 1; i < 3; ++i) {
@@ -95,9 +88,9 @@ FlatTensorIndex(const Index3& index, const Dim3& dims) {
 
 // Convert index to tensor index with dimension.
 template <typename IndexType = int>
-EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Index3
-ConvertTensorIndex(IndexType index, const Dim3& dims) {
-  Index3 tensor_index;
+EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Dim3<IndexType> ConvertTensorIndex(
+    IndexType index, const Dim3<IndexType>& dims) {
+  Dim3<IndexType> tensor_index;
 #pragma unroll
   for (int i = 2; i >= 0; --i) {
     IndexType new_index = index / dims[i];
@@ -107,40 +100,5 @@ ConvertTensorIndex(IndexType index, const Dim3& dims) {
   return tensor_index;
 }
 
-template <typename IntType, bool ceil>
-IntType CeilOrFloor(IntType x, IntType deviser) {
-  PADDLE_ENFORCE_GT(
-      deviser,
-      0,
-      common::errors::InvalidArgument("deviser should be greater than 0, "
-                                      "but received is:%d",
-                                      deviser));
-
-  PADDLE_ENFORCE_GT(
-      x,
-      0,
-      common::errors::InvalidArgument("input should be greater than 0, "
-                                      "but received is:%d",
-                                      x));
-
-  const IntType round_to_zero = x / deviser;
-  const IntType inte_result = round_to_zero * deviser;
-
-  if (ceil) {
-    const bool do_adjustment =
-        (round_to_zero >= 0) && (deviser > 0 && x > inte_result);
-    const IntType adjustment = static_cast<IntType>(do_adjustment);
-    const IntType ceil_val = round_to_zero + adjustment;
-    return ceil_val;
-  } else {
-    const bool do_adjustment =
-        (round_to_zero <= 0) && (deviser > 0 && x < inte_result);
-
-    const IntType adjustment = static_cast<IntType>(do_adjustment);
-    const IntType floor_val = round_to_zero - adjustment;
-    return floor_val;
-  }
-}
-
 }  // namespace funcs
 }  // namespace phi
diff --git a/paddle/phi/kernels/funcs/transpose_function.cu.h b/paddle/phi/kernels/funcs/transpose_function.cu.h
@@ -84,7 +84,7 @@ template <typename T,
           int TileY,
           typename IndexType = int>
 __global__ void TilingSwapDim1And2(const T* __restrict__ input,
-                                   Dim3 input_dims,
+                                   Dim3<IndexType> input_dims,
                                    T* __restrict__ output) {
   assert(blockDim.x == NumThreads);
   assert(blockDim.y == 1);
@@ -104,25 +104,25 @@ __global__ void TilingSwapDim1And2(const T* __restrict__ input,
 
   int x = threadIdx.x;
 
-  Dim3 output_dims = {
+  Dim3<IndexType> output_dims = {
       input_dims[0],
       input_dims[2],
       input_dims[1],
   };
 
   // Align dim to Tiles
-  Dim3 tile_aligned_input_dim = {
+  Dim3<IndexType> tile_aligned_input_dim = {
       input_dims[0],
       (input_dims[1] + TileX - 1) / TileX,
       (input_dims[2] + TileY - 1) / TileY,
   };
 
   // Converts block idx to tile index, each block process a tile
-  Index3 input_block_tile_index =
+  Dim3<IndexType> input_block_tile_index =
       ConvertTensorIndex<IndexType>(blockIdx.x, tile_aligned_input_dim);
 
   // Compute real index align to tile:0, 32, 64...
-  Index3 block_tile_index_in_input = {
+  Dim3<IndexType> block_tile_index_in_input = {
       input_block_tile_index[0],
       input_block_tile_index[1] * TileX,
       input_block_tile_index[2] * TileY,
@@ -179,13 +179,13 @@ __global__ void TilingSwapDim1And2(const T* __restrict__ input,
   __syncthreads();
 
   // Store sm value back to out
-  Index3 output_block_tile_index = {
+  Dim3<IndexType> output_block_tile_index = {
       input_block_tile_index[0],
       input_block_tile_index[2],
       input_block_tile_index[1],
   };
 
-  Index3 block_tile_index_in_output = {
+  Dim3<IndexType> block_tile_index_in_output = {
       output_block_tile_index[0],
       output_block_tile_index[1] * TileY,
       output_block_tile_index[2] * TileX,
@@ -283,7 +283,7 @@ void LaunchNarrowDims2TransposeKernel(const phi::GPUContext& d,
                                       int tile_size_j,
                                       IndexType total_tiles_count,
                                       const T* input,
-                                      const Dim3& input_dims,
+                                      const Dim3<IndexType>& input_dims,
                                       T* output) {
   constexpr int NumThreads = tile_long;
   if (tile_size_i <= tile_long && tile_size_j <= tile_short) {
@@ -308,7 +308,7 @@ struct NarrowDims2TransposeDispatch {
                           int tile_size_j,
                           IndexType total_tiles_count,
                           const T* input,
-                          const Dim3& input_dims,
+                          const Dim3<IndexType>& input_dims,
                           T* output) {
     PADDLE_ENFORCE_EQ(
         (tile_long & (tile_long - 1)),
@@ -373,7 +373,7 @@ struct NarrowDims2TransposeDispatch<
                           int tile_size_j,
                           IndexType total_tiles_count,
                           const T* input,
-                          const Dim3& input_dims,
+                          const Dim3<IndexType>& input_dims,
                           T* output) {
     PADDLE_ENFORCE_EQ(
         (tile_long & (tile_long - 1)),
@@ -423,7 +423,7 @@ struct NarrowDims2TransposeDispatch<
                           int tile_size_j,
                           IndexType total_tiles_count,
                           const T* input,
-                          const Dim3& input_dims,
+                          const Dim3<IndexType>& input_dims,
                           T* output) {
     PADDLE_ENFORCE_EQ(
         (tile_long & (tile_long - 1)),
@@ -447,7 +447,7 @@ struct NarrowDims2TransposeDispatch<
 template <typename T, bool conjugate = false, typename IndexType = int>
 void SwapDim1And2InNarrow(const phi::GPUContext& d,
                           const T* input,
-                          const Dim3& input_dims,
+                          const Dim3<IndexType>& input_dims,
                           T* output,
                           const int kMinTileSize) {
   // First get available tile sizes for the data type requested as backups
@@ -471,13 +471,10 @@ void SwapDim1And2InNarrow(const phi::GPUContext& d,
     // data may not aligned to tile, so some threads wasted, we need
     // to find least wasted threads, which means we need to find tile
     // can split input properly, in another words: num_wasted_threads=0.
-    int num_wasted_threads =
-        input_long_edge -
-        CeilOrFloor<int, false>(input_long_edge, proposed_tile_long_edge) *
-            proposed_tile_long_edge;
+    int num_full_tiles = input_long_edge / proposed_tile_long_edge;
 
-    int num_full_tiles =
-        CeilOrFloor<int, false>(input_long_edge, proposed_tile_long_edge);
+    int num_wasted_threads =
+        input_long_edge - num_full_tiles * proposed_tile_long_edge;
 
     float cost = num_wasted_threads;
 
@@ -509,10 +506,10 @@ void SwapDim1And2InNarrow(const phi::GPUContext& d,
                            : std::min(select_tile_size_j, tile_short_edge);
 
   // Here finally get proper long X short tile size.
-  Dim3 input_dims_aligned = {
+  Dim3<IndexType> input_dims_aligned = {
       input_dims[0],
-      CeilOrFloor<int, true>(input_dims[1], select_tile_size_i),
-      CeilOrFloor<int, true>(input_dims[2], select_tile_size_j),
+      (input_dims[1] + select_tile_size_i - 1) / select_tile_size_i,
+      (input_dims[2] + select_tile_size_j - 1) / select_tile_size_j,
   };
 
   IndexType total_tiles_count = input_dims_aligned[0];
@@ -537,18 +534,18 @@ void SwapDim1And2InNarrow(const phi::GPUContext& d,
 template <typename T, int pos0, int pos1, int pos2, typename IndexType = int>
 __global__ void TransposeSimpleKernel(IndexType nthreads,
                                       const T* __restrict__ input,
-                                      Dim3 input_dims,
+                                      Dim3<IndexType> input_dims,
                                       T* __restrict__ output) {
-  Dim3 output_dims;
+  Dim3<IndexType> output_dims;
   output_dims[pos0] = input_dims[0];
   output_dims[pos1] = input_dims[1];
   output_dims[pos2] = input_dims[2];
 
   CUDA_KERNEL_LOOP_TYPE(output_index, nthreads, IndexType) {
-    Index3 output_tensor_index =
+    Dim3<IndexType> output_tensor_index =
         ConvertTensorIndex<IndexType>(output_index, output_dims);
 
-    Index3 input_tensor_index;
+    Dim3<IndexType> input_tensor_index;
     input_tensor_index[0] = output_tensor_index[pos0];
     input_tensor_index[1] = output_tensor_index[pos1];
     input_tensor_index[2] = output_tensor_index[pos2];
@@ -564,7 +561,7 @@ __global__ void TransposeSimpleKernel(IndexType nthreads,
 template <typename T, typename IndexType = int>
 void SendSwapDim1And2InTranspose(const phi::GPUContext& d,
                                  const T* input,
-                                 const Dim3& input_dims,
+                                 const Dim3<IndexType>& input_dims,
                                  T* output) {
   // Suppose tile size > 16
   static const int kMinTileSize = 16;
@@ -580,10 +577,10 @@ void SendSwapDim1And2InTranspose(const phi::GPUContext& d,
     constexpr int kTileSize = 32;
     constexpr int kNumThreads = 256;
 
-    Dim3 input_dims_aligned = {
+    Dim3<IndexType> input_dims_aligned = {
         input_dims[0],
-        CeilOrFloor<int, true>(input_dims[1], kTileSize),
-        CeilOrFloor<int, true>(input_dims[2], kTileSize),
+        (input_dims[1] + kTileSize - 1) / kTileSize,
+        (input_dims[2] + kTileSize - 1) / kTileSize,
     };
 
     IndexType total_tiles_count = input_dims_aligned[0];
@@ -617,11 +614,11 @@ struct SwapDim1And2InTranspose {
   typedef phi::GPUContext Device;
   void operator()(const Device& d,
                   const T* in,
-                  const std::vector<int>& combined_dims,
+                  const std::vector<int64_t>& combined_dims,
                   T* out) {
-    Dim3 input_dims = {static_cast<int>(combined_dims[0]),
-                       static_cast<int>(combined_dims[1]),
-                       static_cast<int>(combined_dims[2])};
+    Dim3<IndexType> input_dims = {static_cast<IndexType>(combined_dims[0]),
+                                  static_cast<IndexType>(combined_dims[1]),
+                                  static_cast<IndexType>(combined_dims[2])};
     SendSwapDim1And2InTranspose<T, IndexType>(d, in, input_dims, out);
   }
 };
@@ -631,11 +628,11 @@ struct SwapDim0And2InTranspose {
   typedef phi::GPUContext Device;
   void operator()(const Device& d,
                   const T* in,
-                  const std::vector<int>& combined_dims,
+                  const std::vector<int64_t>& combined_dims,
                   T* out) {
-    Dim3 input_dims = {static_cast<int>(combined_dims[0]),
-                       static_cast<int>(combined_dims[1]),
-                       static_cast<int>(combined_dims[2])};
+    Dim3<IndexType> input_dims = {static_cast<IndexType>(combined_dims[0]),
+                                  static_cast<IndexType>(combined_dims[1]),
+                                  static_cast<IndexType>(combined_dims[2])};
 
     IndexType total_size = combined_dims[0];
     total_size *= combined_dims[1];
@@ -653,7 +650,7 @@ struct SwapDim0And2InTranspose {
 inline void CombineTransposeDim3(const DDim& shape,
                                  const std::vector<int>& perm,
                                  std::vector<int>* new_perm,
-                                 std::vector<int>* new_dims) {
+                                 std::vector<int64_t>* new_dims) {
   PADDLE_ENFORCE_EQ(shape.size(),
                     perm.size(),
                     common::errors::InvalidArgument(
@@ -662,7 +659,7 @@ inline void CombineTransposeDim3(const DDim& shape,
                         shape.size(),
                         perm.size()));
 
-  std::vector<int> dim_vec;
+  std::vector<int64_t> dim_vec;
   if (shape.size() == 1) {
     // If input dimension is already 1, no need to combine dim.
     new_perm->resize(1);
@@ -671,7 +668,7 @@ inline void CombineTransposeDim3(const DDim& shape,
   } else {
     int dim_idx = 0;
     std::vector<int> new_dim_pos(shape.size(), -1);
-    std::vector<int> combined_dims(shape.size(), 0);
+    std::vector<int64_t> combined_dims(shape.size(), 0);
 
     int cur_head = perm[0];
     new_dim_pos[cur_head] = 0;
@@ -728,7 +725,7 @@ struct TransposeSimple {
     auto in_data = in.data<T>();
     auto out_data = out->data<T>();
     std::vector<int> new_perm;
-    std::vector<int> new_dims;
+    std::vector<int64_t> new_dims;
     CombineTransposeDim3(in.dims(), perm, &new_perm, &new_dims);
     if (new_perm.size() < 2 || new_perm.size() > 3) return false;
 
