Fix CPU segfault in roi_align when a ROI is non-finite (NaN or ±inf)

[ousamabenyounes]

Summary

roi_align on CPU segfaults the whole Python process when a ROI contains a non-finite value (NaN or ±inf) in its batch index or coordinates. The bilinear bounds check if (y < -1.0 || y > height || ...) returns false for every NaN comparison, so a malformed ROI slips through, gets cast to int (garbage index), and reads out of bounds. ±inf reaches the same cast via roi_end - roi_start = inf - inf = NaN.

What changed vs. the first revision

Two points raised in review, both fixed:

1. std::isnan → std::isfinite. The isnan-only guard let ±inf ROIs through: roi_width = inf - inf = NaN, unclamped when aligned=True, then the same non-finite cast → segfault (and ceil(NaN) with sampling_ratio=-1 also feeds a garbage pre_calc.resize). std::isfinite on the five raw ROI values covers NaN and ±inf at zero extra cost.
2. Guard moved above the batch-index cast. int roi_batch_ind = offset_rois[0]; runs before the old guard and is itself UB on a non-finite value — a UBSan float-cast-overflow build (which the issue reporter runs) flags it even though it doesn't crash on its own. The guard now sits on the raw ROI immediately after offset_rois is taken, before any cast or coordinate math, in both the forward and backward kernels.

A non-finite ROI is skipped, leaving zeros in the output/gradient — the output buffers are zero-initialised (new_zeros), consistent with how the kernel already treats out-of-bounds sample points as empty.

Test verification (RED → GREEN)

test/test_ops.py::TestRoIAlign::test_nonfinite_rois_no_segfault — parametrized over {NaN, +inf, -inf} × column {0..4} × dtype{f32,f64,f16} × sampling_ratio{1,-1} × aligned{T,F} (180 cases), asserting finite, all-zero output and a finite input gradient.

RED — built on the previous isnan-only kernel, the new ±inf cases crash:

$ python -m pytest test/test_ops.py::TestRoIAlign::test_nonfinite_rois_no_segfault -q
collected 180 items
test/test_ops.py ............Fatal Python error: Segmentation fault
...
Segmentation fault (core dumped)   # exit 139

GREEN — with the isfinite fix:

$ python -m pytest test/test_ops.py::TestRoIAlign::test_nonfinite_rois_no_segfault -q
180 passed in 0.60s

No regression in the rest of the op tests:

$ python -m pytest test/test_ops.py::TestRoIAlign -q
235 passed, 199 skipped

Built CPU-only from source (pip install -e . --no-build-isolation), torch 2.13.0+cpu. clang-format and flake8/ufmt clean on the touched files.

Fixes #9273

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/vision/9588

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.