Repository: pytorch
Issue: Problematic handling of NaN and inf in grid_sample, causing segfaults, corrupted CUDA memory, and incorrect results
Labels: high priority,module: crash,module: cuda,triaged,module: interpolation

_This issue is an expansion of the issue reported in https://github.com/pytorch/pytorch/issues/19826.
The discussion there diagnoses the segfault that occurs in the vectorized 2D CPU kernel. This issue covers the wider problematic handling of `NaN` and `inf` in all versions of `grid_sample` kernels. For details on `inf`, see the comment below._

### Summary

The `grid_sample` function does not have proper handling of `NaN` values for in its grid input.

The 2D CPU version segfaults under certain conditions and parameters, as described in https://github.com/pytorch/pytorch/issues/19826, and with simplified examples below.

The other `grid_sample` kernels (3D CPU, and 2D/3D CUDA) do not segfault, but produce incorrect results under certain conditions when the grid contains a `NaN` value.

Proper handling would place a `NaN` in the output for every grid location that has a `NaN`.

### Segmentation fault in the CPU 2D kernel

This is covered and diagnosed by @SsnL at https://github.com/pytorch/pytorch/issues/19826, but I want to provide a simple example to reproduce the segfault behavior, and expand on the exact conditions in which it occurs.

- Here is a simple example to reproduce the segmentation fault:
```python
>>> image = torch.rand(1, 1, 3, 3, device='cpu')
>>> grid = torch.rand(1, 3, 3, 2, device='cpu')
>>> grid[:,1,1,0] = float('nan')
>>> torch.nn.functional.grid_sample(image, grid, padding_mode='border')
Segmentation fault (core dumped)
```

- This segfault does not, however, happen if both components of a grid point are `NaN`.
Example:
```python
>>> image = torch.rand(1, 1, 3, 3, device='cpu')
>>> grid = torch.rand(1, 3, 3, 2, device='cpu')
>>> grid[:,1,1,:] = float('nan')
>>> torch.nn.functional.grid_sample(image, grid, padding_mode='border')
tensor([[[[0.2587, 0.1807, 0.2114],
          [0.1993,    nan, 0.2673],
          [0.2065, 0.1258, 0.2002]]]])
```
which is, in fact, the correct and desired behavior.

- The segfault occurs for padding modes `border` and `reflection`, but not for `zeros` (where it works correctly).


### The CUDA kernels (both 2D and 3D)

The CUDA kernel does not segfault. However, in `border` padding mode, it produces an incorrect result as if the `NaN` value were a `-1`.
Example:
```python
>>> image = torch.arange(9, 0, -1, dtype=torch.float, device='cuda').view(1,1,3,3)
tensor([[[[9., 8., 7.],
          [6., 5., 4.],
          [3., 2., 1.]]]], device='cuda:0')
# set grid to identity. Note: for old versions, drop the align_corners option
>>> grid = torch.nn.functional.affine_grid(torch.tensor([[[1.,0.,0.],[0.,1.,0.]]], device='cuda'), (1,1,3,3), align_corners=True)
>>> grid[:,1,1,0] = float('nan')  # set the x-coordinate of the central grid point to NaN
>>> torch.nn.functional.grid_sample(image, grid, padding_mode='border', align_corners=True)
tensor([[[[9., 8., 7.],
          [6., 6., 4.],
          [3., 2., 1.]]]], device='cuda:0')
>>> grid[:,1,1,:] = float('nan')  # set both coordinates of the central grid point to NaN
>>> torch.nn.functional.grid_sample(image, grid, padding_mode='border', align_corners=True)
tensor([[[[9., 8., 7.],
          [6., 9., 4.],
          [3., 2., 1.]]]], device='cuda:0')
```
Notice the result at the central output pixel. It behaves as if the `NaN` values of the grid were actually `-1`. Unlike the `border` padding mode, however, the `zeros` and `reflection` modes work correctly (produce a `NaN` in that pixel).


### The 3D CPU kernel

- The 3D CPU implementation also does not segfault, but unlike in the CUDA version, a `NaN` value is effectively treated as if it were a `+1` under the `border` padding mode.
```python
>>> image = torch.arange(27, 0, -1, dtype=torch.float, device='cpu').view(1,1,3,3,3)
tensor([[[[[27., 26., 25.],
           [24., 23., 22.],
           [21., 20., 19.]],

          [[18., 17., 16.],
           [15., 14., 13.],
           [12., 11., 10.]],

          [[ 9.,  8.,  7.],
           [ 6.,  5.,  4.],
           [ 3.,  2.,  1.]]]]])
# set grid to identity. Note: for old versions, drop the align_corners option
>>> grid = torch.nn.functional.affine_grid(torch.tensor([[[1.,0.,0.,0.],[0.,1.,0.,0.],[0.,0.,1.,0.]]], device='cpu'), (1,1,3,3,3), align_corners=True)
>>> grid[:,1,1,1,0] = float('nan')  # set the x-coordinate of the central grid point to NaN
>>> torch.nn.functional.grid_sample(image, grid, padding_mode='border', align_corners=True)
tensor([[[[[27., 26., 25.],
           [24., 23., 22.],
           [21., 20., 19.]],

          [[18., 17., 16.],
           [15., 13., 13.],
           [12., 11., 10.]],

          [[ 9.,  8.,  7.],
           [ 6.,  5.,  4.],
           [ 3.,  2.,  1.]]]]])
>>> grid[:,1,1,1,:] = float('nan')  # set all 3 coordinates of the central grid point to NaN
>>> torch.nn.functional.grid_sample(image, grid, padding_mode='border', align_corners=True)
tensor([[[[[27., 26., 25.],
           [24., 23., 22.],
           [21., 20., 19.]],

          [[18., 17., 16.],
           [15.,  1., 13.],
           [12., 11., 10.]],

          [[ 9.,  8.,  7.],
           [ 6.,  5.,  4.],
           [ 3.,  2.,  1.]]]]])
```
Notice the result in the central output pixel is `13.` in the first case (as if the grid there is `[0,0,+1]`), and `1.` in the second case (as if the grid is `[+1,+1,+1]`). As mentioned above, the same thing on CUDA results in a `15.` and a `27.` (as if the grid at that point were `[0,0,-1]` and `[-1,-1,-1]`, respectively).
That output pixel should just be a `NaN`.

- The `zeros` and `reflection` padding modes on 3D CPU always produce a `0.` result in the output wherever there is a `NaN` in the grid. I am not yet sure how to explain this last one.

### Desired behavior

Every pixel in the output for which the corresponding grid point has a `NaN` in one of its components should come out to be a `NaN`.

An alternative behavior (not advocating for this - I'm just presenting it as an option, for completeness) is to fill in a border value wherever there is a `NaN`. For the `zero` padding mode, this would fill in a `0.`. For the `border` and `reflection` padding modes, it's not clear how this would work.

In any case, the behaviors should be standardized across the different kernels.

### (Partial) Diagnoses

The 2D CPU segfault issue is diagnosed in https://github.com/pytorch/pytorch/issues/19826, and I think I have a decent idea of what's going on in the CUDA and 3D CPU `border` mode. The 3D CPU `zeros` and `reflection` modes might need another look at the code to diagnose.

I thought it would be good to write out all these cases explicitly, as I did above, since it helps for reproducing these errors and fixing them.
Once they're fixed, this will also be a good list of test cases to verify that the issues are indeed fixed.

cc @ezyang @gchanan @zou3519 @bdhirsh @jbschlosser @anjali411 @ngimel