[BIT] Fix paddle.unstack api for big tensor (#72931)

cszdrg · web-flow · commit d3496047a415 · 2025-05-30T10:11:40.000+08:00
diff --git a/paddle/phi/infermeta/backward.cc b/paddle/phi/infermeta/backward.cc
@@ -1747,8 +1747,8 @@ void UnStackGradInferMeta(const std::vector<const MetaTensor*>& out_grad,
                         rank));
   if (axis < 0) axis += (rank + 1);
 
-  auto vec = common::vectorize<int>(input_dims[0]);
-  vec.insert(vec.begin() + axis, static_cast<int>(input_dims.size()));
+  auto vec = common::vectorize<int64_t>(input_dims[0]);
+  vec.insert(vec.begin() + axis, static_cast<int64_t>(input_dims.size()));
   x_grad->set_dims(common::make_ddim(vec));
   x_grad->set_dtype(out_grad[0]->dtype());
 }
diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
@@ -6075,7 +6075,7 @@ void UnStackInferMeta(const MetaTensor& x,
             x_dim[axis],
             num));
   }
-  auto vec = common::vectorize<int>(x_dim);
+  auto vec = common::vectorize<int64_t>(x_dim);
   vec.erase(vec.begin() + axis);
   for (size_t i = 0; i < output_count; i++) {
     outs[i]->set_dims(common::make_ddim(vec));
diff --git a/paddle/phi/kernels/funcs/stack_and_unstack.h b/paddle/phi/kernels/funcs/stack_and_unstack.h
@@ -210,7 +210,7 @@ void LaunchUnStackKernel(const Context& ctx,
     constexpr int kWarpSize = 32;
     constexpr int kMaxOut = 16;
 
-    int tid_x = 0, tid_y = 0, bid_x = 0, bid_y = 1;
+    int64_t tid_x = 0, tid_y = 0, bid_x = 0, bid_y = 1;
     if (split_dim < kMaxOut) {
       tid_y = split_dim;
       tid_x =
@@ -219,10 +219,13 @@ void LaunchUnStackKernel(const Context& ctx,
     } else {
       tid_y = kMaxOut;
       tid_x = kWarpSize;
-      bid_y = backends::gpu::DivUp<int>(split_dim, kMaxOut);
+      bid_y = backends::gpu::DivUp<int64_t>(split_dim, kMaxOut);
     }
-    int tile_x_num = backends::gpu::DivUp<int>(out_row, tid_x);
-    bid_x = std::min(tile_x_num, backends::gpu::kMultiDimslimit);
+    int64_t tile_x_num = backends::gpu::DivUp<int64_t>(out_row, tid_x);
+    if (tile_x_num < static_cast<int64_t>(backends::gpu::kMultiDimslimit))
+      bid_x = tile_x_num;
+    else
+      bid_x = backends::gpu::kMultiDimslimit;
     dim3 blocks(tid_x, tid_y, 1);
     dim3 grids(bid_x, bid_y, 1);
 
diff --git a/paddle/phi/kernels/impl/unstack_grad_kernel_impl.h b/paddle/phi/kernels/impl/unstack_grad_kernel_impl.h
@@ -28,13 +28,13 @@ void UnStackGradKernel(const Context &dev_ctx,
                        DenseTensor *x_grad) {
   if (axis < 0) axis += (x[0]->dims().size() + 1);
 
-  int n = static_cast<int>(x.size());
+  int64_t n = static_cast<int64_t>(x.size());
   auto *x_grad_data = dev_ctx.template Alloc<T>(x_grad);
   std::vector<const T *> x_datas(n);
-  for (int i = 0; i < n; i++) x_datas[i] = x[i]->data<T>();
+  for (int64_t i = 0; i < n; i++) x_datas[i] = x[i]->data<T>();
 
-  int pre = 1;
-  int post = 1;
+  int64_t pre = 1;
+  int64_t post = 1;
   auto &dim = x[0]->dims();
   for (auto i = 0; i < axis; ++i) pre *= dim[i];
   for (auto i = axis; i < dim.size(); ++i) post *= dim[i];
@@ -56,8 +56,8 @@ void UnStackGradKernel(const Context &dev_ctx,
 
   size_t x_offset = 0;
   size_t y_offset = 0;
-  for (int i = 0; i < pre; i++) {
-    for (int j = 0; j < n; j++) {
+  for (int64_t i = 0; i < pre; i++) {
+    for (int64_t j = 0; j < n; j++) {
       std::memcpy(
           x_grad_data + y_offset, x_data_arr[j] + x_offset, post * sizeof(T));
       y_offset += post;

Original file line number	Diff line number	Diff line change
`@@ -6075,7 +6075,7 @@ void UnStackInferMeta(const MetaTensor& x,`
`6075`	`6075`	`x_dim[axis],`
`6076`	`6076`	`num));`
`6077`	`6077`	`}`
`6078`		`- auto vec = common::vectorize<int>(x_dim);`
	`6078`	`+ auto vec = common::vectorize<int64_t>(x_dim);`
`6079`	`6079`	`vec.erase(vec.begin() + axis);`
`6080`	`6080`	`for (size_t i = 0; i < output_count; i++) {`
`6081`	`6081`	`outs[i]->set_dims(common::make_ddim(vec));`