Repository: pytorch
Issue: [JIT] JIT Autograd is saving more results than are necessary for bprop like the result of RELU when it is an intermediate
Labels: oncall: jit

## 🐛 Bug
For the operations of RELU + DROPOUT, the intermediate result of the RELU is not needed but it is saved off none-the-less which is an unnecessary performance hit when the goal of fusion is to remove extra reads and writes to global memory.

Here is the Pytorch IR for RELU + Dropout:

For `FusionGroup_0`, the result `%16` does not need to be saved in this instance.

```
graph(%x.1 : Half(*, *, requires_grad=1, device=cuda:0),
      %prob.1 : float):
  %2 : str = prim::Constant[value="dropout probability has to be between 0 and 1, but got {}"]() # /opt/conda/lib/python3.8/site-packages/torch/nn/functional.py:979:25
  %3 : float = prim::Constant[value=1.]() # /opt/conda/lib/python3.8/site-packages/torch/nn/functional.py:978:21
  %4 : float = prim::Constant[value=0.]() # dropout.py:18:25
  %5 : bool = prim::Constant[value=1]() # dropout.py:19:42
  %7 : bool = aten::lt(%prob.1, %4) # /opt/conda/lib/python3.8/site-packages/torch/nn/functional.py:978:7
  %8 : bool = prim::If(%7) # /opt/conda/lib/python3.8/site-packages/torch/nn/functional.py:978:7
    block0():
      -> (%5)
    block1():
      %9 : bool = aten::gt(%prob.1, %3) # /opt/conda/lib/python3.8/site-packages/torch/nn/functional.py:978:17
      -> (%9)
   = prim::If(%8) # /opt/conda/lib/python3.8/site-packages/torch/nn/functional.py:978:4
    block0():
      %10 : str = aten::format(%2, %prob.1) # /opt/conda/lib/python3.8/site-packages/torch/nn/functional.py:979:25
       = prim::RaiseException(%10) # /opt/conda/lib/python3.8/site-packages/torch/nn/functional.py:979:8
      -> ()
    block1():
      -> ()
  %13 : Half(*, *, requires_grad=1, device=cuda:0) = prim::DifferentiableGraph_0(%prob.1, %x.1)
  return (%13)
with prim::DifferentiableGraph_0 = graph(%1 : float,
      %4 : Half(*, *, requires_grad=1, device=cuda:0)):
  %67 : bool = prim::Constant[value=1]()
  %10 : int = prim::Constant[value=1]() # <string>:103:64
  %9 : float = prim::Constant[value=1.]() # <string>:100:18
  %8 : None = prim::Constant()
  %p1m.3 : float = aten::sub(%9, %1) # <string>:100:18
  %22 : float = aten::div(%9, %p1m.3) # <string>:104:57
  %res.1 : Half(*, *, device=cuda:0), %mask.4 : Bool(*, *, device=cuda:0), %70 : Half(*, *, requires_grad=1, device=cuda:0) = prim::FusionGroup_0(%4, %22, %p1m.3)
  return (%res.1, %70, %8, %9, %10, %67, %p1m.3, %mask.4, %p1m.3)
with prim::FusionGroup_0 = graph(%14 : Half(*, *, requires_grad=1, device=cuda:0),
      %1 : float,
      %9 : float):
  %12 : int = prim::Constant[value=1]() # <string>:103:64
  %11 : None = prim::Constant()
  %15 : float = prim::Constant[value=0.]()
  %16 : Half(*, *, requires_grad=1, device=cuda:0) = aten::threshold(%14, %15, %15) # <string>:49:19
  %13 : Half(*, *, device=cuda:0) = aten::rand_like(%16, %11, %11, %11, %11, %12) # <string>:103:27
  %mask.4 : Bool(*, *, device=cuda:0) = aten::lt(%13, %9) # <string>:103:27
  %7 : Half(*, *, device=cuda:0) = aten::type_as(%mask.4, %16) # <string>:104:26
  %5 : Half(*, *, device=cuda:0) = aten::mul(%7, %16) # <string>:104:26
  %res.1 : Half(*, *, device=cuda:0) = aten::mul(%5, %1) # <string>:104:26
  return (%res.1, %mask.4, %16)
```

## To Reproduce

Script to reproduce:

```
import torch
import torch.nn as nn
import torch.nn.functional as F
import argparse

parser = argparse.ArgumentParser(description='Relu-Dropout JIT Test')
parser.add_argument('--bad-jit', action='store_true', help='Use new executor that cause Relu not to fuse with dropout.')
args = parser.parse_args()

assert hasattr(torch._C, '_jit_set_profiling_executor'), "Old JIT behavior doesn't exist!"
if not args.bad_jit :
    torch._C._jit_set_profiling_executor(False)
    torch._C._jit_set_profiling_mode(False)

@torch.jit.script
def jit_relu_dropout(x, prob) :
    # type: (Tensor, float) -> Tensor
    out = F.threshold(x, 0., 0.)
    out = F.dropout(out, p=prob, training=True)
    return out 
                                                                                    
inputs = torch.randn(5120, 1024, dtype=torch.float16, device=torch.device("cuda")).requires_grad_(True)

for x in range(0, 5) :
    outputs = jit_relu_dropout(inputs, 0.1)
    if x == 1:
        print(jit_relu_dropout.graph_for(inputs, 0.1))
```

## Environment

PyTorch version: 1.7.0a0+9093ffa
Is debug build: True
CUDA used to build PyTorch: 11.0
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.4 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.14.0

Python version: 3.8 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration:
GPU 0: Tesla V100-DGXS-16GB
GPU 1: Tesla V100-DGXS-16GB
GPU 2: Tesla V100-DGXS-16GB
GPU 3: Tesla V100-DGXS-16GB

Nvidia driver version: 440.33.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.0.2
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.0.2
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.0.2
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.0.2
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.0.2
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.0.2
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.0.2
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip] numpy==1.19.1
[pip] torch==1.7.0a0
[conda] magma-cuda102             2.5.2                         1    pytorch
[conda] mkl                       2019.1                      144    
[conda] mkl-include               2019.1                      144    
[conda] nomkl                     3.0                           0    
[conda] numpy                     1.19.1           py38h30dfecb_0    
[conda] numpy-base                1.19.1           py38h75fe3a5_0    
[conda] torch                     1.7.0a0                  pypi_0    pypi


cc @gmagogsfm