[API] fix paddle.dist with big tensor

Author: ZhangX-21

paddle/phi/kernels/cpu/p_norm_grad_kernel.cc

@@ -24,21 +24,21 @@ namespace phi {
 
 inline void GetDims(const phi::DDim& dim,
                     int axis,
-                    int* pre,
-                    int* n,
-                    int* post,
+                    int64_t* pre,
+                    int64_t* n,
+                    int64_t* post,
                     bool asvector) {
   *pre = 1;
   *post = 1;
-  *n = static_cast<int>(dim[axis]);
+  *n = dim[axis];
   if (asvector) {
-    *n = static_cast<int>(product(dim));
+    *n = product(dim);
   } else {
     for (int i = 0; i < axis; ++i) {
-      (*pre) *= static_cast<int>(dim[i]);
+      (*pre) *= dim[i];
     }
     for (int i = axis + 1; i < dim.size(); ++i) {
-      (*post) *= static_cast<int>(dim[i]);
+      (*post) *= dim[i];
     }
   }
 }
@@ -64,10 +64,10 @@ void PNormGradKernel(const Context& dev_ctx,
   auto xdim = in_x->dims();
 
   if (axis < 0) axis = xdim.size() + axis;
-  int pre, n, post;
+  int64_t pre, n, post;
   GetDims(xdim, axis, &pre, &n, &post, asvector);
-  Eigen::DSizes<int, 3> shape(pre, n, post);
-  Eigen::DSizes<int, 3> rshape(pre, 1, post);
+  Eigen::DSizes<int64_t, 3> shape(pre, n, post);
+  Eigen::DSizes<int64_t, 3> rshape(pre, 1, post);
 
   auto* place = dev_ctx.eigen_device();
 
@@ -81,8 +81,8 @@ void PNormGradKernel(const Context& dev_ctx,
   auto norm = norm_e.reshape(rshape);
   auto norm_dy = norm_dy_e.reshape(rshape);
 
-  Eigen::DSizes<int, 1> rdim(1);
-  Eigen::DSizes<int, 3> bcast(1, n, 1);
+  Eigen::DSizes<int64_t, 1> rdim(1);
+  Eigen::DSizes<int64_t, 3> bcast(1, n, 1);
 
   if (porder == 0) {
     phi::funcs::SetConstant<Context, T> set_zero;

paddle/phi/kernels/cpu/p_norm_kernel.cc

@@ -26,21 +26,21 @@ namespace phi {
 
 inline void GetDims(const phi::DDim& dim,
                     int axis,
-                    int* pre,
-                    int* n,
-                    int* post,
+                    int64_t* pre,
+                    int64_t* n,
+                    int64_t* post,
                     bool asvector) {
   *pre = 1;
   *post = 1;
-  *n = static_cast<int>(dim[axis]);
+  *n = static_cast<int64_t>(dim[axis]);
   if (asvector) {
-    *n = static_cast<int>(product(dim));
+    *n = static_cast<int64_t>(product(dim));
   } else {
     for (int i = 0; i < axis; ++i) {
-      (*pre) *= static_cast<int>(dim[i]);
+      (*pre) *= static_cast<int64_t>(dim[i]);
     }
     for (int i = axis + 1; i < dim.size(); ++i) {
-      (*post) *= static_cast<int>(dim[i]);
+      (*post) *= static_cast<int64_t>(dim[i]);
     }
   }
 }
@@ -59,7 +59,7 @@ void PNormKernel(const Context& dev_ctx,
 
   auto xdim = in_x->dims();
   if (axis < 0) axis = xdim.size() + axis;
-  int pre = 0, n = 0, post = 0;
+  int64_t pre = 0, n = 0, post = 0;
   GetDims(xdim, axis, &pre, &n, &post, asvector);
 
   if (x.numel() == 0) {
@@ -73,8 +73,8 @@ void PNormKernel(const Context& dev_ctx,
 
   auto* place = dev_ctx.eigen_device();
 
-  Eigen::DSizes<int, 3> shape(pre, n, post);
-  Eigen::DSizes<int, 2> norm_shape(pre, post);
+  Eigen::DSizes<int64_t, 3> shape(pre, n, post);
+  Eigen::DSizes<int64_t, 2> norm_shape(pre, post);
 
   auto x_e = phi::EigenVector<T>::Flatten(*in_x);
   auto norm_e = phi::EigenVector<T>::Flatten(*out);
@@ -86,7 +86,7 @@ void PNormKernel(const Context& dev_ctx,
   // p=inf means the maximum of |xr|
   // p=-inf means the minimum of |xr|
   // otherwise, Lp-norm = pow(sum(pow(|xr|, p)), 1/p)
-  Eigen::DSizes<int, 1> rdim(1);
+  Eigen::DSizes<int64_t, 1> rdim(1);
   if (porder == 0) {
     norm.device(*place) = (xr != xr.constant(0)).template cast<T>().sum(rdim);
   } else if (porder == INFINITY) {

paddle/phi/kernels/funcs/blas/blas_impl.h

@@ -78,31 +78,31 @@ struct CBlas<phi::dtype::bfloat16> {
   }
 
   template <typename... ARGS>
-  static void VADD(int n,
+  static void VADD(int64_t n,
                    const phi::dtype::bfloat16 *x,
                    const phi::dtype::bfloat16 *y,
                    phi::dtype::bfloat16 *z) {
-    for (int i = 0; i < n; ++i) {
+    for (int64_t i = 0; i < n; ++i) {
       z[i] = x[i] + y[i];
     }
   }
 
   template <typename... ARGS>
-  static void VMUL(int n,
+  static void VMUL(int64_t n,
                    const phi::dtype::bfloat16 *x,
                    const phi::dtype::bfloat16 *y,
                    phi::dtype::bfloat16 *z) {
-    for (int i = 0; i < n; ++i) {
+    for (int64_t i = 0; i < n; ++i) {
       z[i] = x[i] * y[i];
     }
   }
 
   template <typename... ARGS>
-  static void VSUB(int n,
+  static void VSUB(int64_t n,
                    const phi::dtype::bfloat16 *x,
                    const phi::dtype::bfloat16 *y,
                    phi::dtype::bfloat16 *z) {
-    for (int i = 0; i < n; ++i) {
+    for (int64_t i = 0; i < n; ++i) {
       z[i] = x[i] - y[i];
     }
   }