[PHI] Preliminary fix for elementwise broadcast int32 shape overflow (#72584)

Enigmatisms · web-flow · commit 88243ff77daf · 2025-05-07T20:16:14.000+08:00
diff --git a/paddle/fluid/operators/elementwise/elementwise_op.h b/paddle/fluid/operators/elementwise/elementwise_op.h
@@ -106,9 +106,9 @@ class ElementwiseOp : public framework::OperatorWithKernel {
                             axis));
       axis = (axis < 0 ? (std::abs(x_dims.size() - y_dims.size()) + axis + 1)
                        : axis);
-      std::vector<int> x_dims_array(max_dim);
-      std::vector<int> y_dims_array(max_dim);
-      std::vector<int> out_dims_array(max_dim);
+      std::vector<int64_t> x_dims_array(max_dim);
+      std::vector<int64_t> y_dims_array(max_dim);
+      std::vector<int64_t> out_dims_array(max_dim);
 #ifdef PADDLE_WITH_DNNL
       // Broadcasting of dims has to be done on Paddle shapes (NHWC)
       // if model is using NHWC and any of shapes in at least 3D
@@ -120,8 +120,8 @@ class ElementwiseOp : public framework::OperatorWithKernel {
       if (should_rotate) {
         // Pick bigger shape and rotate this one
         bool x_over_y = (x_dims.size() > y_dims.size());
-        auto vdims = x_over_y ? common::vectorize<int>(x_dims)
-                              : common::vectorize<int>(y_dims);
+        auto vdims = x_over_y ? common::vectorize<int64_t>(x_dims)
+                              : common::vectorize<int64_t>(y_dims);
         std::rotate(vdims.begin() + 1, vdims.begin() + 2, vdims.end());
         if (x_over_y) {
           x_dims = common::make_ddim(vdims);
diff --git a/paddle/phi/infermeta/binary.cc b/paddle/phi/infermeta/binary.cc
@@ -489,9 +489,9 @@ void CompareRawInferMeta(const MetaTensor& x,
   } else {
     int max_dim = std::max(dim_x.size(), dim_y.size());
     int axis = std::abs(dim_x.size() - dim_y.size());
-    std::vector<int> x_dims_array(max_dim);
-    std::vector<int> y_dims_array(max_dim);
-    std::vector<int> out_dims_array(max_dim);
+    std::vector<int64_t> x_dims_array(max_dim);
+    std::vector<int64_t> y_dims_array(max_dim);
+    std::vector<int64_t> out_dims_array(max_dim);
     funcs::GetBroadcastDimsArrays(dim_x,
                                   dim_y,
                                   x_dims_array.data(),
@@ -543,9 +543,9 @@ void ComplexInferMeta(const MetaTensor& x,
 
     // start align axis
     int axis = std::abs(x_dims.size() - y_dims.size());
-    std::vector<int> x_dims_array(max_dim);
-    std::vector<int> y_dims_array(max_dim);
-    std::vector<int> out_dims_array(max_dim);
+    std::vector<int64_t> x_dims_array(max_dim);
+    std::vector<int64_t> y_dims_array(max_dim);
+    std::vector<int64_t> out_dims_array(max_dim);
     phi::funcs::GetBroadcastDimsArrays(x_dims,
                                        y_dims,
                                        x_dims_array.data(),
@@ -1690,9 +1690,9 @@ void ElementwiseRawInferMeta(const MetaTensor& x,
                           axis));
     axis = (axis < 0 ? (std::abs(x_dims.size() - y_dims.size()) + axis + 1)
                      : axis);
-    std::vector<int> x_dims_array(max_dim);
-    std::vector<int> y_dims_array(max_dim);
-    std::vector<int> out_dims_array(max_dim);
+    std::vector<int64_t> x_dims_array(max_dim);
+    std::vector<int64_t> y_dims_array(max_dim);
+    std::vector<int64_t> out_dims_array(max_dim);
 
 #ifdef PADDLE_WITH_DNNL
     bool should_rotate =
@@ -1703,8 +1703,8 @@ void ElementwiseRawInferMeta(const MetaTensor& x,
     if (should_rotate) {
       // Pick bigger shape and rotate this one
       bool x_over_y = (common::product(x_dims) > common::product(y_dims));
-      auto vdims = x_over_y ? common::vectorize<int>(x_dims)
-                            : common::vectorize<int>(y_dims);
+      auto vdims = x_over_y ? common::vectorize<int64_t>(x_dims)
+                            : common::vectorize<int64_t>(y_dims);
       std::rotate(vdims.begin() + 1, vdims.begin() + 2, vdims.end());
       if (x_over_y) {
         x_dims = common::make_ddim(vdims);
@@ -3141,8 +3141,8 @@ void MatrixRankTolInferMeta(const MetaTensor& x,
                         "The dims of input must be greater than 2"));
 
   if (hermitian) {
-    int rows = static_cast<int>(dim_x[dim_x.size() - 2]);
-    int cols = static_cast<int>(dim_x[dim_x.size() - 1]);
+    int64_t rows = static_cast<int64_t>(dim_x[dim_x.size() - 2]);
+    int64_t cols = static_cast<int64_t>(dim_x[dim_x.size() - 1]);
     PADDLE_ENFORCE_EQ(rows,
                       cols,
                       common::errors::InvalidArgument(
@@ -3155,9 +3155,9 @@ void MatrixRankTolInferMeta(const MetaTensor& x,
   } else {
     int max_dim = std::max(dim_x_batch.size(), dim_tol.size());
     int axis = std::abs(dim_x_batch.size() - dim_tol.size());
-    std::vector<int> x_batch_dims_array(max_dim);
-    std::vector<int> tol_dims_array(max_dim);
-    std::vector<int> out_dims_array(max_dim);
+    std::vector<int64_t> x_batch_dims_array(max_dim);
+    std::vector<int64_t> tol_dims_array(max_dim);
+    std::vector<int64_t> out_dims_array(max_dim);
     phi::funcs::GetBroadcastDimsArrays(dim_x_batch,
                                        dim_tol,
                                        x_batch_dims_array.data(),
diff --git a/paddle/phi/infermeta/multiary.cc b/paddle/phi/infermeta/multiary.cc
@@ -5176,16 +5176,16 @@ void SendUERecvInferMeta(const MetaTensor& x,
 
   // Infer out's shape according to x and e(need broadcasting condition)
   out->set_dtype(x.dtype());
-  auto x_dims1 = common::vectorize<int>(x_dims);
-  auto y_dims1 = common::vectorize<int>(y_dims);
-  std::vector<int> x_dims2(x_dims1.begin() + 1, x_dims1.end());
-  std::vector<int> y_dims2(y_dims1.begin() + 1, y_dims1.end());
+  auto x_dims1 = common::vectorize<int64_t>(x_dims);
+  auto y_dims1 = common::vectorize<int64_t>(y_dims);
+  std::vector<int64_t> x_dims2(x_dims1.begin() + 1, x_dims1.end());
+  std::vector<int64_t> y_dims2(y_dims1.begin() + 1, y_dims1.end());
 
   int max_dim = static_cast<int>(std::max(x_dims2.size(), y_dims2.size()));
   int axis = std::abs(static_cast<int>(x_dims2.size() - y_dims2.size()));
-  std::vector<int> x_dims_array(max_dim);
-  std::vector<int> y_dims_array(max_dim);
-  std::vector<int> out_dims_array(max_dim);
+  std::vector<int64_t> x_dims_array(max_dim);
+  std::vector<int64_t> y_dims_array(max_dim);
+  std::vector<int64_t> out_dims_array(max_dim);
   // Only need to broadcast dimensions other than the 0th dimension.
   phi::funcs::GetBroadcastDimsArrays(common::make_ddim(x_dims2),
                                      common::make_ddim(y_dims2),
@@ -5247,15 +5247,15 @@ void SendUVInferMeta(const MetaTensor& x,
   out->set_dtype(x.dtype());
   auto x_dims = x.dims();
   auto y_dims = y.dims();
-  auto x_dims1 = common::vectorize<int>(x_dims);
-  auto y_dims1 = common::vectorize<int>(y_dims);
-  std::vector<int> x_dims2(x_dims1.begin() + 1, x_dims1.end());
-  std::vector<int> y_dims2(y_dims1.begin() + 1, y_dims1.end());
+  auto x_dims1 = common::vectorize<int64_t>(x_dims);
+  auto y_dims1 = common::vectorize<int64_t>(y_dims);
+  std::vector<int64_t> x_dims2(x_dims1.begin() + 1, x_dims1.end());
+  std::vector<int64_t> y_dims2(y_dims1.begin() + 1, y_dims1.end());
   int max_dim = static_cast<int>(std::max(x_dims2.size(), y_dims2.size()));
   int axis = std::abs(static_cast<int>(x_dims2.size() - y_dims2.size()));
-  std::vector<int> x_dims_array(max_dim);
-  std::vector<int> y_dims_array(max_dim);
-  std::vector<int> out_dims_array(max_dim);
+  std::vector<int64_t> x_dims_array(max_dim);
+  std::vector<int64_t> y_dims_array(max_dim);
+  std::vector<int64_t> out_dims_array(max_dim);
   // Only need to broadcast dimensions other than the 0th dimension.
   phi::funcs::GetBroadcastDimsArrays(common::make_ddim(x_dims2),
                                      common::make_ddim(y_dims2),
diff --git a/paddle/phi/kernels/funcs/common_infer_shape_functions.h b/paddle/phi/kernels/funcs/common_infer_shape_functions.h
@@ -21,9 +21,9 @@ inline phi::DDim BroadcastTwoDims(const phi::DDim &x_dims,
                                   int axis = -1) {
   int max_dim = std::max(x_dims.size(), y_dims.size());
   axis = (axis == -1 ? std::abs(x_dims.size() - y_dims.size()) : axis);
-  std::vector<int> x_dims_array(max_dim);
-  std::vector<int> y_dims_array(max_dim);
-  std::vector<int> out_dims_array(max_dim);
+  std::vector<int64_t> x_dims_array(max_dim);
+  std::vector<int64_t> y_dims_array(max_dim);
+  std::vector<int64_t> out_dims_array(max_dim);
   GetBroadcastDimsArrays(x_dims,
                          y_dims,
                          x_dims_array.data(),
diff --git a/paddle/phi/kernels/funcs/common_shape.h b/paddle/phi/kernels/funcs/common_shape.h
@@ -32,11 +32,12 @@ inline void SetXShape(const DenseTensor &x, DenseTensor *xshape) {
   xshape->ResetLoD(x.meta().legacy_lod);
 }
 
+template <typename T>
 inline void GetBroadcastDimsArrays(const DDim &x_dims,
                                    const DDim &y_dims,
-                                   int *x_dims_array,
-                                   int *y_dims_array,
-                                   int *out_dims_array,
+                                   T *x_dims_array,
+                                   T *y_dims_array,
+                                   T *out_dims_array,
                                    const int max_dim,
                                    const int axis) {
   PADDLE_ENFORCE_GE(
diff --git a/paddle/phi/kernels/funcs/elementwise_base.h b/paddle/phi/kernels/funcs/elementwise_base.h
@@ -274,14 +274,14 @@ template <typename Functor, typename T, typename OutType = T>
 void CommonForwardBroadcastCPU(const DenseTensor &x,
                                const DenseTensor &y,
                                DenseTensor *z,
-                               int *x_dims_array,
-                               int *y_dims_array,
-                               int *out_dims_array,
+                               int64_t *x_dims_array,
+                               int64_t *y_dims_array,
+                               int64_t *out_dims_array,
                                int max_dim,
                                const CPUContext &ctx,
                                Functor func,
                                const bool is_xsize_larger = true) {
-  std::vector<int> index_array(max_dim, 0);
+  std::vector<int64_t> index_array(max_dim, 0);
   const T *x_data = x.data<T>();
   const T *y_data = y.data<T>();
   PADDLE_ENFORCE_NOT_NULL(
@@ -290,8 +290,10 @@ void CommonForwardBroadcastCPU(const DenseTensor &x,
       y_data, errors::InvalidArgument("The input Y should not be empty."));
   OutType *out_data = ctx.Alloc<OutType>(z);
 
-  const int out_size = std::accumulate(
-      out_dims_array, out_dims_array + max_dim, 1, std::multiplies<int>());
+  const int64_t out_size = std::accumulate(out_dims_array,
+                                           out_dims_array + max_dim,
+                                           1ll,
+                                           std::multiplies<int64_t>());
   int x_index, y_index;
   for (int out_index = 0; out_index < out_size; ++out_index) {
     x_index = GetElementwiseIndex(x_dims_array, max_dim, index_array.data());
@@ -331,9 +333,9 @@ void CommonElementwiseBroadcastForward(const CPUContext &dev_ctx,
           "Axis should be less than or equal to %d, but received axis is %d.",
           max_dim,
           axis));
-  std::vector<int> x_dims_array(max_dim);
-  std::vector<int> y_dims_array(max_dim);
-  std::vector<int> out_dims_array(max_dim);
+  std::vector<int64_t> x_dims_array(max_dim);
+  std::vector<int64_t> y_dims_array(max_dim);
+  std::vector<int64_t> out_dims_array(max_dim);
   GetBroadcastDimsArrays(x_dims,
                          y_dims,
                          x_dims_array.data(),
diff --git a/paddle/phi/kernels/funcs/elementwise_utils.h b/paddle/phi/kernels/funcs/elementwise_utils.h
@@ -94,10 +94,11 @@ inline DDim TrimTrailingSingularDims(const DDim &dims) {
   return actual_dims;
 }
 
-inline int GetElementwiseIndex(const int *x_dims_array,
-                               const int max_dim,
-                               const int *index_array) {
-  int index_ = 0;
+template <typename ShapeT = int>
+inline ShapeT GetElementwiseIndex(const ShapeT *x_dims_array,
+                                  const int max_dim,
+                                  const ShapeT *index_array) {
+  ShapeT index_ = 0;
   for (int i = 0; i < max_dim; i++) {
     if (x_dims_array[i] > 1) {
       index_ = index_ * x_dims_array[i] + index_array[i];
@@ -106,9 +107,10 @@ inline int GetElementwiseIndex(const int *x_dims_array,
   return index_;
 }
 
-inline void UpdateElementwiseIndexArray(const int *out_dims_array,
+template <typename ShapeT = int>
+inline void UpdateElementwiseIndexArray(const ShapeT *out_dims_array,
                                         const int max_dim,
-                                        int *index_array) {
+                                        ShapeT *index_array) {
   for (int i = max_dim - 1; i >= 0; --i) {
     ++index_array[i];
     if (index_array[i] >= out_dims_array[i]) {
diff --git a/paddle/phi/kernels/fusion/xpu/add_layernorm_xpu_kernel.cc b/paddle/phi/kernels/fusion/xpu/add_layernorm_xpu_kernel.cc
@@ -25,7 +25,7 @@ namespace fusion {
 static phi::DDim BroadCastInferShape(const DDim x_dims,
                                      const DDim y_dims,
                                      int axis) {
-  std::vector<int> out_dims_array(x_dims.size(), -1);
+  std::vector<int64_t> out_dims_array(x_dims.size(), -1);
   if (x_dims != y_dims) {
     int max_dim = std::max(x_dims.size(), y_dims.size());
     if (x_dims.size() == y_dims.size()) {
@@ -49,8 +49,8 @@ static phi::DDim BroadCastInferShape(const DDim x_dims,
                           axis));
     axis = (axis < 0 ? (std::abs(x_dims.size() - y_dims.size()) + axis + 1)
                      : axis);
-    std::vector<int> x_dims_array(max_dim);
-    std::vector<int> y_dims_array(max_dim);
+    std::vector<int64_t> x_dims_array(max_dim);
+    std::vector<int64_t> y_dims_array(max_dim);
     out_dims_array.resize(max_dim);
     phi::funcs::GetBroadcastDimsArrays(x_dims,
                                        y_dims,
diff --git a/paddle/phi/kernels/impl/graph_message_passing_impl.h b/paddle/phi/kernels/impl/graph_message_passing_impl.h
@@ -87,18 +87,19 @@ inline BroadCastInfo CalcBCastInfo(const phi::DDim& l_dims,
   return binfo;
 }
 
-inline std::vector<int> InferBroadcastShape(const phi::DDim& x_dims,
-                                            const phi::DDim& e_dims,
-                                            const std::string& type = "x") {
-  auto x_dims1 = common::vectorize<int>(x_dims);
-  auto e_dims1 = common::vectorize<int>(e_dims);
-  std::vector<int> x_dims2(x_dims1.begin() + 1, x_dims1.end());
-  std::vector<int> e_dims2(e_dims1.begin() + 1, e_dims1.end());
+template <typename ShapeT = int64_t>
+inline std::vector<ShapeT> InferBroadcastShape(const phi::DDim& x_dims,
+                                               const phi::DDim& e_dims,
+                                               const std::string& type = "x") {
+  auto x_dims1 = common::vectorize<ShapeT>(x_dims);
+  auto e_dims1 = common::vectorize<ShapeT>(e_dims);
+  std::vector<ShapeT> x_dims2(x_dims1.begin() + 1, x_dims1.end());
+  std::vector<ShapeT> e_dims2(e_dims1.begin() + 1, e_dims1.end());
   int max_dim = std::max(x_dims2.size(), e_dims2.size());
   int axis = std::abs(static_cast<int>(x_dims2.size() - e_dims2.size()));
-  std::vector<int> x_dims_array(max_dim);
-  std::vector<int> e_dims_array(max_dim);
-  std::vector<int> out_dims_array(max_dim);
+  std::vector<ShapeT> x_dims_array(max_dim);
+  std::vector<ShapeT> e_dims_array(max_dim);
+  std::vector<ShapeT> out_dims_array(max_dim);
   // Only need to broadcast dimensions other than the 0th dimension.
   phi::funcs::GetBroadcastDimsArrays(common::make_ddim(x_dims2),
                                      common::make_ddim(e_dims2),
