From 9549d13119b833fcb1bed0f779cfa2576a58d466 Mon Sep 17 00:00:00 2001 From: GEAK RDNA Smoke Date: Wed, 26 Aug 2026 04:14:10 -0400 Subject: [PATCH] Guard roi_align/ps_roi_align bilinear sampling against non-finite boxes A non-finite box coordinate makes every comparison in the "inverse elements are out of feature map boundary" check false, so the sample point is neither rejected by the empty branch nor clamped by the y <= 0 / x <= 0 clamps below it. The index is then computed from (int)y with y == NaN, which is undefined and yields INT_MIN on x86, so pos1..pos4 become large negative offsets and the forward kernels read far outside the input buffer: x = torch.zeros(1, 1, 4, 4) b = torch.tensor([[0., nan, nan, nan, nan]]) torchvision.ops.roi_align(x, b, (1, 1), 1.0, 1, True) # SIGSEGV NaN is not the only way in, so filtering the boxes up front would not be enough: [inf, -inf, inf, -inf] makes roi_width NaN via inf - inf, and [-3.4e38, -3.4e38, 3.4e38, 3.4e38] -- entirely finite -- overflows roi_end - roi_start to inf and then computes 0 * inf. Write the guard as !(in range) instead of (out of range) so that a NaN coordinate takes the empty branch. For finite inputs the two forms are exactly complementary; comparing the PreCalc output of both over 20,055,768 sample points gives zero bitwise differences. The backward kernels reach the same conversion but already gate every write on x_low >= 0 && x_high >= 0 && y_low >= 0 && y_high >= 0, which INT_MIN fails, so no memory is corrupted -- this is an out-of-bounds read only. It is also not observable: all four interpolation weights are NaN whenever the coordinate is, so whatever is read is multiplied by NaN. CUDA/MPS did not crash, since float-to-int conversion returns 0 rather than INT_MIN there, but the guard is updated at all ten sites so the four copies stay identical and CPU and GPU agree on the result. Co-Authored-By: Claude Opus 5 (1M context) --- test/test_ops.py | 29 +++++++++++++++++++ .../csrc/ops/cpu/ps_roi_align_kernel.cpp | 10 +++++-- torchvision/csrc/ops/cpu/roi_align_common.h | 5 +++- torchvision/csrc/ops/cpu/roi_align_kernel.cpp | 5 +++- .../csrc/ops/cuda/ps_roi_align_kernel.cu | 10 +++++-- torchvision/csrc/ops/cuda/roi_align_kernel.cu | 10 +++++-- torchvision/csrc/ops/mps/mps_stable_kernels.h | 10 +++++-- 7 files changed, 69 insertions(+), 10 deletions(-) diff --git a/test/test_ops.py b/test/test_ops.py index e8844b47422..9f9c2eb9da3 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -162,6 +162,35 @@ def test_forward(self, device, contiguous, x_dtype, rois_dtype=None, determinist torch.testing.assert_close(gt_y.to(y), y, rtol=tol, atol=tol) + @pytest.mark.parametrize("device", cpu_and_cuda_and_mps()) + @pytest.mark.parametrize( + "coords", + ( + pytest.param([float("nan")] * 4, id="nan"), + pytest.param([float("inf"), float("-inf"), float("inf"), float("-inf")], id="inf"), + # Finite, but roi_end - roi_start overflows to inf, and 0 * inf is nan. + pytest.param([-3.4e38, -3.4e38, 3.4e38, 3.4e38], id="overflow"), + ), + ) + def test_non_finite_boxes(self, device, coords): + # A non-finite coordinate makes every comparison in the bounds check + # false, so the sample point used to be neither rejected nor clamped, + # and the index was then computed from (int)nan -- undefined behaviour, + # INT_MIN on x86 -- reading far outside the input buffer. + pool_size = 5 + # n_channels % (pool_size ** 2) == 0 required for PS operations. + n_channels = 2 * (pool_size**2) + x = torch.rand(1, n_channels, 10, 10, device=device, requires_grad=True) + rois = torch.tensor([[0.0] + coords], dtype=torch.float32, device=device) + + y = self.fn(x, rois, pool_size, pool_size, spatial_scale=1, sampling_ratio=1) + assert y.shape[0] == rois.shape[0] + assert not y.isinf().any() + + y.sum().backward() + assert x.grad.shape == x.shape + assert not x.grad.isinf().any() + @pytest.mark.parametrize("device", cpu_and_cuda()) def test_is_leaf_node(self, device): op_obj = self.make_obj(wrap=True).to(device=device) diff --git a/torchvision/csrc/ops/cpu/ps_roi_align_kernel.cpp b/torchvision/csrc/ops/cpu/ps_roi_align_kernel.cpp index dc1f380f36c..c2f12d55bb5 100644 --- a/torchvision/csrc/ops/cpu/ps_roi_align_kernel.cpp +++ b/torchvision/csrc/ops/cpu/ps_roi_align_kernel.cpp @@ -23,7 +23,10 @@ T bilinear_interpolate( T x, int index /* index for debug only*/) { // deal with cases that inverse elements are out of feature map boundary - if (y < -1.0 || y > height || x < -1.0 || x > width) { + // NB: written as !(in range) rather than (out of range) so that NaN + // coordinates, for which every comparison is false, take this branch + // instead of reaching the (int) casts below. + if (!(y >= -1.0 && y <= height && x >= -1.0 && x <= width)) { // empty return 0; } @@ -166,7 +169,10 @@ void bilinear_interpolate_gradient( int& y_high, int index /* index for debug only*/) { // deal with cases that inverse elements are out of feature map boundary - if (y < -1.0 || y > height || x < -1.0 || x > width) { + // NB: written as !(in range) rather than (out of range) so that NaN + // coordinates, for which every comparison is false, take this branch + // instead of reaching the (int) casts below. + if (!(y >= -1.0 && y <= height && x >= -1.0 && x <= width)) { // empty w1 = w2 = w3 = w4 = 0.; x_low = x_high = y_low = y_high = -1; diff --git a/torchvision/csrc/ops/cpu/roi_align_common.h b/torchvision/csrc/ops/cpu/roi_align_common.h index 03e5084357f..173b2ac6314 100644 --- a/torchvision/csrc/ops/cpu/roi_align_common.h +++ b/torchvision/csrc/ops/cpu/roi_align_common.h @@ -57,7 +57,10 @@ void pre_calc_for_bilinear_interpolate( T x = xx; T y = yy; // deal with: inverse elements are out of feature map boundary - if (y < -1.0 || y > height || x < -1.0 || x > width) { + // NB: written as !(in range) rather than (out of range) so that NaN + // coordinates, for which every comparison is false, take this branch + // instead of reaching the (int) casts below. + if (!(y >= -1.0 && y <= height && x >= -1.0 && x <= width)) { // empty PreCalc pc; pc.pos1 = 0; diff --git a/torchvision/csrc/ops/cpu/roi_align_kernel.cpp b/torchvision/csrc/ops/cpu/roi_align_kernel.cpp index 26dc7423853..d35cd3f550b 100644 --- a/torchvision/csrc/ops/cpu/roi_align_kernel.cpp +++ b/torchvision/csrc/ops/cpu/roi_align_kernel.cpp @@ -130,7 +130,10 @@ void bilinear_interpolate_gradient( int& y_high, int index /* index for debug only*/) { // deal with cases that inverse elements are out of feature map boundary - if (y < -1.0 || y > height || x < -1.0 || x > width) { + // NB: written as !(in range) rather than (out of range) so that NaN + // coordinates, for which every comparison is false, take this branch + // instead of reaching the (int) casts below. + if (!(y >= -1.0 && y <= height && x >= -1.0 && x <= width)) { // empty w1 = w2 = w3 = w4 = 0.; x_low = x_high = y_low = y_high = -1; diff --git a/torchvision/csrc/ops/cuda/ps_roi_align_kernel.cu b/torchvision/csrc/ops/cuda/ps_roi_align_kernel.cu index 74f7e9d6004..97313250cfb 100644 --- a/torchvision/csrc/ops/cuda/ps_roi_align_kernel.cu +++ b/torchvision/csrc/ops/cuda/ps_roi_align_kernel.cu @@ -32,7 +32,10 @@ __device__ T bilinear_interpolate( T x, int index /* index for debug only*/) { // deal with cases that inverse elements are out of feature map boundary - if (y < -1.0 || y > height || x < -1.0 || x > width) { + // NB: written as !(in range) rather than (out of range) so that NaN + // coordinates, for which every comparison is false, take this branch + // instead of reaching the (int) casts below. + if (!(y >= -1.0 && y <= height && x >= -1.0 && x <= width)) { // empty return 0; } @@ -167,7 +170,10 @@ __device__ void bilinear_interpolate_gradient( int& y_high, int index /* index for debug only*/) { // deal with cases that inverse elements are out of feature map boundary - if (y < -1.0 || y > height || x < -1.0 || x > width) { + // NB: written as !(in range) rather than (out of range) so that NaN + // coordinates, for which every comparison is false, take this branch + // instead of reaching the (int) casts below. + if (!(y >= -1.0 && y <= height && x >= -1.0 && x <= width)) { // empty w1 = w2 = w3 = w4 = 0.; x_low = x_high = y_low = y_high = -1; diff --git a/torchvision/csrc/ops/cuda/roi_align_kernel.cu b/torchvision/csrc/ops/cuda/roi_align_kernel.cu index 346b6e4f0e2..99a05c5951c 100644 --- a/torchvision/csrc/ops/cuda/roi_align_kernel.cu +++ b/torchvision/csrc/ops/cuda/roi_align_kernel.cu @@ -32,7 +32,10 @@ __device__ T bilinear_interpolate( T x, int index /* index for debug only*/) { // deal with cases that inverse elements are out of feature map boundary - if (y < -1.0 || y > height || x < -1.0 || x > width) { + // NB: written as !(in range) rather than (out of range) so that NaN + // coordinates, for which every comparison is false, take this branch + // instead of reaching the (int) casts below. + if (!(y >= -1.0 && y <= height && x >= -1.0 && x <= width)) { // empty return 0; } @@ -170,7 +173,10 @@ __device__ void bilinear_interpolate_gradient( int& y_high, int index /* index for debug only*/) { // deal with cases that inverse elements are out of feature map boundary - if (y < -1.0 || y > height || x < -1.0 || x > width) { + // NB: written as !(in range) rather than (out of range) so that NaN + // coordinates, for which every comparison is false, take this branch + // instead of reaching the (int) casts below. + if (!(y >= -1.0 && y <= height && x >= -1.0 && x <= width)) { // empty w1 = w2 = w3 = w4 = 0.; x_low = x_high = y_low = y_high = -1; diff --git a/torchvision/csrc/ops/mps/mps_stable_kernels.h b/torchvision/csrc/ops/mps/mps_stable_kernels.h index 9ced5a95fed..5156f2f924e 100644 --- a/torchvision/csrc/ops/mps/mps_stable_kernels.h +++ b/torchvision/csrc/ops/mps/mps_stable_kernels.h @@ -57,7 +57,10 @@ inline T bilinear_interpolate( T x, uint index /* index for debug only*/) { // deal with cases that inverse elements are out of feature map boundary - if (y < -1.0 || y > height || x < -1.0 || x > width) { + // NB: written as !(in range) rather than (out of range) so that NaN + // coordinates, for which every comparison is false, take this branch + // instead of reaching the (int) casts below. + if (!(y >= -1.0 && y <= height && x >= -1.0 && x <= width)) { // empty return 0; } @@ -164,7 +167,10 @@ inline void bilinear_interpolate_gradient( thread integer_t& y_high, uint index /* index for debug only*/) { // deal with cases that inverse elements are out of feature map boundary - if (y < -1.0 || y > height || x < -1.0 || x > width) { + // NB: written as !(in range) rather than (out of range) so that NaN + // coordinates, for which every comparison is false, take this branch + // instead of reaching the (int) casts below. + if (!(y >= -1.0 && y <= height && x >= -1.0 && x <= width)) { // empty w1 = w2 = w3 = w4 = 0.; x_low = x_high = y_low = y_high = -1;