Fix CPU segfault in roi_align when ROI contains NaN coordinates

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)

torchvision/csrc/ops/cpu/roi_align_common.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <ATen/ATen.h>
+#include <cmath>
 
 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<double>(x)) ||
+              std::isnan(static_cast<double>(y)) || y < -1.0 || y > height ||
+              x < -1.0 || x > width) {
             // empty
             PreCalc<T> pc;
             pc.pos1 = 0;

torchvision/csrc/ops/cpu/roi_align_kernel.cpp

@@ -1,6 +1,8 @@
 #include <ATen/ATen.h>
 #include <torch/library.h>
 
+#include <cmath>
+
 #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<double>(x)) ||
+      std::isnan(static_cast<double>(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;