Guard roi_align/ps_roi_align bilinear sampling against non-finite boxes

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)

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;

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<T> pc;
             pc.pos1 = 0;

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;

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;

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;

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;