From 64b9065dddb93978c7681cf5fab31436dacf7b82 Mon Sep 17 00:00:00 2001 From: Haohongxiang <86215757+haohongxiang@users.noreply.github.com> Date: Tue, 13 Jul 2021 00:26:53 -0500 Subject: [PATCH] Fix gather_op by adding OurOfRangeCheck for param[Index], test=develop (#34096) * Fix gather_op by adding OurOfRangeCheck for param[Index] * Code Optimization --- paddle/fluid/operators/gather.cu.h | 27 +++++++++++++++++++++++---- paddle/fluid/operators/gather.h | 22 +++++++++++++++++++++- 2 files changed, 44 insertions(+), 5 deletions(-) diff --git a/paddle/fluid/operators/gather.cu.h b/paddle/fluid/operators/gather.cu.h index 95cb428abdf34..6469307bc5652 100644 --- a/paddle/fluid/operators/gather.cu.h +++ b/paddle/fluid/operators/gather.cu.h @@ -30,13 +30,20 @@ using platform::DeviceContext; template __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); } } @@ -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]; @@ -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(); @@ -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(); const IndexT* p_index = index.data(); @@ -112,7 +122,7 @@ void GPUGather(const platform::DeviceContext& ctx, const Tensor& src, GatherCUDAKernel<<< grid, block, 0, reinterpret_cast(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 @@ -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) + diff --git a/paddle/fluid/operators/gather.h b/paddle/fluid/operators/gather.h index 8deab709220d7..b7fa4022882b7 100644 --- a/paddle/fluid/operators/gather.h +++ b/paddle/fluid/operators/gather.h @@ -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); } } @@ -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;