Skip to content
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

Jump to bottom

Fix gather_op by adding OutOfRangeCheck for param[Index], test=develop #34096

Merged
merged 2 commits into from
Jul 13, 2021
Merged

Fix gather_op by adding OutOfRangeCheck for param[Index], test=develop #34096

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
aa83a1b
Fix gather_op by adding OurOfRangeCheck for param[Index], test=develop
haohongxiang Jul 12, 2021
3bba806
Code Optimization, test=develop
haohongxiang Jul 13, 2021
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
27 changes: 23 additions & 4 deletions paddle/fluid/operators/gather.cu.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,13 +30,20 @@ using platform::DeviceContext;

template <typename T, typename IndexT = int>
__global__ void GatherCUDAKernel(const T* params, const IndexT* indices,
T* output, size_t index_size,
size_t slice_size) {
T* output, size_t input_size,
size_t index_size, size_t slice_size) {
CUDA_KERNEL_LOOP(i, index_size * slice_size) {
int indices_i = i / slice_size;
int slice_i = i - indices_i * slice_size; // offset inside the slice
IndexT gather_i = indices[indices_i];
IndexT params_i = gather_i * slice_size + slice_i;
PADDLE_ENFORCE(
gather_i >= 0 && gather_i < input_size,
"The index is out of bounds, "
"please check whether the dimensions of index and "
"input meet the requirements. It should "
"be less than [%d] and greater than or equal to 0, but received [%d]",
input_size, gather_i);
*(output + i) = *(params + params_i);
}
}
Expand All @@ -58,7 +65,7 @@ __global__ void GatherNdCUDAKernel(const T* input, const int* input_dims,
"The index is out of bounds, "
"please check whether the dimensions of index and "
"input meet the requirements. It should "
"be less than [%d] and greater or equal to 0, but received [%d]",
"be less than [%d] and greater than or equal to 0, but received [%d]",
input_dims[j], index_value);
gather_i += (index_value * temp);
temp *= input_dims[j];
Expand Down Expand Up @@ -91,6 +98,7 @@ void GPUGather(const platform::DeviceContext& ctx, const Tensor& src,
" the second dimension should be 1."));
}

// index size
int index_size = index.dims()[0];

auto src_dims = src.dims();
Expand All @@ -100,6 +108,8 @@ void GPUGather(const platform::DeviceContext& ctx, const Tensor& src,
// slice size
int slice_size = 1;
for (int i = 1; i < src_dims.size(); ++i) slice_size *= src_dims[i];
// input size
int input_size = src_dims[0] * slice_size;

const T* p_src = src.data<T>();
const IndexT* p_index = index.data<IndexT>();
Expand All @@ -112,7 +122,7 @@ void GPUGather(const platform::DeviceContext& ctx, const Tensor& src,
GatherCUDAKernel<T, IndexT><<<
grid, block, 0,
reinterpret_cast<const platform::CUDADeviceContext&>(ctx).stream()>>>(
p_src, p_index, p_output, index_size, slice_size);
p_src, p_index, p_output, input_size, index_size, slice_size);
}

template <typename DeviceContext, typename T, typename IndexT = int>
Expand Down Expand Up @@ -177,6 +187,15 @@ __global__ void GatherGPUKernel(const T* input, const U* index, T* out,
int next_idx = idx - outer_size * inner_dim_index;
int index_dim_index = next_idx / outer_dim_size;
int index_val = index[index_dim_index];

PADDLE_ENFORCE(
index_val >= 0 && index_val < input_index_dim_size,
"The index is out of bounds, "
"please check whether the dimensions of index and "
"input meet the requirements. It should "
"be less than [%d] and greater than or equal to 0, but received [%d]",
input_index_dim_size, index_val);

int out_dim_index = next_idx - outer_dim_size * index_dim_index;
int input_index =
inner_dim_index * (outer_dim_size * input_index_dim_size) +
Expand Down
22 changes: 21 additions & 1 deletion paddle/fluid/operators/gather.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,11 +67,25 @@ void CPUGather(const platform::DeviceContext& ctx, const Tensor& src,
// slice size
int slice_size = 1;
for (int i = 1; i < src_dims.size(); ++i) slice_size *= src_dims[i];
// input size
int input_size = src_dims[0] * slice_size;

const size_t slice_bytes = slice_size * sizeof(T);

for (int64_t i = 0; i < index_size; ++i) {
IndexT index_ = p_index[i];
PADDLE_ENFORCE_LT(p_index[i], input_size,
platform::errors::OutOfRange(
"The element of Index must be less than the size of "
"input dim size of axis which is %d, but received "
"index element which is %d in the %d index.",
input_size, p_index[i], i));
PADDLE_ENFORCE_GE(p_index[i], 0UL,
platform::errors::OutOfRange(
"The element of Index must be greater than or equal "
"to 0, but received index element which is %d in the "
"%d index.",
p_index[i], i));
memcpy(p_output + i * slice_size, p_src + index_ * slice_size, slice_bytes);
}
}
Expand Down Expand Up @@ -141,11 +155,17 @@ void GatherV2Function(const Tensor* input, const Tensor* index, int axis,
int input_index_dim_size = input_dim[axis_index];
for (int i = 0; i < index_size; i++) {
PADDLE_ENFORCE_LT(index_data[i], input_index_dim_size,
platform::errors::InvalidArgument(
platform::errors::OutOfRange(
"The element of Index must be less than the size of "
"input dim size of axis which is %d, but received "
"index element which is %d in the %d index.",
input_index_dim_size, index_data[i], i));
PADDLE_ENFORCE_GE(index_data[i], 0UL,
platform::errors::OutOfRange(
"The element of Index must be greater than or equal "
"to 0, but received index element which is %d in the "
"%d index.",
index_data[i], i));
}

int inner_dim_size = 1;
Expand Down