Skip to content

[API] fix paddle.diag with big tensor #72638

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
lshpku merged 1 commit into from
May 19, 2025
Merged

[API] fix paddle.diag with big tensor #72638

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
4 changes: 2 additions & 2 deletions paddle/phi/infermeta/unary.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -931,10 +931,10 @@ void DiagInferMeta(const MetaTensor& x,
auto x_dims = x.dims();

if (x_dims.size() <= 1) {
int64_t size_ = (x_dims.size() == 1UL ? x_dims[0] : 1) + std::abs(offset);
int64_t size_ = (x_dims.size() == 1ULL ? x_dims[0] : 1) + std::abs(offset);

@lshpku lshpku May 19, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这里写1ULL和直接写1有什么区别

@ZhangX-21 ZhangX-21 May 19, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这个没有什么区别，只是修改支持unsigned long long int的取值。

out->set_dims({size_, size_});
out->set_dtype(x.dtype());
} else if (x_dims.size() == 2UL) {
} else if (x_dims.size() == 2ULL) {
int64_t size_ = 0;
if (offset >= 0) {
// Note(LutaoChu): Do not use std::min here, otherwise the calculation
Expand Down
20 changes: 10 additions & 10 deletions paddle/phi/kernels/cpu/diag_grad_kernel.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,28 +33,28 @@ void DiagGradKernel(const Context& dev_ctx,
auto dout_dims = out_grad.dims();

if (dx_dims.size() <= 1) {
auto dx_length = (dx_dims.size() == 1 ? dx_dims[0] : int64_t(1));
int dx_stride = 1;
int64_t dx_length = (dx_dims.size() == 1 ? dx_dims[0] : int64_t(1));
int64_t dx_stride = 1;

auto dout_stride_0 = phi::funcs::ComputeStride(0, dout_dims);
auto dout_stride_1 = phi::funcs::ComputeStride(1, dout_dims);
int64_t dout_stride_0 = phi::funcs::ComputeStride(0, dout_dims);
int64_t dout_stride_1 = phi::funcs::ComputeStride(1, dout_dims);
dout_data +=
(offset >= 0 ? offset * dout_stride_1 : -offset * dout_stride_0);

for (int i = 0; i < dx_length; i++) {
for (int64_t i = 0; i < dx_length; i++) {
dx_data[i * dx_stride] = dout_data[i * (dout_stride_0 + dout_stride_1)];
}
} else {
phi::funcs::SetConstant<Context, T> set_padding_value;
set_padding_value(dev_ctx, x_grad, static_cast<T>(0));

int dx_stride_0 = phi::funcs::ComputeStride(0, dx_dims);
int dx_stride_1 = phi::funcs::ComputeStride(1, dx_dims);
auto dout_stride_0 = phi::funcs::ComputeStride(0, dout_dims);
int64_t dx_stride_0 = phi::funcs::ComputeStride(0, dx_dims);
int64_t dx_stride_1 = phi::funcs::ComputeStride(1, dx_dims);
int64_t dout_stride_0 = phi::funcs::ComputeStride(0, dout_dims);
dx_data += (offset >= 0 ? offset * dx_stride_1 : -offset * dx_stride_0);

auto dout_length = dout_dims[0];
for (int i = 0; i < dout_length; i++) {
int64_t dout_length = dout_dims[0];
for (int64_t i = 0; i < dout_length; i++) {
dx_data[i * (dx_stride_0 + dx_stride_1)] = dout_data[i * dout_stride_0];
}
}
Expand Down
16 changes: 8 additions & 8 deletions paddle/phi/kernels/cpu/diag_kernel.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,22 +37,22 @@ void DiagKernel(const Context& dev_ctx,
phi::funcs::SetConstant<Context, T> set_padding_value;
set_padding_value(dev_ctx, out, static_cast<T>(padding_value));

auto x_length = (x_dims.size() == 1UL ? x_dims[0] : int64_t(1));
const int& x_stride = 1;
int64_t x_length = (x_dims.size() == 1ULL ? x_dims[0] : int64_t(1));
const int64_t x_stride = 1;

auto out_stride_0 = phi::funcs::ComputeStride(0, out_dims);
auto out_stride_1 = phi::funcs::ComputeStride(1, out_dims);
int64_t out_stride_0 = phi::funcs::ComputeStride(0, out_dims);
int64_t out_stride_1 = phi::funcs::ComputeStride(1, out_dims);
out_data += (offset >= 0 ? offset * out_stride_1 : -offset * out_stride_0);

for (i = 0; i < x_length; i++) {
out_data[i * (out_stride_0 + out_stride_1)] = x_data[i * x_stride];
}
} else {
auto out_length = out_dims[0];
const int& x_stride_0 = phi::funcs::ComputeStride(0, x_dims);
const int& x_stride_1 = phi::funcs::ComputeStride(1, x_dims);
int64_t out_length = out_dims[0];
const int64_t x_stride_0 = phi::funcs::ComputeStride(0, x_dims);
const int64_t x_stride_1 = phi::funcs::ComputeStride(1, x_dims);

auto out_stride_0 = phi::funcs::ComputeStride(0, out_dims);
int64_t out_stride_0 = phi::funcs::ComputeStride(0, out_dims);
x_data += (offset >= 0 ? offset * x_stride_1 : -offset * x_stride_0);
for (i = 0; i < out_length; i++) {
out_data[i * out_stride_0] = x_data[i * (x_stride_0 + x_stride_1)];
Expand Down
4 changes: 2 additions & 2 deletions paddle/phi/kernels/funcs/diag_functor.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,8 +25,8 @@
namespace phi {
namespace funcs {

inline int ComputeStride(int axis, phi::DDim dims) {
int size = 1;
inline int64_t ComputeStride(int axis, phi::DDim dims) {
int64_t size = 1;
for (int i = axis + 1; i < dims.size(); i++) {
size *= dims[i];
}
Expand Down
48 changes: 24 additions & 24 deletions paddle/phi/kernels/gpu/diag_grad_kernel.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,14 +24,13 @@ namespace phi {
template <typename T>
__global__ void ExtractDiagonalKernel(const T* dout,
T* dx,
std::ptrdiff_t start,
std::ptrdiff_t dx_length,
const std::ptrdiff_t sumStride,
const std::ptrdiff_t xStride) {
for (std::ptrdiff_t idx = blockIdx.x * blockDim.x + threadIdx.x;
idx < dx_length;
int64_t start,
int64_t dx_length,
const int64_t sumStride,
const int64_t xStride) {
for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < dx_length;
idx += gridDim.x * blockDim.x) {
const std::ptrdiff_t outOffset = start + sumStride * idx;
const int64_t outOffset = start + sumStride * idx;
dx[xStride * idx] = dout[outOffset];
}
}
Expand All @@ -40,13 +39,13 @@ __global__ void ExtractDiagonalKernel(const T* dout,
template <typename T>
__global__ void PasteDiagonalKernel(const T* dout,
T* dx,
std::ptrdiff_t start,
std::ptrdiff_t size,
const std::ptrdiff_t sumStride,
const std::ptrdiff_t outStride) {
for (std::ptrdiff_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size;
int64_t start,
int64_t size,
const int64_t sumStride,
const int64_t outStride) {
for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size;
idx += gridDim.x * blockDim.x) {
std::ptrdiff_t xOffset = start + sumStride * idx;
int64_t xOffset = start + sumStride * idx;
dx[xOffset] = dout[outStride * idx];
}
}
Expand Down Expand Up @@ -74,14 +73,14 @@ void DiagGradKernel(const Context& dev_ctx,
};

if (dx_dims.size() <= 1) {
auto dx_length = (dx_dims.size() == 1 ? dx_dims[0] : int64_t(1));
auto size = (offset > 0) ? dx_length + offset : dx_length - offset;
int dx_stride = 1;
int64_t dx_length = (dx_dims.size() == 1ULL ? dx_dims[0] : 1ULL);
int64_t size = (offset > 0ULL) ? dx_length + offset : dx_length - offset;
int64_t dx_stride = 1;
if (size > 0) {
auto dout_stride_0 = phi::funcs::ComputeStride(0, dout_dims);
auto dout_stride_1 = phi::funcs::ComputeStride(1, dout_dims);
auto start =
(offset >= 0 ? offset * dout_stride_1 : -offset * dout_stride_0);
int64_t dout_stride_0 = phi::funcs::ComputeStride(0, dout_dims);
int64_t dout_stride_1 = phi::funcs::ComputeStride(1, dout_dims);
int64_t start =
(offset >= 0ULL ? offset * dout_stride_1 : -offset * dout_stride_0);

std::tuple<int64_t, int64_t> block_grid_size = GetBlockGridSize(size);
ExtractDiagonalKernel<T>
Expand All @@ -99,8 +98,8 @@ void DiagGradKernel(const Context& dev_ctx,
phi::funcs::SetConstant<Context, T> set_padding_value;
set_padding_value(dev_ctx, x_grad, static_cast<T>(0));

int dx_stride_0 = phi::funcs::ComputeStride(0, dx_dims);
int dx_stride_1 = phi::funcs::ComputeStride(1, dx_dims);
int64_t dx_stride_0 = phi::funcs::ComputeStride(0, dx_dims);
int64_t dx_stride_1 = phi::funcs::ComputeStride(1, dx_dims);
int64_t size;
if (offset > 0) {
size = std::min(dx_dims[0], dx_dims[1] - offset);
Expand All @@ -109,8 +108,9 @@ void DiagGradKernel(const Context& dev_ctx,
}

if (size > 0) {
auto start = (offset >= 0 ? offset * dx_stride_1 : -offset * dx_stride_0);
auto dout_stride_0 = phi::funcs::ComputeStride(0, dout_dims);
int64_t start =
(offset >= 0 ? offset * dx_stride_1 : -offset * dx_stride_0);
int64_t dout_stride_0 = phi::funcs::ComputeStride(0, dout_dims);
std::tuple<int64_t, int64_t> block_grid_size = GetBlockGridSize(size);
PasteDiagonalKernel<T><<<std::get<1>(block_grid_size),
std::get<0>(block_grid_size),
Expand Down
47 changes: 23 additions & 24 deletions paddle/phi/kernels/gpu/diag_kernel.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,13 +28,13 @@ namespace phi {
template <typename T>
__global__ void ExtractDiagonalKernel(T* out,
const T* x,
std::ptrdiff_t start,
std::ptrdiff_t size,
const std::ptrdiff_t sumStride,
const std::ptrdiff_t outStride) {
for (std::ptrdiff_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size;
int64_t start,
int64_t size,
const int64_t sumStride,
const int64_t outStride) {
for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size;
idx += gridDim.x * blockDim.x) {
const std::ptrdiff_t xOffset = start + sumStride * idx;
const int64_t xOffset = start + sumStride * idx;
out[outStride * idx] = x[xOffset];
}
}
Expand All @@ -43,14 +43,13 @@ __global__ void ExtractDiagonalKernel(T* out,
template <typename T>
__global__ void PasteDiagonalKernel(T* out,
const T* x,
std::ptrdiff_t start,
std::ptrdiff_t x_length,
const std::ptrdiff_t sumStride,
const std::ptrdiff_t xStride) {
for (std::ptrdiff_t idx = blockIdx.x * blockDim.x + threadIdx.x;
idx < x_length;
int64_t start,
int64_t x_length,
const int64_t sumStride,
const int64_t xStride) {
for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < x_length;

@lshpku lshpku May 19, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这里blockIdx.x * blockDim.x会超过uint32吗？

@ZhangX-21 ZhangX-21 May 19, 2025 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

blockIdx.x * blockDim.x，取值范围不超过uint32, 但是idx < x_length在大Tensor场景下，所以需要采用int64_t的取值范围。

idx += gridDim.x * blockDim.x) {
const std::ptrdiff_t outOffset = start + sumStride * idx;
const int64_t outOffset = start + sumStride * idx;
out[outOffset] = x[xStride * idx];
}
}
Expand Down Expand Up @@ -81,13 +80,13 @@ void DiagKernel(const Context& dev_ctx,
phi::funcs::SetConstant<Context, T> set_padding_value;
set_padding_value(dev_ctx, out, static_cast<T>(padding_value));

auto x_length = (x_dims.size() == 1UL ? x_dims[0] : int64_t(1));
auto size = (offset > 0) ? x_length + offset : x_length - offset;
const int& x_stride = 1;
int64_t x_length = (x_dims.size() == 1ULL ? x_dims[0] : int64_t(1));
int64_t size = (offset > 0) ? x_length + offset : x_length - offset;
const int64_t x_stride = 1;
if (size > 0) {
const auto& out_stride_0 = phi::funcs::ComputeStride(0, out_dims);
const auto& out_stride_1 = phi::funcs::ComputeStride(1, out_dims);
auto start =
const int64_t out_stride_0 = phi::funcs::ComputeStride(0, out_dims);
const int64_t out_stride_1 = phi::funcs::ComputeStride(1, out_dims);
int64_t start =
(offset >= 0 ? offset * out_stride_1 : -offset * out_stride_0);

std::tuple<int64_t, int64_t> block_grid_size = GetBlockGridSize(size);
Expand All @@ -103,8 +102,8 @@ void DiagKernel(const Context& dev_ctx,
x_stride);
}
} else {
const int& x_stride_0 = phi::funcs::ComputeStride(0, x_dims);
const int& x_stride_1 = phi::funcs::ComputeStride(1, x_dims);
const int64_t x_stride_0 = phi::funcs::ComputeStride(0, x_dims);
const int64_t x_stride_1 = phi::funcs::ComputeStride(1, x_dims);

int64_t size;
if (offset > 0) {
Expand All @@ -114,11 +113,11 @@ void DiagKernel(const Context& dev_ctx,
}

if (size > 0) {
auto start = (offset >= 0 ? offset * x_stride_1 : -offset * x_stride_0);
const auto& out_stride_0 = phi::funcs::ComputeStride(0, out_dims);
int64_t start =
(offset >= 0 ? offset * x_stride_1 : -offset * x_stride_0);
const int64_t out_stride_0 = phi::funcs::ComputeStride(0, out_dims);

std::tuple<int64_t, int64_t> block_grid_size = GetBlockGridSize(size);

ExtractDiagonalKernel<T><<<std::get<1>(block_grid_size),
std::get<0>(block_grid_size),
0,
Expand Down