From abf6a69f38b33a011124578a87380d16d81f9bd2 Mon Sep 17 00:00:00 2001 From: nileshpatil6 Date: Sun, 5 Jul 2026 23:48:36 +0530 Subject: [PATCH] fix(roi_align): treat NaN ROI coordinates as out of bounds on CPU roi_align on CPU segfaults when a ROI contains NaN coordinates. The bounds check in pre_calc_for_bilinear_interpolate and bilinear_interpolate_gradient compares y/x against -1 and height/width, but every comparison against NaN evaluates to false, so a NaN box silently falls through the check instead of being treated as empty. The NaN value then gets truncated to an int, which is undefined behavior in C++ and in practice produces a garbage index that is used to read the input feature map out of bounds. This adds explicit isnan checks alongside the existing range checks so NaN coordinates are treated the same as any other out-of-bounds ROI (zero weights, no read). Added a regression test that reproduces the crash with a NaN ROI on CPU. --- test/test_ops.py | 25 +++++++++++++++++++ torchvision/csrc/ops/cpu/roi_align_common.h | 9 ++++++- torchvision/csrc/ops/cpu/roi_align_kernel.cpp | 9 ++++++- 3 files changed, 41 insertions(+), 2 deletions(-) diff --git a/test/test_ops.py b/test/test_ops.py index 8f8756fe675..8b2aa32aba9 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -474,6 +474,31 @@ def expected_fn( def test_boxes_shape(self): self._helper_boxes_shape(ops.roi_align) + @pytest.mark.parametrize("aligned", (True, False)) + @pytest.mark.parametrize("sampling_ratio", (-1, 1)) + def test_nan_roi_coordinates_do_not_crash(self, aligned, sampling_ratio): + # Non-regression test: a ROI with NaN coordinates used to bypass the + # out-of-bounds check in the CPU kernel (every comparison against NaN + # is False), get truncated to an undefined int index, and read out of + # bounds from the input feature map (can segfault under ASAN). + x = torch.zeros(2, 10, 15, 19, dtype=torch.float32) + rois = torch.tensor( + [[0.0, float("nan"), 0.0, float("nan"), 1.0]], + dtype=torch.float32, + ) + + y = ops.roi_align( + x, + rois, + output_size=(7, 7), + spatial_scale=0.1, + sampling_ratio=sampling_ratio, + aligned=aligned, + ) + + assert y.shape == (1, 10, 7, 7) + assert torch.isfinite(y).all() or torch.isnan(y).all() + @pytest.mark.parametrize("aligned", (True, False)) @pytest.mark.parametrize("device", cpu_and_cuda_and_mps()) @pytest.mark.parametrize("x_dtype", (torch.float16, torch.float32, torch.float64)) # , ids=str) diff --git a/torchvision/csrc/ops/cpu/roi_align_common.h b/torchvision/csrc/ops/cpu/roi_align_common.h index e10c67b5b79..1c7bc318846 100644 --- a/torchvision/csrc/ops/cpu/roi_align_common.h +++ b/torchvision/csrc/ops/cpu/roi_align_common.h @@ -1,6 +1,7 @@ #pragma once #include +#include namespace vision { namespace ops { @@ -57,7 +58,13 @@ 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) { + // NaN coordinates must be treated as out of bounds too, since every + // comparison against NaN is false and would otherwise fall through + // this check, get truncated to an undefined int value below, and + // produce an out-of-bounds pos1/pos2/pos3/pos4 index. + if (std::isnan(static_cast(x)) || + std::isnan(static_cast(y)) || 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 e0185da45df..ac1cb00cef8 100644 --- a/torchvision/csrc/ops/cpu/roi_align_kernel.cpp +++ b/torchvision/csrc/ops/cpu/roi_align_kernel.cpp @@ -1,6 +1,8 @@ #include #include +#include + #include "./roi_align_common.h" namespace vision { @@ -123,7 +125,12 @@ 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) { + // NaN coordinates must be treated as out of bounds too, since every + // comparison against NaN is false and would otherwise fall through this + // check and get truncated to an undefined int value below. + if (std::isnan(static_cast(x)) || + std::isnan(static_cast(y)) || 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;