diff --git a/paddle/phi/kernels/gpu/grid_sample_grad_kernel.cu b/paddle/phi/kernels/gpu/grid_sample_grad_kernel.cu index 6e8b12c4b1b900..2b6ceff59afa7f 100644 --- a/paddle/phi/kernels/gpu/grid_sample_grad_kernel.cu +++ b/paddle/phi/kernels/gpu/grid_sample_grad_kernel.cu @@ -121,16 +121,13 @@ ComputePositionsWithMask(T coord, coord = ClipIndexesWithMask(coord, size, &grad_clip); *grad_in = (*grad_in) * grad_clip; } else if (padding_mode == PaddingMode::reflect) { - if (align_corners) { - coord = ReflectIndexesWithMask(coord, 0, 2 * (size - 1), &grad_refl); - } else { - coord = ReflectIndexesWithMask(coord, -1, 2 * size - 1, &grad_refl); - } + coord = align_corners + ? ReflectIndexesWithMask(coord, 0, 2 * (size - 1), &grad_refl) + : ReflectIndexesWithMask(coord, -1, 2 * size - 1, &grad_refl); coord = ClipIndexesWithMask(coord, size, &grad_clip); *grad_in = (*grad_in) * grad_refl * grad_clip; } - - return coord; + return SafeDownGradeToIntRange(coord); } template diff --git a/paddle/phi/kernels/gpu/grid_sample_kernel.cu b/paddle/phi/kernels/gpu/grid_sample_kernel.cu index 3809ae7d5c3388..8499e371d10cf0 100644 --- a/paddle/phi/kernels/gpu/grid_sample_kernel.cu +++ b/paddle/phi/kernels/gpu/grid_sample_kernel.cu @@ -27,16 +27,13 @@ template static __forceinline__ __device__ T Unnormalize(T coord, int size, bool align_corners) { - if (align_corners) { - return ((coord + 1.f) / 2) * (size - 1); - } else { - return ((coord + 1.f) * size - 1) / 2; - } + return align_corners ? ((coord + 1.f) / 2) * (size - 1) + : ((coord + 1.f) * size - 1) / 2; } template static __forceinline__ __device__ T ClipIndexes(T in, int max_value) { - return min(static_cast(max_value), max(in, static_cast(0))); + return min(static_cast(max_value - 1), max(in, static_cast(0))); } template @@ -51,11 +48,7 @@ static __forceinline__ __device__ T ReflectIndexes(T in, in = fabs(in - min); T extra = fmod(in, span); int flips = static_cast(floor(in / span)); - if (flips % 2 == 0) { - return extra + min; - } else { - return span - extra + min; - } + return (flips & 1) ? span - extra + min : extra + min; // cond ? odd : even } template @@ -65,16 +58,13 @@ static __forceinline__ __device__ T ComputePositions(T coord, bool align_corners) { coord = Unnormalize(coord, size, align_corners); if (padding_mode == PaddingMode::border) { - coord = ClipIndexes(coord, size - 1); + coord = ClipIndexes(coord, size); } else if (padding_mode == PaddingMode::reflect) { - if (align_corners) { - coord = ReflectIndexes(coord, 0, 2 * (size - 1)); - } else { - coord = ReflectIndexes(coord, -1, 2 * size - 1); - } - coord = ClipIndexes(coord, size - 1); + coord = align_corners ? ReflectIndexes(coord, 0, 2 * (size - 1)) + : ReflectIndexes(coord, -1, 2 * size - 1); + coord = ClipIndexes(coord, size); } - return coord; + return SafeDownGradeToIntRange(coord); } template diff --git a/paddle/phi/kernels/gpu/grid_sample_utils.h b/paddle/phi/kernels/gpu/grid_sample_utils.h index bd5e859a59d1d2..415305efaa1057 100644 --- a/paddle/phi/kernels/gpu/grid_sample_utils.h +++ b/paddle/phi/kernels/gpu/grid_sample_utils.h @@ -14,6 +14,8 @@ #pragma once +#include + namespace phi { enum class Mode { @@ -21,6 +23,13 @@ enum class Mode { nearest, }; +template +__forceinline__ __device__ T SafeDownGradeToIntRange(T x) { + bool unsafe_cond = + x > INT_MAX - 1 || x < INT_MIN || !::isfinite(static_cast(x)); + return unsafe_cond ? static_cast(-100.0) : x; +} + enum class PaddingMode { zeros, border, reflect }; static __forceinline__ __device__ bool InBounds(int h, int w, int H, int W) {