[PHI] Fix flatten and split kernel for big tensor (PaddlePaddle#72634)

Author: ggggxm

paddle/fluid/primitive/decomp_rule/decomp_rule/composite.h

@@ -909,7 +909,7 @@ Tensor flatten_decomp(const Tensor& x, int start_axis, int end_axis) {
       return reshape<T>(x, x_dim);
     }
 
-    int slice_numel = 1;
+    int64_t slice_numel = 1;
     for (int i = start_axis; i <= end_axis; ++i) {
       slice_numel *= x_dim[i];
     }

paddle/phi/infermeta/unary.cc

@@ -1617,7 +1617,7 @@ void FlattenInferMeta(const MetaTensor& x,
   }
 
   int64_t outer = 1;
-  std::vector<int32_t> out_shape;
+  std::vector<int64_t> out_shape;
   out_shape.reserve(in_dims_size - stop_axis + start_axis + 1);
 
   for (int i = 0; i < start_axis; ++i) {
@@ -4380,7 +4380,7 @@ void SplitInferMeta(const MetaTensor& x,
         axis_value == -1) {  // NOLINT
       out_dims = std::vector<phi::DDim>(
           sections_data.size(),
-          common::make_ddim(std::vector<int>(x.dims().size(), -1)));
+          common::make_ddim(std::vector<int64_t>(x.dims().size(), -1)));
     } else {
       out_dims = std::vector<phi::DDim>(sections_data.size(), x.dims());
     }
@@ -4403,7 +4403,7 @@ void SplitInferMeta(const MetaTensor& x,
     const int unknow_dim_val = -1;
     int unknow_dim_idx = -1;
     int num_of_unknow = 0;
-    int sum_of_section = 0;
+    int64_t sum_of_section = 0;
 
     for (int i = 0; i < static_cast<int>(sections_data.size()); ++i) {
       sections_vec.push_back(sections_data[i]);
@@ -4412,7 +4412,7 @@ void SplitInferMeta(const MetaTensor& x,
         num_of_unknow++;
         unknow_dim_idx = i;
       } else {
-        sum_of_section += static_cast<int>(sections_data[i]);
+        sum_of_section += static_cast<int64_t>(sections_data[i]);
       }
     }