[PHI] fix paddle.diag with big tensor (#72638)

ZhangX-21 · web-flow · commit 7940a6154380 · 2025-05-19T10:43:28.000+08:00
diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
@@ -931,10 +931,10 @@ void DiagInferMeta(const MetaTensor& x,
   auto x_dims = x.dims();
 
   if (x_dims.size() <= 1) {
-    int64_t size_ = (x_dims.size() == 1UL ? x_dims[0] : 1) + std::abs(offset);
+    int64_t size_ = (x_dims.size() == 1ULL ? x_dims[0] : 1) + std::abs(offset);
     out->set_dims({size_, size_});
     out->set_dtype(x.dtype());
-  } else if (x_dims.size() == 2UL) {
+  } else if (x_dims.size() == 2ULL) {
     int64_t size_ = 0;
     if (offset >= 0) {
       // Note(LutaoChu): Do not use std::min here, otherwise the calculation
diff --git a/paddle/phi/kernels/cpu/diag_grad_kernel.cc b/paddle/phi/kernels/cpu/diag_grad_kernel.cc
@@ -33,28 +33,28 @@ void DiagGradKernel(const Context& dev_ctx,
   auto dout_dims = out_grad.dims();
 
   if (dx_dims.size() <= 1) {
-    auto dx_length = (dx_dims.size() == 1 ? dx_dims[0] : int64_t(1));
-    int dx_stride = 1;
+    int64_t dx_length = (dx_dims.size() == 1 ? dx_dims[0] : int64_t(1));
+    int64_t dx_stride = 1;
 
-    auto dout_stride_0 = phi::funcs::ComputeStride(0, dout_dims);
-    auto dout_stride_1 = phi::funcs::ComputeStride(1, dout_dims);
+    int64_t dout_stride_0 = phi::funcs::ComputeStride(0, dout_dims);
+    int64_t dout_stride_1 = phi::funcs::ComputeStride(1, dout_dims);
     dout_data +=
         (offset >= 0 ? offset * dout_stride_1 : -offset * dout_stride_0);
 
-    for (int i = 0; i < dx_length; i++) {
+    for (int64_t i = 0; i < dx_length; i++) {
       dx_data[i * dx_stride] = dout_data[i * (dout_stride_0 + dout_stride_1)];
     }
   } else {
     phi::funcs::SetConstant<Context, T> set_padding_value;
     set_padding_value(dev_ctx, x_grad, static_cast<T>(0));
 
-    int dx_stride_0 = phi::funcs::ComputeStride(0, dx_dims);
-    int dx_stride_1 = phi::funcs::ComputeStride(1, dx_dims);
-    auto dout_stride_0 = phi::funcs::ComputeStride(0, dout_dims);
+    int64_t dx_stride_0 = phi::funcs::ComputeStride(0, dx_dims);
+    int64_t dx_stride_1 = phi::funcs::ComputeStride(1, dx_dims);
+    int64_t dout_stride_0 = phi::funcs::ComputeStride(0, dout_dims);
     dx_data += (offset >= 0 ? offset * dx_stride_1 : -offset * dx_stride_0);
 
-    auto dout_length = dout_dims[0];
-    for (int i = 0; i < dout_length; i++) {
+    int64_t dout_length = dout_dims[0];
+    for (int64_t i = 0; i < dout_length; i++) {
       dx_data[i * (dx_stride_0 + dx_stride_1)] = dout_data[i * dout_stride_0];
     }
   }
diff --git a/paddle/phi/kernels/cpu/diag_kernel.cc b/paddle/phi/kernels/cpu/diag_kernel.cc
@@ -37,22 +37,22 @@ void DiagKernel(const Context& dev_ctx,
     phi::funcs::SetConstant<Context, T> set_padding_value;
     set_padding_value(dev_ctx, out, static_cast<T>(padding_value));
 
-    auto x_length = (x_dims.size() == 1UL ? x_dims[0] : int64_t(1));
-    const int& x_stride = 1;
+    int64_t x_length = (x_dims.size() == 1ULL ? x_dims[0] : int64_t(1));
+    const int64_t x_stride = 1;
 
-    auto out_stride_0 = phi::funcs::ComputeStride(0, out_dims);
-    auto out_stride_1 = phi::funcs::ComputeStride(1, out_dims);
+    int64_t out_stride_0 = phi::funcs::ComputeStride(0, out_dims);
+    int64_t out_stride_1 = phi::funcs::ComputeStride(1, out_dims);
     out_data += (offset >= 0 ? offset * out_stride_1 : -offset * out_stride_0);
 
     for (i = 0; i < x_length; i++) {
       out_data[i * (out_stride_0 + out_stride_1)] = x_data[i * x_stride];
     }
   } else {
-    auto out_length = out_dims[0];
-    const int& x_stride_0 = phi::funcs::ComputeStride(0, x_dims);
-    const int& x_stride_1 = phi::funcs::ComputeStride(1, x_dims);
+    int64_t out_length = out_dims[0];
+    const int64_t x_stride_0 = phi::funcs::ComputeStride(0, x_dims);
+    const int64_t x_stride_1 = phi::funcs::ComputeStride(1, x_dims);
 
-    auto out_stride_0 = phi::funcs::ComputeStride(0, out_dims);
+    int64_t out_stride_0 = phi::funcs::ComputeStride(0, out_dims);
     x_data += (offset >= 0 ? offset * x_stride_1 : -offset * x_stride_0);
     for (i = 0; i < out_length; i++) {
       out_data[i * out_stride_0] = x_data[i * (x_stride_0 + x_stride_1)];
diff --git a/paddle/phi/kernels/funcs/diag_functor.h b/paddle/phi/kernels/funcs/diag_functor.h
@@ -25,8 +25,8 @@
 namespace phi {
 namespace funcs {
 
-inline int ComputeStride(int axis, phi::DDim dims) {
-  int size = 1;
+inline int64_t ComputeStride(int axis, phi::DDim dims) {
+  int64_t size = 1;
   for (int i = axis + 1; i < dims.size(); i++) {
     size *= dims[i];
   }
diff --git a/paddle/phi/kernels/gpu/diag_grad_kernel.cu b/paddle/phi/kernels/gpu/diag_grad_kernel.cu
@@ -24,14 +24,13 @@ namespace phi {
 template <typename T>
 __global__ void ExtractDiagonalKernel(const T* dout,
                                       T* dx,
-                                      std::ptrdiff_t start,
-                                      std::ptrdiff_t dx_length,
-                                      const std::ptrdiff_t sumStride,
-                                      const std::ptrdiff_t xStride) {
-  for (std::ptrdiff_t idx = blockIdx.x * blockDim.x + threadIdx.x;
-       idx < dx_length;
+                                      int64_t start,
+                                      int64_t dx_length,
+                                      const int64_t sumStride,
+                                      const int64_t xStride) {
+  for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < dx_length;
        idx += gridDim.x * blockDim.x) {
-    const std::ptrdiff_t outOffset = start + sumStride * idx;
+    const int64_t outOffset = start + sumStride * idx;
     dx[xStride * idx] = dout[outOffset];
   }
 }
@@ -40,13 +39,13 @@ __global__ void ExtractDiagonalKernel(const T* dout,
 template <typename T>
 __global__ void PasteDiagonalKernel(const T* dout,
                                     T* dx,
-                                    std::ptrdiff_t start,
-                                    std::ptrdiff_t size,
-                                    const std::ptrdiff_t sumStride,
-                                    const std::ptrdiff_t outStride) {
-  for (std::ptrdiff_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size;
+                                    int64_t start,
+                                    int64_t size,
+                                    const int64_t sumStride,
+                                    const int64_t outStride) {
+  for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size;
        idx += gridDim.x * blockDim.x) {
-    std::ptrdiff_t xOffset = start + sumStride * idx;
+    int64_t xOffset = start + sumStride * idx;
     dx[xOffset] = dout[outStride * idx];
   }
 }
@@ -74,14 +73,14 @@ void DiagGradKernel(const Context& dev_ctx,
   };
 
   if (dx_dims.size() <= 1) {
-    auto dx_length = (dx_dims.size() == 1 ? dx_dims[0] : int64_t(1));
-    auto size = (offset > 0) ? dx_length + offset : dx_length - offset;
-    int dx_stride = 1;
+    int64_t dx_length = (dx_dims.size() == 1ULL ? dx_dims[0] : 1ULL);
+    int64_t size = (offset > 0ULL) ? dx_length + offset : dx_length - offset;
+    int64_t dx_stride = 1;
     if (size > 0) {
-      auto dout_stride_0 = phi::funcs::ComputeStride(0, dout_dims);
-      auto dout_stride_1 = phi::funcs::ComputeStride(1, dout_dims);
-      auto start =
-          (offset >= 0 ? offset * dout_stride_1 : -offset * dout_stride_0);
+      int64_t dout_stride_0 = phi::funcs::ComputeStride(0, dout_dims);
+      int64_t dout_stride_1 = phi::funcs::ComputeStride(1, dout_dims);
+      int64_t start =
+          (offset >= 0ULL ? offset * dout_stride_1 : -offset * dout_stride_0);
 
       std::tuple<int64_t, int64_t> block_grid_size = GetBlockGridSize(size);
       ExtractDiagonalKernel<T>
@@ -99,8 +98,8 @@ void DiagGradKernel(const Context& dev_ctx,
     phi::funcs::SetConstant<Context, T> set_padding_value;
     set_padding_value(dev_ctx, x_grad, static_cast<T>(0));
 
-    int dx_stride_0 = phi::funcs::ComputeStride(0, dx_dims);
-    int dx_stride_1 = phi::funcs::ComputeStride(1, dx_dims);
+    int64_t dx_stride_0 = phi::funcs::ComputeStride(0, dx_dims);
+    int64_t dx_stride_1 = phi::funcs::ComputeStride(1, dx_dims);
     int64_t size;
     if (offset > 0) {
       size = std::min(dx_dims[0], dx_dims[1] - offset);
@@ -109,8 +108,9 @@ void DiagGradKernel(const Context& dev_ctx,
     }
 
     if (size > 0) {
-      auto start = (offset >= 0 ? offset * dx_stride_1 : -offset * dx_stride_0);
-      auto dout_stride_0 = phi::funcs::ComputeStride(0, dout_dims);
+      int64_t start =
+          (offset >= 0 ? offset * dx_stride_1 : -offset * dx_stride_0);
+      int64_t dout_stride_0 = phi::funcs::ComputeStride(0, dout_dims);
       std::tuple<int64_t, int64_t> block_grid_size = GetBlockGridSize(size);
       PasteDiagonalKernel<T><<<std::get<1>(block_grid_size),
                                std::get<0>(block_grid_size),
diff --git a/paddle/phi/kernels/gpu/diag_kernel.cu b/paddle/phi/kernels/gpu/diag_kernel.cu
@@ -28,13 +28,13 @@ namespace phi {
 template <typename T>
 __global__ void ExtractDiagonalKernel(T* out,
                                       const T* x,
-                                      std::ptrdiff_t start,
-                                      std::ptrdiff_t size,
-                                      const std::ptrdiff_t sumStride,
-                                      const std::ptrdiff_t outStride) {
-  for (std::ptrdiff_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size;
+                                      int64_t start,
+                                      int64_t size,
+                                      const int64_t sumStride,
+                                      const int64_t outStride) {
+  for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size;
        idx += gridDim.x * blockDim.x) {
-    const std::ptrdiff_t xOffset = start + sumStride * idx;
+    const int64_t xOffset = start + sumStride * idx;
     out[outStride * idx] = x[xOffset];
   }
 }
@@ -43,14 +43,13 @@ __global__ void ExtractDiagonalKernel(T* out,
 template <typename T>
 __global__ void PasteDiagonalKernel(T* out,
                                     const T* x,
-                                    std::ptrdiff_t start,
-                                    std::ptrdiff_t x_length,
-                                    const std::ptrdiff_t sumStride,
-                                    const std::ptrdiff_t xStride) {
-  for (std::ptrdiff_t idx = blockIdx.x * blockDim.x + threadIdx.x;
-       idx < x_length;
+                                    int64_t start,
+                                    int64_t x_length,
+                                    const int64_t sumStride,
+                                    const int64_t xStride) {
+  for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < x_length;
        idx += gridDim.x * blockDim.x) {
-    const std::ptrdiff_t outOffset = start + sumStride * idx;
+    const int64_t outOffset = start + sumStride * idx;
     out[outOffset] = x[xStride * idx];
   }
 }
@@ -81,13 +80,13 @@ void DiagKernel(const Context& dev_ctx,
     phi::funcs::SetConstant<Context, T> set_padding_value;
     set_padding_value(dev_ctx, out, static_cast<T>(padding_value));
 
-    auto x_length = (x_dims.size() == 1UL ? x_dims[0] : int64_t(1));
-    auto size = (offset > 0) ? x_length + offset : x_length - offset;
-    const int& x_stride = 1;
+    int64_t x_length = (x_dims.size() == 1ULL ? x_dims[0] : int64_t(1));
+    int64_t size = (offset > 0) ? x_length + offset : x_length - offset;
+    const int64_t x_stride = 1;
     if (size > 0) {
-      const auto& out_stride_0 = phi::funcs::ComputeStride(0, out_dims);
-      const auto& out_stride_1 = phi::funcs::ComputeStride(1, out_dims);
-      auto start =
+      const int64_t out_stride_0 = phi::funcs::ComputeStride(0, out_dims);
+      const int64_t out_stride_1 = phi::funcs::ComputeStride(1, out_dims);
+      int64_t start =
           (offset >= 0 ? offset * out_stride_1 : -offset * out_stride_0);
 
       std::tuple<int64_t, int64_t> block_grid_size = GetBlockGridSize(size);
@@ -103,8 +102,8 @@ void DiagKernel(const Context& dev_ctx,
                                                    x_stride);
     }
   } else {
-    const int& x_stride_0 = phi::funcs::ComputeStride(0, x_dims);
-    const int& x_stride_1 = phi::funcs::ComputeStride(1, x_dims);
+    const int64_t x_stride_0 = phi::funcs::ComputeStride(0, x_dims);
+    const int64_t x_stride_1 = phi::funcs::ComputeStride(1, x_dims);
 
     int64_t size;
     if (offset > 0) {
@@ -114,11 +113,11 @@ void DiagKernel(const Context& dev_ctx,
     }
 
     if (size > 0) {
-      auto start = (offset >= 0 ? offset * x_stride_1 : -offset * x_stride_0);
-      const auto& out_stride_0 = phi::funcs::ComputeStride(0, out_dims);
+      int64_t start =
+          (offset >= 0 ? offset * x_stride_1 : -offset * x_stride_0);
+      const int64_t out_stride_0 = phi::funcs::ComputeStride(0, out_dims);
 
       std::tuple<int64_t, int64_t> block_grid_size = GetBlockGridSize(size);
-
       ExtractDiagonalKernel<T><<<std::get<1>(block_grid_size),
                                  std::get<0>(block_grid_size),
                                  0,
