Fix int32 overflow issues for large tensor support in paddle/phi/kernels/impl (#76107) (#76276)

* Fix int32 overflow in svd_grad and conv kernel impl

* fix

Co-authored-by: Zhan Rongrui <[email protected]>

Author: zhengshengning

paddle/phi/kernels/impl/accuracy_check_kernel_impl.h

@@ -143,12 +143,12 @@ __global__ void AccuracyCheckCUDAKernel(const T* in_data,
                                         const double rtol,
                                         const double atol,
                                         bool equal_nan,
-                                        int num,
+                                        int64_t num,
                                         bool* out_data) {
-  unsigned int idx = threadIdx.x + blockIdx.x * blockDim.x;
+  int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
   bool val;
   using MPType = typename phi::dtype::MPTypeTrait<T>::Type;
-  for (int i = idx; i < num; i += blockDim.x * gridDim.x) {
+  for (int64_t i = idx; i < num; i += blockDim.x * gridDim.x) {
     const double a = static_cast<MPType>(in_data[i]);
     const double b = static_cast<MPType>(other_data[i]);
     if (isnan(a) || isnan(b)) {
@@ -173,11 +173,11 @@ __global__ void AccuracyCheckCUDAKernel<phi::complex64>(
     const double rtol,
     const double atol,
     bool equal_nan,
-    int num,
+    int64_t num,
     bool* out_data) {
-  unsigned int idx = threadIdx.x + blockIdx.x * blockDim.x;
+  int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
   bool val;
-  for (int i = idx; i < num; i += blockDim.x * gridDim.x) {
+  for (int64_t i = idx; i < num; i += blockDim.x * gridDim.x) {
     const phi::complex64 a = in_data[i];
     const phi::complex64 b = other_data[i];
     if (isnan(a) || isnan(b)) {
@@ -203,11 +203,11 @@ __global__ void AccuracyCheckCUDAKernel<phi::complex128>(
     const double rtol,
     const double atol,
     bool equal_nan,
-    int num,
+    int64_t num,
     bool* out_data) {
-  unsigned int idx = threadIdx.x + blockIdx.x * blockDim.x;
+  int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
   bool val;
-  for (int i = idx; i < num; i += blockDim.x * gridDim.x) {
+  for (int64_t i = idx; i < num; i += blockDim.x * gridDim.x) {
     const phi::complex128 a = in_data[i];
     const phi::complex128 b = other_data[i];
     if (isnan(a) || isnan(b)) {
@@ -236,12 +236,12 @@ struct AccuracyCheckFunctor<phi::GPUContext, T> {
                   const double atol,
                   bool equal_nan,
                   DenseTensor* output) {
-    int num = in.numel();
+    int64_t num = in.numel();
     const T* in_data = in.data<T>();
     const T* other_data = other.data<T>();
     bool* out_data = dev_ctx.template Alloc<bool>(output);
     int block = 1024;
-    int grid = (block - 1 + num) / block;
+    int64_t grid = (block - 1 + num) / block;
     grid = (grid > block) ? block : grid;
 #ifdef PADDLE_WITH_HIP
     hipMemset(out_data, true, num * sizeof(bool));

paddle/phi/kernels/impl/conv_grad_kernel_impl.h

@@ -85,7 +85,7 @@ void ConvGradKernel(const Context& dev_ctx,
   UpdatePaddingAndDilation<int>(
       &paddings, &dilations, padding_algorithm, in_data_dims, strides, ksize);
 
-  const int batch_size = static_cast<int>(transformed_input.dims()[0]);
+  const int64_t batch_size = transformed_input.dims()[0];
 
   // filter_shape_vec: {k_o, k_i, k_h, k_w} or {k_o, k_i, k_d, k_h, k_w}
   std::vector<int64_t> filter_shape_vec(common::vectorize(filter.dims()));
@@ -125,8 +125,8 @@ void ConvGradKernel(const Context& dev_ctx,
 
   // convolution backward input operator:  gemm + col2im(or col2vol)
   // convolution backward weight operator: im2col(or vol2col) + gemm
-  int in_step = static_cast<int>(transformed_input.dims()[1]) / groups;
-  int out_step = static_cast<int>(transformed_output_grad.dims()[1]) / groups;
+  int64_t in_step = transformed_input.dims()[1] / groups;
+  int64_t out_step = transformed_output_grad.dims()[1] / groups;
 
   bool is_expand = IsExpand(filter_shape_vec, strides, paddings, dilations);
 
@@ -163,7 +163,7 @@ void ConvGradKernel(const Context& dev_ctx,
     phi::funcs::Col2ImFunctor<phi::funcs::ColFormat::kCFO, Context, T> col2im;
     phi::funcs::Col2VolFunctor<Context, T> col2vol;
 
-    for (int i = 0; i < batch_size; i++) {
+    for (int64_t i = 0; i < batch_size; i++) {
       DenseTensor out_grad_batch =
           transformed_output_grad.Slice(i, i + 1).Resize(output_matrix_shape);
       DenseTensor in_grad_batch =
@@ -327,7 +327,7 @@ void ConvGradGradKernel(const Context& dev_ctx,
   UpdatePaddingAndDilation(
       &paddings, &dilations, padding_algorithm, in_data_dims, strides, ksize);
 
-  const int batch_size = static_cast<int>(transformed_X.dims()[0]);
+  const int64_t batch_size = transformed_X.dims()[0];
   std::vector<int64_t> filter_shape_vec(common::vectorize(W.dims()));
   std::vector<int64_t> output_shape_vec(
       common::vectorize(transformed_dY.dims()));
@@ -354,8 +354,8 @@ void ConvGradGradKernel(const Context& dev_ctx,
       transformed_dY.dims()[1],
       transformed_dY.numel() /
           (transformed_dY.dims()[0] * transformed_dY.dims()[1])};
-  int in_step = static_cast<int>(transformed_X.dims()[1]) / groups;
-  int out_step = static_cast<int>(transformed_dY.dims()[1]) / groups;
+  int64_t in_step = transformed_X.dims()[1] / groups;
+  int64_t out_step = transformed_dY.dims()[1] / groups;
 
   bool is_expand = IsExpand(filter_shape_vec, strides, paddings, dilations);
   DenseTensor col;
@@ -394,7 +394,7 @@ void ConvGradGradKernel(const Context& dev_ctx,
     phi::funcs::Col2ImFunctor<phi::funcs::ColFormat::kCFO, Context, T> col2im;
     phi::funcs::Col2VolFunctor<Context, T> col2vol;
 
-    for (int i = 0; i < batch_size; i++) {
+    for (int64_t i = 0; i < batch_size; i++) {
       DenseTensor dy_batch =
           transformed_dY.Slice(i, i + 1).Resize(output_matrix_shape);
       DenseTensor dx_batch = transformed_dX.Slice(i, i + 1).Resize(input_shape);

paddle/phi/kernels/impl/conv_kernel_impl.h

@@ -76,7 +76,7 @@ void ConvKernelImpl(const Context& dev_ctx,
   UpdatePaddingAndDilation(
       &paddings, &dilations, padding_algorithm, in_data_dims, strides, ksize);
 
-  const int batch_size = static_cast<int>(transformed_input.dims()[0]);
+  const int64_t batch_size = transformed_input.dims()[0];
 
   // filter_shape_vec:
   // {k_o, k_i, k_h, k_w} or {k_o, k_i, k_d, k_h, k_w}
@@ -137,14 +137,14 @@ void ConvKernelImpl(const Context& dev_ctx,
           (transformed_output.dims()[0] * transformed_output.dims()[1])};
 
   // convolution operator: im2col(or vol2col) + gemm
-  int in_step = static_cast<int>(transformed_input.dims()[1]) / groups;
-  int out_step = static_cast<int>(transformed_output.dims()[1]) / groups;
+  int64_t in_step = transformed_input.dims()[1] / groups;
+  int64_t out_step = transformed_output.dims()[1] / groups;
 
   phi::funcs::Im2ColFunctor<phi::funcs::ColFormat::kCFO, Context, T> im2col;
   phi::funcs::Vol2ColFunctor<Context, T> vol2col;
 
   auto blas = phi::funcs::GetBlas<Context, T>(dev_ctx);
-  for (int i = 0; i < batch_size; i++) {
+  for (int64_t i = 0; i < batch_size; i++) {
     DenseTensor in_batch =
         transformed_input.Slice(i, i + 1).Resize(in_matrix_shape);
     DenseTensor out_batch =

paddle/phi/kernels/impl/elementwise_grad_kernel_impl.h

@@ -259,7 +259,7 @@ void ComputeDDoutWithoutBroadcast(const CPUContext& dev_ctx UNUSED,
   auto* y_data = y.data<T>();
   auto* out_data = out.data<T>();
   auto* ddout_data = ddout->data<T>();
-  for (int i = 0; i < out_numel; i++) {
+  for (int64_t i = 0; i < out_numel; i++) {
     ddout_data[i] = dout_op(ddx_data[i], ddy_data[i], y_data[i], out_data[i]);
   }
 }
@@ -283,7 +283,7 @@ void ComputeDDoutWithBroadcast(const CPUContext& dev_ctx UNUSED,
   auto* out_data = out.data<T>();
   auto* ddout_data = ddout->data<T>();
   std::vector<int> index_array(max_dim, 0);
-  for (int i = 0; i < out_numel; i++) {
+  for (int64_t i = 0; i < out_numel; i++) {
     int x_index = phi::funcs::GetElementwiseIndex(
         x_dims_array, max_dim, index_array.data());
     int y_index = phi::funcs::GetElementwiseIndex(
@@ -381,9 +381,9 @@ __global__ void ComputeDDoutWithoutBroadcastGPUKernel(const T* ddx_data,
                                                       const T* y_data,
                                                       const T* out_data,
                                                       T* ddout_data,
-                                                      int numel,
+                                                      int64_t numel,
                                                       DDout_OP dout_op) {
-  int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  int64_t tid = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
   if (tid >= numel) return;
   ddout_data[tid] =
       dout_op(ddx_data[tid], ddy_data[tid], y_data[tid], out_data[tid]);
@@ -418,16 +418,16 @@ __global__ void ComputeDDoutWithBroadcastGPUKernel(
     const T* y_data,
     const T* out_data,
     T* ddout_data,
-    int numel,
+    int64_t numel,
     const CudaIntArray x_dims_array,
     const CudaIntArray y_dims_array,
     const CudaIntArray out_dims_array,
     const int max_dim,
     DDout_OP dout_op) {
-  int tid = threadIdx.x + blockIdx.x * blockDim.x;
+  int64_t tid = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
   if (tid >= numel) return;
-  int x_index = 0, y_index = 0, x_index_prod = 1, y_index_prod = 1,
-      out_index = tid, dim_index;
+  int64_t x_index = 0, y_index = 0, x_index_prod = 1, y_index_prod = 1,
+          out_index = tid, dim_index;
   for (int64_t i = max_dim - 1; i >= 0; i--) {
     if (out_index == 0) break;
     dim_index = out_index % out_dims_array[i];

paddle/phi/kernels/impl/fold_grad_kernel_impl.h

@@ -38,7 +38,7 @@ void FoldGradKernel(const Context& dev_ctx,
   if (!x_grad) return;
 
   const auto& x_dims = x_grad->dims();
-  const int batch_size = static_cast<int>(x_dims[0]);
+  const int64_t batch_size = x_dims[0];
 
   int output_height = (output_sizes[0] + 2 * paddings[0] -
                        (dilations[0] * (kernel_sizes[0] - 1) + 1)) /
@@ -49,8 +49,8 @@ void FoldGradKernel(const Context& dev_ctx,
                          strides[1] +
                      1;
 
-  int n_input_plane = x_dims[1];
-  int n_output_plane = n_input_plane / (kernel_sizes[0] * kernel_sizes[1]);
+  int64_t n_input_plane = x_dims[1];
+  int64_t n_output_plane = n_input_plane / (kernel_sizes[0] * kernel_sizes[1]);
 
   DDim out_shape =
       common::make_ddim({n_output_plane, output_sizes[0], output_sizes[1]});
@@ -59,7 +59,7 @@ void FoldGradKernel(const Context& dev_ctx,
 
   phi::funcs::Im2ColFunctor<phi::funcs::ColFormat::kCFO, Context, T> im2col;
 
-  for (int i = 0; i < batch_size; i++) {
+  for (int64_t i = 0; i < batch_size; i++) {
     DenseTensor out_grad_batch = out_grad.Slice(i, i + 1).Resize(out_shape);
     DenseTensor x_grad_batch =
         x_grad->Slice(i, i + 1).Resize(input_matrix_shape);

paddle/phi/kernels/impl/fold_kernel_impl.h

@@ -33,7 +33,7 @@ void FoldKernel(const Context& dev_ctx,
                 const std::vector<int>& paddings,
                 const std::vector<int>& dilations,
                 DenseTensor* out) {
-  const int batch_size = static_cast<int>(x.dims()[0]);
+  const int64_t batch_size = x.dims()[0];
   dev_ctx.template Alloc<T>(out);
 
   phi::funcs::Col2ImFunctor<phi::funcs::ColFormat::kCFO, Context, T> col2im;
@@ -48,8 +48,8 @@ void FoldKernel(const Context& dev_ctx,
                          strides[1] +
                      1;
 
-  int n_input_plane = x_dims[1];
-  int n_output_plane = n_input_plane / (kernel_sizes[0] * kernel_sizes[1]);
+  int64_t n_input_plane = x_dims[1];
+  int64_t n_output_plane = n_input_plane / (kernel_sizes[0] * kernel_sizes[1]);
 
   DDim output_shape =
       common::make_ddim({n_output_plane, output_sizes[0], output_sizes[1]});
@@ -60,7 +60,7 @@ void FoldKernel(const Context& dev_ctx,
   phi::funcs::SetConstant<Context, T> set_zero;
   set_zero(dev_ctx, out, static_cast<T>(0));
 
-  for (int i = 0; i < batch_size; i++) {
+  for (int64_t i = 0; i < batch_size; i++) {
     DenseTensor out_batch =
         out->Slice(i, i + 1).Resize(output_shape);  // im size=3
 

paddle/phi/kernels/impl/frame_grad_kernel_impl.h

@@ -29,9 +29,10 @@ void FrameGradKernel(const Context& dev_ctx,
   dev_ctx.template Alloc<T>(dx);
   const size_t dout_rank = dout.dims().size();
   const size_t dx_rank = dx->dims().size();
-  const int n_frames =
+  const int64_t n_frames =
       (axis == 0) ? dout.dims()[0] : dout.dims()[dout_rank - 1];
-  const int seq_length = (axis == 0) ? dx->dims()[0] : dx->dims()[dx_rank - 1];
+  const int64_t seq_length =
+      (axis == 0) ? dx->dims()[0] : dx->dims()[dx_rank - 1];
   DenseTensor dout_tmp = dout;
 
   DDim preserved_dims;

paddle/phi/kernels/impl/frame_kernel_impl.h

@@ -28,8 +28,9 @@ void FrameKernel(const Context& dev_ctx,
   dev_ctx.template Alloc<T>(out);
   const size_t x_rank = x.dims().size();
   const size_t out_rank = out->dims().size();
-  const int n_frames = (axis == 0) ? out->dims()[0] : out->dims()[out_rank - 1];
-  const int seq_length = (axis == 0) ? x.dims()[0] : x.dims()[x_rank - 1];
+  const int64_t n_frames =
+      (axis == 0) ? out->dims()[0] : out->dims()[out_rank - 1];
+  const int64_t seq_length = (axis == 0) ? x.dims()[0] : x.dims()[x_rank - 1];
   // When the number of input dims is larger than 2, it needs to copy
   // from x to resize input into 2d and output into 3d. Moreover, output
   // dims will be restored at the last step.

paddle/phi/kernels/impl/gumbel_softmax_grad_kernel_impl.h

@@ -30,7 +30,8 @@ void GumbelSoftmaxGradKernel(const Context& dev_ctx,
                              DenseTensor* dx) {
   const int rank = dx->dims().size();
   axis = funcs::CanonicalAxis(axis, rank);
-  int axis_dim = dx->dims()[axis];
+  int64_t axis_dim = dx->dims()[axis];
+
   // allocate memory on device.
 
   dev_ctx.template Alloc<T>(dx);
@@ -44,6 +45,19 @@ void GumbelSoftmaxGradKernel(const Context& dev_ctx,
     return;
   }
 
+  // TODO(large-tensor): Softmax functor implementation still uses int for
+  // dimensions. Need to update Softmax functor to support dimensions >
+  // INT32_MAX.
+  PADDLE_ENFORCE_LE(
+      axis_dim,
+      std::numeric_limits<int>::max(),
+      common::errors::InvalidArgument(
+          "The axis dimension (%ld) exceeds the maximum value that int can "
+          "represent (%d). GumbelSoftmax gradient operation does not support "
+          "such large tensors yet.",
+          axis_dim,
+          std::numeric_limits<int>::max()));
+
   const int size_to_axis = funcs::SizeToAxis(axis, dx->dims());
   const int size_from_axis = funcs::SizeFromAxis(axis, dx->dims());
   DenseTensor dx_2d(*dx), out_2d(out), dout_2d(dout);

paddle/phi/kernels/impl/gumbel_softmax_kernel_impl.h

@@ -14,6 +14,7 @@
 
 #pragma once
 
+#include <iostream>
 #include <random>
 
 #include "paddle/phi/core/dense_tensor.h"
@@ -52,7 +53,7 @@ void GumbelSoftmaxKernelHelper(const Context& dev_ctx,
                                DenseTensor* out) {
   const int rank = x.dims().size();
   axis = funcs::CanonicalAxis(axis, rank);
-  int axis_dim = x.dims()[axis];
+  int64_t axis_dim = x.dims()[axis];
 
   PADDLE_ENFORCE_GT(temperature,
                     0,
@@ -73,6 +74,19 @@ void GumbelSoftmaxKernelHelper(const Context& dev_ctx,
     return;
   }
 
+  // TODO(large-tensor): Softmax functor implementation still uses int for
+  // dimensions. Need to update Softmax functor to support dimensions >
+  // INT32_MAX.
+  PADDLE_ENFORCE_LE(
+      axis_dim,
+      std::numeric_limits<int>::max(),
+      common::errors::InvalidArgument(
+          "The axis dimension (%ld) exceeds the maximum value that int can "
+          "represent (%d). GumbelSoftmax operation does not support such "
+          "large tensors yet.",
+          axis_dim,
+          std::numeric_limits<int>::max()));
+
   const int size_to_axis = funcs::SizeToAxis(axis, x.dims());
   const int size_from_axis = funcs::SizeFromAxis(axis, x.dims());
   DenseTensor x_noise_2d, out_2d(*out);