Skip to content

Commit

Permalink
Fix gather_op by adding OurOfRangeCheck for param[Index], test=develop (
Browse files Browse the repository at this point in the history 
#34096)

* Fix gather_op by adding OurOfRangeCheck for param[Index]

* Code Optimization
  • Loading branch information
@haohongxiang
haohongxiang authored Jul 13, 2021
1 parent d8343f4 commit 64b9065
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 5 deletions.
27 changes: 23 additions & 4 deletions paddle/fluid/operators/gather.cu.h
View file
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
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

1 comment on commit 64b9065

@haohongxiang
Copy link
Contributor Author

Choose a reason for hiding this comment

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

Spelling mistake : the fixed contents should be "OutOfRangeCheck", not "OurOfRangeCheck".

Please sign in to comment.