[API] fix paddle.cross with big tensor

Author: ZhangX-21

paddle/phi/kernels/cpu/cross_grad_kernel.cc

@@ -74,13 +74,13 @@ void CrossGradKernel(const Context &dev_ctx,
                                 "But received: Input(X/Y).dims() == [%s].",
                                 input_x_dims));
   }
-  auto outer_loops = 1;
-  for (auto i = 0; i < dim; i++) {
-    outer_loops *= static_cast<int>(input_x_dims[i]);
+  int64_t outer_loops = 1;
+  for (int64_t i = 0; i < dim; i++) {
+    outer_loops *= input_x_dims[i];
   }
-  auto slice_size = 1;
-  for (auto i = dim + 1; i < input_x_dims.size(); i++) {
-    slice_size *= static_cast<int>(input_x_dims[i]);
+  int64_t slice_size = 1;
+  for (int64_t i = dim + 1; i < input_x_dims.size(); i++) {
+    slice_size *= input_x_dims[i];
   }
 
   int64_t numel = x.numel();
@@ -111,12 +111,12 @@ void CrossGradKernel(const Context &dev_ctx,
   dev_ctx.template Alloc<T>(output_x_grad);
   dev_ctx.template Alloc<T>(output_y_grad);
 
-  for (auto i = 0; i < outer_loops; i++) {
-    for (auto j = 0; j < 3; j++) {
-      auto dst_pos = (3 * i + j) * slice_size;
-      auto in_pos1 = (3 * i + ((j + 1) % 3)) * slice_size;
-      auto in_pos2 = (3 * i + ((j + 2) % 3)) * slice_size;
-      for (auto k = 0; k < slice_size; k++) {
+  for (int64_t i = 0; i < outer_loops; i++) {
+    for (int64_t j = 0; j < 3; j++) {
+      int64_t dst_pos = (3 * i + j) * slice_size;
+      int64_t in_pos1 = (3 * i + ((j + 1) % 3)) * slice_size;
+      int64_t in_pos2 = (3 * i + ((j + 2) % 3)) * slice_size;
+      for (int64_t k = 0; k < slice_size; k++) {
         out_dx_vec[dst_pos + k] =
             input_dout_vec[in_pos2 + k] * input_y_vec[in_pos1 + k] -
             input_dout_vec[in_pos1 + k] * input_y_vec[in_pos2 + k];

paddle/phi/kernels/cpu/cross_kernel.cc

@@ -75,13 +75,13 @@ void CrossKernel(const Context& dev_ctx,
     dev_ctx.template Alloc<T>(output);
     return;
   }
-  auto outer_loops = 1;
+  int64_t outer_loops = 1;
   for (auto i = 0; i < dim; i++) {
-    outer_loops *= static_cast<int>(input_x_dims[i]);
+    outer_loops *= input_x_dims[i];
   }
-  auto slice_size = 1;
+  int64_t slice_size = 1;
   for (auto i = dim + 1; i < input_x_dims.size(); i++) {
-    slice_size *= static_cast<int>(input_x_dims[i]);
+    slice_size *= input_x_dims[i];
   }
 
   std::vector<T> input_x_vec, input_y_vec;
@@ -91,13 +91,13 @@ void CrossKernel(const Context& dev_ctx,
 
   dev_ctx.template Alloc<T>(output);
 
-  for (auto i = 0; i < outer_loops; i++) {
-    for (auto j = 0; j < 3; j++) {
-      auto dst_pos = (3 * i + j) * slice_size;
-      auto in_pos1 = (3 * i + ((j + 1) % 3)) * slice_size;
-      auto in_pos2 = (3 * i + ((j + 2) % 3)) * slice_size;
+  for (int64_t i = 0; i < outer_loops; i++) {
+    for (int64_t j = 0; j < 3; j++) {
+      int64_t dst_pos = (3 * i + j) * slice_size;
+      int64_t in_pos1 = (3 * i + ((j + 1) % 3)) * slice_size;
+      int64_t in_pos2 = (3 * i + ((j + 2) % 3)) * slice_size;
 
-      for (auto k = 0; k < slice_size; k++) {
+      for (int64_t k = 0; k < slice_size; k++) {
         out_vec[dst_pos + k] =
             input_x_vec[in_pos1 + k] * input_y_vec[in_pos2 + k] -
             input_x_vec[in_pos2 + k] * input_y_vec[in_pos1 + k];

paddle/phi/kernels/gpu/cross_grad_kernel.cu

@@ -25,20 +25,21 @@
 namespace phi {
 
 template <typename T>
-__global__ void CrossGrad(const T* x,
-                          const T* y,
-                          const T* out,
-                          T* out_dx,
-                          T* out_dy,
-                          const int64_t stride,
-                          const int64_t N,
-                          phi::funcs::IndexCalculator<int> index_calculator) {
+__global__ void CrossGrad(
+    const T* x,
+    const T* y,
+    const T* out,
+    T* out_dx,
+    T* out_dy,
+    const int64_t stride,
+    const int64_t N,
+    phi::funcs::IndexCalculator<int64_t> index_calculator) {
   CUDA_KERNEL_LOOP(i, N) {
     int64_t offset = index_calculator(i);
 
-    auto pos0 = offset + 0 * stride;
-    auto pos1 = offset + 1 * stride;
-    auto pos2 = offset + 2 * stride;
+    int64_t pos0 = offset + 0 * stride;
+    int64_t pos1 = offset + 1 * stride;
+    int64_t pos2 = offset + 2 * stride;
 
     using MPType = typename phi::dtype::MPTypeTrait<T>::Type;
 
@@ -168,7 +169,7 @@ void CrossGradKernel(const Context& dev_ctx,
   const auto* input_out_grad_data = input_out_grad.data<T>();
   auto* output_x_grad_data = dev_ctx.template Alloc<T>(x_grad);
   auto* output_y_grad_data = dev_ctx.template Alloc<T>(y_grad);
-  auto index_calculator = phi::funcs::IndexCalculator<int>(
+  auto index_calculator = phi::funcs::IndexCalculator<int64_t>(
       merged_dims.size() - 1, cal_dims, left_strides, full_strides);
 
   backends::gpu::GpuLaunchConfig config =

paddle/phi/kernels/gpu/cross_kernel.cu

@@ -29,13 +29,13 @@ __global__ void Cross(const T* x,
                       T* out,
                       const int64_t stride,
                       const int64_t N,
-                      phi::funcs::IndexCalculator<int> index_calculator) {
+                      phi::funcs::IndexCalculator<int64_t> index_calculator) {
   CUDA_KERNEL_LOOP(i, N) {
     int64_t offset = index_calculator(i);
 
-    auto pos0 = offset + 0 * stride;
-    auto pos1 = offset + 1 * stride;
-    auto pos2 = offset + 2 * stride;
+    int64_t pos0 = offset + 0 * stride;
+    int64_t pos1 = offset + 1 * stride;
+    int64_t pos2 = offset + 2 * stride;
 
     using MPType = typename phi::dtype::MPTypeTrait<T>::Type;
 
@@ -149,7 +149,7 @@ void CrossKernel(const Context& dev_ctx,
   const auto* input_x_data = input_x.data<T>();
   const auto* input_y_data = input_y.data<T>();
   auto* out_data = dev_ctx.template Alloc<T>(out);
-  auto index_calculator = phi::funcs::IndexCalculator<int>(
+  auto index_calculator = phi::funcs::IndexCalculator<int64_t>(
       merged_dims.size() - 1, cal_dims, left_strides, full_strides);
 
   int64_t numel = x.numel();