Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
May 7, 2025
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7a5d4f6
[PHI] Preliminary fix for elementwise broadcast int32 shape overflow
Enigmatisms May 6, 2025
6c8ca0c
[PHI] Preliminary fix for elementwise broadcast int32 shape overflow
Enigmatisms May 7, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions paddle/fluid/operators/elementwise/elementwise_op.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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);
Expand Down
32 changes: 16 additions & 16 deletions paddle/phi/infermeta/binary.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down Expand Up @@ -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(),
Expand Down Expand Up @@ -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 =
Expand All @@ -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);
Expand Down Expand Up @@ -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(
Expand All @@ -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(),
Expand Down
28 changes: 14 additions & 14 deletions paddle/phi/infermeta/multiary.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down Expand Up @@ -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),
Expand Down
6 changes: 3 additions & 3 deletions paddle/phi/kernels/funcs/common_infer_shape_functions.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down
7 changes: 4 additions & 3 deletions paddle/phi/kernels/funcs/common_shape.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
20 changes: 11 additions & 9 deletions paddle/phi/kernels/funcs/elementwise_base.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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());
Expand Down Expand Up @@ -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(),
Expand Down
14 changes: 8 additions & 6 deletions paddle/phi/kernels/funcs/elementwise_utils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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];
Expand All @@ -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]) {
Expand Down
6 changes: 3 additions & 3 deletions paddle/phi/kernels/fusion/xpu/add_layernorm_xpu_kernel.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()) {
Expand All @@ -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,
Expand Down
21 changes: 11 additions & 10 deletions paddle/phi/kernels/impl/graph_message_passing_impl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down
Loading