Repository: pytorch
Issue: [JIT] torch.jit.optimized_execution(True) greatly slows down some operations in PyTorch 1.8.0 
Labels: oncall: jit

## 🐛 Bug

In PyTorch 1.8.0 JIT recompiles some functions every time if input tensor changes its content (not the shape).

## To Reproduce
If I run the following code
```
optimize = True
device = 'cuda:0'
num_runs = 5


import time

import torch


def func(mask: torch.Tensor):
    H, W = mask.size()

    tensor = torch.zeros([H, W], device=mask.device)
    masked_view = tensor[mask]
    output = torch.stack([masked_view, masked_view + W + 1], dim=1)

    return output


jit_func = torch.jit.script(func)


def get_random_mask():
    mask = torch.randint(2, size=[1000, 1000], dtype=torch.bool, device=device)

    return mask


with torch.jit.optimized_execution(optimize):
    times = []
    for i in range(num_runs):
        mask = get_random_mask()

        torch.cuda.synchronize(device)
        start = time.perf_counter()

        _ = jit_func(mask)

        torch.cuda.synchronize(device)
        elapsed_time = time.perf_counter() - start

        times.append(elapsed_time)

print(f'PyTorch version: {torch.__version__}')
print(f'Optimized execution: {optimize}')
print(f"Times:")
print("\n".join([f"{x:.4f} sec." for x in times]))
```
I got the following results:
```
PyTorch version: 1.8.0
Optimized execution: False
Times:
0.0007 sec.
0.0002 sec.
0.0002 sec.
0.0002 sec.
0.0002 sec.

PyTorch version: 1.8.0
Optimized execution: True
Times:
0.0402 sec.
0.1237 sec.
0.1194 sec.
0.1202 sec.
0.1204 sec.

PyTorch version: 1.7.1+cu110
Optimized execution: True
Times:
0.0024 sec.
0.1230 sec.
0.0003 sec.
0.0002 sec.
0.0002 sec.

PyTorch version: 1.7.1+cu110
Optimized execution: False
Times:
0.0007 sec.
0.0003 sec.
0.0002 sec.
0.0002 sec.
0.0002 sec.

```

Evidently, PyTorch 1.8.0 recompiles this function for every new random mask, even though its shape is unchanged. 

## Expected behavior

JIT should not recompile this function for each new mask.

## Environment

PyTorch version: 1.8.0
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: 6.0.0-1ubuntu2 (tags/RELEASE_600/final)
CMake version: version 3.10.2

Python version: 3.8 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: 10.1.243
GPU models and configuration: GPU 0: GeForce RTX 2080 SUPER
Nvidia driver version: 460.32.03
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.20.1
[pip3] pytorch-lightning==1.1.4
[pip3] torch==1.8.0
[pip3] torchvision==0.9.0
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               11.1.1               h6406543_8    conda-forge
[conda] mkl                       2020.2                      256  
[conda] mkl-service               2.3.0            py38he904b0f_0  
[conda] mkl_fft                   1.3.0            py38h54f3939_0  
[conda] mkl_random                1.1.1            py38h0573a6f_0  
[conda] numpy                     1.20.1                   pypi_0    pypi
[conda] pytorch                   1.8.0           py3.8_cuda11.1_cudnn8.0.5_0    pytorch
[conda] pytorch-lightning         1.1.4                    pypi_0    pypi
[conda] torchvision               0.9.0                py38_cu111    pytorch



cc @gmagogsfm